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

Author: bratpiorka

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

@@ -478,10 +478,9 @@ struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
     // 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);
+    // Writes to the pointer are UB w/ __grid_constant__, but param space is
+    // read-only on CUDA so we need to force a copy.
+    return PI.setAborted(&SI);
   }
 
   void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
@@ -529,10 +528,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 +589,9 @@ 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.
+  if ((IsGridConstant || 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 +609,14 @@ 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.