[llvm] 0d51fec - [NVPTX] PerformSELECTShiftCombine drops high bits of a wide guarded shift amount (#201165)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Jun 2 11:49:21 PDT 2026
Author: Justin Lebar
Date: 2026-06-02T11:49:16-07:00
New Revision: 0d51fec35520f39f9b2f55989f50d694244c7582
URL: https://github.com/llvm/llvm-project/commit/0d51fec35520f39f9b2f55989f50d694244c7582
DIFF: https://github.com/llvm/llvm-project/commit/0d51fec35520f39f9b2f55989f50d694244c7582.diff
LOG: [NVPTX] PerformSELECTShiftCombine drops high bits of a wide guarded shift amount (#201165)
LLVM shifts produce poison if you shift greater than the width of the
operand. But PTX shifts clamp the shift amount:
> shl/shr: Shift amounts greater than the register width N are clamped
> to N.
NVPTXISelLowering looks for shl/shr which guard against an out-of-range
shift and lower these to an unguarded PTX shift:
define i64 @f(i64 %x, i64 %shift) {
%cmp = icmp ult i64 %shift, 64
%shl = shl i64 %x, %shift
%sel = select i1 %cmp, i64 %shl, i64 0
ret i64 %sel
}
In PTX shifts the shift amount is always i32, whereas in LLVM the shift
amount is the same type as the "shiftee". Therefore in the i64 case,
this lowering is only sound if we know that the upper 32 bits of the
shift amount are all zero.
Added:
Modified:
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
llvm/test/CodeGen/NVPTX/shift-opt.ll
Removed:
################################################################################
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 75b8e736251d8..7854309c42e5d 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -6719,13 +6719,12 @@ static SDValue PerformEXTRACTCombine(SDNode *N,
/// Into:
/// (NVPTXISD::SRL_CLAMP x, shift_amt) or (NVPTXISD::SHL_CLAMP x, shift_amt)
///
-/// These patterns arise from C/C++ code like `shift >= 32 ? 0 : x >> shift`
-/// which guards against undefined behavior. PTX shr/shl instructions clamp
-/// shift amounts >= BitWidth to produce 0 for logical shifts, making the
-/// guard redundant.
+/// These patterns arise from code like `s >= 32 ? 0 : x >> s`. In LLVM,
+/// over-shifting a value results in poison, but PTX shr/shl instructions clamp
+/// the shift amount to BitWidth, making the guard redundant.
///
-/// Note: We only handle SRL and SHL, not SRA, because arithmetic right
-/// shifts could produce 0 or -1 when shift >= BitWidth.
+/// Note: We only handle SRL and SHL, not SRA, because arithmetic right shifts
+/// can produce 0 or -1 when shift >= BitWidth.
/// Note: We don't handle uge or ule. These don't appear because of
/// canonicalization.
static SDValue PerformSELECTShiftCombine(SDNode *N,
@@ -6761,11 +6760,21 @@ static SDValue PerformSELECTShiftCombine(SDNode *N,
if (!MatchedUGT && !MatchedULT)
return SDValue();
+ // In LLVM IR, the shift amount and the value-to-be-shifted are the same
+ // type, whereas in PTX the shift amount is always i32. Therefore when
+ // shifting types larger than i32, we can only do this transformation if we
+ // know that the upper bits of the shift amount are known zero.
+ SDValue ClampAmt = ShiftOp.getOperand(1);
+ unsigned ClampAmtBits = ClampAmt.getValueSizeInBits();
+ if (ShiftAmt.getValueSizeInBits() > ClampAmtBits &&
+ DCI.DAG.computeKnownBits(ShiftAmt).countMaxActiveBits() > ClampAmtBits)
+ return SDValue();
+
// Return a clamp shift operation, which has the same semantics as PTX shift.
unsigned ClampOpc = ShiftOp.getOpcode() == ISD::SRL ? NVPTXISD::SRL_CLAMP
: NVPTXISD::SHL_CLAMP;
return DCI.DAG.getNode(ClampOpc, SDLoc(N), ShiftOp.getValueType(),
- ShiftOp.getOperand(0), ShiftOp.getOperand(1));
+ ShiftOp.getOperand(0), ClampAmt);
}
static SDValue PerformVSELECTCombine(SDNode *N,
diff --git a/llvm/test/CodeGen/NVPTX/shift-opt.ll b/llvm/test/CodeGen/NVPTX/shift-opt.ll
index 554839436f451..e8855c6e3b795 100644
--- a/llvm/test/CodeGen/NVPTX/shift-opt.ll
+++ b/llvm/test/CodeGen/NVPTX/shift-opt.ll
@@ -326,14 +326,18 @@ define i32 @test_guarded_i32_ult(i32 %x, i32 %shift) {
define i64 @test_guarded_i64_ugt(i64 %x, i64 %shift) {
; CHECK-LABEL: test_guarded_i64_ugt(
; CHECK: {
+; CHECK-NEXT: .reg .pred %p<2>;
; CHECK-NEXT: .reg .b32 %r<2>;
-; CHECK-NEXT: .reg .b64 %rd<3>;
+; CHECK-NEXT: .reg .b64 %rd<5>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.b64 %rd1, [test_guarded_i64_ugt_param_0];
-; CHECK-NEXT: ld.param.b32 %r1, [test_guarded_i64_ugt_param_1];
-; CHECK-NEXT: shr.u64 %rd2, %rd1, %r1;
-; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
+; CHECK-NEXT: ld.param.b64 %rd2, [test_guarded_i64_ugt_param_1];
+; CHECK-NEXT: setp.gt.u64 %p1, %rd2, 63;
+; CHECK-NEXT: cvt.u32.u64 %r1, %rd2;
+; CHECK-NEXT: shr.u64 %rd3, %rd1, %r1;
+; CHECK-NEXT: selp.b64 %rd4, 0, %rd3, %p1;
+; CHECK-NEXT: st.param.b64 [func_retval0], %rd4;
; CHECK-NEXT: ret;
%cmp = icmp ugt i64 %shift, 63
%shr = lshr i64 %x, %shift
@@ -345,14 +349,18 @@ define i64 @test_guarded_i64_ugt(i64 %x, i64 %shift) {
define i64 @test_guarded_i64_ult(i64 %x, i64 %shift) {
; CHECK-LABEL: test_guarded_i64_ult(
; CHECK: {
+; CHECK-NEXT: .reg .pred %p<2>;
; CHECK-NEXT: .reg .b32 %r<2>;
-; CHECK-NEXT: .reg .b64 %rd<3>;
+; CHECK-NEXT: .reg .b64 %rd<5>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.b64 %rd1, [test_guarded_i64_ult_param_0];
-; CHECK-NEXT: ld.param.b32 %r1, [test_guarded_i64_ult_param_1];
-; CHECK-NEXT: shr.u64 %rd2, %rd1, %r1;
-; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
+; CHECK-NEXT: ld.param.b64 %rd2, [test_guarded_i64_ult_param_1];
+; CHECK-NEXT: setp.lt.u64 %p1, %rd2, 64;
+; CHECK-NEXT: cvt.u32.u64 %r1, %rd2;
+; CHECK-NEXT: shr.u64 %rd3, %rd1, %r1;
+; CHECK-NEXT: selp.b64 %rd4, %rd3, 0, %p1;
+; CHECK-NEXT: st.param.b64 [func_retval0], %rd4;
; CHECK-NEXT: ret;
%cmp = icmp ult i64 %shift, 64
%shr = lshr i64 %x, %shift
@@ -494,14 +502,18 @@ define i32 @test_guarded_i32_ult_shl(i32 %x, i32 %shift) {
define i64 @test_guarded_i64_ugt_shl(i64 %x, i64 %shift) {
; CHECK-LABEL: test_guarded_i64_ugt_shl(
; CHECK: {
+; CHECK-NEXT: .reg .pred %p<2>;
; CHECK-NEXT: .reg .b32 %r<2>;
-; CHECK-NEXT: .reg .b64 %rd<3>;
+; CHECK-NEXT: .reg .b64 %rd<5>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.b64 %rd1, [test_guarded_i64_ugt_shl_param_0];
-; CHECK-NEXT: ld.param.b32 %r1, [test_guarded_i64_ugt_shl_param_1];
-; CHECK-NEXT: shl.b64 %rd2, %rd1, %r1;
-; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
+; CHECK-NEXT: ld.param.b64 %rd2, [test_guarded_i64_ugt_shl_param_1];
+; CHECK-NEXT: setp.gt.u64 %p1, %rd2, 63;
+; CHECK-NEXT: cvt.u32.u64 %r1, %rd2;
+; CHECK-NEXT: shl.b64 %rd3, %rd1, %r1;
+; CHECK-NEXT: selp.b64 %rd4, 0, %rd3, %p1;
+; CHECK-NEXT: st.param.b64 [func_retval0], %rd4;
; CHECK-NEXT: ret;
%cmp = icmp ugt i64 %shift, 63
%shl = shl i64 %x, %shift
@@ -513,14 +525,18 @@ define i64 @test_guarded_i64_ugt_shl(i64 %x, i64 %shift) {
define i64 @test_guarded_i64_ult_shl(i64 %x, i64 %shift) {
; CHECK-LABEL: test_guarded_i64_ult_shl(
; CHECK: {
+; CHECK-NEXT: .reg .pred %p<2>;
; CHECK-NEXT: .reg .b32 %r<2>;
-; CHECK-NEXT: .reg .b64 %rd<3>;
+; CHECK-NEXT: .reg .b64 %rd<5>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.b64 %rd1, [test_guarded_i64_ult_shl_param_0];
-; CHECK-NEXT: ld.param.b32 %r1, [test_guarded_i64_ult_shl_param_1];
-; CHECK-NEXT: shl.b64 %rd2, %rd1, %r1;
-; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
+; CHECK-NEXT: ld.param.b64 %rd2, [test_guarded_i64_ult_shl_param_1];
+; CHECK-NEXT: setp.lt.u64 %p1, %rd2, 64;
+; CHECK-NEXT: cvt.u32.u64 %r1, %rd2;
+; CHECK-NEXT: shl.b64 %rd3, %rd1, %r1;
+; CHECK-NEXT: selp.b64 %rd4, %rd3, 0, %p1;
+; CHECK-NEXT: st.param.b64 [func_retval0], %rd4;
; CHECK-NEXT: ret;
%cmp = icmp ult i64 %shift, 64
%shl = shl i64 %x, %shift
@@ -550,3 +566,26 @@ define i64 @test_guarded_i64_ult_shl_
diff erent_shift(i64 %x, i64 %shift1, i64 %s
%sel = select i1 %cmp, i64 %shl, i64 0
ret i64 %sel
}
+
+; The shift amount is zext'd from i32, so its high 32 bits are known zero and
+; the guard's icmp agrees with the narrowed clamp amount: the combine SHOULD
+; fire, lowering to a bare clamp shift with no setp/selp guard.
+; (select (ult zext(s), 64), (shl x, zext(s)), 0) --> clamp shl
+define i64 @test_guarded_i64_shl_ult_zext_amt(i64 %x, i32 %s) {
+; CHECK-LABEL: test_guarded_i64_shl_ult_zext_amt(
+; CHECK: {
+; CHECK-NEXT: .reg .b32 %r<2>;
+; CHECK-NEXT: .reg .b64 %rd<3>;
+; CHECK-EMPTY:
+; CHECK-NEXT: // %bb.0:
+; CHECK-NEXT: ld.param.b64 %rd1, [test_guarded_i64_shl_ult_zext_amt_param_0];
+; CHECK-NEXT: ld.param.b32 %r1, [test_guarded_i64_shl_ult_zext_amt_param_1];
+; CHECK-NEXT: shl.b64 %rd2, %rd1, %r1;
+; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
+; CHECK-NEXT: ret;
+ %shift = zext i32 %s to i64
+ %cmp = icmp ult i64 %shift, 64
+ %shl = shl i64 %x, %shift
+ %sel = select i1 %cmp, i64 %shl, i64 0
+ ret i64 %sel
+}
More information about the llvm-commits
mailing list