Skip to content

[NVPTX] Fix bug in sign of bfe folding #130862

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Mar 12, 2025
Merged

[NVPTX] Fix bug in sign of bfe folding #130862

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
96436ff
pre-commit tests -- use update_llc_test_checks.py, add some more cases
AlexMaclean Mar 11, 2025
f5ad399
[NVPTX] Fix bug in sign of bfe folding
AlexMaclean Mar 11, 2025
70ffef2
address comment comment
AlexMaclean Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2039,6 +2039,11 @@ bool NVPTXDAGToDAGISel::tryBFE(SDNode *N) {
Val = AndLHS;
Start = CurDAG->getTargetConstant(ShiftAmt, DL, MVT::i32);
Len = CurDAG->getTargetConstant(NumBits, DL, MVT::i32);

// If we have a arithmetic right shift, we need to use the signed bfe
// variant only if the MSB is copied as part of the mask
Copy link
Member

@Artem-B Artem-B Mar 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could use some editing. The comment explains what we do, but not why.

E.g. something along the lines of "if pre-shift AND includes the sign bit in the bitfield, we must use signed BFE to replicate that bit during bitfield extraction. If the sign bit is not part of the mask, unsigned BFE will zero out upper bits of the result"

if (N->getOpcode() == ISD::SRA)
IsSigned = (ShiftAmt + NumBits) == Val.getValueSizeInBits();
} else if (LHS->getOpcode() == ISD::SHL) {
// Here, we have a pattern like:
//
Expand Down
200 changes: 177 additions & 23 deletions llvm/test/CodeGen/NVPTX/bfe.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,67 +1,221 @@
; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_20 | FileCheck %s
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefixes=CHECK,CHECK-O3
; RUN: llc < %s -mtriple=nvptx64 -mcpu=sm_20 -O0 | FileCheck %s --check-prefixes=CHECK,CHECK-O0
; RUN: %if ptxas %{ llc < %s -mtriple=nvptx64 -mcpu=sm_20 | %ptxas-verify %}
; RUN: %if ptxas %{ llc < %s -mtriple=nvptx64 -mcpu=sm_20 -O0 | %ptxas-verify %}

target triple = "nvptx64-nvidia-cuda"

; CHECK: bfe0
define i32 @bfe0(i32 %a) {
; CHECK: bfe.u32 %r{{[0-9]+}}, %r{{[0-9]+}}, 4, 4
; CHECK-NOT: shr
; CHECK-NOT: and
; CHECK-LABEL: bfe0(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [bfe0_param_0];
; CHECK-NEXT: bfe.u32 %r2, %r1, 4, 4;
; CHECK-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-NEXT: ret;
%val0 = ashr i32 %a, 4
%val1 = and i32 %val0, 15
ret i32 %val1
}

; CHECK: bfe1
define i32 @bfe1(i32 %a) {
; CHECK: bfe.u32 %r{{[0-9]+}}, %r{{[0-9]+}}, 3, 3
; CHECK-NOT: shr
; CHECK-NOT: and
; CHECK-LABEL: bfe1(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [bfe1_param_0];
; CHECK-NEXT: bfe.u32 %r2, %r1, 3, 3;
; CHECK-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-NEXT: ret;
%val0 = ashr i32 %a, 3
%val1 = and i32 %val0, 7
ret i32 %val1
}

; CHECK: bfe2
define i32 @bfe2(i32 %a) {
; CHECK: bfe.u32 %r{{[0-9]+}}, %r{{[0-9]+}}, 5, 3
; CHECK-NOT: shr
; CHECK-NOT: and
; CHECK-LABEL: bfe2(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [bfe2_param_0];
; CHECK-NEXT: bfe.u32 %r2, %r1, 5, 3;
; CHECK-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-NEXT: ret;
%val0 = ashr i32 %a, 5
%val1 = and i32 %val0, 7
ret i32 %val1
}

; CHECK-LABEL: no_bfe_on_32bit_overflow
define i32 @no_bfe_on_32bit_overflow(i32 %a) {
; CHECK-NOT: bfe.u32 %r{{[0-9]+}}, %r{{[0-9]+}}, 31, 4
; CHECK-LABEL: no_bfe_on_32bit_overflow(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<4>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [no_bfe_on_32bit_overflow_param_0];
; CHECK-NEXT: shr.s32 %r2, %r1, 31;
; CHECK-NEXT: and.b32 %r3, %r2, 15;
; CHECK-NEXT: st.param.b32 [func_retval0], %r3;
; CHECK-NEXT: ret;
%val0 = ashr i32 %a, 31
%val1 = and i32 %val0, 15
ret i32 %val1
}

; CHECK-LABEL: no_bfe_on_32bit_overflow_shr_and_pair
define i32 @no_bfe_on_32bit_overflow_shr_and_pair(i32 %a) {
; CHECK: shr.s32 %r{{[0-9]+}}, %r{{[0-9]+}}, 31
; CHECK: and.b32 %r{{[0-9]+}}, %r{{[0-9]+}}, 15
; CHECK-LABEL: no_bfe_on_32bit_overflow_shr_and_pair(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<4>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [no_bfe_on_32bit_overflow_shr_and_pair_param_0];
; CHECK-NEXT: shr.s32 %r2, %r1, 31;
; CHECK-NEXT: and.b32 %r3, %r2, 15;
; CHECK-NEXT: st.param.b32 [func_retval0], %r3;
; CHECK-NEXT: ret;
%val0 = ashr i32 %a, 31
%val1 = and i32 %val0, 15
ret i32 %val1
}

; CHECK-LABEL: no_bfe_on_64bit_overflow
define i64 @no_bfe_on_64bit_overflow(i64 %a) {
; CHECK-NOT: bfe.u64 %rd{{[0-9]+}}, %rd{{[0-9]+}}, 63, 3
; CHECK-LABEL: no_bfe_on_64bit_overflow(
; CHECK: {
; CHECK-NEXT: .reg .b64 %rd<4>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u64 %rd1, [no_bfe_on_64bit_overflow_param_0];
; CHECK-NEXT: shr.s64 %rd2, %rd1, 63;
; CHECK-NEXT: and.b64 %rd3, %rd2, 7;
; CHECK-NEXT: st.param.b64 [func_retval0], %rd3;
; CHECK-NEXT: ret;
%val0 = ashr i64 %a, 63
%val1 = and i64 %val0, 7
ret i64 %val1
}

; CHECK-LABEL: no_bfe_on_64bit_overflow_shr_and_pair
define i64 @no_bfe_on_64bit_overflow_shr_and_pair(i64 %a) {
; CHECK: shr.s64 %rd{{[0-9]+}}, %rd{{[0-9]+}}, 63
; CHECK: and.b64 %rd{{[0-9]+}}, %rd{{[0-9]+}}, 7
; CHECK-LABEL: no_bfe_on_64bit_overflow_shr_and_pair(
; CHECK: {
; CHECK-NEXT: .reg .b64 %rd<4>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u64 %rd1, [no_bfe_on_64bit_overflow_shr_and_pair_param_0];
; CHECK-NEXT: shr.s64 %rd2, %rd1, 63;
; CHECK-NEXT: and.b64 %rd3, %rd2, 7;
; CHECK-NEXT: st.param.b64 [func_retval0], %rd3;
; CHECK-NEXT: ret;
%val0 = ashr i64 %a, 63
%val1 = and i64 %val0, 7
ret i64 %val1
}

define i32 @bfe_ashr_signed_32(i32 %x) {
; CHECK-O3-LABEL: bfe_ashr_signed_32(
; CHECK-O3: {
; CHECK-O3-NEXT: .reg .b32 %r<3>;
; CHECK-O3-EMPTY:
; CHECK-O3-NEXT: // %bb.0:
; CHECK-O3-NEXT: ld.param.u16 %r1, [bfe_ashr_signed_32_param_0+2];
; CHECK-O3-NEXT: bfe.s32 %r2, %r1, 4, 12;
Comment on lines +124 to +125
Copy link
Member

@Artem-B Artem-B Mar 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just a drive-by observation. No action needed for the patch.

This may actually produce worse SASS than the unoptimized case.
Assuming that the parameter is passed in register, ptxas may end up right-shifting the input value, before applying BFE, while unoptimized case will just run BFE directly.
Or it could all be a wash if ptxas is smart enough to collapse shifts into BFEs.

; CHECK-O3-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-O3-NEXT: ret;
;
; CHECK-O0-LABEL: bfe_ashr_signed_32(
; CHECK-O0: {
; CHECK-O0-NEXT: .reg .b32 %r<3>;
; CHECK-O0-EMPTY:
; CHECK-O0-NEXT: // %bb.0:
; CHECK-O0-NEXT: ld.param.u32 %r1, [bfe_ashr_signed_32_param_0];
; CHECK-O0-NEXT: bfe.s32 %r2, %r1, 20, 12;
; CHECK-O0-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-O0-NEXT: ret;
%and = and i32 %x, -65536
%shr = ashr exact i32 %and, 20
ret i32 %shr
}

define i32 @bfe_ashr_unsigned_32(i32 %x) {
; CHECK-LABEL: bfe_ashr_unsigned_32(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u32 %r1, [bfe_ashr_unsigned_32_param_0];
; CHECK-NEXT: bfe.u32 %r2, %r1, 5, 6;
; CHECK-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-NEXT: ret;
%and = and i32 %x, 2047
%shr = ashr exact i32 %and, 5
ret i32 %shr
}

define i64 @bfe_ashr_signed_64(i64 %x) {
; CHECK-LABEL: bfe_ashr_signed_64(
; CHECK: {
; CHECK-NEXT: .reg .b64 %rd<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u64 %rd1, [bfe_ashr_signed_64_param_0];
; CHECK-NEXT: bfe.s64 %rd2, %rd1, 16, 48;
; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
; CHECK-NEXT: ret;
%and = and i64 %x, -65536
%shr = ashr exact i64 %and, 16
ret i64 %shr
}

define i64 @bfe_ashr_unsigned_64(i64 %x) {
; CHECK-LABEL: bfe_ashr_unsigned_64(
; CHECK: {
; CHECK-NEXT: .reg .b64 %rd<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.u64 %rd1, [bfe_ashr_unsigned_64_param_0];
; CHECK-NEXT: bfe.u64 %rd2, %rd1, 5, 6;
; CHECK-NEXT: st.param.b64 [func_retval0], %rd2;
; CHECK-NEXT: ret;
%and = and i64 %x, 2047
%shr = ashr exact i64 %and, 5
ret i64 %shr
}

define i32 @bfe3(i128 %a) {
; CHECK-LABEL: bfe3(
; CHECK: {
; CHECK-NEXT: .reg .b32 %r<3>;
; CHECK-NEXT: .reg .b64 %rd<3>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.v2.u64 {%rd1, %rd2}, [bfe3_param_0];
; CHECK-NEXT: cvt.u32.u64 %r1, %rd1;
; CHECK-NEXT: bfe.s32 %r2, %r1, 15, 17;
; CHECK-NEXT: st.param.b32 [func_retval0], %r2;
; CHECK-NEXT: ret;
%trunc = trunc i128 %a to i32
%and = and i32 %trunc, -32768
%shr = ashr exact i32 %and, 15
ret i32 %shr
}

define i64 @bfe4(i128 %a) {
; CHECK-LABEL: bfe4(
; CHECK: {
; CHECK-NEXT: .reg .b64 %rd<4>;
; CHECK-EMPTY:
; CHECK-NEXT: // %bb.0:
; CHECK-NEXT: ld.param.v2.u64 {%rd1, %rd2}, [bfe4_param_0];
; CHECK-NEXT: bfe.s64 %rd3, %rd1, 17, 47;
; CHECK-NEXT: st.param.b64 [func_retval0], %rd3;
; CHECK-NEXT: ret;
%trunc = trunc i128 %a to i64
%and = and i64 %trunc, -131072
%shr = ashr exact i64 %and, 17
ret i64 %shr
}