[SYCL][Fusion] Adapt internalization to opaque pointers. (#10882)

Fixes `sycl/test-e2e/KernelFusion/internalize_array_wrapper.cpp` which
was temporarily deactivated in #9828.

Switching to opaque pointers meant that accesses to the data structure
in the testcase were no longer represented as one "deep" GEP, but rather
a series of GEPs into the different nesting levels of the structure.
Rather than bailing out, we now keep track of GEPs that index into an
aggregate object, and exclude their users from the modulo-`LocalSize`
rewriting (the pointer types still have to be changed to the desired
target address space, though).

In addition, perform `-simplifycfg` before internalization in order to
get rid of unreachable code which, in this particular case, contained
Phi nodes on the candidate accessor.

---------

Signed-off-by: Julian Oppermann <[email protected]>

Author: jopperm

sycl-fusion/jit-compiler/lib/fusion/FusionPipeline.cpp

@@ -105,6 +105,9 @@ FusionPipeline::runFusionPasses(Module &Mod, SYCLModuleInfo &InputInfo,
     // Ideally, the static compiler should have performed that job.
     const unsigned FlatAddressSpace = getFlatAddressSpace(Mod);
     FPM.addPass(InferAddressSpacesPass(FlatAddressSpace));
+    // Run CFG simplification to prevent unreachable code from obscuring
+    // internalization opportunities.
+    FPM.addPass(SimplifyCFGPass{});
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
   }
   // Run dataflow internalization and runtime constant propagation.

sycl-fusion/passes/internalization/Internalization.cpp

@@ -58,10 +58,10 @@ struct SYCLInternalizerImpl {
   ///   - Promote the function and call the new function instead,
   ///   keeping the original function.
   /// - The value appears in a load/store operation: Do nothing
-  void promoteValue(Value *Val, std::size_t LocalSize) const;
+  void promoteValue(Value *Val, std::size_t LocalSize, bool InAggregate) const;
 
   void promoteGEPI(GetElementPtrInst *GEPI, const Value *Val,
-                   std::size_t LocalSize) const;
+                   std::size_t LocalSize, bool InAggregate) const;
 
   void promoteCall(CallBase *C, const Value *Val, std::size_t LocalSize) const;
 
@@ -81,23 +81,26 @@ struct SYCLInternalizerImpl {
   ///
   /// Check that an value can be promoted.
   /// For GEP and Call instructions, delegate to the specific implementations.
+  /// \p InAggregate indicates that at least one GEP instruction addressing into
+  /// an aggregate object was encountered, hence \p Val no longer represents a
+  /// pure offset computation on the original candidate argument.
   /// For address-space casts, pointer-to-int conversions and unknown users,
   /// return an error.
-  Error canPromoteValue(Value *Val, size_t LocalSize) const;
+  Error canPromoteValue(Value *Val, size_t LocalSize, bool InAggregate) const;
 
   ///
-  /// Check that the operand of a GEP can be promoted.
-  /// If the GEP uses more than one index, return an error.
-  /// Otherwise, check if the GEP itself can be promoted in its users.
+  /// Check that the operand of a GEP can be promoted to its users, and
+  /// propagate whether it represents a pointer into an aggregate object.
   Error canPromoteGEP(GetElementPtrInst *GEPI, const Value *Val,
-                      size_t LocalSize) const;
+                      size_t LocalSize, bool InAggregate) const;
 
   ///
   /// Check if operand to a function call can be promoted.
-  /// If the function returns a pointer, return an error.
-  /// Otherwise, check if the corresponding formal parameter can be promoted in
-  /// the function body.
-  Error canPromoteCall(CallBase *C, const Value *Val, size_t LocalSize) const;
+  /// If the function returns a pointer, or the operand points into an aggregate
+  /// object, return an error. Otherwise, check if the corresponding formal
+  /// parameter can be promoted in the function body.
+  Error canPromoteCall(CallBase *C, const Value *Val, size_t LocalSize,
+                       bool InAggregate) const;
 
   Error checkArgsPromotable(Function *F,
                             SmallVectorImpl<size_t> &PromoteArgSizes) const;
@@ -212,7 +215,8 @@ getUsagesInternalization(const User *U, const Value *V, std::size_t LocalSize) {
 }
 
 Error SYCLInternalizerImpl::canPromoteCall(CallBase *C, const Value *Val,
-                                           size_t LocalSize) const {
+                                           size_t LocalSize,
+                                           bool InAggregate) const {
   if (isa<PointerType>(C->getType())) {
     // With opaque pointers, we do not have the necessary information to compare
     // the element-type of the pointer returned by the function and the element
@@ -222,6 +226,12 @@ Error SYCLInternalizerImpl::canPromoteCall(CallBase *C, const Value *Val,
         inconvertibleErrorCode(),
         "It is not safe to promote a called function which returns a pointer.");
   }
+  if (InAggregate) {
+    return createStringError(
+        inconvertibleErrorCode(),
+        "Promotion of a pointer into an aggregate object to a called function "
+        "is currently not supported.");
+  }
 
   SmallVector<size_t> InternInfo = getUsagesInternalization(C, Val, LocalSize);
   assert(!InternInfo.empty() && "Value must be used at least once");
@@ -232,27 +242,23 @@ Error SYCLInternalizerImpl::canPromoteCall(CallBase *C, const Value *Val,
 }
 
 Error SYCLInternalizerImpl::canPromoteGEP(GetElementPtrInst *GEPI,
-                                          const Value *Val,
-                                          size_t LocalSize) const {
+                                          const Value *Val, size_t LocalSize,
+                                          bool InAggregate) const {
   if (cast<PointerType>(GEPI->getType())->getAddressSpace() == AS) {
     // If the GEPI is already using the correct address-space, no change is
     // required.
     return Error::success();
   }
-  if (GEPI->getNumIndices() != 1 &&
-      std::any_of(GEPI->user_begin(), GEPI->user_end(), [](const auto *User) {
-        return isa<GetElementPtrInst>(User);
-      })) {
-    return createStringError(inconvertibleErrorCode(),
-                             "Only one index expected in source of "
-                             "promotable GEP instruction pointer argument");
-  }
-  // Recurse to check all users of the GEP.
-  return canPromoteValue(GEPI, LocalSize);
+  // Recurse to check all users of the GEP. We are either already in
+  // `InAggregate` mode, or inspect the current instruction. Recall that a GEP's
+  // first index is used to step through the base pointer, whereas any
+  // additional indices represent addressing into an aggregrate type.
+  return canPromoteValue(GEPI, LocalSize,
+                         InAggregate || GEPI->getNumIndices() >= 2);
 }
 
-Error SYCLInternalizerImpl::canPromoteValue(Value *Val,
-                                            size_t LocalSize) const {
+Error SYCLInternalizerImpl::canPromoteValue(Value *Val, size_t LocalSize,
+                                            bool InAggregate) const {
   for (auto *U : Val->users()) {
     auto *I = dyn_cast<Instruction>(U);
     if (!I) {
@@ -272,13 +278,14 @@ Error SYCLInternalizerImpl::canPromoteValue(Value *Val,
     case Instruction::Call:
     case Instruction::Invoke:
     case Instruction::CallBr:
-      if (auto Err = canPromoteCall(cast<CallBase>(I), Val, LocalSize)) {
+      if (auto Err =
+              canPromoteCall(cast<CallBase>(I), Val, LocalSize, InAggregate)) {
         return Err;
       }
       break;
     case Instruction::GetElementPtr:
-      if (auto Err =
-              canPromoteGEP(cast<GetElementPtrInst>(I), Val, LocalSize)) {
+      if (auto Err = canPromoteGEP(cast<GetElementPtrInst>(I), Val, LocalSize,
+                                   InAggregate)) {
         return Err;
       }
       break;
@@ -316,7 +323,7 @@ Error SYCLInternalizerImpl::checkArgsPromotable(
       PromoteArgSizes[Index] = 0;
       continue;
     }
-    if (auto Err = canPromoteValue(Arg, LocalSize)) {
+    if (auto Err = canPromoteValue(Arg, LocalSize, /*InAggregate=*/false)) {
       // Set the local size to 0 to indicate that this argument should not be
       // promoted.
       PromoteArgSizes[Index] = 0;
@@ -360,30 +367,37 @@ void SYCLInternalizerImpl::promoteCall(CallBase *C, const Value *Val,
 }
 
 void SYCLInternalizerImpl::promoteGEPI(GetElementPtrInst *GEPI,
-                                       const Value *Val,
-                                       std::size_t LocalSize) const {
+                                       const Value *Val, std::size_t LocalSize,
+                                       bool InAggregate) const {
   // Not PointerType is unreachable. Other case is catched in caller.
   if (cast<PointerType>(GEPI->getType())->getAddressSpace() != AS) {
-    remapIndices(GEPI, LocalSize);
+    if (!InAggregate)
+      remapIndices(GEPI, LocalSize);
     auto *ValTy = cast<PointerType>(Val->getType());
     GEPI->mutateType(PointerType::getWithSamePointeeType(
         cast<PointerType>(GEPI->getType()), ValTy->getAddressSpace()));
-    return promoteValue(GEPI, LocalSize);
+    // Recurse to promote to all users of the GEP. We are either already in
+    // `InAggregate` mode, or inspect the current instruction. Recall that a
+    // GEP's first index is used to step through the base pointer, whereas any
+    // additional indices represent addressing into an aggregrate type.
+    return promoteValue(GEPI, LocalSize,
+                        InAggregate || GEPI->getNumIndices() >= 2);
   }
 }
 
-void SYCLInternalizerImpl::promoteValue(Value *Val,
-                                        std::size_t LocalSize) const {
+void SYCLInternalizerImpl::promoteValue(Value *Val, std::size_t LocalSize,
+                                        bool InAggregate) const {
   for (auto *U : Val->users()) {
     auto *I = cast<Instruction>(U);
     switch (I->getOpcode()) {
     case Instruction::Call:
     case Instruction::Invoke:
     case Instruction::CallBr:
+      assert(!InAggregate);
       promoteCall(cast<CallBase>(I), Val, LocalSize);
       break;
     case Instruction::GetElementPtr:
-      promoteGEPI(cast<GetElementPtrInst>(I), Val, LocalSize);
+      promoteGEPI(cast<GetElementPtrInst>(I), Val, LocalSize, InAggregate);
       break;
     case Instruction::Load:
     case Instruction::Store:
@@ -535,7 +549,7 @@ Function *SYCLInternalizerImpl::promoteFunctionArgs(
     if (CreateAllocas) {
       Arg = replaceByNewAlloca(cast<Argument>(Arg), AS, LocalSize);
     }
-    promoteValue(Arg, LocalSize);
+    promoteValue(Arg, LocalSize, /*InAggregate=*/false);
   }
 
   TargetInfo.updateAddressSpaceMetadata(NewF, PromoteToLocal, AS);

sycl-fusion/test/internalization/abort-promote-call.ll

@@ -0,0 +1,36 @@
+; RUN: opt -load-pass-plugin %shlibdir/SYCLKernelFusion%shlibext \
+; RUN: -passes=sycl-internalization --sycl-info-path %S/abort-kernel-info.yaml -S %s | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v24:32:32-v32:32:32-v48:64:64-v64:64:64-v96:128:128-v128:128:128-v192:256:256-v256:256:256-v512:512:512-v1024:1024:1024"
+target triple = "spir64-unknown-unknown"
+
+; Function Attrs: noinline
+define void @fun(ptr addrspace(1) %arg) #0 {
+  ret void
+}
+
+%struct = type { i32, i32, i32 }
+
+; CHECK-LABEL: define {{[^@]+}}@fused_0
+; CHECK-SAME: (ptr addrspace(1) align 4 %[[ACC:.*]])
+define spir_kernel void @fused_0(ptr addrspace(1) align 4 %acc) !kernel_arg_addr_space !12 !kernel_arg_access_qual !13 !kernel_arg_type !14 !kernel_arg_type_qual !15 !kernel_arg_base_type !14 !kernel_arg_name !16 !sycl.kernel.promote !17 !sycl.kernel.promote.localsize !18 {
+; Scenario: Test private internalization is not performed when pointers into
+; aggregate object are passed to function calls.
+
+; CHECK-NOT:  alloca [1 x %struct]
+  %gep1 = getelementptr %struct, ptr addrspace(1) %acc, i64 17
+  %gep2 = getelementptr %struct, ptr addrspace(1) %gep1, i64 0, i32 2
+  call void @fun(ptr addrspace(1) %gep2)
+  store i32 42, ptr addrspace(1) %gep2
+  ret void
+}
+
+attributes #0 = { noinline }
+
+!12 = !{i32 1}
+!13 = !{!"none"}
+!14 = !{!"ptr"}
+!15 = !{!""}
+!16 = !{!"acc"}
+!17 = !{!"private"}
+!18 = !{i64 1}