[llvm] i[AMDGPU] Sink async DMA out of s_cbranch_execz then-blocks (PR #196374)
via llvm-commits
llvm-commits at lists.llvm.org
Thu May 7 10:18:18 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-amdgpu
Author: Vigneshwar Jayakumar (VigneshwarJ)
<details>
<summary>Changes</summary>
LLVM lowers a divergent branch around global_load_async_to_lds /
global_store_async_from_lds with s_cbranch_execz, fully-masked waves skip the DMA entirely, forcing software-pipelined kernels to use conservative waitcnts.
Sink each async DMA into the join block, immediately before the SI_END_CF EXEC restore. EXEC at the sunk slot is the same masked value, so per-lane behavior is unchanged, but async_waitcnts can be pipelined.
---
Full diff: https://github.com/llvm/llvm-project/pull/196374.diff
2 Files Affected:
- (modified) llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp (+136)
- (added) llvm/test/CodeGen/AMDGPU/sink-async-dma-out-of-execz.ll (+216)
``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp b/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
index 9cc86e84407b1..1efc15233f080 100644
--- a/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
+++ b/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
@@ -70,6 +70,11 @@ static cl::opt<bool>
RemoveRedundantEndcf("amdgpu-remove-redundant-endcf",
cl::init(true), cl::ReallyHidden);
+static cl::opt<bool> EnableSinkAsyncDMAOutOfExecz(
+ "amdgpu-sink-async-dma-out-of-execz", cl::init(true), cl::ReallyHidden,
+ cl::desc("Sink async DMA out of EXECZ-skipped then-blocks for "
+ "ASYNCcnt determinism"));
+
namespace {
class SILowerControlFlow {
@@ -132,6 +137,10 @@ class SILowerControlFlow {
// Remove redundant SI_END_CF instructions.
void optimizeEndCf();
+ // Sink async DMA ops out of EXECZ-skippable then-blocks for ASYNCcnt
+ // determinism.
+ bool sinkAsyncDMAOutOfExeczBlocks(MachineFunction &MF);
+
public:
SILowerControlFlow(const GCNSubtarget *ST, LiveIntervals *LIS,
LiveVariables *LV, MachineDominatorTree *MDT,
@@ -654,6 +663,128 @@ void SILowerControlFlow::optimizeEndCf() {
}
}
+static MachineInstr *findExecRestoreAtBlockStart(MachineBasicBlock &JoinBB,
+ const SIRegisterInfo *TRI) {
+ for (MachineInstr &MI : JoinBB) {
+ if (MI.isMetaInstruction() || MI.isDebugInstr())
+ continue;
+ return MI.modifiesRegister(AMDGPU::EXEC, TRI) ? &MI : nullptr;
+ }
+ return nullptr;
+}
+
+bool SILowerControlFlow::sinkAsyncDMAOutOfExeczBlocks(MachineFunction &MF) {
+ bool Changed = false;
+
+ for (MachineBasicBlock &MBB : MF) {
+ // Header must end with an S_CBRANCH_EXECZ to JoinBB.
+ auto BrIt = llvm::find_if(MBB.terminators(), [](const MachineInstr &MI) {
+ return MI.getOpcode() == AMDGPU::S_CBRANCH_EXECZ;
+ });
+ if (BrIt == MBB.terminators().end())
+ continue;
+ MachineBasicBlock *JoinBB = BrIt->getOperand(0).getMBB();
+
+ if (MBB.succ_size() != 2)
+ continue;
+ MachineBasicBlock *S0 = *MBB.succ_begin();
+ MachineBasicBlock *S1 = *std::next(MBB.succ_begin());
+ MachineBasicBlock *ThenBB = (S0 == JoinBB) ? S1 : S0;
+ if (ThenBB == JoinBB || ThenBB->succ_size() != 1 ||
+ *ThenBB->succ_begin() != JoinBB)
+ continue;
+
+ // Sink before the EXEC restore at the top of JoinBB.
+ MachineInstr *ExecRestore = findExecRestoreAtBlockStart(*JoinBB, TRI);
+ if (!ExecRestore)
+ continue;
+
+ SmallVector<MachineInstr *, 4> ToSink;
+ bool Eligible = true;
+ for (MachineInstr &TMI : *ThenBB) {
+ if (TMI.isMetaInstruction() || TMI.isDebugInstr() || TMI.isTerminator())
+ continue;
+ if (TII->hasUnwantedEffectsWhenEXECEmpty(TMI)) {
+ Eligible = false;
+ break;
+ }
+ if (SIInstrInfo::usesASYNC_CNT(TMI)) {
+ ToSink.push_back(&TMI);
+ continue;
+ }
+ // Bail on stores that could alias the sunk DMAs.
+ if (TMI.mayStore()) {
+ Eligible = false;
+ break;
+ }
+ // Bail on non-invariant loads that could alias the sunk DMAs.
+ if (TMI.mayLoad() && !TMI.isDereferenceableInvariantLoad()) {
+ Eligible = false;
+ break;
+ }
+ }
+ if (!Eligible || ToSink.empty())
+ continue;
+
+ // Restore dominance for ThenBB-defined operands of the sunk DMAs by
+ // inserting a header-side IMPLICIT_DEF.
+ MachineBasicBlock::iterator FirstTerm = MBB.getFirstTerminator();
+ DebugLoc DL = FirstTerm->getDebugLoc();
+ SmallSet<Register, 4> NeedsImpDef;
+ for (MachineInstr *DmaMI : ToSink) {
+ for (const MachineOperand &MO : DmaMI->uses()) {
+ if (!MO.isReg() || !MO.isUse() || !MO.getReg().isVirtual())
+ continue;
+ Register Reg = MO.getReg();
+ MachineInstr *Def = MRI->getUniqueVRegDef(Reg);
+ if (!Def || Def->getParent() != ThenBB || is_contained(ToSink, Def))
+ continue;
+ if (!NeedsImpDef.insert(Reg).second)
+ continue;
+ MachineInstr *ImpDef = BuildMI(
+ MBB, FirstTerm, DL, TII->get(TargetOpcode::IMPLICIT_DEF), Reg);
+ if (LIS)
+ LIS->InsertMachineInstrInMaps(*ImpDef);
+ RecomputeRegs.insert(Reg);
+ }
+ }
+
+ // Sink each DMA before ExecRestore and propagate phys-reg uses to
+ // JoinBB live-ins.
+ for (MachineInstr *DmaMI : ToSink) {
+ LLVM_DEBUG(dbgs() << "Sinking async DMA out of execz then-block: "
+ << *DmaMI);
+ if (LIS)
+ LIS->RemoveMachineInstrFromMaps(*DmaMI);
+ JoinBB->splice(ExecRestore->getIterator(), ThenBB, DmaMI->getIterator());
+ if (LIS)
+ LIS->InsertMachineInstrInMaps(*DmaMI);
+ for (const MachineOperand &MO : DmaMI->operands()) {
+ if (!MO.isReg() || !MO.isUse())
+ continue;
+ Register Reg = MO.getReg();
+ if (Reg.isVirtual()) {
+ if (LIS)
+ RecomputeRegs.insert(Reg);
+ // Vregs defined outside ThenBB are now live through ThenBB to reach
+ // the sunk DMA in JoinBB. Update LiveVariables AliveBlocks.
+ if (LV) {
+ MachineInstr *Def = MRI->getUniqueVRegDef(Reg);
+ if (Def && Def->getParent() != ThenBB && !is_contained(ToSink, Def))
+ LV->getVarInfo(Reg).AliveBlocks.set(ThenBB->getNumber());
+ }
+ } else if (Reg != LMC.ExecReg) {
+ JoinBB->addLiveIn(Reg);
+ }
+ }
+ }
+ JoinBB->sortUniqueLiveIns();
+ Changed = true;
+ }
+
+ return Changed;
+}
+
MachineBasicBlock *SILowerControlFlow::process(MachineInstr &MI) {
MachineBasicBlock &MBB = *MI.getParent();
MachineBasicBlock::iterator I(MI);
@@ -829,6 +960,11 @@ bool SILowerControlFlow::run(MachineFunction &MF) {
}
}
+ // Sink async DMA out of s_cbranch_execz then-blocks for ASYNCcnt
+ // determinism. Must run after the main loop so the SI_END_CF anchor exists.
+ if (EnableSinkAsyncDMAOutOfExecz)
+ Changed |= sinkAsyncDMAOutOfExeczBlocks(MF);
+
optimizeEndCf();
if (LIS && Changed) {
diff --git a/llvm/test/CodeGen/AMDGPU/sink-async-dma-out-of-execz.ll b/llvm/test/CodeGen/AMDGPU/sink-async-dma-out-of-execz.ll
new file mode 100644
index 0000000000000..5f9351e63b8b2
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/sink-async-dma-out-of-execz.ll
@@ -0,0 +1,216 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=amdgcn -mcpu=gfx1250 -verify-machineinstrs < %s | FileCheck -check-prefix=GFX1250 %s
+
+; SILowerControlFlow sinks async DMA out of s_cbranch_execz-guarded then-blocks
+; into the join-block, immediately before the s_or_b32 EXEC restore. This
+; makes ASYNCcnt deterministic regardless of whether the wave skipped the
+; block, which is required for correct waitcnt counting in software-pipelined
+; loops.
+
+
+; Basic shape: async DMA sinks into %join right before s_or_b32 exec.
+define amdgpu_ps void @async_store_simple(ptr addrspace(1) inreg %gaddr, ptr addrspace(3) %laddr, i32 %bound) {
+; GFX1250-LABEL: async_store_simple:
+; GFX1250: ; %bb.0: ; %entry
+; GFX1250-NEXT: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT: s_mov_b32 s2, exec_lo
+; GFX1250-NEXT: v_cmpx_lt_i32_e64 s0, v1
+; GFX1250-NEXT: ; implicit-def: $vgpr1
+; GFX1250-NEXT: ; %bb.1: ; %do_store
+; GFX1250-NEXT: v_mov_b32_e32 v1, 0
+; GFX1250-NEXT: ; %bb.2: ; %join
+; GFX1250-NEXT: global_store_async_from_lds_b128 v1, v0, s[0:1]
+; GFX1250-NEXT: s_or_b32 exec_lo, exec_lo, s2
+; GFX1250-NEXT: s_wait_asynccnt 0x0
+; GFX1250-NEXT: s_endpgm
+entry:
+ %tid = tail call i32 @llvm.amdgcn.workitem.id.x()
+ %cmp = icmp slt i32 %tid, %bound
+ br i1 %cmp, label %do_store, label %skip
+
+do_store:
+ tail call void @llvm.amdgcn.global.store.async.from.lds.b128(ptr addrspace(1) %gaddr, ptr addrspace(3) %laddr, i32 0, i32 0)
+ br label %join
+
+skip:
+ br label %join
+
+join:
+ tail call void @llvm.amdgcn.s.wait.asynccnt(i16 0)
+ ret void
+}
+
+; SWP-style pattern: address compute lives inside the masked block, so the
+; sink inserts header-side IMPLICIT_DEFs for the DMA's address vregs to
+; satisfy SSA dominance on the EXECZ-skip path.
+ at global_smem = external addrspace(3) global [0 x i8], align 16
+
+define amdgpu_kernel void @async_load_pipelined(ptr addrspace(1) readonly %src, i32 %M, i32 %N) {
+; GFX1250-LABEL: async_load_pipelined:
+; GFX1250: ; %bb.0: ; %entry
+; GFX1250-NEXT: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT: s_load_b64 s[0:1], s[4:5], 0x2c nv
+; GFX1250-NEXT: v_and_b32_e32 v3, 0x7f, v0
+; GFX1250-NEXT: v_bfe_u32 v2, v0, 3, 7
+; GFX1250-NEXT: s_wait_kmcnt 0x0
+; GFX1250-NEXT: s_delay_alu instid0(VALU_DEP_2) | instskip(NEXT) | instid1(VALU_DEP_2)
+; GFX1250-NEXT: v_cmp_gt_i32_e32 vcc_lo, s0, v3
+; GFX1250-NEXT: v_cmp_gt_i32_e64 s0, s1, v2
+; GFX1250-NEXT: s_and_b32 s1, vcc_lo, s0
+; GFX1250-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
+; GFX1250-NEXT: s_and_saveexec_b32 s0, s1
+; GFX1250-NEXT: ; implicit-def: $vgpr4
+; GFX1250-NEXT: ; implicit-def: $vgpr0_vgpr1
+; GFX1250-NEXT: s_cbranch_execz .LBB1_2
+; GFX1250-NEXT: ; %bb.1: ; %do_load
+; GFX1250-NEXT: s_load_b64 s[2:3], s[4:5], 0x24 nv
+; GFX1250-NEXT: v_dual_mov_b32 v1, 0 :: v_dual_lshlrev_b32 v0, 3, v3
+; GFX1250-NEXT: v_lshl_add_u32 v3, v3, 7, global_smem at abs32@lo
+; GFX1250-NEXT: s_wait_kmcnt 0x0
+; GFX1250-NEXT: s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_1) | instid1(VALU_DEP_1)
+; GFX1250-NEXT: v_add_nc_u64_e32 v[4:5], s[2:3], v[0:1]
+; GFX1250-NEXT: v_lshlrev_b32_e32 v0, 3, v2
+; GFX1250-NEXT: v_add_nc_u64_e32 v[0:1], v[4:5], v[0:1]
+; GFX1250-NEXT: v_lshl_add_u32 v4, v2, 4, v3
+; GFX1250-NEXT: .LBB1_2: ; %join
+; GFX1250-NEXT: global_load_async_to_lds_b128 v4, v[0:1], off
+; GFX1250-NEXT: s_or_b32 exec_lo, exec_lo, s0
+; GFX1250-NEXT: s_wait_asynccnt 0x0
+; GFX1250-NEXT: s_endpgm
+entry:
+ %tid = tail call i32 @llvm.amdgcn.workitem.id.x()
+ %row = and i32 %tid, 127
+ %col = lshr i32 %tid, 3
+ %cmp_row = icmp slt i32 %row, %M
+ %cmp_col = icmp slt i32 %col, %N
+ %mask = and i1 %cmp_row, %cmp_col
+ %lds_row_off = shl nuw nsw i32 %row, 7
+ %lds_col_off = shl nuw nsw i32 %col, 4
+ %lds_base = getelementptr inbounds nuw i8, ptr addrspace(3) @global_smem, i32 %lds_row_off
+ %lds_addr = getelementptr inbounds nuw i8, ptr addrspace(3) %lds_base, i32 %lds_col_off
+ %off_row = sext i32 %row to i64
+ %src_row = getelementptr [8 x i8], ptr addrspace(1) %src, i64 %off_row
+ %off_col = sext i32 %col to i64
+ %src_addr = getelementptr [8 x i8], ptr addrspace(1) %src_row, i64 %off_col
+ br i1 %mask, label %do_load, label %skip
+
+do_load:
+ tail call void @llvm.amdgcn.global.load.async.to.lds.b128(ptr addrspace(1) %src_addr, ptr addrspace(3) %lds_addr, i32 0, i32 0)
+ br label %join
+
+skip:
+ br label %join
+
+join:
+ tail call void @llvm.amdgcn.s.wait.asynccnt(i16 0)
+ ret void
+}
+
+; Negative test: an explicit s_wait_asynccnt inside the then-block disqualifies
+; sinking. Skipping the wait on the EXECZ path would observe a different
+; ASYNCcnt value, so the eligibility scan must bail and the branch must stay.
+define amdgpu_ps void @async_load_then_block_has_waitcnt(ptr addrspace(1) inreg %gaddr, ptr addrspace(3) %laddr, i32 %bound) {
+; GFX1250-LABEL: async_load_then_block_has_waitcnt:
+; GFX1250: ; %bb.0: ; %entry
+; GFX1250-NEXT: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT: s_mov_b32 s2, exec_lo
+; GFX1250-NEXT: v_cmpx_lt_i32_e64 s0, v1
+; GFX1250-NEXT: s_cbranch_execz .LBB2_2
+; GFX1250-NEXT: ; %bb.1: ; %do_load
+; GFX1250-NEXT: v_mov_b32_e32 v1, 0
+; GFX1250-NEXT: global_load_async_to_lds_b128 v0, v1, s[0:1]
+; GFX1250-NEXT: s_wait_asynccnt 0x0
+; GFX1250-NEXT: .LBB2_2: ; %join
+; GFX1250-NEXT: s_endpgm
+entry:
+ %tid = tail call i32 @llvm.amdgcn.workitem.id.x()
+ %cmp = icmp slt i32 %tid, %bound
+ br i1 %cmp, label %do_load, label %skip
+
+do_load:
+ tail call void @llvm.amdgcn.global.load.async.to.lds.b128(ptr addrspace(1) %gaddr, ptr addrspace(3) %laddr, i32 0, i32 0)
+ tail call void @llvm.amdgcn.s.wait.asynccnt(i16 0)
+ br label %join
+
+skip:
+ br label %join
+
+join:
+ ret void
+}
+
+; Negative test: a non-invariant DS load reading from the LDS region the
+; async load writes to would read stale data if the DMA were sunk past it.
+; The mayLoad bail must fire.
+define amdgpu_ps void @async_load_then_block_has_aliasing_lds_load(ptr addrspace(1) inreg %gaddr, ptr addrspace(3) %laddr, ptr addrspace(1) %out, i32 %bound) {
+; GFX1250-LABEL: async_load_then_block_has_aliasing_lds_load:
+; GFX1250: ; %bb.0: ; %entry
+; GFX1250-NEXT: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT: v_dual_mov_b32 v5, v2 :: v_dual_mov_b32 v4, v1
+; GFX1250-NEXT: s_mov_b32 s2, exec_lo
+; GFX1250-NEXT: v_cmpx_lt_i32_e64 s0, v3
+; GFX1250-NEXT: s_cbranch_execz .LBB3_2
+; GFX1250-NEXT: ; %bb.1: ; %do_load
+; GFX1250-NEXT: v_mov_b32_e32 v1, 0
+; GFX1250-NEXT: global_load_async_to_lds_b128 v0, v1, s[0:1]
+; GFX1250-NEXT: ds_load_b32 v0, v0
+; GFX1250-NEXT: s_wait_dscnt 0x0
+; GFX1250-NEXT: global_store_b32 v[4:5], v0, off
+; GFX1250-NEXT: .LBB3_2: ; %join
+; GFX1250-NEXT: s_wait_xcnt 0x0
+; GFX1250-NEXT: s_or_b32 exec_lo, exec_lo, s2
+; GFX1250-NEXT: s_wait_asynccnt 0x0
+; GFX1250-NEXT: s_endpgm
+entry:
+ %tid = tail call i32 @llvm.amdgcn.workitem.id.x()
+ %cmp = icmp slt i32 %tid, %bound
+ br i1 %cmp, label %do_load, label %skip
+
+do_load:
+ tail call void @llvm.amdgcn.global.load.async.to.lds.b128(ptr addrspace(1) %gaddr, ptr addrspace(3) %laddr, i32 0, i32 0)
+ %val = load i32, ptr addrspace(3) %laddr
+ store i32 %val, ptr addrspace(1) %out
+ br label %join
+
+skip:
+ br label %join
+
+join:
+ tail call void @llvm.amdgcn.s.wait.asynccnt(i16 0)
+ ret void
+}
+
+; Negative test: a regular store to the same global address the async load
+define amdgpu_ps void @async_load_then_block_has_aliasing_global_store(ptr addrspace(1) %gaddr, ptr addrspace(3) %laddr, i32 %bound) {
+; GFX1250-LABEL: async_load_then_block_has_aliasing_global_store:
+; GFX1250: ; %bb.0: ; %entry
+; GFX1250-NEXT: s_setreg_imm32_b32 hwreg(HW_REG_WAVE_MODE, 25, 1), 1 ; msbs: dst=0 src0=0 src1=0 src2=0
+; GFX1250-NEXT: v_cmp_lt_i32_e32 vcc_lo, s0, v3
+; GFX1250-NEXT: s_and_saveexec_b32 s0, vcc_lo
+; GFX1250-NEXT: s_cbranch_execz .LBB4_2
+; GFX1250-NEXT: ; %bb.1: ; %do_load
+; GFX1250-NEXT: global_load_async_to_lds_b128 v2, v[0:1], off
+; GFX1250-NEXT: v_mov_b32_e32 v2, 0
+; GFX1250-NEXT: global_store_b32 v[0:1], v2, off
+; GFX1250-NEXT: .LBB4_2: ; %join
+; GFX1250-NEXT: s_wait_xcnt 0x0
+; GFX1250-NEXT: s_or_b32 exec_lo, exec_lo, s0
+; GFX1250-NEXT: s_wait_asynccnt 0x0
+; GFX1250-NEXT: s_endpgm
+entry:
+ %tid = tail call i32 @llvm.amdgcn.workitem.id.x()
+ %cmp = icmp slt i32 %tid, %bound
+ br i1 %cmp, label %do_load, label %skip
+
+do_load:
+ tail call void @llvm.amdgcn.global.load.async.to.lds.b128(ptr addrspace(1) %gaddr, ptr addrspace(3) %laddr, i32 0, i32 0)
+ store i32 0, ptr addrspace(1) %gaddr
+ br label %join
+
+skip:
+ br label %join
+
+join:
+ tail call void @llvm.amdgcn.s.wait.asynccnt(i16 0)
+ ret void
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/196374
More information about the llvm-commits
mailing list