[llvm] [AMDGPU][LDS] Fix dynamic LDS interaction with "amdgpu-no-lds-kernel-id" (PR #107092)

[via llvm-commits]

Juan Manuel Martinez =?utf-8?q?Caamaño?= <juamarti at amd.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/107092 at github.com>


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Juan Manuel Martinez Caamaño (jmmartinez)

<details>
<summary>Changes</summary>

Dynamic lds and Table lds both use the amdgpu_lds_kernel_id intrinsic.
Kernels and functons that make an indirect use of this should not have the
"amdgpu-no-lds-kernel-id" attribute.

For the later, this was done. For the dynamic lds case, this was
missing. This patch fixes it.

Pre commit of test in https://github.com/llvm/llvm-project/pull/107091

---
Full diff: https://github.com/llvm/llvm-project/pull/107092.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp (+10-7) 
- (modified) llvm/test/CodeGen/AMDGPU/lower-module-lds-zero-size-arr.ll (+14-3) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
index 38f0b6dda1997b..2a4441d70429c3 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp
@@ -1010,13 +1010,6 @@ class AMDGPULowerModuleLDS {
           M, TableLookupVariablesOrdered, OrderedKernels, KernelToReplacement);
       replaceUsesInInstructionsWithTableLookup(M, TableLookupVariablesOrdered,
                                                LookupTable);
-
-      // Strip amdgpu-no-lds-kernel-id from all functions reachable from the
-      // kernel. We may have inferred this wasn't used prior to the pass.
-      //
-      // TODO: We could filter out subgraphs that do not access LDS globals.
-      for (Function *F : KernelsThatAllocateTableLDS)
-        removeFnAttrFromReachable(CG, F, {"amdgpu-no-lds-kernel-id"});
     }
 
     DenseMap<Function *, GlobalVariable *> KernelToCreatedDynamicLDS =
@@ -1024,6 +1017,16 @@ class AMDGPULowerModuleLDS {
                                  KernelsThatIndirectlyAllocateDynamicLDS,
                                  DynamicVariables, OrderedKernels);
 
+    // Strip amdgpu-no-lds-kernel-id from all functions reachable from the
+    // kernel. We may have inferred this wasn't used prior to the pass.
+    // TODO: We could filter out subgraphs that do not access LDS globals.
+    for (auto *KernelSet : {&KernelsThatIndirectlyAllocateDynamicLDS,
+                            &KernelsThatAllocateTableLDS})
+      for (Function *F : *KernelSet) {
+        removeFnAttrFromReachable(CG, F, {"amdgpu-no-lds-kernel-id"});
+        errs() << "bug code triggered!\n";
+      }
+
     // All kernel frames have been allocated. Calculate and record the
     // addresses.
     {
diff --git a/llvm/test/CodeGen/AMDGPU/lower-module-lds-zero-size-arr.ll b/llvm/test/CodeGen/AMDGPU/lower-module-lds-zero-size-arr.ll
index da1d23f1496cf3..6d3928e5b4285a 100644
--- a/llvm/test/CodeGen/AMDGPU/lower-module-lds-zero-size-arr.ll
+++ b/llvm/test/CodeGen/AMDGPU/lower-module-lds-zero-size-arr.ll
@@ -1,4 +1,4 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 3
 ; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds < %s | FileCheck %s
 
 ; This is an extension and should be rejected by the front-end in most cases.
@@ -6,7 +6,11 @@
 
 @Var0 = linkonce_odr hidden local_unnamed_addr addrspace(3) global [0 x float] poison
 
-define void @fn(float %val, i32 %idx) {
+;.
+; CHECK: @llvm.amdgcn.kernelA.dynlds = external addrspace(3) global [0 x i8], align 4, !absolute_symbol [[META0:![0-9]+]]
+; CHECK: @llvm.amdgcn.dynlds.offset.table = internal addrspace(4) constant [1 x i32] [i32 ptrtoint (ptr addrspace(3) @llvm.amdgcn.kernelA.dynlds to i32)]
+;.
+define void @fn(float %val, i32 %idx) #0 {
 ; CHECK-LABEL: define void @fn(
 ; CHECK-SAME: float [[VAL:%.*]], i32 [[IDX:%.*]]) {
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.amdgcn.lds.kernel.id()
@@ -22,7 +26,7 @@ define void @fn(float %val, i32 %idx) {
   ret void
 }
 
-define amdgpu_kernel void @kernelA(float %val, i32 %idx) {
+define amdgpu_kernel void @kernelA(float %val, i32 %idx) #0 {
 ; CHECK-LABEL: define amdgpu_kernel void @kernelA(
 ; CHECK-SAME: float [[VAL:%.*]], i32 [[IDX:%.*]]) !llvm.amdgcn.lds.kernel.id [[META1:![0-9]+]] {
 ; CHECK-NEXT:    call void @llvm.donothing() [ "ExplicitUse"(ptr addrspace(3) @llvm.amdgcn.kernelA.dynlds) ]
@@ -32,6 +36,13 @@ define amdgpu_kernel void @kernelA(float %val, i32 %idx) {
   tail call void @fn(float %val, i32 %idx)
   ret void
 }
+
+attributes #0 = { "amdgpu-no-lds-kernel-id" }
+
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { nocallback nofree nosync nounwind willreturn memory(none) }
+; CHECK: attributes #[[ATTR1:[0-9]+]] = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
 ;.
+; CHECK: [[META0]] = !{i32 0, i32 1}
 ; CHECK: [[META1]] = !{i32 0}
 ;.

``````````

</details>


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