[amdgpu] Increase alignment of all LDS variables

[JonChesterfield]

Summary

Currently the superalign option only increases the alignment of
variables that are moved into the module.lds block. Change that to all LDS
variables. Also only increase the alignment once, instead of once per function.

[JonChesterfield]

• add missing chunk to diff

[JonChesterfield]

This will increase the LDS used by some cases (specifically, padding around variables that aren't moved into the module.lds block). I assume that's expected to be a net win. Internal CI running at 611966.

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

229

Somewhat misleading diff. processUsedLDS has had the if (SuperAlignLDSGlobals) {} block copied out into a helper function but is otherwise unchanged.

llvm/test/CodeGen/AMDGPU/lower-kernel-lds-super-align.ll

1–2

preference for explicit here, means the test keeps working if we change the default on the pass

[arsenm]

In D115488#3184510, @JonChesterfield wrote:

This will increase the LDS used by some cases (specifically, padding around variables that aren't moved into the module.lds block). I assume that's expected to be a net win. Internal CI running at 611966.

If this is supposed to be the guaranteed to work in all cases pass, isn't that an issue? Also it's probably not a net win if the increase pushes over an occupancy window

[JonChesterfield]

I thought overaligning these would be a loss in general when it was proposed (hence the flag) but haven't done the benchmarking. Kind of taking it on faith that someone else has since the feature was implemented. epsdb passed fwiw.

To be more precise, I'm not convinced this makes much real world difference as most variables are naturally aligned anyway but it looks like a mistake that we only change the subset of LDS variables that are already slow to access and allocated in many places.

Also totally happy with deleting this feature or defaulting it to off. The point of this exercise is to get it out of the 'lowering' LDS pass and into a to be written 'optimisation' LDS pass, where this is the (minor?) semantic half of that change and moving it to a separate pass would then be non-functional.

[rampitec]

LGTM

[JonChesterfield]

Cool. If it hits us in a benchmark there's the chicken switch.

[JonChesterfield]

• rebase

llvm/lib/Target/AMDGPU/AMDGPULowerModuleLDSPass.cpp

public:
   static char ID;
 
   AMDGPULowerModuleLDS() : ModulePass(ID) {
     initializeAMDGPULowerModuleLDSPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnModule(Module &M) override {
     UsedList = getUsedList(M);
-
-    bool Changed = processUsedLDS(M);
+    bool Changed = superAlignLDSGlobals(M);
+    Changed |= processUsedLDS(M);
 
     for (Function &F : M.functions()) {
       if (F.isDeclaration())
         continue;
 
       // Only lower compute kernels' LDS.
       if (!AMDGPU::isKernel(F.getCallingConv()))
         continue;
       Changed |= processUsedLDS(M, &F);
     }
 
     UsedList.clear();
     return Changed;
   }
 
 private:
-  bool processUsedLDS(Module &M, Function *F = nullptr) {
-    LLVMContext &Ctx = M.getContext();
+  // Increase the alignment of LDS globals if necessary to maximise the chance
+  // that we can use aligned LDS instructions to access them.
+  static bool superAlignLDSGlobals(Module &M) {
     const DataLayout &DL = M.getDataLayout();
+    bool Changed = false;
+    if (!SuperAlignLDSGlobals) {
+      return Changed;
+    }
 
-    // Find variables to move into new struct instance
-    std::vector<GlobalVariable *> FoundLocalVars =
-        AMDGPU::findVariablesToLower(M, F);
-
-    if (FoundLocalVars.empty()) {
-      // No variables to rewrite, no changes made.
-      return false;
+    for (auto &GV : M.globals()) {
+      if (GV.getType()->getPointerAddressSpace() != AMDGPUAS::LOCAL_ADDRESS) {
+        // Only changing alignment of LDS variables
+        continue;
+      }
+      if (!GV.hasInitializer()) {
+        // cuda/hip extern __shared__ variable, leave alignment alone
+        continue;
       }
 
-    // Increase the alignment of LDS globals if necessary to maximise the chance
-    // that we can use aligned LDS instructions to access them.
-    if (SuperAlignLDSGlobals) {
-      for (auto *GV : FoundLocalVars) {
-        Align Alignment = AMDGPU::getAlign(DL, GV);
-        TypeSize GVSize = DL.getTypeAllocSize(GV->getValueType());
+      Align Alignment = AMDGPU::getAlign(DL, &GV);
+      TypeSize GVSize = DL.getTypeAllocSize(GV.getValueType());
 
-        if (GVSize > 8) {
-          // We might want to use a b96 or b128 load/store
-          Alignment = std::max(Alignment, Align(16));
-        } else if (GVSize > 4) {
-          // We might want to use a b64 load/store
-          Alignment = std::max(Alignment, Align(8));
-        } else if (GVSize > 2) {
-          // We might want to use a b32 load/store
-          Alignment = std::max(Alignment, Align(4));
-        } else if (GVSize > 1) {
-          // We might want to use a b16 load/store
-          Alignment = std::max(Alignment, Align(2));
-        }
 
-        GV->setAlignment(Alignment);
-      }
+      if (Alignment != AMDGPU::getAlign(DL, &GV)) {
+        Changed = true;
+        GV.setAlignment(Alignment);
+      }
+    }
+    return Changed;
+
+  bool processUsedLDS(Module &M, Function *F = nullptr) {
+    LLVMContext &Ctx = M.getContext();
+    const DataLayout &DL = M.getDataLayout();
+
+    // Find variables to move into new struct instance
+    std::vector<GlobalVariable *> FoundLocalVars =
+        AMDGPU::findVariablesToLower(M, F);
+
+    if (FoundLocalVars.empty()) {
+      // No variables to rewrite, no changes made.
+      return false;
     }
 
     SmallVector<OptimizedStructLayoutField, 8> LayoutFields;
     LayoutFields.reserve(FoundLocalVars.size());
     for (GlobalVariable *GV : FoundLocalVars) {
       OptimizedStructLayoutField F(GV, DL.getTypeAllocSize(GV->getValueType()),
                                    AMDGPU::getAlign(DL, GV));
       LayoutFields.emplace_back(F);

llvm/test/CodeGen/AMDGPU/lower-kernel-lds-super-align.ll

-; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_ON %s
-; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_ON %s
+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds --amdgpu-super-align-lds-globals=true < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_ON %s
+; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-super-align-lds-globals=true < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_ON %s
 ; RUN: opt -S -mtriple=amdgcn-- -amdgpu-lower-module-lds --amdgpu-super-align-lds-globals=false < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_OFF %s
 ; RUN: opt -S -mtriple=amdgcn-- -passes=amdgpu-lower-module-lds --amdgpu-super-align-lds-globals=false < %s | FileCheck --check-prefixes=CHECK,SUPER-ALIGN_OFF %s
 
 ; CHECK: %llvm.amdgcn.kernel.k1.lds.t = type { [32 x i8] }
 ; CHECK: %llvm.amdgcn.kernel.k2.lds.t = type { i16, [2 x i8], i16 }
 ; CHECK: %llvm.amdgcn.kernel.k3.lds.t = type { [32 x i64], [32 x i32] }
 ; CHECK: %llvm.amdgcn.kernel.k4.lds.t = type { [2 x i32 addrspace(3)*] }
 
+; SUPER-ALIGN_ON: @lds.unused = addrspace(3) global i32 undef, align 4
+; SUPER-ALIGN_OFF: @lds.unused = addrspace(3) global i32 undef, align 2
[e-mail] = addrspace(3) global i32 undef, align 2
+
 ; CHECK-NOT: @lds.1
 @lds.1 = internal unnamed_addr addrspace(3) global [32 x i8] undef, align 1
 
 ; SUPER-ALIGN_ON: @llvm.amdgcn.kernel.k1.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k1.lds.t undef, align 16
 ; SUPER-ALIGN_OFF: @llvm.amdgcn.kernel.k1.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k1.lds.t undef, align 1
 
 ; CHECK: @llvm.amdgcn.kernel.k2.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k2.lds.t undef, align 4
 ; SUPER-ALIGN_ON:  @llvm.amdgcn.kernel.k3.lds = internal addrspace(3) global %llvm.amdgcn.kernel.k3.lds.t undef, align 16