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

[Abhinav Garg via llvm-commits]

================
@@ -1304,6 +1305,69 @@ 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());
+
+    // After lowering Dst is used by another instruction, need to erase old MI
+    // to avoid hitting multiple Dst assert since we use CSE builder
+    B.setInsertPt(*MI.getParent(), std::next(MI.getIterator()));
+    MI.eraseFromParent();
+
+    const RegisterBank *SizeBank = MRI.getRegBank(AllocSize);
+    if (SizeBank != 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);
+    LLT IntPtrTy = LLT::scalar(PtrTy.getSizeInBits());
+    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, IntPtrTy}, AllocSize, WaveSize).getReg(0);
+    }
+    auto OldSP = B.buildCopy({SgprRB, PtrTy}, SPReg);
+
+    if (Alignment > TFI.getStackAlign()) {
+      const uint64_t EffectiveAlignment =
+          HasFlatScratch ? Alignment.value()
+                         : (Alignment.value() << WavefrontSizeLog2);
+      const uint64_t StackAlignMask = EffectiveAlignment - 1;
+      auto Tmp1 =
+          B.buildPtrAdd({SgprRB, PtrTy}, OldSP,
+                        B.buildConstant({SgprRB, IntPtrTy}, StackAlignMask));
+      auto MaskReg = B.buildConstant(
+          {SgprRB, IntPtrTy},
+          maskTrailingZeros<uint64_t>(HasFlatScratch ? Log2(Alignment)
+                                                     : Log2(Alignment) +
+                                                           WavefrontSizeLog2));
+      B.buildPtrMask(Dst, Tmp1, MaskReg);
+    } else {
+      B.buildCopy(Dst, OldSP);
+    }
----------------
abhigargrepo wrote:

Accepted

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