[llvm-branch-commits] [llvm] [AMDGPU][SwLowerLDS] Lower non-kernels with LDS instructions and assign kernel IDs to their callers (PR #228043)

[via llvm-branch-commits]

================
@@ -251,35 +253,92 @@ SetVector<Function *> AMDGPUSwLowerLDS::getOrderedIndirectLDSAccessingKernels(
   return OrderedKernels;
 }
 
-void AMDGPUSwLowerLDS::getNonKernelsWithLDSArguments(const CallGraph &CG) {
-  // Among the kernels accessing LDS, get list of
-  // Non-kernels to which a call is made and a ptr
-  // to addrspace(3) is passed as argument.
-  for (auto &K : FuncLDSAccessInfo.KernelToLDSParametersMap) {
-    Function *Func = K.first;
-    const CallGraphNode *CGN = CG[Func];
-    if (!CGN)
+// LDS memory operations and casts between flat and LDS.
+static bool isLDSMemoryInstruction(const Instruction &Inst) {
+  if (auto *LI = dyn_cast<LoadInst>(&Inst))
+    return LI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+  if (auto *SI = dyn_cast<StoreInst>(&Inst))
+    return SI->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+  if (auto *RMW = dyn_cast<AtomicRMWInst>(&Inst))
+    return RMW->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+  if (auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&Inst))
+    return XCHG->getPointerAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+  if (auto *ASC = dyn_cast<AddrSpaceCastInst>(&Inst)) {
+    unsigned SrcAS = ASC->getSrcAddressSpace();
+    unsigned DstAS = ASC->getDestAddressSpace();
+    return (SrcAS == AMDGPUAS::LOCAL_ADDRESS &&
+            DstAS == AMDGPUAS::FLAT_ADDRESS) ||
+           (SrcAS == AMDGPUAS::FLAT_ADDRESS &&
+            DstAS == AMDGPUAS::LOCAL_ADDRESS);
+  }
+  if (auto *MI = dyn_cast<AnyMemIntrinsic>(&Inst)) {
+    if (MI->getDestAddressSpace() == AMDGPUAS::LOCAL_ADDRESS)
+      return true;
+    if (auto *MTI = dyn_cast<AnyMemTransferInst>(MI))
+      return MTI->getSourceAddressSpace() == AMDGPUAS::LOCAL_ADDRESS;
+  }
+  return false;
+}
+
+void AMDGPUSwLowerLDS::getNonKernelsWithLDSInstructions(const CallGraph &CG) {
+  // Non-kernels with LDS instructions find the malloc buffer through the base
+  // table, so every kernel reaching them needs a kernel ID, even one without
+  // LDS.
+  DenseMap<Function *, bool> HasLDSInstructionsCache;
+  auto HasLDSInstructions = [&](Function *F) {
+    auto [It, Inserted] = HasLDSInstructionsCache.try_emplace(F, false);
+    if (Inserted) {
+      It->second = any_of(instructions(*F), isLDSMemoryInstruction);
+      if (It->second)
+        FuncLDSAccessInfo.NonKernelsWithLDSInstructions.insert(F);
+    }
+    return It->second;
+  };
+
+  // Indirect calls may reach any address taken non-kernel.
+  SmallVector<Function *> AddressTakenFuncs;
+  for (Function &F : M)
----------------
skc7 wrote:

The walk is over the `CallGraph` itself. For indirect calls, we can't use its external node because it also includes every non-local function, not just address taken functions. AFAIU, `getTransitiveUsesOfLDS` uses the same address taken set. 

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