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

[Pablo Reble via llvm-commits]

================
@@ -777,17 +919,78 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
       MI->eraseFromParent();
     } else if (AddrSpaceCastInst *ASC = dyn_cast<AddrSpaceCastInst>(Inst)) {
       Value *AIOperand = ASC->getPointerOperand();
-      Value *Replacement =
-          getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, AIOperand);
+      Value *Replacement = TranslatePtr(AIOperand);
+      if (!Replacement)
+        continue;
       Value *NewAI = IRB.CreateAddrSpaceCast(Replacement, ASC->getType());
       // Note: No need to add the instruction to AsanInfo instructions to be
       // instrumented list. FLAT_ADDRESS ptr would have been already
       // instrumented by asan pass prior to this pass.
       ASC->replaceAllUsesWith(NewAI);
       ASC->eraseFromParent();
+      MaybeDeadPtrs.push_back(AIOperand);
     } else
       report_fatal_error("Unimplemented LDS lowering instruction");
   }
+
+  // Drop the now-dead flat->local casts and local GEPs.
+  RecursivelyDeleteTriviallyDeadInstructionsPermissive(MaybeDeadPtrs);
+}
+
+Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
+                                                         Value *LDSPtr) {
+  // If LDSPtr is lowered storage reached via a flat argument, i.e.
+  // addrspacecast(flat->local) optionally followed by local GEPs, return the
+  // equivalent flat pointer into the global backing buffer. Otherwise nullptr,
+  // so the caller falls back to the base/offset-table translation.
+  auto It = FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.find(Func);
+  if (It == FuncLDSAccessInfo.NonKernelsWithLDSFlatArg.end())
+    return nullptr;
+  const SmallPtrSet<Argument *, 4> &LDSArgs = It->second;
+
+  // Peel off local GEPs sitting between the cast and the access.
+  SmallVector<GEPOperator *, 4> GEPs;
+  Value *Cur = LDSPtr;
+  while (auto *GEP = dyn_cast<GEPOperator>(Cur)) {
+    if (GEP->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
+      return nullptr;
+    GEPs.push_back(GEP);
+    Cur = GEP->getPointerOperand();
+  }
+
+  auto *ASC = dyn_cast<Operator>(Cur);
+  if (!ASC || ASC->getOpcode() != Instruction::AddrSpaceCast)
+    return nullptr;
+  Value *FlatSrc = ASC->getOperand(0);
+  if (FlatSrc->getType()->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
+    return nullptr;
----------------
reble wrote:

```suggestion
  auto *ASC = dyn_cast<AddrSpaceCastOperator>(Cur);
  if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
      ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
    return nullptr;
  Value *FlatSrc = ASC->getPointerOperand();
```
Using `AddrSpaceCastOperator` instead of manual handling for consistency?

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