[AMDGPU] Support divergent sized dynamic alloca #121148

easyonaadit · 2024-12-26T12:59:04Z

Currently, AMDGPU backend can handle uniform-sized dynamic allocas.
This patch extends support for divergent-sized dynamic allocas.
When the size argument of a dynamic alloca is divergent,
a wave-wide reduction is performed to get the required stack space.
@llvm.amdgcn.wave.reduce.umax is used to perform the
wave reduction.

Dynamic allocas are not completely supported yet,
as the stack is not properly restored on function exit.
This patch doesn't attempt to address the aforementioned issue.

Note: Compiler already Zero-Extends or Truncates all other
types(of alloca size arg) to i32.

Currently, AMDGPU backend supports static sized allocas. This patch extends support for truly dynamic sized allocas. The value representing dynamic alloca size can be divergent which needs a wave-wide reduction to get the stack space. @llvm.amdgcn.wave.reduce.umax is used to perform the wave reduction. Note: Compiler already Zero-Extends or Truncates all other types(of alloca size arg) to i32.

llvmbot · 2024-12-26T12:59:39Z

@llvm/pr-subscribers-llvm-globalisel

@llvm/pr-subscribers-backend-amdgpu

Author: Aaditya (easyonaadit)

Changes

Currently, AMDGPU backend supports static sized allocas.
This patch extends support for truly dynamic sized allocas.
The value representing dynamic alloca size can be divergent
which needs a wave-wide reduction to get the stack space.
@llvm.amdgcn.wave.reduce.umax is used to perform the
wave reduction.

Note: Compiler already Zero-Extends or Truncates all other
types(of alloca size arg) to i32.

Patch is 98.72 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/121148.diff

6 Files Affected:

(modified) llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp (+8-3)
(modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+34-26)
(modified) llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll (+255-23)
(modified) llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-dyn-stackalloc.mir (+129)
(modified) llvm/test/CodeGen/AMDGPU/dynamic_stackalloc.ll (+1287-41)
(added) llvm/test/CodeGen/AMDGPU/dynamic_stackalloc_wave_reduce.ll (+170)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
index d94c400ad14225..bdc737958b0e96 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp
@@ -1190,9 +1190,14 @@ bool AMDGPURegisterBankInfo::applyMappingDynStackAlloc(
 
   const RegisterBank *SizeBank = getRegBank(AllocSize, MRI, *TRI);
 
-  // TODO: Need to emit a wave reduction to get the maximum size.
-  if (SizeBank != &AMDGPU::SGPRRegBank)
-    return false;
+  if (SizeBank != &AMDGPU::SGPRRegBank) {
+    auto WaveReduction =
+        B.buildIntrinsic(Intrinsic::amdgcn_wave_reduce_umax,
+                         {LLT::scalar(MRI.getType(AllocSize).getSizeInBits())})
+            .addUse(AllocSize)
+            .addImm(0);
+    AllocSize = WaveReduction.getReg(0);
+  }
 
   LLT PtrTy = MRI.getType(Dst);
   LLT IntPtrTy = LLT::scalar(PtrTy.getSizeInBits());
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 58b061f5c1af0d..dad743fd00c2ca 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -4017,8 +4017,9 @@ SDValue SITargetLowering::LowerCall(CallLoweringInfo &CLI,
 }
 
 // This is similar to the default implementation in ExpandDYNAMIC_STACKALLOC,
-// except for stack growth direction(default: downwards, AMDGPU: upwards) and
-// applying the wave size scale to the increment amount.
+// except for:
+// 1. stack growth direction(default: downwards, AMDGPU: upwards), and
+// 2. scale size where, scale = wave-reduction(alloca-size) * wave-size
 SDValue SITargetLowering::lowerDYNAMIC_STACKALLOCImpl(SDValue Op,
                                                       SelectionDAG &DAG) const {
   const MachineFunction &MF = DAG.getMachineFunction();
@@ -4026,20 +4027,16 @@ SDValue SITargetLowering::lowerDYNAMIC_STACKALLOCImpl(SDValue Op,
 
   SDLoc dl(Op);
   EVT VT = Op.getValueType();
-  SDValue Tmp1 = Op;
-  SDValue Tmp2 = Op.getValue(1);
-  SDValue Tmp3 = Op.getOperand(2);
-  SDValue Chain = Tmp1.getOperand(0);
-
+  SDValue Chain = Op.getOperand(0);
   Register SPReg = Info->getStackPtrOffsetReg();
 
   // Chain the dynamic stack allocation so that it doesn't modify the stack
   // pointer when other instructions are using the stack.
   Chain = DAG.getCALLSEQ_START(Chain, 0, 0, dl);
 
-  SDValue Size = Tmp2.getOperand(1);
+  SDValue Size = Op.getValue(1).getOperand(1);
   SDValue BaseAddr = DAG.getCopyFromReg(Chain, dl, SPReg, VT);
-  Align Alignment = cast<ConstantSDNode>(Tmp3)->getAlignValue();
+  Align Alignment = cast<ConstantSDNode>(Op.getOperand(2))->getAlignValue();
 
   const TargetFrameLowering *TFL = Subtarget->getFrameLowering();
   assert(TFL->getStackGrowthDirection() == TargetFrameLowering::StackGrowsUp &&
@@ -4057,30 +4054,41 @@ SDValue SITargetLowering::lowerDYNAMIC_STACKALLOCImpl(SDValue Op,
                            DAG.getSignedConstant(-ScaledAlignment, dl, VT));
   }
 
-  SDValue ScaledSize = DAG.getNode(
-      ISD::SHL, dl, VT, Size,
-      DAG.getConstant(Subtarget->getWavefrontSizeLog2(), dl, MVT::i32));
-
-  SDValue NewSP = DAG.getNode(ISD::ADD, dl, VT, BaseAddr, ScaledSize); // Value
+  assert(Size.getValueType() == MVT::i32 && "Size must be 32-bit");
+  SDValue NewSP;
+  if (isa<ConstantSDNode>(Op.getOperand(1))) {
+    // for constant sized alloca, scale alloca size by wave-size
+    SDValue ScaledSize = DAG.getNode(
+        ISD::SHL, dl, VT, Size,
+        DAG.getConstant(Subtarget->getWavefrontSizeLog2(), dl, MVT::i32));
+    NewSP = DAG.getNode(ISD::ADD, dl, VT, BaseAddr, ScaledSize); // Value
+  } else {
+    // for dynamic sized alloca, perform wave-wide reduction to get max of
+    // alloca size(divergent) and then scale it by wave-size
+    SDValue WaveReduction =
+        DAG.getTargetConstant(Intrinsic::amdgcn_wave_reduce_umax, dl, MVT::i32);
+    Size = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, dl, MVT::i32, WaveReduction,
+                       Size, DAG.getConstant(0, dl, MVT::i32));
+    SDValue ScaledSize = DAG.getNode(
+        ISD::SHL, dl, VT, Size,
+        DAG.getConstant(Subtarget->getWavefrontSizeLog2(), dl, MVT::i32));
+    NewSP =
+        DAG.getNode(ISD::ADD, dl, VT, BaseAddr, ScaledSize); // Value in vgpr.
+    SDValue ReadFirstLaneID =
+        DAG.getTargetConstant(Intrinsic::amdgcn_readfirstlane, dl, MVT::i32);
+    NewSP = DAG.getNode(ISD::INTRINSIC_WO_CHAIN, dl, MVT::i32, ReadFirstLaneID,
+                        NewSP);
+  }
 
   Chain = DAG.getCopyToReg(Chain, dl, SPReg, NewSP); // Output chain
-  Tmp2 = DAG.getCALLSEQ_END(Chain, 0, 0, SDValue(), dl);
+  SDValue CallSeqEnd = DAG.getCALLSEQ_END(Chain, 0, 0, SDValue(), dl);
 
-  return DAG.getMergeValues({BaseAddr, Tmp2}, dl);
+  return DAG.getMergeValues({BaseAddr, CallSeqEnd}, dl);
 }
 
 SDValue SITargetLowering::LowerDYNAMIC_STACKALLOC(SDValue Op,
                                                   SelectionDAG &DAG) const {
-  // We only handle constant sizes here to allow non-entry block, static sized
-  // allocas. A truly dynamic value is more difficult to support because we
-  // don't know if the size value is uniform or not. If the size isn't uniform,
-  // we would need to do a wave reduction to get the maximum size to know how
-  // much to increment the uniform stack pointer.
-  SDValue Size = Op.getOperand(1);
-  if (isa<ConstantSDNode>(Size))
-    return lowerDYNAMIC_STACKALLOCImpl(Op, DAG); // Use "generic" expansion.
-
-  return AMDGPUTargetLowering::LowerDYNAMIC_STACKALLOC(Op, DAG);
+  return lowerDYNAMIC_STACKALLOCImpl(Op, DAG); // Use "generic" expansion.
 }
 
 SDValue SITargetLowering::LowerSTACKSAVE(SDValue Op, SelectionDAG &DAG) const {
diff --git a/llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll b/llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll
index aefcad491073fc..a76a125f0580b0 100644
--- a/llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll
+++ b/llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll
@@ -1,10 +1,38 @@
-; RUN: not llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -global-isel-abort=2 -pass-remarks-missed="gisel.*" -verify-machineinstrs=0 -o /dev/null 2>&1 %s | FileCheck -check-prefix=ERR %s
-
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 1 (in function: kernel_dynamic_stackalloc_vgpr_align4)
-; ERR-NEXT: warning: Instruction selection used fallback path for kernel_dynamic_stackalloc_vgpr_align4
-; ERR-NEXT: error: <unknown>:0:0: in function kernel_dynamic_stackalloc_vgpr_align4 void (ptr addrspace(1)): unsupported dynamic alloca
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -global-isel -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefix=GFX9 %s
 
 define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_align4(ptr addrspace(1) %ptr) {
+; GFX9-LABEL: kernel_dynamic_stackalloc_vgpr_align4:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_load_dwordx2 s[4:5], s[8:9], 0x0
+; GFX9-NEXT:    v_lshlrev_b32_e32 v0, 2, v0
+; GFX9-NEXT:    s_add_u32 s0, s0, s17
+; GFX9-NEXT:    s_addc_u32 s1, s1, 0
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX9-NEXT:    global_load_dword v0, v0, s[4:5]
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s33, 0
+; GFX9-NEXT:    s_movk_i32 s32, 0x400
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:  .LBB0_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB0_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_mov_b32 s4, s32
+; GFX9-NEXT:    s_lshl_b32 s5, s6, 6
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x7b
+; GFX9-NEXT:    v_mov_b32_e32 v1, s4
+; GFX9-NEXT:    s_add_u32 s32, s4, s5
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_endpgm
   %id = call i32 @llvm.amdgcn.workitem.id.x()
   %gep = getelementptr i32, ptr addrspace(1) %ptr, i32 %id
   %n = load i32, ptr addrspace(1) %gep
@@ -13,11 +41,38 @@ define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_align4(ptr addrspace(1
   ret void
 }
 
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 1 (in function: kernel_dynamic_stackalloc_vgpr_default_align)
-; ERR-NEXT: warning: Instruction selection used fallback path for kernel_dynamic_stackalloc_vgpr_default_align
-; ERR-NEXT: error: <unknown>:0:0: in function kernel_dynamic_stackalloc_vgpr_default_align void (ptr addrspace(1)): unsupported dynamic alloca
-
 define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_default_align(ptr addrspace(1) %ptr) {
+; GFX9-LABEL: kernel_dynamic_stackalloc_vgpr_default_align:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_load_dwordx2 s[4:5], s[8:9], 0x0
+; GFX9-NEXT:    v_lshlrev_b32_e32 v0, 2, v0
+; GFX9-NEXT:    s_add_u32 s0, s0, s17
+; GFX9-NEXT:    s_addc_u32 s1, s1, 0
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX9-NEXT:    global_load_dword v0, v0, s[4:5]
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s33, 0
+; GFX9-NEXT:    s_movk_i32 s32, 0x400
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:  .LBB1_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB1_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_mov_b32 s4, s32
+; GFX9-NEXT:    s_lshl_b32 s5, s6, 6
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x7b
+; GFX9-NEXT:    v_mov_b32_e32 v1, s4
+; GFX9-NEXT:    s_add_u32 s32, s4, s5
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_endpgm
   %id = call i32 @llvm.amdgcn.workitem.id.x()
   %gep = getelementptr i32, ptr addrspace(1) %ptr, i32 %id
   %n = load i32, ptr addrspace(1) %gep
@@ -25,11 +80,40 @@ define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_default_align(ptr addr
   store volatile i32 123, ptr addrspace(5) %alloca
   ret void
 }
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 64 (in function: kernel_dynamic_stackalloc_vgpr_align64)
-; ERR-NEXT: warning: Instruction selection used fallback path for kernel_dynamic_stackalloc_vgpr_align64
-; ERR-NEXT: error: <unknown>:0:0: in function kernel_dynamic_stackalloc_vgpr_align64 void (ptr addrspace(1)): unsupported dynamic alloca
 
 define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_align64(ptr addrspace(1) %ptr) {
+; GFX9-LABEL: kernel_dynamic_stackalloc_vgpr_align64:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_load_dwordx2 s[4:5], s[8:9], 0x0
+; GFX9-NEXT:    v_lshlrev_b32_e32 v0, 2, v0
+; GFX9-NEXT:    s_add_u32 s0, s0, s17
+; GFX9-NEXT:    s_addc_u32 s1, s1, 0
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX9-NEXT:    global_load_dword v0, v0, s[4:5]
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s33, 0
+; GFX9-NEXT:    s_movk_i32 s32, 0x1000
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:  .LBB2_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB2_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_add_u32 s5, s32, 0xfff
+; GFX9-NEXT:    s_and_b32 s5, s5, 0xfffff000
+; GFX9-NEXT:    s_lshl_b32 s4, s6, 6
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x7b
+; GFX9-NEXT:    v_mov_b32_e32 v1, s5
+; GFX9-NEXT:    s_add_u32 s32, s5, s4
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_endpgm
   %id = call i32 @llvm.amdgcn.workitem.id.x()
   %gep = getelementptr i32, ptr addrspace(1) %ptr, i32 %id
   %n = load i32, ptr addrspace(1) %gep
@@ -38,35 +122,183 @@ define amdgpu_kernel void @kernel_dynamic_stackalloc_vgpr_align64(ptr addrspace(
   ret void
 }
 
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 1 (in function: func_dynamic_stackalloc_vgpr_align4)
-; ERR-NEXT: warning: Instruction selection used fallback path for func_dynamic_stackalloc_vgpr_align4
-; ERR-NEXT: error: <unknown>:0:0: in function func_dynamic_stackalloc_vgpr_align4 void (i32): unsupported dynamic alloca
-
 define void @func_dynamic_stackalloc_vgpr_align4(i32 %n) {
+; GFX9-LABEL: func_dynamic_stackalloc_vgpr_align4:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    s_mov_b32 s9, s33
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_mov_b32 s33, s32
+; GFX9-NEXT:    s_addk_i32 s32, 0x400
+; GFX9-NEXT:  .LBB3_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB3_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_mov_b32 s4, s32
+; GFX9-NEXT:    s_lshl_b32 s5, s6, 6
+; GFX9-NEXT:    s_add_u32 s32, s4, s5
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x1c8
+; GFX9-NEXT:    v_mov_b32_e32 v1, s4
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_addk_i32 s32, 0xfc00
+; GFX9-NEXT:    s_mov_b32 s33, s9
+; GFX9-NEXT:    s_setpc_b64 s[30:31]
   %alloca = alloca i32, i32 %n, align 4, addrspace(5)
   store volatile i32 456, ptr addrspace(5) %alloca
   ret void
 }
 
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 1 (in function: func_dynamic_stackalloc_vgpr_default_align)
-; ERR-NEXT: warning: Instruction selection used fallback path for func_dynamic_stackalloc_vgpr_default_align
-; ERR-NEXT: error: <unknown>:0:0: in function func_dynamic_stackalloc_vgpr_default_align void (i32): unsupported dynamic alloca
-
 define void @func_dynamic_stackalloc_vgpr_default_align(i32 %n) {
+; GFX9-LABEL: func_dynamic_stackalloc_vgpr_default_align:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    s_mov_b32 s9, s33
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_mov_b32 s33, s32
+; GFX9-NEXT:    s_addk_i32 s32, 0x400
+; GFX9-NEXT:  .LBB4_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB4_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_mov_b32 s4, s32
+; GFX9-NEXT:    s_lshl_b32 s5, s6, 6
+; GFX9-NEXT:    s_add_u32 s32, s4, s5
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x1c8
+; GFX9-NEXT:    v_mov_b32_e32 v1, s4
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_addk_i32 s32, 0xfc00
+; GFX9-NEXT:    s_mov_b32 s33, s9
+; GFX9-NEXT:    s_setpc_b64 s[30:31]
   %alloca = alloca i32, i32 %n, addrspace(5)
   store volatile i32 456, ptr addrspace(5) %alloca
   ret void
 }
-; ERR: remark: <unknown>:0:0: cannot select: %{{[0-9]+}}:sreg_32(p5) = G_DYN_STACKALLOC %{{[0-9]+}}:vgpr(s32), 64 (in function: func_dynamic_stackalloc_vgpr_align64)
-; ERR-NEXT: warning: Instruction selection used fallback path for func_dynamic_stackalloc_vgpr_align64
-; ERR-NEXT: error: <unknown>:0:0: in function func_dynamic_stackalloc_vgpr_align64 void (i32): unsupported dynamic alloca
 
 define void @func_dynamic_stackalloc_vgpr_align64(i32 %n) {
+; GFX9-LABEL: func_dynamic_stackalloc_vgpr_align64:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    s_mov_b32 s9, s33
+; GFX9-NEXT:    s_add_i32 s33, s32, 0xfc0
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_and_b32 s33, s33, 0xfffff000
+; GFX9-NEXT:    s_addk_i32 s32, 0x2000
+; GFX9-NEXT:  .LBB5_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB5_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_add_u32 s5, s32, 0xfff
+; GFX9-NEXT:    s_lshl_b32 s4, s6, 6
+; GFX9-NEXT:    s_and_b32 s5, s5, 0xfffff000
+; GFX9-NEXT:    s_add_u32 s32, s5, s4
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x1c8
+; GFX9-NEXT:    v_mov_b32_e32 v1, s5
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_addk_i32 s32, 0xe000
+; GFX9-NEXT:    s_mov_b32 s33, s9
+; GFX9-NEXT:    s_setpc_b64 s[30:31]
   %alloca = alloca i32, i32 %n, align 64, addrspace(5)
   store volatile i32 456, ptr addrspace(5) %alloca
   ret void
 }
 
+define void @func_dynamic_stackalloc_non_standard_size_i64(i64 %n) {
+; GFX9-LABEL: func_dynamic_stackalloc_non_standard_size_i64:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    s_mov_b32 s9, s33
+; GFX9-NEXT:    s_add_i32 s33, s32, 0xfc0
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_and_b32 s33, s33, 0xfffff000
+; GFX9-NEXT:    s_addk_i32 s32, 0x2000
+; GFX9-NEXT:  .LBB6_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB6_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_add_u32 s5, s32, 0xfff
+; GFX9-NEXT:    s_lshl_b32 s4, s6, 6
+; GFX9-NEXT:    s_and_b32 s5, s5, 0xfffff000
+; GFX9-NEXT:    s_add_u32 s32, s5, s4
+; GFX9-NEXT:    v_mov_b32_e32 v0, 0x1c8
+; GFX9-NEXT:    v_mov_b32_e32 v1, s5
+; GFX9-NEXT:    buffer_store_dword v0, v1, s[0:3], 0 offen
+; GFX9-NEXT:    s_waitcnt vmcnt(0)
+; GFX9-NEXT:    s_addk_i32 s32, 0xe000
+; GFX9-NEXT:    s_mov_b32 s33, s9
+; GFX9-NEXT:    s_setpc_b64 s[30:31]
+  %alloca = alloca i32, i64 %n, align 64, addrspace(5)
+  store volatile i32 456, ptr addrspace(5) %alloca
+  ret void
+}
+
+define void @func_dynamic_stackalloc_non_standard_size_i16(i16 %n) {
+; GFX9-LABEL: func_dynamic_stackalloc_non_standard_size_i16:
+; GFX9:       ; %bb.0:
+; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
+; GFX9-NEXT:    v_and_b32_e32 v0, 0xffff, v0
+; GFX9-NEXT:    v_lshl_add_u32 v0, v0, 2, 15
+; GFX9-NEXT:    s_mov_b32 s9, s33
+; GFX9-NEXT:    s_add_i32 s33, s32, 0xfc0
+; GFX9-NEXT:    v_and_b32_e32 v0, -16, v0
+; GFX9-NEXT:    s_mov_b64 s[4:5], exec
+; GFX9-NEXT:    s_mov_b32 s6, 0
+; GFX9-NEXT:    s_and_b32 s33, s33, 0xfffff000
+; GFX9-NEXT:    s_addk_i32 s32, 0x2000
+; GFX9-NEXT:  .LBB7_1: ; =>This Inner Loop Header: Depth=1
+; GFX9-NEXT:    s_ff1_i32_b64 s7, s[4:5]
+; GFX9-NEXT:    v_readlane_b32 s8, v0, s7
+; GFX9-NEXT:    s_bitset0_b64 s[4:5], s7
+; GFX9-NEXT:    s_max_u32 s6, s6, s8
+; GFX9-NEXT:    s_cmp_lg_u64 s[4:5], 0
+; GFX9-NEXT:    s_cbranch_scc1 .LBB7_1
+; GFX9-NEXT:  ; %bb.2:
+; GFX9-NEXT:    s_add_u32 s5, s32, 0xfff
+; GFX9-NEXT:    s_lshl_b32 s4, s6, 6
+; GFX9-NEXT:    s_and_b32 s5, s5, 0xfffff000
+; GF...
[truncated]

arsenm · 2024-12-26T13:04:57Z

Title is misleading, this is attempting to only handle dynamic allocation with divergent sizes. We still have the SP is not restored issue with dynamic allocate, in general

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

github-actions · 2024-12-27T06:05:06Z

✅ With the latest revision this PR passed the C/C++ code formatter.

pravinjagtap · 2024-12-27T06:34:18Z

We still have the SP is not restored issue with dynamic allocate, in general

Will take up this in next patch.

pravinjagtap

LGTM, wait for others.
Title could be "[AMDGPU] Support divergent sized allocation"?

llvm/test/CodeGen/AMDGPU/dynamic_stackalloc_wave_reduce.ll

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

pravinjagtap · 2025-01-02T13:52:43Z

Ping @arsenm @s-barannikov

jayfoad · 2025-01-02T13:54:33Z

Currently, AMDGPU backend supports static sized allocas.
This patch extends support for truly dynamic sized allocas.

Update this description as well as the title.

pravinjagtap · 2025-01-02T14:00:51Z

Currently, AMDGPU backend supports static sized allocas.
This patch extends support for truly dynamic sized allocas.

Update this description as well as the title.

The title is already updated by matt.

arsenm · 2025-01-02T14:10:42Z

The title is already updated by matt.

I updated the title but the longer description is still not accurate

easyonaadit · 2025-01-02T15:07:42Z

Currently, AMDGPU backend supports static sized allocas.
This patch extends support for truly dynamic sized allocas.

Update this description as well as the title.

Done. I have tried to explicitly mention what is not supported, hopefully its no longer ambiguous.

llvm/test/CodeGen/AMDGPU/dynamic_stackalloc.ll

s-barannikov

a couple of nits

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

easyonaadit · 2025-01-02T16:40:16Z

The title is already updated by matt.

I updated the title but the longer description is still not accurate

Is the description better suited now?

Currently, AMDGPU backend can handle uniform-sized dynamic allocas. This patch extends support for divergent-sized dynamic allocas. When the size argument of a dynamic alloca is divergent, a wave-wide reduction is performed to get the required stack space. `@llvm.amdgcn.wave.reduce.umax` is used to perform the wave reduction. Dynamic allocas are not completely supported yet, as the stack is not properly restored on function exit. This patch doesn't attempt to address the aforementioned issue. Note: Compiler already Zero-Extends or Truncates all other types(of alloca size arg) to i32.

llvmbot added backend:AMDGPU llvm:globalisel labels Dec 26, 2024

pravinjagtap requested review from jayfoad, arsenm, cdevadas and pravinjagtap December 26, 2024 13:04

pravinjagtap requested a review from s-barannikov December 26, 2024 13:05

arsenm reviewed Dec 26, 2024

View reviewed changes

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp Outdated Show resolved Hide resolved

llvm/lib/Target/AMDGPU/SIISelLowering.cpp Outdated Show resolved Hide resolved

pravinjagtap requested a review from nikic December 26, 2024 13:07

s-barannikov reviewed Dec 26, 2024

View reviewed changes

llvm/lib/Target/AMDGPU/SIISelLowering.cpp Outdated Show resolved Hide resolved

Review comments

6a6fe25

easyonaadit force-pushed the support-dynamically-sized-allocas branch from 0377bb6 to 6a6fe25 Compare December 27, 2024 06:25

easyonaadit changed the title ~~[AMDGPU] Support dynamically sized allocas~~ [AMDGPU] Support divergent sized allocas Dec 27, 2024

pravinjagtap approved these changes Dec 27, 2024

View reviewed changes

easyonaadit changed the title ~~[AMDGPU] Support divergent sized allocas~~ [AMDGPU] Support divergent sized allocation Dec 27, 2024

arsenm changed the title ~~[AMDGPU] Support divergent sized allocation~~ [AMDGPU] Support divergent sized dynamic alloca Dec 27, 2024

arsenm reviewed Dec 27, 2024

View reviewed changes

llvm/test/CodeGen/AMDGPU/dynamic_stackalloc_wave_reduce.ll Outdated Show resolved Hide resolved

llvm/test/CodeGen/AMDGPU/GlobalISel/dynamic-alloca-divergent.ll Outdated Show resolved Hide resolved

Merged SDAG and G-ISEL tests

f03f744

Removed redundant test cases

80d7439

arsenm reviewed Jan 2, 2025

View reviewed changes

llvm/test/CodeGen/AMDGPU/dynamic_stackalloc.ll Outdated Show resolved Hide resolved

Review Comments

a4519a8

s-barannikov reviewed Jan 2, 2025

View reviewed changes

llvm/lib/Target/AMDGPU/SIISelLowering.cpp Outdated Show resolved Hide resolved

llvm/lib/Target/AMDGPU/SIISelLowering.cpp Outdated Show resolved Hide resolved

Review comments

ce257c7

arsenm approved these changes Jan 6, 2025

View reviewed changes

arsenm merged commit 0bd1c87 into llvm:main Jan 6, 2025
8 checks passed

pravinjagtap mentioned this pull request Feb 10, 2025

[AMDGPU] Extend wave reduce intrinsics for i32 type #126469

Merged

[AMDGPU] Support divergent sized dynamic alloca #121148

[AMDGPU] Support divergent sized dynamic alloca #121148

Uh oh!

Conversation

easyonaadit commented Dec 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Dec 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

arsenm commented Dec 26, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Dec 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pravinjagtap commented Dec 27, 2024

Uh oh!

pravinjagtap left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

pravinjagtap commented Jan 2, 2025

Uh oh!

jayfoad commented Jan 2, 2025

Uh oh!

pravinjagtap commented Jan 2, 2025

Uh oh!

arsenm commented Jan 2, 2025

Uh oh!

easyonaadit commented Jan 2, 2025

Uh oh!

Uh oh!

s-barannikov left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

easyonaadit commented Jan 2, 2025

Uh oh!

Uh oh!

Uh oh!

easyonaadit commented Dec 26, 2024 •

edited

Loading

llvmbot commented Dec 26, 2024 •

edited

Loading

github-actions bot commented Dec 27, 2024 •

edited

Loading