[llvm] [RISCV] Add additional conditions for signed icmp cases when canonicalizing masks for VL-predicate (PR #218068)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-risc-v

Author: Min-Yih Hsu (mshockwave)

<details>
<summary>Changes</summary>

Fixes #<!-- -->217500 

Previously in #<!-- -->214877 we added a mask canonicalization to make it easier for the mask-consuming instructions to use VL-predicated instead. But the transformation will be incorrect if the icmp against boundary value is signed. Specifically, the problem arises if `%base + %offset` is negative. This patch fixes this issue by guarding `%base + %offset` to ensure this expression (1) never overflow, and (2) is non-negative. 

Alive2 Proof: https://alive2.llvm.org/ce/z/L2M_LL

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


2 Files Affected:

- (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+16) 
- (modified) llvm/test/CodeGen/RISCV/rvv/masked-load-vl-predicatable.ll (+27-2) 


``````````diff
diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index 9406b697882d7..0f0bf048956ba 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -19216,6 +19216,22 @@ static SDValue canonicalizeMaskForVLPredicate(EVT MaskVT, SDValue LHS,
   APInt Offset =
       IsSigned ? Start.sextOrTrunc(LenWidth) : Start.zextOrTrunc(LenWidth);
 
+  if (IsSigned) {
+    // Two additional conditions:
+    // 1. Offset + BaseIndex never overflow
+    if (!Offset.isZero() && (DAG.ComputeNumSignBits(BaseIndex) <= 1 ||
+                             Offset.getNumSignBits() <= 1))
+      return SDValue();
+
+    // 2. Offset + BaseIndex has to be non-negative
+    auto BaseIndexKB = DAG.computeKnownBits(BaseIndex);
+    auto OffsetKB = KnownBits::makeConstant(Offset);
+    if (!KnownBits::add(BaseIndexKB, OffsetKB, /*NSW=*/true,
+                        /*NUW=*/false)
+             .isNonNegative())
+      return SDValue();
+  }
+
   unsigned MinOpc = IsSigned ? ISD::SMIN : ISD::UMIN;
   SDValue NewStepVector = DAG.getStepVector(DL, LHS.getValueType());
   BaseIndex = DAG.getNode(
diff --git a/llvm/test/CodeGen/RISCV/rvv/masked-load-vl-predicatable.ll b/llvm/test/CodeGen/RISCV/rvv/masked-load-vl-predicatable.ll
index 6eeb9839418df..a0036d7b02040 100644
--- a/llvm/test/CodeGen/RISCV/rvv/masked-load-vl-predicatable.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/masked-load-vl-predicatable.ll
@@ -92,6 +92,7 @@ define <16 x float> @masked_load_signed_cmp(ptr %p, i32 signext %n, i32 signext
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 16, e32, m4, ta, ma
 ; CHECK-NEXT:    vmv.v.i v8, 0
+; CHECK-NEXT:    srliw a2, a2, 3
 ; CHECK-NEXT:    min a2, a1, a2
 ; CHECK-NEXT:    sub a1, a1, a2
 ; CHECK-NEXT:    li a2, 16
@@ -99,7 +100,8 @@ define <16 x float> @masked_load_signed_cmp(ptr %p, i32 signext %n, i32 signext
 ; CHECK-NEXT:    vsetvli zero, a1, e32, m4, tu, ma
 ; CHECK-NEXT:    vle32.v v8, (a0)
 ; CHECK-NEXT:    ret
-  %11 = insertelement <16 x i32> poison, i32 %iv, i32 0
+  %base = lshr i32 %iv, 3
+  %11 = insertelement <16 x i32> poison, i32 %base, i32 0
   %12 = shufflevector <16 x i32> %11, <16 x i32> poison, <16 x i32> zeroinitializer
   %13 = or disjoint <16 x i32> %12, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
   %14 = insertelement <16 x i32> poison, i32 %n, i32 0
@@ -114,6 +116,7 @@ define <16 x float> @masked_load_sgt(ptr %p, i32 signext %n, i32 signext %iv) no
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 16, e32, m4, ta, ma
 ; CHECK-NEXT:    vmv.v.i v8, 0
+; CHECK-NEXT:    srliw a2, a2, 3
 ; CHECK-NEXT:    min a2, a1, a2
 ; CHECK-NEXT:    sub a1, a1, a2
 ; CHECK-NEXT:    li a2, 16
@@ -121,7 +124,8 @@ define <16 x float> @masked_load_sgt(ptr %p, i32 signext %n, i32 signext %iv) no
 ; CHECK-NEXT:    vsetvli zero, a1, e32, m4, tu, ma
 ; CHECK-NEXT:    vle32.v v8, (a0)
 ; CHECK-NEXT:    ret
-  %11 = insertelement <16 x i32> poison, i32 %iv, i32 0
+  %base = lshr i32 %iv, 3
+  %11 = insertelement <16 x i32> poison, i32 %base, i32 0
   %12 = shufflevector <16 x i32> %11, <16 x i32> poison, <16 x i32> zeroinitializer
   %13 = or disjoint <16 x i32> %12, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
   %14 = insertelement <16 x i32> poison, i32 %n, i32 0
@@ -253,3 +257,24 @@ define <16 x float> @negative_masked_load_wrong_step_vector(ptr %p, i32 signext
   %19 = tail call <16 x float> @llvm.masked.load(ptr %p, <16 x i1> %16, <16 x float> zeroinitializer)
   ret <16 x float> %19
 }
+
+; We can't to do this transformation if we cannot gaurantee %base + %offset to always be non-negative.
+define <16 x float> @negative_masked_load_signed_cmp(ptr %p, i32 signext %n, i32 signext %base) nounwind {
+; CHECK-LABEL: negative_masked_load_signed_cmp:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 16, e32, m4, ta, mu
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vor.vx v8, v8, a2
+; CHECK-NEXT:    vmslt.vx v0, v8, a1
+; CHECK-NEXT:    vmv.v.i v8, 0
+; CHECK-NEXT:    vle32.v v8, (a0), v0.t
+; CHECK-NEXT:    ret
+  %11 = insertelement <16 x i32> poison, i32 %base, i32 0
+  %12 = shufflevector <16 x i32> %11, <16 x i32> poison, <16 x i32> zeroinitializer
+  %13 = or disjoint <16 x i32> %12, <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %14 = insertelement <16 x i32> poison, i32 %n, i32 0
+  %15 = shufflevector <16 x i32> %14, <16 x i32> poison, <16 x i32> zeroinitializer
+  %16 = icmp slt <16 x i32> %13, %15
+  %19 = tail call <16 x float> @llvm.masked.load(ptr %p, <16 x i1> %16, <16 x float> zeroinitializer)
+  ret <16 x float> %19
+}

``````````

</details>


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