TODO [SYCL] Fix circular dependency in byval parameter handling in NVPTXLowerArgs

Author: bratpiorka

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

@@ -79,15 +79,15 @@
 //
 //    define void @foo({i32*, i32*}* byval %input) {
 //      %b_param = addrspacecat ptr %input to ptr addrspace(101)
-//      %b_ptr = getelementptr {ptr, ptr}, ptr addrspace(101) %b_param, i64 0, i32 1
-//      %b = load ptr, ptr addrspace(101) %b_ptr
-//      %b_global = addrspacecast ptr %b to ptr addrspace(1)
-//      ; use %b_generic
+//      %b_ptr = getelementptr {ptr, ptr}, ptr addrspace(101) %b_param, i64 0,
+//      i32 1 %b = load ptr, ptr addrspace(101) %b_ptr %b_global = addrspacecast
+//      ptr %b to ptr addrspace(1) ; use %b_generic
 //    }
 //
-//    Create a local copy of kernel byval parameters used in a way that *might* mutate
-//    the parameter, by storing it in an alloca. Mutations to "grid_constant" parameters
-//    are undefined behaviour, and don't require local copies.
+//    Create a local copy of kernel byval parameters used in a way that *might*
+//    mutate the parameter, by storing it in an alloca. Mutations to
+//    "grid_constant" parameters are undefined behaviour, and don't require
+//    local copies.
 //
 //    define void @foo(ptr byval(%struct.s) align 4 %input) {
 //       store i32 42, ptr %input
@@ -124,11 +124,11 @@
 //
 //    define void @foo(ptr byval(%struct.s) %input) {
 //      %input1 = addrspacecast ptr %input to ptr addrspace(101)
-//      ; the following intrinsic converts pointer to generic. We don't use an addrspacecast
-//      ; to prevent generic -> param -> generic from getting cancelled out
-//      %input1.gen = call ptr @llvm.nvvm.ptr.param.to.gen.p0.p101(ptr addrspace(101) %input1)
-//      %call = call i32 @escape(ptr %input1.gen)
-//      ret void
+//      ; the following intrinsic converts pointer to generic. We don't use an
+//      addrspacecast ; to prevent generic -> param -> generic from getting
+//      cancelled out %input1.gen = call ptr
+//      @llvm.nvvm.ptr.param.to.gen.p0.p101(ptr addrspace(101) %input1) %call =
+//      call i32 @escape(ptr %input1.gen) ret void
 //    }
 //
 // TODO: merge this pass with NVPTXInferAddressSpaces so that other passes don't
@@ -475,13 +475,18 @@ struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
   }
 
   void visitStoreInst(StoreInst &SI) {
-    // Storing the pointer escapes it.
-    if (U->get() == SI.getValueOperand())
-      return PI.setEscapedAndAborted(&SI);
-    // Writes to the pointer are UB w/ __grid_constant__, but do not force a
-    // copy.
-    if (!IsGridConstant)
-      return PI.setAborted(&SI);
+    // Storing the pointer value (as opposed to storing through it) escapes it.
+    // For grid_constant params, this is allowed - we can pass the generic
+    // pointer. For non-grid-constant, this requires a copy.
+    if (U->get() == SI.getValueOperand()) {
+      if (IsGridConstant)
+        return PI.setEscaped(&SI);
+      else
+        return PI.setEscapedAndAborted(&SI);
+    }
+    // Writes through the pointer to param space are UB w/ __grid_constant__,
+    // and param space is read-only on CUDA, so we need to force a copy.
+    return PI.setAborted(&SI);
   }
 
   void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
@@ -529,10 +534,22 @@ void copyByValParam(Function &F, Argument &Arg) {
   // the use of the byval parameter with this alloca instruction.
   AllocA->setAlignment(
       Arg.getParamAlign().value_or(DL.getPrefTypeAlign(StructType)));
-  Arg.replaceAllUsesWith(AllocA);
 
+  // CRITICAL: Must create ArgInParam BEFORE replacing uses of Arg.
+  // createNVVMInternalAddrspaceWrap needs to use Arg as an operand.
   CallInst *ArgInParam = createNVVMInternalAddrspaceWrap(IRB, Arg);
 
+  // Replace all uses of Arg with AllocA, except the use in ArgInParam.
+  // We can't use replaceAllUsesWith because it would replace ArgInParam's
+  // operand too, creating a circular dependency.
+  SmallVector<Use *, 8> UsesToReplace;
+  for (Use &U : Arg.uses()) {
+    if (U.getUser() != ArgInParam)
+      UsesToReplace.push_back(&U);
+  }
+  for (Use *U : UsesToReplace)
+    U->set(AllocA);
+
   // Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
   // addrspacecast preserves alignment.  Since params are constant, this load
   // is definitely not volatile.
@@ -578,10 +595,12 @@ static void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg) {
 
   // We can't access byval arg directly and need a pointer. on sm_70+ we have
   // ability to take a pointer to the argument without making a local copy.
-  // However, we're still not allowed to write to it. If the user specified
-  // `__grid_constant__` for the argument, we'll consider escaped pointer as
-  // read-only.
-  if (IsGridConstant || (HasCvtaParam && ArgUseIsReadOnly)) {
+  // However, param space is read-only, so we can only use this optimization
+  // if the argument is not written to.
+  // Grid constant params can escape (pointer passed to functions) but cannot
+  // have stores. Non-grid-constant params must be fully read-only.
+  if ((IsGridConstant && !PI.isAborted()) ||
+      (HasCvtaParam && ArgUseIsReadOnly)) {
     LLVM_DEBUG(dbgs() << "Using non-copy pointer to " << *Arg << "\n");
     // Replace all argument pointer uses (which might include a device function
     // call) with a cast to the generic address space using cvta.param
@@ -599,10 +618,15 @@ static void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg) {
         ParamSpaceArg, IRB.getPtrTy(ADDRESS_SPACE_GENERIC),
         Arg->getName() + ".gen");
 
-    Arg->replaceAllUsesWith(GenericArg);
-
-    // Do not replace Arg in the cast to param space
-    ParamSpaceArg->setOperand(0, Arg);
+    // Collect all uses of Arg except the one in ParamSpaceArg, then replace
+    // them
+    SmallVector<Use *> UsesToReplace;
+    for (Use &U : Arg->uses()) {
+      if (U.getUser() != ParamSpaceArg)
+        UsesToReplace.push_back(&U);
+    }
+    for (Use *U : UsesToReplace)
+      U->set(GenericArg);
   } else
     copyByValParam(*Func, *Arg);
 }

sycl/test-e2e/FreeFunctionKernels/structs_with_special_types_as_kernel_paramters.cpp

@@ -1,9 +1,6 @@
 // RUN: %{build} -o %t.out
 // RUN: %{run} %t.out
 
-// XFAIL: target-nvidia
-// XFAIL-TRACKER: https://github.com/intel/llvm/issues/20908
-
 // This test verifies whether struct that contains either sycl::local_accesor or
 // sycl::accessor can be used with free function kernels extension.