[AMDGPULowerModuleLDSPass] Kernels do not always have entries in KernelToReplacement map

[jmmartinez]

Summary

The crash was a result of the accesses in lines 512 and 1176. If the kernel was not present in the map, a default initialized entry was created.

Fixes SWDEV-404491

[arsenm]

Does this also fix https://github.com/llvm/llvm-project/issues/63338

[jmmartinez]

In D154946#4488685, @arsenm wrote:

Does this also fix https://github.com/llvm/llvm-project/issues/63338

Yes, seems like the same bug

[JonChesterfield]

Thanks for the patch! On reading through I think there's something more fundamentally broken here. Please could you hold fire on this for a moment while I dig through it?

There's some refactoring in line which seems worth having asap - the extra const, the spurious copy deleted - would you prefer I post a diff with those parts pulled out or you post one for me to accept?

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

486–487

I don't follow this chunk. Previously count != 0 followed by lookup. New find != end followed by using the previous value. Is there a functional change I'm missing here?

1178–1179

Nervous about this. There's a bunch of equality relationships between structures in this pass and the IR generated, probably not adequately documented.

Some number of kernels get identified as needing LDS related stuff and those get put in the OrderedKernels vector. There's an intrinsic which amounts to an index into that vector which persists beyond this pass. If KernelToReplacement is missing an entry for one of those kernels, I've botched the control flow elsewhere in this function.

[jmmartinez]

In D154946#4488779, @JonChesterfield wrote:

Thanks for the patch! On reading through I think there's something more fundamentally broken here. Please could you hold fire on this for a moment while I dig through it?
There's some refactoring in line which seems worth having asap - the extra const, the spurious copy deleted - would you prefer I post a diff with those parts pulled out or you post one for me to accept?

No problem! I will move it into a separate patch.

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

486–487

I'll move it out in another patch. It is a consequence of adding the const.

With const, we cannot use operator[] since it inserts a default initialized element if none exist (the count was guarding against it).

1178–1179

If I'm not wrong, this happens since lowerKernelScopeStructVariables doesn't add replacements for Dynamic LDS variables.

Kernels accesing only Dynamic LDS still need an id for the lookup table., even if they don't have any static LDS kernel variable.

(I hope I'm not mixing concepts)

[jmmartinez]

• rebased

[JonChesterfield]

OK, I think I've got the root cause.

The slow path lowering for normal variables is a table lookup indexed by kernel id, which is an integer assigned somewhat arbitrarily and later made available through the lds intrinsic (involves dedicating a sgpr to the i32 value).

The lowering for dynamic variables is a table lookup indexed by the same kernel id because I didn't think it was worth burning a second sgpr on it.

The pass successfully assigns integers to kernels that use either of these variables. assignLDSKernelIDToEachKernel takes two sets, combines them, sorts the result.

Where it then goes wrong is the lookup table construction assumes every kernel with an index will have a corresponding element, which was true before dynamic variable lowering was introduced. Now however, if a kernel uses dynamic variables and no table lookup static ones, it'll have an assigned index but the corresponding table entry needs to be poison.

Lowerdynamicldsvariables handles this correctly because it was written at the same time as breaking that invariant.

Right fix is going to involve poison at some positions in the lookup table, probably similarly to buildLookupTable. Likely to look like a hybrid of this patch and D154972.

[arsenm]

I assume this is redundant now?

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

static Constant *getAddressesOfVariablesInKernel(
     // TODO: Drop the ptrtoint conversion
 
     Type *I32 = Type::getInt32Ty(Ctx);
 
     ArrayType *KernelOffsetsType = ArrayType::get(I32, Variables.size());
 
     SmallVector<Constant *> Elements;
     for (size_t i = 0; i < Variables.size(); i++) {
       GlobalVariable *GV = Variables[i];
       auto ConstantGepIt = LDSVarsToConstantGEP.find(GV);
       if (ConstantGepIt != LDSVarsToConstantGEP.end()) {
         auto elt = ConstantExpr::getPtrToInt(ConstantGepIt->second, I32);
         Elements.push_back(elt);
       } else {
         Elements.push_back(PoisonValue::get(I32));
       }
     }
     return ConstantArray::get(KernelOffsetsType, Elements);
   }
 
   static GlobalVariable *buildLookupTable(
       Module &M, ArrayRef<GlobalVariable *> Variables,
       ArrayRef<Function *> kernels,
-      DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) {
+      const DenseMap<Function *, LDSVariableReplacement> &KernelToReplacement) {
     if (Variables.empty()) {
       return nullptr;
     }
     LLVMContext &Ctx = M.getContext();
 
     const size_t NumberVariables = Variables.size();
     const size_t NumberKernels = kernels.size();
 
     ArrayType *KernelOffsetsType =
         ArrayType::get(Type::getInt32Ty(Ctx), NumberVariables);
 
     ArrayType *AllKernelsOffsetsType =
         ArrayType::get(KernelOffsetsType, NumberKernels);
 
-    std::vector<Constant *> overallConstantExprElts(NumberKernels);
+    std::vector<Constant *> overallConstantExprElts(
+        NumberKernels, PoisonValue::get(KernelOffsetsType));
     for (size_t i = 0; i < NumberKernels; i++) {
-      LDSVariableReplacement Replacement = KernelToReplacement[kernels[i]];
+      Function *Kernel = kernels[i];
+      auto Replacement = KernelToReplacement.find(Kernel);
+      if (Replacement == KernelToReplacement.end())
+        continue;
       overallConstantExprElts[i] = getAddressesOfVariablesInKernel(
-          Ctx, Variables, Replacement.LDSVarsToConstantGEP);
+          Ctx, Variables, Replacement->second.LDSVarsToConstantGEP);
     }
 
     Constant *init =
         ConstantArray::get(AllKernelsOffsetsType, overallConstantExprElts);
 
     return new GlobalVariable(
         M, AllKernelsOffsetsType, true, GlobalValue::InternalLinkage, init,
         "llvm.amdgcn.lds.offset.table", nullptr, GlobalValue::NotThreadLocal,
bool runOnModule(Module &M) override {
 
     const DenseSet<Function *> KernelsThatIndirectlyAllocateDynamicLDS =
         kernelsThatIndirectlyAccessAnyOfPassedVariables(M, LDSUsesInfo,
                                                         DynamicVariables);
 
     GlobalVariable *MaybeModuleScopeStruct = lowerModuleScopeStructVariables(
         M, ModuleScopeVariables, KernelsThatAllocateModuleLDS);
 
-    DenseMap<Function *, LDSVariableReplacement> KernelToReplacement =
+    const DenseMap<Function *, LDSVariableReplacement> KernelToReplacement =
         lowerKernelScopeStructVariables(M, LDSUsesInfo, ModuleScopeVariables,
                                         KernelsThatAllocateModuleLDS,
                                         MaybeModuleScopeStruct);
 
     // Lower zero cost accesses to the kernel instances just created
     for (auto &GV : KernelAccessVariables) {
       auto &funcs = LDSToKernelsThatNeedToAccessItIndirectly[GV];
       assert(funcs.size() == 1); // Only one kernel can access it
-      LDSVariableReplacement Replacement =
-          KernelToReplacement[*(funcs.begin())];
+      const LDSVariableReplacement &Replacement =
+          KernelToReplacement.at(*funcs.begin());
 
       DenseSet<GlobalVariable *> Vec;
       Vec.insert(GV);
 
       // TODO: Looks like a latent bug, Replacement may not be marked
       // UsedByKernel here
       replaceLDSVariablesWithStruct(M, Vec, Replacement, [](Use &U) {
         return isa<Instruction>(U.getUser());
       });
     }
 
     // The ith element of this vector is kernel id i
     std::vector<Function *> OrderedKernels =
         assignLDSKernelIDToEachKernel(&M, KernelsThatAllocateTableLDS,
                                       KernelsThatIndirectlyAllocateDynamicLDS);
 
     if (!KernelsThatAllocateTableLDS.empty()) {
       LLVMContext &Ctx = M.getContext();
       IRBuilder<> Builder(Ctx);
 
-      for (size_t i = 0; i < OrderedKernels.size(); i++) {
-        markUsedByKernel(Builder, OrderedKernels[i],
-                         KernelToReplacement[OrderedKernels[i]].SGV);
+      for (Function *Kernel : OrderedKernels) {
+        auto Replacement = KernelToReplacement.find(Kernel);
+        if (Replacement == KernelToReplacement.end())
+          continue;
+        markUsedByKernel(Builder, Kernel, Replacement->second.SGV);
       }
 
       // The order must be consistent between lookup table and accesses to
       // lookup table
       std::vector<GlobalVariable *> TableLookupVariablesOrdered(
           TableLookupVariables.begin(), TableLookupVariables.end());
       llvm::sort(TableLookupVariablesOrdered.begin(),
                  TableLookupVariablesOrdered.end(),

llvm/test/CodeGen/AMDGPU/lower-module-lds-all-indirect-accesses.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-lower-module-lds-strategy=hybrid < %s | FileCheck %s
+
+@A = external addrspace(3) global [8 x ptr]
+@B = external addrspace(3) global [0 x i32]
+
+define amdgpu_kernel void @kernel_0() {
+; CHECK-LABEL: define amdgpu_kernel void @kernel_0() !llvm.amdgcn.lds.kernel.id !1 {
+; CHECK-NEXT:    call void @llvm.donothing() [ "ExplicitUse"(ptr addrspace(3) @llvm.amdgcn.kernel.kernel_0.lds) ]
+; CHECK-NEXT:    call void @call_store_A()
+; CHECK-NEXT:    ret void
+;
+  call void @call_store_A()
+  ret void
+}
+
+define amdgpu_kernel void @kernel_1() {
+; CHECK-LABEL: define amdgpu_kernel void @kernel_1() !llvm.amdgcn.lds.kernel.id !2 {
+; CHECK-NEXT:    call void @llvm.donothing() [ "ExplicitUse"(ptr addrspace(3) @llvm.amdgcn.kernel_1.dynlds) ]
+; CHECK-NEXT:    [[PTR:%.*]] = call ptr @get_B_ptr()
+; CHECK-NEXT:    ret void
+;
+  %ptr = call ptr @get_B_ptr()
+  ret void
+}
+
+define amdgpu_kernel void @kernel_2() {
+; CHECK-LABEL: define amdgpu_kernel void @kernel_2() !llvm.amdgcn.lds.kernel.id !3 {
+; CHECK-NEXT:    call void @llvm.donothing() [ "ExplicitUse"(ptr addrspace(3) @llvm.amdgcn.kernel.kernel_2.lds) ]
+; CHECK-NEXT:    call void @store_A()
+; CHECK-NEXT:    ret void
+;
+  call void @store_A()
+  ret void
+}
+
+define amdgpu_kernel void @kernel_3() {
+; CHECK-LABEL: define amdgpu_kernel void @kernel_3() !llvm.amdgcn.lds.kernel.id !4 {
+; CHECK-NEXT:    call void @llvm.donothing() [ "ExplicitUse"(ptr addrspace(3) @llvm.amdgcn.kernel_3.dynlds) ]
+; CHECK-NEXT:    [[PTR:%.*]] = call ptr @get_B_ptr()
+; CHECK-NEXT:    ret void
+;
+  %ptr = call ptr @get_B_ptr()
+  ret void
+}
+
+define private void @call_store_A() {
+; CHECK-LABEL: define private void @call_store_A() {
+; CHECK-NEXT:    call void @store_A()
+; CHECK-NEXT:    ret void
+;
+  call void @store_A()
+  ret void
+}
+
+define private void @store_A() {
+; CHECK-LABEL: define private void @store_A() {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.amdgcn.lds.kernel.id()
+; CHECK-NEXT:    [[A:%.*]] = getelementptr inbounds [4 x [1 x i32]], ptr addrspace(4) @llvm.amdgcn.lds.offset.table, i32 0, i32 [[TMP1]], i32 0
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr addrspace(4) [[A]], align 4
+; CHECK-NEXT:    [[A1:%.*]] = inttoptr i32 [[TMP2]] to ptr addrspace(3)
+; CHECK-NEXT:    [[TMP3:%.*]] = addrspacecast ptr addrspace(3) [[A1]] to ptr
+; CHECK-NEXT:    store ptr [[TMP3]], ptr null, align 8
+; CHECK-NEXT:    ret void
+;
+  store ptr addrspacecast (ptr addrspace(3) @A to ptr), ptr null
+  ret void
+}
+
+define private ptr @get_B_ptr() {
+; CHECK-LABEL: define private ptr @get_B_ptr() {
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.amdgcn.lds.kernel.id()
+; CHECK-NEXT:    [[B:%.*]] = getelementptr inbounds [4 x i32], ptr addrspace(4) @llvm.amdgcn.dynlds.offset.table, i32 0, i32 [[TMP1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr addrspace(4) [[B]], align 4
+; CHECK-NEXT:    [[B1:%.*]] = inttoptr i32 [[TMP2]] to ptr addrspace(3)
+; CHECK-NEXT:    [[TMP3:%.*]] = addrspacecast ptr addrspace(3) [[B1]] to ptr
+; CHECK-NEXT:    ret ptr [[TMP3]]
+;
+  ret ptr addrspacecast (ptr addrspace(3) @B to ptr)
+}