[llvm] [AMDGPU] Fix a wave.reduce crash on a move of a non-immediate (PR #228819)
Mian Miftah via llvm-commits
llvm-commits at lists.llvm.org
Sat Oct 3 21:15:12 PDT 2026
https://github.com/mmiftahx created https://github.com/llvm/llvm-project/pull/228819
lowerWaveReduce checks whether a uniform source is the constant 1 by
looking at the move that defines it (#212755), and reads the move's
operand with getImm() without checking that it is an immediate. A source
defined by a move of a register, a frame index or a global address hits
the assertion in getImm(). An add reduction of a uniform umin's result
does this, since the umin is lowered to a move of its source.
Use getImmOrMaterializedImm, which returns nothing unless the operand is
an immediate.
>From b6d104c6105a2c57375d97943891dcb0b5f632b9 Mon Sep 17 00:00:00 2001
From: mmiftahx <mmiftah.duna at gmail.com>
Date: Thu, 1 Oct 2026 21:38:10 -0500
Subject: [PATCH] [AMDGPU] Fix a wave.reduce crash on a move of a non-immediate
lowerWaveReduce checks whether a uniform source is the constant 1 by
looking at the move that defines it (#212755), and reads the move's
operand with getImm() without checking that it is an immediate. A source
defined by a move of a register, a frame index or a global address hits
the assertion in getImm(). An add reduction of a uniform umin's result
does this, since the umin is lowered to a move of its source.
Use getImmOrMaterializedImm, which returns nothing unless the operand is
an immediate.
---
llvm/lib/Target/AMDGPU/SIISelLowering.cpp | 58 ++++------
.../CodeGen/AMDGPU/wave-reduce-move-source.ll | 109 ++++++++++++++++++
2 files changed, 133 insertions(+), 34 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/wave-reduce-move-source.ll
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 4012992d3f1d6a3..eefcaf9c5c9120c 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -6128,21 +6128,18 @@ static MachineBasicBlock *lowerWaveReduce(MachineInstr &MI,
.addImm(1)
.setOperandDead(3); // Dead scc
// Check if Src is a known identity constant.
- MachineInstr *SrcDef = MRI.getVRegDef(SrcReg);
- if (SrcDef && SrcDef->isMoveImmediate()) {
- int64_t Imm = SrcDef->getOperand(1).getImm();
- if (Imm == 1) { // 1 * parity(exec) = parity(exec)
- if (Opc == AMDGPU::S_XOR_B32) {
- BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstReg)
- .addReg(ParityRegister);
- } else {
- Register DstHi =
- MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
- BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstHi).addImm(0);
- BuildRegSequence(BB, MI, DstReg, ParityRegister, DstHi);
- }
- break;
+ if (TII->getImmOrMaterializedImm(MRI, SrcReg) == 1) {
+ // 1 * parity(exec) = parity(exec)
+ if (Opc == AMDGPU::S_XOR_B32) {
+ BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstReg)
+ .addReg(ParityRegister);
+ } else {
+ Register DstHi =
+ MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
+ BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstHi).addImm(0);
+ BuildRegSequence(BB, MI, DstReg, ParityRegister, DstHi);
}
+ break;
}
if (Opc == AMDGPU::S_XOR_B32) {
BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MUL_I32), DstReg)
@@ -6179,14 +6176,11 @@ static MachineBasicBlock *lowerWaveReduce(MachineInstr &MI,
}
case AMDGPU::S_ADD_I32: {
// Check if Src is a known identity constant.
- MachineInstr *SrcDef = MRI.getVRegDef(SrcReg);
- if (SrcDef && SrcDef->isMoveImmediate()) {
- int64_t Imm = SrcDef->getOperand(1).getImm();
- if (Imm == 1) { // 1 * bitcount(exec) = bitcount(exec)
- BuildMI(BB, MI, DL, TII->get(AMDGPU::COPY), DstReg)
- .addReg(NewAccumulator->getOperand(0).getReg());
- break;
- }
+ if (TII->getImmOrMaterializedImm(MRI, SrcReg) == 1) {
+ // 1 * bitcount(exec) = bitcount(exec)
+ BuildMI(BB, MI, DL, TII->get(AMDGPU::COPY), DstReg)
+ .addReg(NewAccumulator->getOperand(0).getReg());
+ break;
}
BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MUL_I32), DstReg)
.addReg(SrcReg)
@@ -6211,18 +6205,14 @@ static MachineBasicBlock *lowerWaveReduce(MachineInstr &MI,
auto [Op1L, Op1H] = ExtractSubRegs(MI, MI.getOperand(1),
MRI.getRegClass(SrcReg), ST, MRI);
// Check if Src is a known identity constant.
- MachineInstr *SrcDef = MRI.getVRegDef(SrcReg);
- if (SrcDef && SrcDef->isMoveImmediate()) {
- int64_t Imm = SrcDef->getOperand(1).getImm();
- if (Imm == 1 && Opc == AMDGPU::S_ADD_U64_PSEUDO) {
- // 1 * bitcount(exec) = bitcount(exec)
- Register DstHi =
- MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
- BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstHi).addImm(0);
- BuildRegSequence(BB, MI, DstReg,
- NewAccumulator->getOperand(0).getReg(), DstHi);
- break;
- }
+ if (Opc == AMDGPU::S_ADD_U64_PSEUDO &&
+ TII->getImmOrMaterializedImm(MRI, SrcReg) == 1) {
+ // 1 * bitcount(exec) = bitcount(exec)
+ Register DstHi = MRI.createVirtualRegister(&AMDGPU::SReg_32RegClass);
+ BuildMI(BB, MI, DL, TII->get(AMDGPU::S_MOV_B32), DstHi).addImm(0);
+ BuildRegSequence(BB, MI, DstReg,
+ NewAccumulator->getOperand(0).getReg(), DstHi);
+ break;
}
if (Opc == AMDGPU::S_SUB_U64_PSEUDO) {
BuildMI(BB, MI, DL, TII->get(AMDGPU::S_SUB_I32), NegatedValLo)
diff --git a/llvm/test/CodeGen/AMDGPU/wave-reduce-move-source.ll b/llvm/test/CodeGen/AMDGPU/wave-reduce-move-source.ll
new file mode 100644
index 000000000000000..d68acc96c43f22e
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/wave-reduce-move-source.ll
@@ -0,0 +1,109 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=amdgpu12.50-amd-amdpal -global-isel=0 < %s | FileCheck %s
+; RUN: llc -mtriple=amdgpu12.50-amd-amdpal -global-isel=1 < %s | FileCheck %s
+
+; Uniform wave.reduce sources defined by a move of something other than an
+; immediate: a register (the uniform lowering of wave.reduce.umin and
+; wave.reduce.and), a frame index and a global address.
+; gfx1250 with PAL moves the 64-bit address of @gv with one s_mov_b64.
+
+ at gv = addrspace(4) global i32 0
+
+define i32 @add_of_umin(i32 inreg %x) {
+; CHECK-LABEL: add_of_umin:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_wait_loadcnt_dscnt 0x0
+; CHECK-NEXT: s_wait_kmcnt 0x0
+; CHECK-NEXT: s_mov_b32 s1, exec_lo
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_bcnt1_i32_b32 s1, s1
+; CHECK-NEXT: s_mul_i32 s0, s0, s1
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
+; CHECK-NEXT: v_mov_b32_e32 v0, s0
+; CHECK-NEXT: s_set_pc_i64 s[30:31]
+ %m = call i32 @llvm.amdgcn.wave.reduce.umin.i32(i32 %x, i32 1)
+ %r = call i32 @llvm.amdgcn.wave.reduce.add.i32(i32 %m, i32 1)
+ ret i32 %r
+}
+
+define i32 @add_frame_index() {
+; CHECK-LABEL: add_frame_index:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_wait_loadcnt_dscnt 0x0
+; CHECK-NEXT: s_wait_kmcnt 0x0
+; CHECK-NEXT: s_mov_b32 s0, exec_lo
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_bcnt1_i32_b32 s0, s0
+; CHECK-NEXT: s_mul_i32 s0, s0, s32
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
+; CHECK-NEXT: v_mov_b32_e32 v0, s0
+; CHECK-NEXT: s_set_pc_i64 s[30:31]
+ %a = alloca i32, align 4, addrspace(5)
+ %p = ptrtoint ptr addrspace(5) %a to i32
+ %r = call i32 @llvm.amdgcn.wave.reduce.add.i32(i32 %p, i32 1)
+ ret i32 %r
+}
+
+define i32 @xor_frame_index() {
+; CHECK-LABEL: xor_frame_index:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_wait_loadcnt_dscnt 0x0
+; CHECK-NEXT: s_wait_kmcnt 0x0
+; CHECK-NEXT: s_mov_b32 s0, exec_lo
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_bcnt1_i32_b32 s0, s0
+; CHECK-NEXT: s_and_b32 s0, s0, 1
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_mul_i32 s0, s0, s32
+; CHECK-NEXT: v_mov_b32_e32 v0, s0
+; CHECK-NEXT: s_set_pc_i64 s[30:31]
+ %a = alloca i32, align 4, addrspace(5)
+ %p = ptrtoint ptr addrspace(5) %a to i32
+ %r = call i32 @llvm.amdgcn.wave.reduce.xor.i32(i32 %p, i32 1)
+ ret i32 %r
+}
+
+define i64 @add_global_address_i64() {
+; CHECK-LABEL: add_global_address_i64:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_wait_loadcnt_dscnt 0x0
+; CHECK-NEXT: s_wait_kmcnt 0x0
+; CHECK-NEXT: s_mov_b32 s2, exec_lo
+; CHECK-NEXT: s_mov_b64 s[0:1], gv at abs64
+; CHECK-NEXT: s_bcnt1_i32_b32 s2, s2
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(SKIP_3) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_mul_hi_u32 s3, s0, s2
+; CHECK-NEXT: s_mul_i32 s1, s1, s2
+; CHECK-NEXT: s_mul_i32 s0, s0, s2
+; CHECK-NEXT: s_add_co_u32 s1, s3, s1
+; CHECK-NEXT: v_dual_mov_b32 v0, s0 :: v_dual_mov_b32 v1, s1
+; CHECK-NEXT: s_set_pc_i64 s[30:31]
+ %p = ptrtoint ptr addrspace(4) @gv to i64
+ %r = call i64 @llvm.amdgcn.wave.reduce.add.i64(i64 %p, i32 1)
+ ret i64 %r
+}
+
+define i64 @sub_of_and_i64(i64 inreg %x) {
+; CHECK-LABEL: sub_of_and_i64:
+; CHECK: ; %bb.0:
+; CHECK-NEXT: s_wait_loadcnt_dscnt 0x0
+; CHECK-NEXT: s_wait_kmcnt 0x0
+; CHECK-NEXT: s_mov_b32 s2, exec_lo
+; CHECK-NEXT: s_mov_b64 s[0:1], s[0:1]
+; CHECK-NEXT: s_bcnt1_i32_b32 s2, s2
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
+; CHECK-NEXT: s_sub_co_i32 s2, 0, s2
+; CHECK-NEXT: s_ashr_i32 s3, s2, 31
+; CHECK-NEXT: s_mul_hi_u32 s4, s0, s2
+; CHECK-NEXT: s_mul_i32 s1, s1, s2
+; CHECK-NEXT: s_mul_i32 s3, s0, s3
+; CHECK-NEXT: s_add_co_u32 s1, s4, s1
+; CHECK-NEXT: s_mul_i32 s0, s0, s2
+; CHECK-NEXT: s_add_co_u32 s1, s1, s3
+; CHECK-NEXT: s_delay_alu instid0(SALU_CYCLE_1)
+; CHECK-NEXT: v_dual_mov_b32 v0, s0 :: v_dual_mov_b32 v1, s1
+; CHECK-NEXT: s_set_pc_i64 s[30:31]
+ %m = call i64 @llvm.amdgcn.wave.reduce.and.i64(i64 %x, i32 1)
+ %r = call i64 @llvm.amdgcn.wave.reduce.sub.i64(i64 %m, i32 1)
+ ret i64 %r
+}
More information about the llvm-commits
mailing list