[AMDGPU][GISEL] Adding new reg bank select rules for G_DYN_STACKALLOC (llvm#200369)

Add register bank selection and legalization support for
G_DYN_STACKALLOC in the new RegBankLegalize framework.

Author: abhigargrepo

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeHelper.cpp

@@ -19,6 +19,7 @@
 #include "AMDGPURegisterBankInfo.h"
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
+#include "SIMachineFunctionInfo.h"
 #include "llvm/CodeGen/GlobalISel/GenericMachineInstrs.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
@@ -1304,6 +1305,61 @@ bool RegBankLegalizeHelper::lower(MachineInstr &MI,
     }
     return true;
   }
+  case DynStackAlloc: {
+    const auto &TFI = *ST.getFrameLowering();
+    // Guard in case the stack growth direction ever changes with scratch
+    // instructions.
+    assert(TFI.getStackGrowthDirection() == TargetFrameLowering::StackGrowsUp &&
+           "Stack grows upwards for AMDGPU");
+
+    Register Dst = MI.getOperand(0).getReg();
+    Register AllocSize = MI.getOperand(1).getReg();
+    Align Alignment = assumeAligned(MI.getOperand(2).getImm());
+
+    // Erase before building new instrs to avoid hitting multiple Dst assert
+    // with CSE.
+    B.setInsertPt(*MI.getParent(), std::next(MI.getIterator()));
+    MI.eraseFromParent();
+
+    if (MRI.getRegBank(AllocSize) != SgprRB) {
+      auto WaveReduction =
+          B.buildIntrinsic(Intrinsic::amdgcn_wave_reduce_umax, {SgprRB_S32})
+              .addUse(AllocSize)
+              .addImm(0);
+      AllocSize = WaveReduction.getReg(0);
+    }
+
+    LLT PtrTy = MRI.getType(Dst);
+    assert(PtrTy.getSizeInBits() == 32 &&
+           "Expected 32-bit pointer for stack allocation");
+    const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
+    Register SPReg = Info->getStackPtrOffsetReg();
+
+    // When using flat-scratch, the stack offset is unscaled.
+    const bool HasFlatScratch = ST.hasFlatScratchEnabled();
+    const unsigned WavefrontSizeLog2 = ST.getWavefrontSizeLog2();
+
+    Register AdjustedSize = AllocSize;
+    if (!HasFlatScratch) {
+      auto WaveSize = B.buildConstant(SgprRB_S32, WavefrontSizeLog2);
+      AdjustedSize = B.buildShl(SgprRB_S32, AllocSize, WaveSize).getReg(0);
+    }
+    if (Alignment > TFI.getStackAlign()) {
+      const uint64_t EffectiveAlignment =
+          Alignment.value() << (HasFlatScratch ? 0 : WavefrontSizeLog2);
+      auto OldSP = B.buildCopy({SgprRB, PtrTy}, SPReg);
+      auto Tmp1 =
+          B.buildPtrAdd({SgprRB, PtrTy}, OldSP,
+                        B.buildConstant(SgprRB_S32, EffectiveAlignment - 1));
+      uint64_t Mask = maskTrailingZeros<uint64_t>(Log2_64(EffectiveAlignment));
+      B.buildPtrMask(Dst, Tmp1, B.buildConstant(SgprRB_S32, Mask));
+    } else {
+      B.buildCopy(Dst, SPReg);
+    }
+    auto PtrAdd = B.buildPtrAdd({SgprRB, PtrTy}, Dst, AdjustedSize);
+    B.buildCopy(SPReg, PtrAdd);
+    return true;
+  }
   case WidenLoad: {
     LLT DstTy = MRI.getType(MI.getOperand(0).getReg());
     if (DstTy == S96)

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.cpp

@@ -1333,6 +1333,10 @@ RegBankLegalizeRules::RegBankLegalizeRules(const GCNSubtarget &_ST,
       .Any({{UniP3}, {{SgprP3}, {SgprP3, Sgpr32}}})
       .Any({{DivP3}, {{VgprP3}, {VgprP3, Vgpr32}}});
 
+  addRulesForGOpcs({G_DYN_STACKALLOC})
+      .Any({{UniP5, UniS32}, {{SgprP5}, {Sgpr32}, DynStackAlloc}})
+      .Any({{UniP5, DivS32}, {{SgprP5}, {Vgpr32}, DynStackAlloc}});
+
   addRulesForGOpcs({G_ABS}, Standard)
       .Uni(S16, {{Sgpr32Trunc}, {Sgpr32SExt}})
       .Div(S16, {{Vgpr16}, {Vgpr16}, AbsToNegMax})

llvm/lib/Target/AMDGPU/AMDGPURegBankLegalizeRules.h

@@ -339,7 +339,8 @@ enum LoweringMethodID {
   InsVecEltToSel,
   InsVecEltTo32,
   AbsToNegMax,
-  AbsToS32
+  AbsToS32,
+  DynStackAlloc
 };
 
 enum FastRulesTypes {

llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-uniform.ll

@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck --check-prefix=GFX9 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1010 < %s | FileCheck --check-prefix=GFX10 %s
-; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1100 < %s | FileCheck --check-prefix=GFX11 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck --check-prefix=GFX9 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1010 < %s | FileCheck --check-prefix=GFX10 %s
+; RUN: llc -global-isel -new-reg-bank-select -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1100 < %s | FileCheck --check-prefix=GFX11 %s
 
 @gv = external addrspace(4) constant i32
 

llvm/test/CodeGen/AMDGPU/GlobalISel/non-entry-alloca.ll

@@ -1,6 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -global-isel -amdgpu-load-store-vectorizer=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,DEFAULTSIZE %s
-; RUN: llc -global-isel -amdgpu-load-store-vectorizer=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -amdgpu-assume-dynamic-stack-object-size=1024 < %s | FileCheck -check-prefixes=GCN,ASSUME1024 %s
+; RUN: llc -global-isel -new-reg-bank-select -amdgpu-load-store-vectorizer=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,DEFAULTSIZE %s
+; RUN: llc -global-isel -new-reg-bank-select -amdgpu-load-store-vectorizer=0 -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -amdgpu-assume-dynamic-stack-object-size=1024 < %s | FileCheck -check-prefixes=GCN,ASSUME1024 %s
 
 ; FIXME: Generated test checks do not check metadata at the end of the
 ; function, so this also includes manually added checks.