[llvm] [AMDGPU] Fix LDS access via flat pointer argument in amdgpu-sw-lower-lds (PR #209842)

[via llvm-commits]

================
@@ -287,6 +303,99 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
   }
 }
 
+// True if flat pointer V is derived from an LDS-carrying argument in LDSArgs,
+// looking through GEP/bitcast/phi/select; Visited breaks phi cycles. When
+// AcceptLocalCast is set, a flat<-local addrspacecast also counts as an origin.
+// Detection (getNonKernelsWithLDSFlatArguments) sets it, since that cast is how
+// LDS first enters a flat pointer; the rewrite
+// (getFlatPtrForRoundTripLDSAccess) leaves it clear so only argument origins
+// qualify.
+static bool flatPtrDerivesFromLDS(Value *V,
+                                  const SmallPtrSetImpl<Argument *> &LDSArgs,
+                                  SmallPtrSetImpl<Value *> &Visited,
+                                  bool AcceptLocalCast) {
+  if (!V->getType()->isPointerTy() || !Visited.insert(V).second)
+    return false;
+  if (auto *A = dyn_cast<Argument>(V))
+    return LDSArgs.contains(A);
+  auto *Op = dyn_cast<Operator>(V);
+  if (!Op)
+    return false;
+  switch (Op->getOpcode()) {
+  case Instruction::AddrSpaceCast: {
+    if (!AcceptLocalCast)
+      return false;
+    Value *Src = Op->getOperand(0);
+    if (Src->getType()->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
+      return true;
+    return flatPtrDerivesFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
+  }
+  case Instruction::GetElementPtr:
+  case Instruction::BitCast:
+    return flatPtrDerivesFromLDS(Op->getOperand(0), LDSArgs, Visited,
+                                 AcceptLocalCast);
+  case Instruction::PHI:
+    for (Value *In : cast<PHINode>(Op)->incoming_values())
+      if (flatPtrDerivesFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
+        return true;
+    return false;
+  case Instruction::Select: {
+    auto *SI = cast<SelectInst>(Op);
+    return flatPtrDerivesFromLDS(SI->getTrueValue(), LDSArgs, Visited,
+                                 AcceptLocalCast) ||
+           flatPtrDerivesFromLDS(SI->getFalseValue(), LDSArgs, Visited,
+                                 AcceptLocalCast);
+  }
+  default:
+    return false;
+  }
+}
+
+void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
+  // A kernel or non-kernel may pass lowered LDS storage to a non-kernel as a
+  // flat pointer instead of addrspace(3). Record, per callee, which flat
+  // parameters carry LDS. Must run before lowering rewrites the call-site
+  // casts.
+  auto &FlatArgMap = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg;
+  bool Changed = true;
+  while (Changed) {
+    Changed = false;
+    for (Function &F : M) {
+      if (F.isDeclaration())
+        continue;
+      SmallPtrSet<Argument *, 4> CallerLDSArgs;
+      if (auto It = FlatArgMap.find(&F); It != FlatArgMap.end())
+        CallerLDSArgs = It->second;
+      for (BasicBlock &BB : F) {
+        for (Instruction &I : BB) {
+          auto *CB = dyn_cast<CallBase>(&I);
+          if (!CB)
+            continue;
+          Function *Callee = CB->getCalledFunction();
+          if (!Callee || Callee->isDeclaration() || AMDGPU::isKernel(*Callee))
+            continue;
+          unsigned NumArgs =
+              std::min(CB->arg_size(), (unsigned)Callee->arg_size());
----------------
skc7 wrote:

Thanks for feedback. Updated in latest patch.

https://github.com/llvm/llvm-project/pull/209842