Skip to content

Commit c0fc9c7

Browse files
petar-avramovicpetar-avramovic
committed
AMDGPU/GlobalISel: Fix inst-selection of ballot
Both input and output of ballot are lane-masks: result is lane-mask with 'S32/S64 LLT and SGPR bank' input is lane-mask with 'S1 LLT and VCC reg bank'. Ballot copies bits from input lane-mask for all active lanes and puts 0 for inactive lanes. GlobalISel did not set 0 in result for inactive lanes for non-constant input.
1 parent af7aa22 commit c0fc9c7

File tree

63 files changed

+797
-306
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

63 files changed

+797
-306
lines changed

‎llvm/docs/AMDGPUUsage.rst‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1369,6 +1369,12 @@ The AMDGPU backend implements the following LLVM IR intrinsics.
13691369
sign-extended from the width of the underlying PC hardware register even on
13701370
processors where the s_getpc_b64 instruction returns a zero-extended value.
13711371

1372+
llvm.amdgcn.ballot Returns a bitfield(i32 or i64) containing the result of its i1 argument
1373+
in all active lanes, and zero in all inactive lanes.
1374+
Provides a way to convert i1 in LLVM IR to i32 or i64 lane mask - bitfield
1375+
used by hardware to control active lanes when used in EXEC register.
1376+
For example, ballot(i1 true) return EXEC mask.
1377+
13721378
============================================== ==========================================================
13731379

13741380
.. TODO::

‎llvm/include/llvm/IR/IntrinsicsAMDGPU.td‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2085,6 +2085,8 @@ def int_amdgcn_fcmp :
20852085
[IntrNoMem, IntrConvergent,
20862086
ImmArg<ArgIndex<2>>, IntrWillReturn, IntrNoCallback, IntrNoFree]>;
20872087

2088+
// Returns a bitfield(i32 or i64) containing the result of its i1 argument
2089+
// in all active lanes, and zero in all inactive lanes.
20882090
def int_amdgcn_ballot :
20892091
Intrinsic<[llvm_anyint_ty], [llvm_i1_ty],
20902092
[IntrNoMem, IntrConvergent, IntrWillReturn, IntrNoCallback, IntrNoFree]>;

‎llvm/lib/Target/AMDGPU/AMDGPUInstructionSelector.cpp‎

Lines changed: 112 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1429,34 +1429,129 @@ bool AMDGPUInstructionSelector::selectBallot(MachineInstr &I) const {
14291429
std::optional<ValueAndVReg> Arg =
14301430
getIConstantVRegValWithLookThrough(I.getOperand(2).getReg(), *MRI);
14311431

1432-
const auto BuildCopy = [&](Register SrcReg) {
1433-
if (Size == STI.getWavefrontSize()) {
1434-
BuildMI(*BB, &I, DL, TII.get(AMDGPU::COPY), DstReg)
1435-
.addReg(SrcReg);
1436-
return;
1432+
const auto getCmpInput = [&]() -> MachineInstr * {
1433+
MachineInstr *SrcMI = getDefIgnoringCopies(I.getOperand(2).getReg(), *MRI);
1434+
// Try to fold sgpr compare.
1435+
if (SrcMI->getOpcode() == AMDGPU::G_TRUNC)
1436+
SrcMI = MRI->getVRegDef(SrcMI->getOperand(1).getReg());
1437+
1438+
if (SrcMI->getOpcode() == AMDGPU::G_ICMP ||
1439+
SrcMI->getOpcode() == AMDGPU::G_FCMP)
1440+
return SrcMI;
1441+
return nullptr;
1442+
};
1443+
1444+
const auto FoldCmp = [&](Register Dst, MachineInstr *CmpMI) -> bool {
1445+
// Fold ballot of a compare. Active lanes when the ballot is executed need
1446+
// to also be active when the compare is executed for this fold to be
1447+
// correct. If an inactive lane on compare becomes active for the ballot,
1448+
// divergent control flow is involved. The compare was in a divergent branch
1449+
// and needs to go through phi before being used by the ballot. The ballot
1450+
// is in a block that merged control flow. Using the compare directly in the
1451+
// ballot implies that active lanes for the ballot are a subset of active
1452+
// lanes for the compare.
1453+
auto Pred = (CmpInst::Predicate)CmpMI->getOperand(1).getPredicate();
1454+
Register Op0 = CmpMI->getOperand(2).getReg();
1455+
Register Op1 = CmpMI->getOperand(3).getReg();
1456+
unsigned OpSize = MRI->getType(Op0).getSizeInBits();
1457+
const TargetRegisterClass *VgprRC = TRI.getVGPRClassForBitWidth(OpSize);
1458+
1459+
int CmpOpcode = getV_CMPOpcode(Pred, OpSize, *Subtarget);
1460+
if (CmpOpcode == -1)
1461+
return false;
1462+
1463+
MachineInstr *Cmp;
1464+
unsigned Op0Idx, Op1Idx;
1465+
if (CmpMI->getOpcode() == AMDGPU::G_ICMP) {
1466+
Cmp =
1467+
BuildMI(*BB, &I, DL, TII.get(CmpOpcode), Dst).addReg(Op0).addReg(Op1);
1468+
Op0Idx = 1;
1469+
Op1Idx = 2;
1470+
} else {
1471+
// fcmp compares have modifiers
1472+
Cmp = BuildMI(*BB, &I, DL, TII.get(CmpOpcode), Dst)
1473+
.addImm(0)
1474+
.addReg(Op0)
1475+
.addImm(0)
1476+
.addReg(Op1)
1477+
.addImm(0);
1478+
Op0Idx = 2;
1479+
Op1Idx = 4;
14371480
}
14381481

1439-
// If emitting a i64 ballot in wave32, fill the upper bits with zeroes.
1440-
Register HiReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
1441-
BuildMI(*BB, &I, DL, TII.get(AMDGPU::S_MOV_B32), HiReg).addImm(0);
1442-
BuildMI(*BB, &I, DL, TII.get(AMDGPU::REG_SEQUENCE), DstReg)
1443-
.addReg(SrcReg)
1482+
return constrainOperandRegClass(*MF, TRI, *MRI, TII, RBI, *Cmp, *VgprRC,
1483+
Cmp->getOperand(Op0Idx)) &&
1484+
constrainOperandRegClass(*MF, TRI, *MRI, TII, RBI, *Cmp, *VgprRC,
1485+
Cmp->getOperand(Op1Idx)) &&
1486+
constrainSelectedInstRegOperands(*Cmp, TII, TRI, RBI);
1487+
};
1488+
1489+
const auto BuildAnd = [&](unsigned Opcode, Register Dst, Register Src,
1490+
Register Exec) -> bool {
1491+
auto And = BuildMI(*BB, &I, DL, TII.get(Opcode), Dst)
1492+
.addReg(Src)
1493+
.addReg(Exec)
1494+
.setOperandDead(3); // Dead scc
1495+
return constrainSelectedInstRegOperands(*And, TII, TRI, RBI);
1496+
};
1497+
1498+
const auto BuildREG_SEQUENCE = [&](Register Dst, Register Lo, Register Hi) {
1499+
BuildMI(*BB, &I, DL, TII.get(AMDGPU::REG_SEQUENCE), Dst)
1500+
.addReg(Lo)
14441501
.addImm(AMDGPU::sub0)
1445-
.addReg(HiReg)
1502+
.addReg(Hi)
14461503
.addImm(AMDGPU::sub1);
14471504
};
14481505

14491506
if (Arg) {
1450-
const int64_t Value = Arg->Value.getSExtValue();
1507+
const int64_t Value = Arg->Value.getZExtValue();
14511508
if (Value == 0) {
1509+
// DstReg(32or64) = S_MOV 0
14521510
unsigned Opcode = Is64 ? AMDGPU::S_MOV_B64 : AMDGPU::S_MOV_B32;
14531511
BuildMI(*BB, &I, DL, TII.get(Opcode), DstReg).addImm(0);
1454-
} else if (Value == -1) // all ones
1455-
BuildCopy(IsWave32 ? AMDGPU::EXEC_LO : AMDGPU::EXEC);
1456-
else
1512+
} else if (Value == 1) {
1513+
if (Size == STI.getWavefrontSize()) {
1514+
// DstReg(32or64) = COPY EXEC
1515+
BuildMI(*BB, &I, DL, TII.get(AMDGPU::COPY), DstReg)
1516+
.addReg(TRI.getExec());
1517+
} else {
1518+
// DstReg(64) = REG_SEQUENCE EXEC_LO, 0
1519+
Register HiReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
1520+
BuildMI(*BB, &I, DL, TII.get(AMDGPU::S_MOV_B32), HiReg).addImm(0);
1521+
BuildREG_SEQUENCE(DstReg, TRI.getExec(), HiReg);
1522+
}
1523+
} else
14571524
return false;
1458-
} else
1459-
BuildCopy(I.getOperand(2).getReg());
1525+
} else {
1526+
Register SrcReg = I.getOperand(2).getReg();
1527+
if (Size == STI.getWavefrontSize()) {
1528+
if (MachineInstr *Cmp = getCmpInput()) {
1529+
// DstReg(32or64) = V_CMP...
1530+
if (!FoldCmp(DstReg, Cmp))
1531+
return false;
1532+
} else {
1533+
// DstReg(32or64) = AND SrcReg, EXEC
1534+
unsigned AndOpc = IsWave32 ? AMDGPU::S_AND_B32 : AMDGPU::S_AND_B64;
1535+
if (!BuildAnd(AndOpc, DstReg, SrcReg, TRI.getExec()))
1536+
return false;
1537+
}
1538+
} else {
1539+
Register LoReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
1540+
if (MachineInstr *Cmp = getCmpInput()) {
1541+
// LoReg(32) = V_CMP...
1542+
if (!FoldCmp(LoReg, Cmp))
1543+
return false;
1544+
} else {
1545+
// LoReg(32) = AND SrcReg, EXEC
1546+
if (!BuildAnd(AMDGPU::S_AND_B32, LoReg, SrcReg, AMDGPU::EXEC_LO))
1547+
return false;
1548+
}
1549+
// DstReg(64) = REG_SEQUENCE (LoReg(32), EXEC_LO), 0
1550+
Register HiReg = MRI->createVirtualRegister(&AMDGPU::SReg_32RegClass);
1551+
BuildMI(*BB, &I, DL, TII.get(AMDGPU::S_MOV_B32), HiReg).addImm(0);
1552+
BuildREG_SEQUENCE(DstReg, LoReg, HiReg);
1553+
}
1554+
}
14601555

14611556
I.eraseFromParent();
14621557
return true;

‎llvm/test/CodeGen/AMDGPU/GlobalISel/image-waterfall-loop-O0.ll‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ define <4 x float> @waterfall_loop(<8 x i32> %vgpr_srd) {
138138
; CHECK-NEXT: s_and_b32 s4, s4, s5
139139
; CHECK-NEXT: v_cmp_eq_u64_e64 s5, s[6:7], v[1:2]
140140
; CHECK-NEXT: s_and_b32 s4, s4, s5
141+
; CHECK-NEXT: s_and_b32 s4, s4, exec_lo
141142
; CHECK-NEXT: s_and_saveexec_b32 s4, s4
142143
; CHECK-NEXT: v_writelane_b32 v0, s4, 13
143144
; CHECK-NEXT: s_or_saveexec_b32 s21, -1

‎llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.ballot.i32.ll‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@ define amdgpu_cs i32 @non_compare(i32 %x) {
3333
; CHECK-LABEL: non_compare:
3434
; CHECK: ; %bb.0:
3535
; CHECK-NEXT: v_and_b32_e32 v0, 1, v0
36-
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, v0
36+
; CHECK-NEXT: v_cmp_ne_u32_e32 vcc_lo, 0, v0
37+
; CHECK-NEXT: s_and_b32 s0, vcc_lo, exec_lo
3738
; CHECK-NEXT: ; return to shader part epilog
3839
%trunc = trunc i32 %x to i1
3940
%ballot = call i32 @llvm.amdgcn.ballot.i32(i1 %trunc)
@@ -89,7 +90,8 @@ define amdgpu_cs i32 @branch_divergent_ballot_ne_zero_non_compare(i32 %v) {
8990
; CHECK: ; %bb.0:
9091
; CHECK-NEXT: v_and_b32_e32 v0, 1, v0
9192
; CHECK-NEXT: v_cmp_ne_u32_e32 vcc_lo, 0, v0
92-
; CHECK-NEXT: s_cmp_eq_u32 vcc_lo, 0
93+
; CHECK-NEXT: s_and_b32 s0, vcc_lo, exec_lo
94+
; CHECK-NEXT: s_cmp_eq_u32 s0, 0
9395
; CHECK-NEXT: s_cbranch_scc1 .LBB7_2
9496
; CHECK-NEXT: ; %bb.1: ; %true
9597
; CHECK-NEXT: s_mov_b32 s0, 42
@@ -113,6 +115,7 @@ define amdgpu_cs i32 @branch_uniform_ballot_ne_zero_non_compare(i32 inreg %v) {
113115
; CHECK: ; %bb.0:
114116
; CHECK-NEXT: s_and_b32 s0, 1, s0
115117
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
118+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
116119
; CHECK-NEXT: s_cmp_eq_u32 s0, 0
117120
; CHECK-NEXT: s_cbranch_scc1 .LBB8_2
118121
; CHECK-NEXT: ; %bb.1: ; %true
@@ -137,7 +140,8 @@ define amdgpu_cs i32 @branch_divergent_ballot_eq_zero_non_compare(i32 %v) {
137140
; CHECK: ; %bb.0:
138141
; CHECK-NEXT: v_and_b32_e32 v0, 1, v0
139142
; CHECK-NEXT: v_cmp_ne_u32_e32 vcc_lo, 0, v0
140-
; CHECK-NEXT: s_cmp_lg_u32 vcc_lo, 0
143+
; CHECK-NEXT: s_and_b32 s0, vcc_lo, exec_lo
144+
; CHECK-NEXT: s_cmp_lg_u32 s0, 0
141145
; CHECK-NEXT: s_cbranch_scc0 .LBB9_2
142146
; CHECK-NEXT: ; %bb.1: ; %false
143147
; CHECK-NEXT: s_mov_b32 s0, 33
@@ -161,6 +165,7 @@ define amdgpu_cs i32 @branch_uniform_ballot_eq_zero_non_compare(i32 inreg %v) {
161165
; CHECK: ; %bb.0:
162166
; CHECK-NEXT: s_and_b32 s0, 1, s0
163167
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
168+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
164169
; CHECK-NEXT: s_cmp_lg_u32 s0, 0
165170
; CHECK-NEXT: s_cbranch_scc0 .LBB10_2
166171
; CHECK-NEXT: ; %bb.1: ; %false
@@ -206,10 +211,7 @@ false:
206211
define amdgpu_cs i32 @branch_uniform_ballot_ne_zero_compare(i32 inreg %v) {
207212
; CHECK-LABEL: branch_uniform_ballot_ne_zero_compare:
208213
; CHECK: ; %bb.0:
209-
; CHECK-NEXT: s_cmp_lt_u32 s0, 12
210-
; CHECK-NEXT: s_cselect_b32 s0, 1, 0
211-
; CHECK-NEXT: s_and_b32 s0, 1, s0
212-
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
214+
; CHECK-NEXT: v_cmp_lt_u32_e64 s0, s0, 12
213215
; CHECK-NEXT: s_cmp_eq_u32 s0, 0
214216
; CHECK-NEXT: s_cbranch_scc1 .LBB12_2
215217
; CHECK-NEXT: ; %bb.1: ; %true
@@ -255,10 +257,7 @@ false:
255257
define amdgpu_cs i32 @branch_uniform_ballot_eq_zero_compare(i32 inreg %v) {
256258
; CHECK-LABEL: branch_uniform_ballot_eq_zero_compare:
257259
; CHECK: ; %bb.0:
258-
; CHECK-NEXT: s_cmp_lt_u32 s0, 12
259-
; CHECK-NEXT: s_cselect_b32 s0, 1, 0
260-
; CHECK-NEXT: s_and_b32 s0, 1, s0
261-
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
260+
; CHECK-NEXT: v_cmp_lt_u32_e64 s0, s0, 12
262261
; CHECK-NEXT: s_cmp_lg_u32 s0, 0
263262
; CHECK-NEXT: s_cbranch_scc0 .LBB14_2
264263
; CHECK-NEXT: ; %bb.1: ; %false
@@ -284,6 +283,7 @@ define amdgpu_cs i32 @branch_divergent_ballot_ne_zero_and(i32 %v1, i32 %v2) {
284283
; CHECK-NEXT: v_cmp_gt_u32_e32 vcc_lo, 12, v0
285284
; CHECK-NEXT: v_cmp_lt_u32_e64 s0, 34, v1
286285
; CHECK-NEXT: s_and_b32 s0, vcc_lo, s0
286+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
287287
; CHECK-NEXT: s_cmp_eq_u32 s0, 0
288288
; CHECK-NEXT: s_cbranch_scc1 .LBB15_2
289289
; CHECK-NEXT: ; %bb.1: ; %true
@@ -315,6 +315,7 @@ define amdgpu_cs i32 @branch_uniform_ballot_ne_zero_and(i32 inreg %v1, i32 inreg
315315
; CHECK-NEXT: s_and_b32 s0, s0, s1
316316
; CHECK-NEXT: s_and_b32 s0, 1, s0
317317
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
318+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
318319
; CHECK-NEXT: s_cmp_eq_u32 s0, 0
319320
; CHECK-NEXT: s_cbranch_scc1 .LBB16_2
320321
; CHECK-NEXT: ; %bb.1: ; %true
@@ -342,6 +343,7 @@ define amdgpu_cs i32 @branch_divergent_ballot_eq_zero_and(i32 %v1, i32 %v2) {
342343
; CHECK-NEXT: v_cmp_gt_u32_e32 vcc_lo, 12, v0
343344
; CHECK-NEXT: v_cmp_lt_u32_e64 s0, 34, v1
344345
; CHECK-NEXT: s_and_b32 s0, vcc_lo, s0
346+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
345347
; CHECK-NEXT: s_cmp_lg_u32 s0, 0
346348
; CHECK-NEXT: s_cbranch_scc0 .LBB17_2
347349
; CHECK-NEXT: ; %bb.1: ; %false
@@ -373,6 +375,7 @@ define amdgpu_cs i32 @branch_uniform_ballot_eq_zero_and(i32 inreg %v1, i32 inreg
373375
; CHECK-NEXT: s_and_b32 s0, s0, s1
374376
; CHECK-NEXT: s_and_b32 s0, 1, s0
375377
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
378+
; CHECK-NEXT: s_and_b32 s0, s0, exec_lo
376379
; CHECK-NEXT: s_cmp_lg_u32 s0, 0
377380
; CHECK-NEXT: s_cbranch_scc0 .LBB18_2
378381
; CHECK-NEXT: ; %bb.1: ; %false
@@ -397,10 +400,7 @@ false:
397400
define amdgpu_cs i32 @branch_uniform_ballot_sgt_N_compare(i32 inreg %v) {
398401
; CHECK-LABEL: branch_uniform_ballot_sgt_N_compare:
399402
; CHECK: ; %bb.0:
400-
; CHECK-NEXT: s_cmp_lt_u32 s0, 12
401-
; CHECK-NEXT: s_cselect_b32 s0, 1, 0
402-
; CHECK-NEXT: s_and_b32 s0, 1, s0
403-
; CHECK-NEXT: v_cmp_ne_u32_e64 s0, 0, s0
403+
; CHECK-NEXT: v_cmp_lt_u32_e64 s0, s0, 12
404404
; CHECK-NEXT: s_cmp_le_i32 s0, 22
405405
; CHECK-NEXT: s_cbranch_scc1 .LBB19_2
406406
; CHECK-NEXT: ; %bb.1: ; %true

0 commit comments

Comments
 (0)