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

via llvm-commits llvm-commits at lists.llvm.org
Fri Jul 31 04:39:32 PDT 2026


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

>From c256325bb70344fd067978a14ce77aa8f7d04af7 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Wed, 15 Jul 2026 22:57:15 +0530
Subject: [PATCH 1/4] [AMDGPU] Fix LDS access via flat pointer argument in
 amdgpu-sw-lower-lds

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp   | 82 +++++++++++++++++++
 .../amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll  | 73 +++++++++++++++++
 2 files changed, 155 insertions(+)
 create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 2b78094837c68..e2668272ea29c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -94,6 +94,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
@@ -204,6 +205,7 @@ class AMDGPUSwLowerLDS {
   void lowerNonKernelLDSAccesses(Function *Func,
                                  SetVector<GlobalVariable *> &LDSGlobals,
                                  NonKernelLDSParameters &NKLDSParams);
+  void lowerFlatToLocalRoundTrips(Function *Func);
   void
   updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
                                 Value *HiddenDynLDSSize,
@@ -790,6 +792,80 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
   }
 }
 
+void AMDGPUSwLowerLDS::lowerFlatToLocalRoundTrips(Function *Func) {
+  // After SW LDS lowering, storage that used to live in LDS now lives in global
+  // memory, but it may still be handed to a callee as a flat (generic) pointer.
+  // Such a callee can re-derive an LDS pointer with a round-trip pattern:
+  //   %p = getelementptr ..., ptr %flat_arg        ; flat, backed by global mem
+  //   %c = addrspacecast ptr %p to ptr addrspace(3)
+  //   load/store ... ptr addrspace(3) %c
+  // The addrspace(3) access would then hit dead LDS. Redo the access through
+  // the flat source pointer, which still points at the global backing buffer.
+  SmallVector<Instruction *> ToErase;
+  SmallPtrSet<AddrSpaceCastInst *, 8> Casts;
+  for (BasicBlock &BB : *Func) {
+    for (Instruction &Inst : BB) {
+      Value *Ptr = nullptr;
+      if (auto *LI = dyn_cast<LoadInst>(&Inst))
+        Ptr = LI->getPointerOperand();
+      else if (auto *SI = dyn_cast<StoreInst>(&Inst))
+        Ptr = SI->getPointerOperand();
+      else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst))
+        Ptr = RMW->getPointerOperand();
+      else if (auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst))
+        Ptr = XCHG->getPointerOperand();
+      else
+        continue;
+
+      auto *ASC = dyn_cast<AddrSpaceCastInst>(Ptr);
+      if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
+          ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
+        continue;
+
+      // Only collapse pointers that originate from a flat function argument,
+      // i.e. storage passed in from a caller. Genuine in-function LDS pointers
+      // are lowered through the base/offset table machinery instead.
+      Value *FlatPtr = ASC->getPointerOperand();
+      if (!isa<Argument>(getUnderlyingObject(FlatPtr)))
+        continue;
+
+      IRB.SetInsertPoint(&Inst);
+      Instruction *NewInst = nullptr;
+      if (auto *LI = dyn_cast<LoadInst>(&Inst)) {
+        NewInst = IRB.CreateLoad(LI->getType(), FlatPtr, LI->getProperties());
+      } else if (auto *SI = dyn_cast<StoreInst>(&Inst)) {
+        NewInst = IRB.CreateStore(SI->getValueOperand(), FlatPtr,
+                                  SI->getProperties());
+      } else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst)) {
+        auto *NewRMW = IRB.CreateAtomicRMW(
+            RMW->getOperation(), FlatPtr, RMW->getValOperand(), RMW->getAlign(),
+            RMW->getOrdering(), RMW->getSyncScopeID());
+        NewRMW->setVolatile(RMW->isVolatile());
+        NewInst = NewRMW;
+      } else {
+        auto *XCHG = cast<AtomicCmpXchgInst>(&Inst);
+        auto *NewXCHG = IRB.CreateAtomicCmpXchg(
+            FlatPtr, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
+            XCHG->getAlign(), XCHG->getSuccessOrdering(),
+            XCHG->getFailureOrdering(), XCHG->getSyncScopeID());
+        NewXCHG->setVolatile(XCHG->isVolatile());
+        NewInst = NewXCHG;
+      }
+      // Flat access into the global backing buffer must be sanitized.
+      AsanInfo.Instructions.insert(NewInst);
+      Inst.replaceAllUsesWith(NewInst);
+      ToErase.push_back(&Inst);
+      Casts.insert(ASC);
+    }
+  }
+  for (Instruction *I : ToErase)
+    I->eraseFromParent();
+  // Drop the now-dead round-trip casts.
+  for (AddrSpaceCastInst *ASC : Casts)
+    if (ASC->use_empty())
+      ASC->eraseFromParent();
+}
+
 void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
   auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
   Type *Int64Ty = IRB.getInt64Ty();
@@ -1323,6 +1399,12 @@ bool AMDGPUSwLowerLDS::run() {
   if (!Changed)
     return Changed;
 
+  // Fix up callees that still access lowered storage through the
+  // flat-arg -> addrspace(3) round-trip pattern.
+  for (Function &F : M)
+    if (!F.isDeclaration())
+      lowerFlatToLocalRoundTrips(&F);
+
   for (auto &GV : make_early_inc_range(M.globals())) {
     if (AMDGPU::isLDSVariableToLower(GV)) {
       // probably want to remove from used lists
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
new file mode 100644
index 0000000000000..646371a5bde24
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
@@ -0,0 +1,73 @@
+; RUN: opt < %s -passes=amdgpu-sw-lower-lds -amdgpu-asan-instrument-lds=false -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s
+
+; A non-kernel receives lowered LDS storage as a flat (generic) pointer and
+; re-derives an LDS pointer via a flat -> addrspace(3) round-trip. After SW LDS
+; lowering the storage lives in global memory, so the addrspace(3) access would
+; hit dead LDS. The round-trip must be collapsed into a flat access on the
+; incoming argument, which still points at the global backing buffer. Check the
+; store, load, atomicrmw and cmpxchg cases.
+
+ at lds_var = internal addrspace(3) global [64 x i32] poison, align 4
+
+define void @store_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @store_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT:     addrspace(3)
+; CHECK:         store i32 42, ptr [[STORAGE]], align 4
+; CHECK-NOT:     addrspace(3)
+; CHECK:         ret void
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  store i32 42, ptr addrspace(3) %cast, align 4
+  ret void
+}
+
+define i32 @load_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @load_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT:     addrspace(3)
+; CHECK:         [[VAL:%.*]] = load i32, ptr [[STORAGE]], align 4
+; CHECK-NOT:     addrspace(3)
+; CHECK:         ret i32 [[VAL]]
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  %val = load i32, ptr addrspace(3) %cast, align 4
+  ret i32 %val
+}
+
+define i32 @atomicrmw_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @atomicrmw_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT:     addrspace(3)
+; CHECK:         [[OLD:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
+; CHECK-NOT:     addrspace(3)
+; CHECK:         ret i32 [[OLD]]
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  %old = atomicrmw add ptr addrspace(3) %cast, i32 1 seq_cst, align 4
+  ret i32 %old
+}
+
+define i32 @cmpxchg_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define i32 @cmpxchg_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
+; CHECK-NOT:     addrspace(3)
+; CHECK:         [[RES:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
+; CHECK-NOT:     addrspace(3)
+; CHECK:         [[VAL:%.*]] = extractvalue { i32, i1 } [[RES]], 0
+; CHECK:         ret i32 [[VAL]]
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  %res = cmpxchg ptr addrspace(3) %cast, i32 0, i32 1 seq_cst seq_cst, align 4
+  %val = extractvalue { i32, i1 } %res, 0
+  ret i32 %val
+}
+
+define amdgpu_kernel void @kernel() sanitize_address {
+; CHECK-LABEL: define amdgpu_kernel void @kernel(
+  %flat = addrspacecast ptr addrspace(3) @lds_var to ptr
+  call void @store_case(ptr %flat)
+  %l = call i32 @load_case(ptr %flat)
+  %a = call i32 @atomicrmw_case(ptr %flat)
+  %c = call i32 @cmpxchg_case(ptr %flat)
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"nosanitize_address", i32 1}

>From 25d57b23b33ff21483d4ec5e3c322f02bf1a22c5 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 17 Jul 2026 21:47:54 +0530
Subject: [PATCH 2/4] update

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp   | 325 +++++++++++++-----
 .../amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll  | 243 +++++++++++--
 2 files changed, 444 insertions(+), 124 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index e2668272ea29c..25c82883bd419 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -81,6 +81,15 @@
 //    The Offset into the base "SW LDS" is obtained from
 //    corresponding element in offset table. With this information, replacement
 //    value is obtained.
+//
+// Replacement of LDS accessed via flat pointer arguments:
+//    An LDS pointer can reach a non-kernel as a flat (generic) pointer
+//    argument, e.g. after an addrspacecast of an LDS global to flat. A
+//    fixed-point call-graph walk identifies the functions and argument
+//    indices that carry such LDS-derived flat pointers. Memory operations
+//    that reach LDS through these arguments are redirected to the relocated
+//    global memory by reconstructing an equivalent flat pointer, instead of
+//    accessing the now-dead hardware LDS.
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
@@ -94,7 +103,6 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/DomTreeUpdater.h"
-#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
@@ -104,10 +112,13 @@
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/MDBuilder.h"
+#include "llvm/IR/Operator.h"
 #include "llvm/IR/ReplaceConstant.h"
+#include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 #include <algorithm>
@@ -170,6 +181,7 @@ struct FunctionsAndLDSAccess {
   SetVector<Function *> NonKernelsWithLDSArgument;
   SetVector<GlobalVariable *> AllNonKernelLDSAccess;
   FunctionVariableMap NonKernelToLDSAccessMap;
+  DenseMap<Function *, SmallPtrSet<Argument *, 4>> NonKernelsWithLDSFlatArg;
 };
 
 class AMDGPUSwLowerLDS {
@@ -179,6 +191,7 @@ class AMDGPUSwLowerLDS {
   bool run();
   void getUsesOfLDSByNonKernels();
   void getNonKernelsWithLDSArguments(const CallGraph &CG);
+  void getNonKernelsWithLDSFlatArguments();
   SetVector<Function *>
   getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
   SetVector<GlobalVariable *>
@@ -205,7 +218,8 @@ class AMDGPUSwLowerLDS {
   void lowerNonKernelLDSAccesses(Function *Func,
                                  SetVector<GlobalVariable *> &LDSGlobals,
                                  NonKernelLDSParameters &NKLDSParams);
-  void lowerFlatToLocalRoundTrips(Function *Func);
+  void lowerNonKernelLDSFlatArgAccesses(Function *Func);
+  Value *getFlatPtrForRoundTripLDSAccess(Function *Func, Value *LDSPtr);
   void
   updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
                                 Value *HiddenDynLDSSize,
@@ -289,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());
+          for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
+            Argument *CalleeArg = Callee->getArg(ArgNo);
+            Type *ArgTy = CalleeArg->getType();
+            if (!ArgTy->isPointerTy() ||
+                ArgTy->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
+              continue;
+            Value *Actual = CB->getArgOperand(ArgNo);
+            SmallPtrSet<Value *, 8> Visited;
+            if (!flatPtrDerivesFromLDS(Actual, CallerLDSArgs, Visited,
+                                       /*AcceptLocalCast=*/true))
+              continue;
+            if (FlatArgMap[Callee].insert(CalleeArg).second)
+              Changed = true;
+          }
+        }
+      }
+    }
+  }
+}
+
 void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
   for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
     if (!AMDGPU::isLDSVariableToLower(*GV))
@@ -690,31 +797,51 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
     SetVector<Instruction *> &LDSInstructions) {
   LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
                     << Func->getName());
+  // Map an LDS pointer to its global-memory equivalent: a flat-argument access
+  // uses the flat source pointer, otherwise the base/offset table. nullptr
+  // means neither applies, so the operation is left unchanged.
+  auto TranslatePtr = [&](Value *LDSPtr) -> Value * {
+    if (Value *Flat = getFlatPtrForRoundTripLDSAccess(Func, LDSPtr))
+      return Flat;
+    if (!LoadMallocPtr)
+      return nullptr;
+    return getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LDSPtr);
+  };
+  // Old LDS pointer operands to clean up. Flat-arg round-trips leave the cast
+  // and local GEPs dead; the base/offset path keeps them live via ptrtoint, so
+  // cleanup is a no-op there. WeakTrackingVH tolerates chains freed
+  // recursively.
+  SmallVector<WeakTrackingVH, 8> MaybeDeadPtrs;
   for (Instruction *Inst : LDSInstructions) {
     IRB.SetInsertPoint(Inst);
     if (LoadInst *LI = dyn_cast<LoadInst>(Inst)) {
       Value *LIOperand = LI->getPointerOperand();
-      Value *Replacement =
-          getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, LIOperand);
+      Value *Replacement = TranslatePtr(LIOperand);
+      if (!Replacement)
+        continue;
       LoadInst *NewLI =
           IRB.CreateLoad(LI->getType(), Replacement, LI->getProperties());
       AsanInfo.Instructions.insert(NewLI);
       LI->replaceAllUsesWith(NewLI);
       LI->eraseFromParent();
+      MaybeDeadPtrs.push_back(LIOperand);
     } else if (StoreInst *SI = dyn_cast<StoreInst>(Inst)) {
       Value *SIOperand = SI->getPointerOperand();
-      Value *Replacement =
-          getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, SIOperand);
+      Value *Replacement = TranslatePtr(SIOperand);
+      if (!Replacement)
+        continue;
       StoreInst *NewSI = IRB.CreateStore(SI->getValueOperand(), Replacement,
                                          SI->getProperties());
       AsanInfo.Instructions.insert(NewSI);
       SI->replaceAllUsesWith(NewSI);
       SI->eraseFromParent();
+      MaybeDeadPtrs.push_back(SIOperand);
     } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(Inst)) {
       Value *RMWPtrOperand = RMW->getPointerOperand();
       Value *RMWValOperand = RMW->getValOperand();
-      Value *Replacement =
-          getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, RMWPtrOperand);
+      Value *Replacement = TranslatePtr(RMWPtrOperand);
+      if (!Replacement)
+        continue;
       AtomicRMWInst *NewRMW = IRB.CreateAtomicRMW(
           RMW->getOperation(), Replacement, RMWValOperand, RMW->getAlign(),
           RMW->getOrdering(), RMW->getSyncScopeID());
@@ -722,10 +849,12 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
       AsanInfo.Instructions.insert(NewRMW);
       RMW->replaceAllUsesWith(NewRMW);
       RMW->eraseFromParent();
+      MaybeDeadPtrs.push_back(RMWPtrOperand);
     } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(Inst)) {
       Value *XCHGPtrOperand = XCHG->getPointerOperand();
-      Value *Replacement =
-          getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, XCHGPtrOperand);
+      Value *Replacement = TranslatePtr(XCHGPtrOperand);
+      if (!Replacement)
+        continue;
       AtomicCmpXchgInst *NewXCHG = IRB.CreateAtomicCmpXchg(
           Replacement, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
           XCHG->getAlign(), XCHG->getSuccessOrdering(),
@@ -734,10 +863,16 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
       AsanInfo.Instructions.insert(NewXCHG);
       XCHG->replaceAllUsesWith(NewXCHG);
       XCHG->eraseFromParent();
+      MaybeDeadPtrs.push_back(XCHGPtrOperand);
     } else if (AnyMemIntrinsic *MI = dyn_cast<AnyMemIntrinsic>(Inst)) {
-      Value *NewDest = MI->getRawDest();
-      if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
-        NewDest = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, NewDest);
+      Value *OldRawDest = MI->getRawDest();
+      Value *NewDest = OldRawDest;
+      if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
+        NewDest = TranslatePtr(NewDest);
+        if (!NewDest)
+          continue;
+        MaybeDeadPtrs.push_back(OldRawDest);
+      }
       CallInst *NewMI = nullptr;
       if (AnyMemSetInst *MSI = dyn_cast<AnyMemSetInst>(MI)) {
         if (MI->isAtomic()) {
@@ -750,9 +885,14 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
                                    cast<MemSetInst>(MI)->isVolatile());
         }
       } else if (AnyMemTransferInst *MTI = dyn_cast<AnyMemTransferInst>(MI)) {
-        Value *NewSrc = MTI->getRawSource();
-        if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
-          NewSrc = getTranslatedGlobalMemoryPtrOfLDS(LoadMallocPtr, NewSrc);
+        Value *OldRawSrc = MTI->getRawSource();
+        Value *NewSrc = OldRawSrc;
+        if (MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS) {
+          NewSrc = TranslatePtr(NewSrc);
+          if (!NewSrc)
+            continue;
+          MaybeDeadPtrs.push_back(OldRawSrc);
+        }
         if (MI->isAtomic()) {
           if (MI->getIntrinsicID() ==
               Intrinsic::memmove_element_unordered_atomic) {
@@ -779,91 +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");
   }
-}
-
-void AMDGPUSwLowerLDS::lowerFlatToLocalRoundTrips(Function *Func) {
-  // After SW LDS lowering, storage that used to live in LDS now lives in global
-  // memory, but it may still be handed to a callee as a flat (generic) pointer.
-  // Such a callee can re-derive an LDS pointer with a round-trip pattern:
-  //   %p = getelementptr ..., ptr %flat_arg        ; flat, backed by global mem
-  //   %c = addrspacecast ptr %p to ptr addrspace(3)
-  //   load/store ... ptr addrspace(3) %c
-  // The addrspace(3) access would then hit dead LDS. Redo the access through
-  // the flat source pointer, which still points at the global backing buffer.
-  SmallVector<Instruction *> ToErase;
-  SmallPtrSet<AddrSpaceCastInst *, 8> Casts;
-  for (BasicBlock &BB : *Func) {
-    for (Instruction &Inst : BB) {
-      Value *Ptr = nullptr;
-      if (auto *LI = dyn_cast<LoadInst>(&Inst))
-        Ptr = LI->getPointerOperand();
-      else if (auto *SI = dyn_cast<StoreInst>(&Inst))
-        Ptr = SI->getPointerOperand();
-      else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst))
-        Ptr = RMW->getPointerOperand();
-      else if (auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst))
-        Ptr = XCHG->getPointerOperand();
-      else
-        continue;
 
-      auto *ASC = dyn_cast<AddrSpaceCastInst>(Ptr);
-      if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
-          ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
-        continue;
+  // Drop the now-dead flat->local casts and local GEPs.
+  RecursivelyDeleteTriviallyDeadInstructionsPermissive(MaybeDeadPtrs);
+}
 
-      // Only collapse pointers that originate from a flat function argument,
-      // i.e. storage passed in from a caller. Genuine in-function LDS pointers
-      // are lowered through the base/offset table machinery instead.
-      Value *FlatPtr = ASC->getPointerOperand();
-      if (!isa<Argument>(getUnderlyingObject(FlatPtr)))
-        continue;
+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();
+  }
 
-      IRB.SetInsertPoint(&Inst);
-      Instruction *NewInst = nullptr;
-      if (auto *LI = dyn_cast<LoadInst>(&Inst)) {
-        NewInst = IRB.CreateLoad(LI->getType(), FlatPtr, LI->getProperties());
-      } else if (auto *SI = dyn_cast<StoreInst>(&Inst)) {
-        NewInst = IRB.CreateStore(SI->getValueOperand(), FlatPtr,
-                                  SI->getProperties());
-      } else if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst)) {
-        auto *NewRMW = IRB.CreateAtomicRMW(
-            RMW->getOperation(), FlatPtr, RMW->getValOperand(), RMW->getAlign(),
-            RMW->getOrdering(), RMW->getSyncScopeID());
-        NewRMW->setVolatile(RMW->isVolatile());
-        NewInst = NewRMW;
-      } else {
-        auto *XCHG = cast<AtomicCmpXchgInst>(&Inst);
-        auto *NewXCHG = IRB.CreateAtomicCmpXchg(
-            FlatPtr, XCHG->getCompareOperand(), XCHG->getNewValOperand(),
-            XCHG->getAlign(), XCHG->getSuccessOrdering(),
-            XCHG->getFailureOrdering(), XCHG->getSyncScopeID());
-        NewXCHG->setVolatile(XCHG->isVolatile());
-        NewInst = NewXCHG;
-      }
-      // Flat access into the global backing buffer must be sanitized.
-      AsanInfo.Instructions.insert(NewInst);
-      Inst.replaceAllUsesWith(NewInst);
-      ToErase.push_back(&Inst);
-      Casts.insert(ASC);
-    }
+  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;
+
+  // The flat source must trace back to an LDS-carrying arg.
+  SmallPtrSet<Value *, 8> Visited;
+  if (!flatPtrDerivesFromLDS(FlatSrc, LDSArgs, Visited,
+                             /*AcceptLocalCast=*/false))
+    return nullptr;
+
+  // Re-apply the peeled GEPs in the flat address space.
+  Value *Flat = FlatSrc;
+  for (GEPOperator *GEP : reverse(GEPs)) {
+    SmallVector<Value *, 4> Indices(GEP->idx_begin(), GEP->idx_end());
+    Flat = IRB.CreateGEP(GEP->getSourceElementType(), Flat, Indices, "",
+                         GEP->getNoWrapFlags());
   }
-  for (Instruction *I : ToErase)
-    I->eraseFromParent();
-  // Drop the now-dead round-trip casts.
-  for (AddrSpaceCastInst *ASC : Casts)
-    if (ASC->use_empty())
-      ASC->eraseFromParent();
+  return Flat;
+}
+
+void AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
+  // Lower a non-kernel that only reaches lowered LDS through a flat argument.
+  // No base/offset table is needed: the flat pointer already points at the
+  // global backing buffer, so the shared translation reuses the flat source.
+  SetVector<Instruction *> LDSInstructions;
+  getLDSMemoryInstructions(Func, LDSInstructions);
+  if (LDSInstructions.empty())
+    return;
+  translateLDSMemoryOperationsToGlobalMemory(Func, /*LoadMallocPtr=*/nullptr,
+                                             LDSInstructions);
 }
 
 void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
@@ -1334,6 +1461,11 @@ bool AMDGPUSwLowerLDS::run() {
   // Get address sanitizer scale.
   initAsanInfo();
 
+  // Discover non-kernels that receive lowered LDS via a flat pointer argument.
+  // This must happen before lowering rewrites the call-site address-space
+  // casts.
+  getNonKernelsWithLDSFlatArguments();
+
   for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
     Function *Func = K.first;
     auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
@@ -1396,15 +1528,20 @@ bool AMDGPUSwLowerLDS::run() {
     Changed = true;
   }
 
+  // Lower non-kernels that reach lowered LDS only through a flat argument.
+  // Functions also handled above already had these accesses translated there.
+  for (auto &K : FuncLDSAccessInfo.NonKernelsWithLDSFlatArg) {
+    Function *Func = K.first;
+    if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(Func) ||
+        FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(Func))
+      continue;
+    lowerNonKernelLDSFlatArgAccesses(Func);
+    Changed = true;
+  }
+
   if (!Changed)
     return Changed;
 
-  // Fix up callees that still access lowered storage through the
-  // flat-arg -> addrspace(3) round-trip pattern.
-  for (Function &F : M)
-    if (!F.isDeclaration())
-      lowerFlatToLocalRoundTrips(&F);
-
   for (auto &GV : make_early_inc_range(M.globals())) {
     if (AMDGPU::isLDSVariableToLower(GV)) {
       // probably want to remove from used lists
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
index 646371a5bde24..f69ebe54cb694 100644
--- a/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-sw-lower-lds-flat-ptr-arg-asan.ll
@@ -1,21 +1,18 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
 ; RUN: opt < %s -passes=amdgpu-sw-lower-lds -amdgpu-asan-instrument-lds=false -S -mtriple=amdgcn-amd-amdhsa | FileCheck %s
 
-; A non-kernel receives lowered LDS storage as a flat (generic) pointer and
-; re-derives an LDS pointer via a flat -> addrspace(3) round-trip. After SW LDS
-; lowering the storage lives in global memory, so the addrspace(3) access would
-; hit dead LDS. The round-trip must be collapsed into a flat access on the
-; incoming argument, which still points at the global backing buffer. Check the
-; store, load, atomicrmw and cmpxchg cases.
-
 @lds_var = internal addrspace(3) global [64 x i32] poison, align 4
 
+declare void @llvm.memset.p3.i64(ptr addrspace(3), i8, i64, i1)
+declare void @llvm.memcpy.p3.p1.i64(ptr addrspace(3), ptr addrspace(1), i64, i1)
+declare void @llvm.memcpy.p1.p3.i64(ptr addrspace(1), ptr addrspace(3), i64, i1)
+
 define void @store_case(ptr %storage) sanitize_address {
 ; CHECK-LABEL: define void @store_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT:     addrspace(3)
-; CHECK:         store i32 42, ptr [[STORAGE]], align 4
-; CHECK-NOT:     addrspace(3)
-; CHECK:         ret void
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT:    store i32 42, ptr [[STORAGE]], align 4
+; CHECK-NEXT:    ret void
+;
   %cast = addrspacecast ptr %storage to ptr addrspace(3)
   store i32 42, ptr addrspace(3) %cast, align 4
   ret void
@@ -23,11 +20,10 @@ define void @store_case(ptr %storage) sanitize_address {
 
 define i32 @load_case(ptr %storage) sanitize_address {
 ; CHECK-LABEL: define i32 @load_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT:     addrspace(3)
-; CHECK:         [[VAL:%.*]] = load i32, ptr [[STORAGE]], align 4
-; CHECK-NOT:     addrspace(3)
-; CHECK:         ret i32 [[VAL]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr [[STORAGE]], align 4
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
   %cast = addrspacecast ptr %storage to ptr addrspace(3)
   %val = load i32, ptr addrspace(3) %cast, align 4
   ret i32 %val
@@ -35,11 +31,10 @@ define i32 @load_case(ptr %storage) sanitize_address {
 
 define i32 @atomicrmw_case(ptr %storage) sanitize_address {
 ; CHECK-LABEL: define i32 @atomicrmw_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT:     addrspace(3)
-; CHECK:         [[OLD:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
-; CHECK-NOT:     addrspace(3)
-; CHECK:         ret i32 [[OLD]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = atomicrmw add ptr [[STORAGE]], i32 1 seq_cst, align 4
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
   %cast = addrspacecast ptr %storage to ptr addrspace(3)
   %old = atomicrmw add ptr addrspace(3) %cast, i32 1 seq_cst, align 4
   ret i32 %old
@@ -47,25 +42,213 @@ define i32 @atomicrmw_case(ptr %storage) sanitize_address {
 
 define i32 @cmpxchg_case(ptr %storage) sanitize_address {
 ; CHECK-LABEL: define i32 @cmpxchg_case(
-; CHECK-SAME: ptr [[STORAGE:%.*]]){{.*}} {
-; CHECK-NOT:     addrspace(3)
-; CHECK:         [[RES:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
-; CHECK-NOT:     addrspace(3)
-; CHECK:         [[VAL:%.*]] = extractvalue { i32, i1 } [[RES]], 0
-; CHECK:         ret i32 [[VAL]]
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = cmpxchg ptr [[STORAGE]], i32 0, i32 1 seq_cst seq_cst, align 4
+; CHECK-NEXT:    [[VAL:%.*]] = extractvalue { i32, i1 } [[TMP1]], 0
+; CHECK-NEXT:    ret i32 [[VAL]]
+;
   %cast = addrspacecast ptr %storage to ptr addrspace(3)
   %res = cmpxchg ptr addrspace(3) %cast, i32 0, i32 1 seq_cst seq_cst, align 4
   %val = extractvalue { i32, i1 } %res, 0
   ret i32 %val
 }
 
-define amdgpu_kernel void @kernel() sanitize_address {
+; GEP applied in the local address space after the cast must be rebuilt in flat.
+define void @gep_chain_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @gep_chain_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[STORAGE]], i64 5
+; CHECK-NEXT:    store i32 7, ptr [[TMP1]], align 4
+; CHECK-NEXT:    ret void
+;
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  %elem = getelementptr i32, ptr addrspace(3) %cast, i64 5
+  store i32 7, ptr addrspace(3) %elem, align 4
+  ret void
+}
+
+; A flat LDS argument forwarded through a non-kernel must be handled in the
+; callee too.
+define void @forwarded_callee(ptr %p) sanitize_address {
+; CHECK-LABEL: define void @forwarded_callee(
+; CHECK-SAME: ptr [[P:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    store i32 9, ptr [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+  %cast = addrspacecast ptr %p to ptr addrspace(3)
+  store i32 9, ptr addrspace(3) %cast, align 4
+  ret void
+}
+
+define void @forwarder(ptr %q) sanitize_address {
+; CHECK-LABEL: define void @forwarder(
+; CHECK-SAME: ptr [[Q:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    call void @forwarded_callee(ptr [[Q]])
+; CHECK-NEXT:    ret void
+;
+  call void @forwarded_callee(ptr %q)
+  ret void
+}
+
+; Memory intrinsics with the lowered LDS as the destination must be redirected
+; through the flat argument.
+define void @memset_case(ptr %storage) sanitize_address {
+; CHECK-LABEL: define void @memset_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr [[STORAGE]], i8 0, i64 32, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  call void @llvm.memset.p3.i64(ptr addrspace(3) %cast, i8 0, i64 32, i1 false)
+  ret void
+}
+
+define void @memcpy_dst_case(ptr %storage, ptr addrspace(1) %src) sanitize_address {
+; CHECK-LABEL: define void @memcpy_dst_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], ptr addrspace(1) [[SRC:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    call void @llvm.memcpy.p0.p1.i64(ptr [[STORAGE]], ptr addrspace(1) [[SRC]], i64 32, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  call void @llvm.memcpy.p3.p1.i64(ptr addrspace(3) %cast, ptr addrspace(1) %src, i64 32, i1 false)
+  ret void
+}
+
+; ... and with the lowered LDS as the source.
+define void @memcpy_src_case(ptr %storage, ptr addrspace(1) %dst) sanitize_address {
+; CHECK-LABEL: define void @memcpy_src_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], ptr addrspace(1) [[DST:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) [[DST]], ptr [[STORAGE]], i64 32, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %cast = addrspacecast ptr %storage to ptr addrspace(3)
+  call void @llvm.memcpy.p1.p3.i64(ptr addrspace(1) %dst, ptr addrspace(3) %cast, i64 32, i1 false)
+  ret void
+}
+
+; A flat pointer merged with a select before the cast must still be recognized
+; as an LDS-carrying flat argument (both operands trace to %storage).
+define void @select_case(ptr %storage, i1 %c) sanitize_address {
+; CHECK-LABEL: define void @select_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], i1 [[C:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:    [[G:%.*]] = getelementptr i32, ptr [[STORAGE]], i64 1
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C]], ptr [[STORAGE]], ptr [[G]]
+; CHECK-NEXT:    store i32 5, ptr [[SEL]], align 4
+; CHECK-NEXT:    ret void
+;
+  %g = getelementptr i32, ptr %storage, i64 1
+  %sel = select i1 %c, ptr %storage, ptr %g
+  %cast = addrspacecast ptr %sel to ptr addrspace(3)
+  store i32 5, ptr addrspace(3) %cast, align 4
+  ret void
+}
+
+; Loop-carried flat pointer: the phi lives in the flat address space and the
+; addrspace(3) pointer is re-derived each iteration.
+define void @loop_case(ptr %storage, i32 %n) sanitize_address {
+; CHECK-LABEL: define void @loop_case(
+; CHECK-SAME: ptr [[STORAGE:%.*]], i32 [[N:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[P:%.*]] = phi ptr [ [[STORAGE]], %[[ENTRY]] ], [ [[P_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    store i32 [[I]], ptr [[P]], align 4
+; CHECK-NEXT:    [[P_NEXT]] = getelementptr i32, ptr [[P]], i64 1
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i32 [[I_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[COND]], label %[[LOOP]], label %[[EXIT:.*]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %loop ]
+  %p = phi ptr [ %storage, %entry ], [ %p.next, %loop ]
+  %cast = addrspacecast ptr %p to ptr addrspace(3)
+  store i32 %i, ptr addrspace(3) %cast, align 4
+  %p.next = getelementptr i32, ptr %p, i64 1
+  %i.next = add i32 %i, 1
+  %cond = icmp slt i32 %i.next, %n
+  br i1 %cond, label %loop, label %exit
+exit:
+  ret void
+}
+
+define amdgpu_kernel void @kernel(i1 %c, i32 %n, ptr addrspace(1) %buf) sanitize_address {
 ; CHECK-LABEL: define amdgpu_kernel void @kernel(
+; CHECK-SAME: i1 [[C:%.*]], i32 [[N:%.*]], ptr addrspace(1) [[BUF:%.*]]) #[[ATTR3:[0-9]+]] {
+; CHECK-NEXT:  [[WID:.*]]:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.amdgcn.workitem.id.y()
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.amdgcn.workitem.id.z()
+; CHECK-NEXT:    [[TMP3:%.*]] = or i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = or i32 [[TMP3]], [[TMP2]]
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp eq i32 [[TMP4]], 0
+; CHECK-NEXT:    br i1 [[TMP5]], label %[[MALLOC:.*]], label %[[BB18:.*]]
+; CHECK:       [[MALLOC]]:
+; CHECK-NEXT:    [[TMP6:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE:%.*]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 0), align 4
+; CHECK-NEXT:    [[TMP7:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 2), align 4
+; CHECK-NEXT:    [[TMP8:%.*]] = add i32 [[TMP6]], [[TMP7]]
+; CHECK-NEXT:    [[TMP9:%.*]] = zext i32 [[TMP8]] to i64
+; CHECK-NEXT:    [[TMP10:%.*]] = call ptr @llvm.returnaddress.p0(i32 0)
+; CHECK-NEXT:    [[TMP11:%.*]] = ptrtoint ptr [[TMP10]] to i64
+; CHECK-NEXT:    [[TMP12:%.*]] = call i64 @__asan_malloc_impl(i64 [[TMP9]], i64 [[TMP11]])
+; CHECK-NEXT:    [[TMP13:%.*]] = inttoptr i64 [[TMP12]] to ptr addrspace(1)
+; CHECK-NEXT:    store ptr addrspace(1) [[TMP13]], ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, align 8
+; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP13]], i64 8
+; CHECK-NEXT:    [[TMP15:%.*]] = ptrtoint ptr addrspace(1) [[TMP14]] to i64
+; CHECK-NEXT:    call void @__asan_poison_region(i64 [[TMP15]], i64 24)
+; CHECK-NEXT:    [[TMP16:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP13]], i64 288
+; CHECK-NEXT:    [[TMP17:%.*]] = ptrtoint ptr addrspace(1) [[TMP16]] to i64
+; CHECK-NEXT:    call void @__asan_poison_region(i64 [[TMP17]], i64 64)
+; CHECK-NEXT:    br label %[[BB18]]
+; CHECK:       [[BB18]]:
+; CHECK-NEXT:    [[XYZCOND:%.*]] = phi i1 [ false, %[[WID]] ], [ true, %[[MALLOC]] ]
+; CHECK-NEXT:    call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    [[TMP19:%.*]] = load ptr addrspace(1), ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, align 8
+; CHECK-NEXT:    [[TMP20:%.*]] = load i32, ptr addrspace(1) getelementptr inbounds ([[LLVM_AMDGCN_SW_LDS_KERNEL_MD_TYPE]], ptr addrspace(1) @llvm.amdgcn.sw.lds.kernel.md, i32 0, i32 1, i32 0), align 4
+; CHECK-NEXT:    [[TMP21:%.*]] = getelementptr inbounds i8, ptr addrspace(3) @llvm.amdgcn.sw.lds.kernel, i32 [[TMP20]]
+; CHECK-NEXT:    [[TMP22:%.*]] = ptrtoint ptr addrspace(3) [[TMP21]] to i32
+; CHECK-NEXT:    [[TMP23:%.*]] = getelementptr inbounds i8, ptr addrspace(1) [[TMP19]], i32 [[TMP22]]
+; CHECK-NEXT:    [[TMP24:%.*]] = addrspacecast ptr addrspace(1) [[TMP23]] to ptr
+; CHECK-NEXT:    call void @store_case(ptr [[TMP24]])
+; CHECK-NEXT:    [[L:%.*]] = call i32 @load_case(ptr [[TMP24]])
+; CHECK-NEXT:    [[A:%.*]] = call i32 @atomicrmw_case(ptr [[TMP24]])
+; CHECK-NEXT:    [[X:%.*]] = call i32 @cmpxchg_case(ptr [[TMP24]])
+; CHECK-NEXT:    call void @gep_chain_case(ptr [[TMP24]])
+; CHECK-NEXT:    call void @forwarder(ptr [[TMP24]])
+; CHECK-NEXT:    call void @memset_case(ptr [[TMP24]])
+; CHECK-NEXT:    call void @memcpy_dst_case(ptr [[TMP24]], ptr addrspace(1) [[BUF]])
+; CHECK-NEXT:    call void @memcpy_src_case(ptr [[TMP24]], ptr addrspace(1) [[BUF]])
+; CHECK-NEXT:    call void @select_case(ptr [[TMP24]], i1 [[C]])
+; CHECK-NEXT:    call void @loop_case(ptr [[TMP24]], i32 [[N]])
+; CHECK-NEXT:    br label %[[CONDFREE:.*]]
+; CHECK:       [[CONDFREE]]:
+; CHECK-NEXT:    call void @llvm.amdgcn.s.barrier()
+; CHECK-NEXT:    br i1 [[XYZCOND]], label %[[FREE:.*]], label %[[END:.*]]
+; CHECK:       [[FREE]]:
+; CHECK-NEXT:    [[TMP25:%.*]] = call ptr @llvm.returnaddress.p0(i32 0)
+; CHECK-NEXT:    [[TMP26:%.*]] = ptrtoint ptr [[TMP25]] to i64
+; CHECK-NEXT:    [[TMP27:%.*]] = ptrtoint ptr addrspace(1) [[TMP19]] to i64
+; CHECK-NEXT:    call void @__asan_free_impl(i64 [[TMP27]], i64 [[TMP26]])
+; CHECK-NEXT:    br label %[[END]]
+; CHECK:       [[END]]:
+; CHECK-NEXT:    ret void
+;
   %flat = addrspacecast ptr addrspace(3) @lds_var to ptr
   call void @store_case(ptr %flat)
   %l = call i32 @load_case(ptr %flat)
   %a = call i32 @atomicrmw_case(ptr %flat)
-  %c = call i32 @cmpxchg_case(ptr %flat)
+  %x = call i32 @cmpxchg_case(ptr %flat)
+  call void @gep_chain_case(ptr %flat)
+  call void @forwarder(ptr %flat)
+  call void @memset_case(ptr %flat)
+  call void @memcpy_dst_case(ptr %flat, ptr addrspace(1) %buf)
+  call void @memcpy_src_case(ptr %flat, ptr addrspace(1) %buf)
+  call void @select_case(ptr %flat, i1 %c)
+  call void @loop_case(ptr %flat, i32 %n)
   ret void
 }
 

>From 6e56fa3a874da74147771e97c092a1793627fc68 Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Thu, 30 Jul 2026 20:41:22 +0530
Subject: [PATCH 3/4] changes as per review feedback

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 63 +++++++++++----------
 1 file changed, 33 insertions(+), 30 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 25c82883bd419..69d89b8862230 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -303,17 +303,21 @@ 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) {
+/// Returns true if the flat pointer \p V is derived from an LDS-carrying
+/// argument in \p LDSArgs, looking through GEP/bitcast/phi/select.
+///
+/// \param V The flat pointer value to test.
+/// \param LDSArgs Arguments known to carry LDS-backed storage.
+/// \param Visited Set of already-visited values, used to break phi cycles.
+/// \param AcceptLocalCast If set, a flat<-local addrspacecast also counts as
+///        an origin. Detection (getNonKernelsWithLDSFlatArguments) sets this,
+///        since that cast is how LDS first enters a flat pointer; the
+///        rewrite (getFlatPtrForRoundTripLDSAccess) leaves it clear so only
+///        argument origins qualify.
+static bool isFlatPtrDerivedFromLDS(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))
@@ -328,23 +332,23 @@ static bool flatPtrDerivesFromLDS(Value *V,
     Value *Src = Op->getOperand(0);
     if (Src->getType()->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
       return true;
-    return flatPtrDerivesFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
+    return isFlatPtrDerivedFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
   }
   case Instruction::GetElementPtr:
   case Instruction::BitCast:
-    return flatPtrDerivesFromLDS(Op->getOperand(0), LDSArgs, Visited,
-                                 AcceptLocalCast);
+    return isFlatPtrDerivedFromLDS(Op->getOperand(0), LDSArgs, Visited,
+                                   AcceptLocalCast);
   case Instruction::PHI:
     for (Value *In : cast<PHINode>(Op)->incoming_values())
-      if (flatPtrDerivesFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
+      if (isFlatPtrDerivedFromLDS(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);
+    return isFlatPtrDerivedFromLDS(SI->getTrueValue(), LDSArgs, Visited,
+                                   AcceptLocalCast) ||
+           isFlatPtrDerivedFromLDS(SI->getFalseValue(), LDSArgs, Visited,
+                                   AcceptLocalCast);
   }
   default:
     return false;
@@ -374,8 +378,8 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
           Function *Callee = CB->getCalledFunction();
           if (!Callee || Callee->isDeclaration() || AMDGPU::isKernel(*Callee))
             continue;
-          unsigned NumArgs =
-              std::min(CB->arg_size(), (unsigned)Callee->arg_size());
+          unsigned NumArgs = std::min(
+              CB->arg_size(), static_cast<unsigned>(Callee->arg_size()));
           for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
             Argument *CalleeArg = Callee->getArg(ArgNo);
             Type *ArgTy = CalleeArg->getType();
@@ -384,8 +388,8 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSFlatArguments() {
               continue;
             Value *Actual = CB->getArgOperand(ArgNo);
             SmallPtrSet<Value *, 8> Visited;
-            if (!flatPtrDerivesFromLDS(Actual, CallerLDSArgs, Visited,
-                                       /*AcceptLocalCast=*/true))
+            if (!isFlatPtrDerivedFromLDS(Actual, CallerLDSArgs, Visited,
+                                         /*AcceptLocalCast=*/true))
               continue;
             if (FlatArgMap[Callee].insert(CalleeArg).second)
               Changed = true;
@@ -958,17 +962,16 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
     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)
+  auto *ASC = dyn_cast<AddrSpaceCastOperator>(Cur);
+  if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
+      ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
     return nullptr;
+  Value *FlatSrc = ASC->getPointerOperand();
 
   // The flat source must trace back to an LDS-carrying arg.
   SmallPtrSet<Value *, 8> Visited;
-  if (!flatPtrDerivesFromLDS(FlatSrc, LDSArgs, Visited,
-                             /*AcceptLocalCast=*/false))
+  if (!isFlatPtrDerivedFromLDS(FlatSrc, LDSArgs, Visited,
+                               /*AcceptLocalCast=*/false))
     return nullptr;
 
   // Re-apply the peeled GEPs in the flat address space.

>From d27d924bd05ad42ddfcc83634ffee991afd0d90d Mon Sep 17 00:00:00 2001
From: skc7 <Krishna.Sankisa at amd.com>
Date: Fri, 31 Jul 2026 17:01:26 +0530
Subject: [PATCH 4/4] [AMDGPU] Lower LDS flat round trips without provenance
 analysis

---
 llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp | 193 +++++---------------
 1 file changed, 42 insertions(+), 151 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
index 69d89b8862230..ec2d4a24586a6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSwLowerLDS.cpp
@@ -82,14 +82,13 @@
 //    corresponding element in offset table. With this information, replacement
 //    value is obtained.
 //
-// Replacement of LDS accessed via flat pointer arguments:
-//    An LDS pointer can reach a non-kernel as a flat (generic) pointer
-//    argument, e.g. after an addrspacecast of an LDS global to flat. A
-//    fixed-point call-graph walk identifies the functions and argument
-//    indices that carry such LDS-derived flat pointers. Memory operations
-//    that reach LDS through these arguments are redirected to the relocated
-//    global memory by reconstructing an equivalent flat pointer, instead of
-//    accessing the now-dead hardware LDS.
+// Replacement of LDS accessed via a flat round trip:
+//    An LDS pointer can reach a non-kernel as a flat (generic) pointer, e.g.
+//    after an addrspacecast of an LDS global to flat that the callee casts
+//    back to addrspace(3). Such a flat<->local round trip is collapsed into a
+//    direct flat access into the relocated global memory, instead of accessing
+//    the now-dead hardware LDS. This is valid regardless of the pointer's
+//    origin, so no provenance analysis is needed.
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
@@ -181,7 +180,6 @@ struct FunctionsAndLDSAccess {
   SetVector<Function *> NonKernelsWithLDSArgument;
   SetVector<GlobalVariable *> AllNonKernelLDSAccess;
   FunctionVariableMap NonKernelToLDSAccessMap;
-  DenseMap<Function *, SmallPtrSet<Argument *, 4>> NonKernelsWithLDSFlatArg;
 };
 
 class AMDGPUSwLowerLDS {
@@ -191,7 +189,6 @@ class AMDGPUSwLowerLDS {
   bool run();
   void getUsesOfLDSByNonKernels();
   void getNonKernelsWithLDSArguments(const CallGraph &CG);
-  void getNonKernelsWithLDSFlatArguments();
   SetVector<Function *>
   getOrderedIndirectLDSAccessingKernels(SetVector<Function *> &Kernels);
   SetVector<GlobalVariable *>
@@ -218,8 +215,8 @@ class AMDGPUSwLowerLDS {
   void lowerNonKernelLDSAccesses(Function *Func,
                                  SetVector<GlobalVariable *> &LDSGlobals,
                                  NonKernelLDSParameters &NKLDSParams);
-  void lowerNonKernelLDSFlatArgAccesses(Function *Func);
-  Value *getFlatPtrForRoundTripLDSAccess(Function *Func, Value *LDSPtr);
+  bool lowerNonKernelLDSFlatArgAccesses(Function *Func);
+  Value *getFlatPtrForRoundTripLDSAccess(Value *LDSPtr);
   void
   updateMallocSizeForDynamicLDS(Function *Func, Value **CurrMallocSize,
                                 Value *HiddenDynLDSSize,
@@ -303,103 +300,6 @@ void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
   }
 }
 
-/// Returns true if the flat pointer \p V is derived from an LDS-carrying
-/// argument in \p LDSArgs, looking through GEP/bitcast/phi/select.
-///
-/// \param V The flat pointer value to test.
-/// \param LDSArgs Arguments known to carry LDS-backed storage.
-/// \param Visited Set of already-visited values, used to break phi cycles.
-/// \param AcceptLocalCast If set, a flat<-local addrspacecast also counts as
-///        an origin. Detection (getNonKernelsWithLDSFlatArguments) sets this,
-///        since that cast is how LDS first enters a flat pointer; the
-///        rewrite (getFlatPtrForRoundTripLDSAccess) leaves it clear so only
-///        argument origins qualify.
-static bool isFlatPtrDerivedFromLDS(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 isFlatPtrDerivedFromLDS(Src, LDSArgs, Visited, AcceptLocalCast);
-  }
-  case Instruction::GetElementPtr:
-  case Instruction::BitCast:
-    return isFlatPtrDerivedFromLDS(Op->getOperand(0), LDSArgs, Visited,
-                                   AcceptLocalCast);
-  case Instruction::PHI:
-    for (Value *In : cast<PHINode>(Op)->incoming_values())
-      if (isFlatPtrDerivedFromLDS(In, LDSArgs, Visited, AcceptLocalCast))
-        return true;
-    return false;
-  case Instruction::Select: {
-    auto *SI = cast<SelectInst>(Op);
-    return isFlatPtrDerivedFromLDS(SI->getTrueValue(), LDSArgs, Visited,
-                                   AcceptLocalCast) ||
-           isFlatPtrDerivedFromLDS(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(), static_cast<unsigned>(Callee->arg_size()));
-          for (unsigned ArgNo = 0; ArgNo < NumArgs; ++ArgNo) {
-            Argument *CalleeArg = Callee->getArg(ArgNo);
-            Type *ArgTy = CalleeArg->getType();
-            if (!ArgTy->isPointerTy() ||
-                ArgTy->getPointerAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
-              continue;
-            Value *Actual = CB->getArgOperand(ArgNo);
-            SmallPtrSet<Value *, 8> Visited;
-            if (!isFlatPtrDerivedFromLDS(Actual, CallerLDSArgs, Visited,
-                                         /*AcceptLocalCast=*/true))
-              continue;
-            if (FlatArgMap[Callee].insert(CalleeArg).second)
-              Changed = true;
-          }
-        }
-      }
-    }
-  }
-}
-
 void AMDGPUSwLowerLDS::getUsesOfLDSByNonKernels() {
   for (GlobalVariable *GV : FuncLDSAccessInfo.AllNonKernelLDSAccess) {
     if (!AMDGPU::isLDSVariableToLower(*GV))
@@ -801,11 +701,11 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
     SetVector<Instruction *> &LDSInstructions) {
   LLVM_DEBUG(dbgs() << "Translating LDS memory operations to global memory : "
                     << Func->getName());
-  // Map an LDS pointer to its global-memory equivalent: a flat-argument access
-  // uses the flat source pointer, otherwise the base/offset table. nullptr
-  // means neither applies, so the operation is left unchanged.
+  // Map an LDS pointer to its global-memory equivalent: a flat<->local round
+  // trip collapses to its flat source pointer, otherwise use the base/offset
+  // table. nullptr means neither applies, so the operation is left unchanged.
   auto TranslatePtr = [&](Value *LDSPtr) -> Value * {
-    if (Value *Flat = getFlatPtrForRoundTripLDSAccess(Func, LDSPtr))
+    if (Value *Flat = getFlatPtrForRoundTripLDSAccess(LDSPtr))
       return Flat;
     if (!LoadMallocPtr)
       return nullptr;
@@ -941,18 +841,17 @@ void AMDGPUSwLowerLDS::translateLDSMemoryOperationsToGlobalMemory(
   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.
+Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Value *LDSPtr) {
+  // If LDSPtr is an addrspacecast(flat->local) optionally followed by local
+  // GEPs, return the equivalent pointer with the GEPs re-applied in the flat
+  // address space; otherwise nullptr so the caller falls back to the
+  // base/offset table.
+  //
+  // No provenance check is needed: collapsing a flat<->local round trip into a
+  // direct flat access is always valid. When the flat pointer is in the LDS
+  // aperture the flat access reaches the same memory as the local one; when it
+  // is not (e.g. it points into the relocated global buffer) the local cast was
+  // already undefined, so the flat access is strictly better.
   SmallVector<GEPOperator *, 4> GEPs;
   Value *Cur = LDSPtr;
   while (auto *GEP = dyn_cast<GEPOperator>(Cur)) {
@@ -966,16 +865,9 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
   if (!ASC || ASC->getSrcAddressSpace() != AMDGPUAS::FLAT_ADDRESS ||
       ASC->getDestAddressSpace() != AMDGPUAS::LOCAL_ADDRESS)
     return nullptr;
-  Value *FlatSrc = ASC->getPointerOperand();
-
-  // The flat source must trace back to an LDS-carrying arg.
-  SmallPtrSet<Value *, 8> Visited;
-  if (!isFlatPtrDerivedFromLDS(FlatSrc, LDSArgs, Visited,
-                               /*AcceptLocalCast=*/false))
-    return nullptr;
 
   // Re-apply the peeled GEPs in the flat address space.
-  Value *Flat = FlatSrc;
+  Value *Flat = ASC->getPointerOperand();
   for (GEPOperator *GEP : reverse(GEPs)) {
     SmallVector<Value *, 4> Indices(GEP->idx_begin(), GEP->idx_end());
     Flat = IRB.CreateGEP(GEP->getSourceElementType(), Flat, Indices, "",
@@ -984,16 +876,17 @@ Value *AMDGPUSwLowerLDS::getFlatPtrForRoundTripLDSAccess(Function *Func,
   return Flat;
 }
 
-void AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
-  // Lower a non-kernel that only reaches lowered LDS through a flat argument.
-  // No base/offset table is needed: the flat pointer already points at the
-  // global backing buffer, so the shared translation reuses the flat source.
+bool AMDGPUSwLowerLDS::lowerNonKernelLDSFlatArgAccesses(Function *Func) {
+  // Lower a non-kernel that reaches lowered LDS only through a flat round trip
+  // (addrspacecast flat->local). No base/offset table is needed: collapsing the
+  // round trip yields a direct flat access into the relocated buffer.
   SetVector<Instruction *> LDSInstructions;
   getLDSMemoryInstructions(Func, LDSInstructions);
   if (LDSInstructions.empty())
-    return;
+    return false;
   translateLDSMemoryOperationsToGlobalMemory(Func, /*LoadMallocPtr=*/nullptr,
                                              LDSInstructions);
+  return true;
 }
 
 void AMDGPUSwLowerLDS::poisonRedzones(Function *Func, Value *MallocPtr) {
@@ -1464,11 +1357,6 @@ bool AMDGPUSwLowerLDS::run() {
   // Get address sanitizer scale.
   initAsanInfo();
 
-  // Discover non-kernels that receive lowered LDS via a flat pointer argument.
-  // This must happen before lowering rewrites the call-site address-space
-  // casts.
-  getNonKernelsWithLDSFlatArguments();
-
   for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
     Function *Func = K.first;
     auto &LDSParams = FuncLDSAccessInfo.KernelToLDSParametersMap[Func];
@@ -1531,15 +1419,18 @@ bool AMDGPUSwLowerLDS::run() {
     Changed = true;
   }
 
-  // Lower non-kernels that reach lowered LDS only through a flat argument.
-  // Functions also handled above already had these accesses translated there.
-  for (auto &K : FuncLDSAccessInfo.NonKernelsWithLDSFlatArg) {
-    Function *Func = K.first;
-    if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(Func) ||
-        FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(Func))
+  // Lower any remaining non-kernel that reaches lowered LDS through a flat
+  // round trip (addrspacecast flat->local). Functions handled above already had
+  // these accesses translated there, so skip them. No detection is needed:
+  // collapsing the round trip is valid regardless of the pointer's origin.
+  for (Function &F : M) {
+    if (F.isDeclaration() || AMDGPU::isKernel(F))
       continue;
-    lowerNonKernelLDSFlatArgAccesses(Func);
-    Changed = true;
+    if (FuncLDSAccessInfo.NonKernelToLDSAccessMap.contains(&F) ||
+        FuncLDSAccessInfo.NonKernelsWithLDSArgument.contains(&F))
+      continue;
+    if (lowerNonKernelLDSFlatArgAccesses(&F))
+      Changed = true;
   }
 
   if (!Changed)



More information about the llvm-commits mailing list