[llvm] [NVPTX] Fix BFE selection for i64 values (PR #228685)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: wenhu1024

<details>
<summary>Changes</summary>

I meet a problem concerning BFE selection：

`tryBFE` currently uses the shift amount type width when checking the BFE start position. Since the shift amount is i32 even when the source value is i64, valid extracts starting at bit 32 or above are rejected.

Use the source value width instead, allowing these patterns to be selected as `bfe.u64`.

For example, an extract from bit 40:

```llvm
%v = xor i64 %x, %y
%s = lshr i64 %v, 40
%r = and i64 %s, 255
```

can now be lowered to `bfe.u64` instead of `shr.u64` + `and.b64`.



---
Full diff: https://github.com/llvm/llvm-project/pull/228685.diff


2 Files Affected:

- (modified) llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp (+1-1) 
- (modified) llvm/test/CodeGen/NVPTX/bfe.ll (+30-1) 


``````````diff
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
index 186bfa185e135a..b048362d985117 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp
@@ -1816,7 +1816,7 @@ bool NVPTXDAGToDAGISel::tryBFE(SDNode *N) {
         uint64_t StartVal = StartConst->getZExtValue();
         // How many "good" bits do we have left?  "good" is defined here as bits
         // that exist in the original value, not shifted in.
-        int64_t GoodBits = Start.getValueSizeInBits() - StartVal;
+        int64_t GoodBits = Val.getValueSizeInBits() - StartVal;
         if (NumBits > GoodBits) {
           // Do not handle the case where bits have been shifted in. In theory
           // we could handle this, but the cost is likely higher than just
diff --git a/llvm/test/CodeGen/NVPTX/bfe.ll b/llvm/test/CodeGen/NVPTX/bfe.ll
index 644bf3606e8f64..515529cc0461b1 100644
--- a/llvm/test/CodeGen/NVPTX/bfe.ll
+++ b/llvm/test/CodeGen/NVPTX/bfe.ll
@@ -51,6 +51,36 @@ define i32 @bfe2(i32 %a) {
   ret i32 %val1
 }
 
+define i64 @bfe_i64_start_above_32(i64 %a, i64 %b) {
+; CHECK-O3-LABEL: bfe_i64_start_above_32(
+; CHECK-O3:       {
+; CHECK-O3-NEXT:    .reg .b64 %rd<5>;
+; CHECK-O3-EMPTY:
+; CHECK-O3-NEXT:  // %bb.0:
+; CHECK-O3-NEXT:    ld.param.b64 %rd1, [bfe_i64_start_above_32_param_0];
+; CHECK-O3-NEXT:    ld.param.b64 %rd2, [bfe_i64_start_above_32_param_1];
+; CHECK-O3-NEXT:    xor.b64 %rd3, %rd1, %rd2;
+; CHECK-O3-NEXT:    bfe.u64 %rd4, %rd3, 40, 8;
+; CHECK-O3-NEXT:    st.param.b64 [func_retval0], %rd4;
+; CHECK-O3-NEXT:    ret;
+;
+; CHECK-O0-LABEL: bfe_i64_start_above_32(
+; CHECK-O0:       {
+; CHECK-O0-NEXT:    .reg .b64 %rd<5>;
+; CHECK-O0-EMPTY:
+; CHECK-O0-NEXT:  // %bb.0:
+; CHECK-O0-NEXT:    ld.param.b64 %rd2, [bfe_i64_start_above_32_param_1];
+; CHECK-O0-NEXT:    ld.param.b64 %rd1, [bfe_i64_start_above_32_param_0];
+; CHECK-O0-NEXT:    xor.b64 %rd3, %rd1, %rd2;
+; CHECK-O0-NEXT:    bfe.u64 %rd4, %rd3, 40, 8;
+; CHECK-O0-NEXT:    st.param.b64 [func_retval0], %rd4;
+; CHECK-O0-NEXT:    ret;
+  %x = xor i64 %a, %b
+  %val0 = lshr i64 %x, 40
+  %val1 = and i64 %val0, 255
+  ret i64 %val1
+}
+
 define i32 @no_bfe_on_32bit_overflow(i32 %a) {
 ; CHECK-LABEL: no_bfe_on_32bit_overflow(
 ; CHECK:       {
@@ -218,4 +248,3 @@ define i64 @bfe4(i128 %a) {
   %shr = ashr exact i64 %and, 17
   ret i64 %shr
 }
-

``````````

</details>


https://github.com/llvm/llvm-project/pull/228685