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

[Krzysztof Drewniak 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;
----------------
krzysz00 wrote:

I take it the destination is already checked by the if above?

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