[llvm] Only check assertions that were meant to apply to the normal case of non-splat vector SREM expansion when we aren't hitting the special case. (PR #86238)

[Owen Anderson via llvm-commits]

https://github.com/resistor updated https://github.com/llvm/llvm-project/pull/86238

>From 640e70e0c57462d96519af0dcec0180cbbd1ae81 Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor at mac.com>
Date: Thu, 21 Mar 2024 22:36:08 -0500
Subject: [PATCH 1/2] Only check assertions that were meant to apply to the
 normal case of non-splat vector SREM expansion when we aren't hitting the
 special case.

Fixes https://github.com/llvm/llvm-project/issues/84830
Introduced in https://github.com/llvm/llvm-project/pull/82706
---
 llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp | 10 +++++-----
 llvm/test/CodeGen/AArch64/srem-vec-crash.ll      | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/srem-vec-crash.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index da29b1d5b312f8..8be03b66e155f6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -6916,6 +6916,11 @@ TargetLowering::prepareSREMEqFold(EVT SETCCVT, SDValue REMNode,
     // Q = floor((2 * A) / (2^K))
     APInt Q = (2 * A).udiv(APInt::getOneBitSet(W, K));
 
+    assert(APInt::getAllOnes(SVT.getSizeInBits()).ugt(A) &&
+           "We are expecting that A is always less than all-ones for SVT");
+    assert(APInt::getAllOnes(ShSVT.getSizeInBits()).ugt(K) &&
+           "We are expecting that K is always less than all-ones for ShSVT");
+
     // If D was a power of two, apply the alternate constant derivation.
     if (D0.isOne()) {
       // A = 2^(W-1)
@@ -6924,11 +6929,6 @@ TargetLowering::prepareSREMEqFold(EVT SETCCVT, SDValue REMNode,
       Q = APInt::getAllOnes(W - K).zext(W);
     }
 
-    assert(APInt::getAllOnes(SVT.getSizeInBits()).ugt(A) &&
-           "We are expecting that A is always less than all-ones for SVT");
-    assert(APInt::getAllOnes(ShSVT.getSizeInBits()).ugt(K) &&
-           "We are expecting that K is always less than all-ones for ShSVT");
-
     // If the divisor is 1 the result can be constant-folded. Likewise, we
     // don't care about INT_MIN lanes, those can be set to undef if appropriate.
     if (D.isOne()) {
diff --git a/llvm/test/CodeGen/AArch64/srem-vec-crash.ll b/llvm/test/CodeGen/AArch64/srem-vec-crash.ll
new file mode 100644
index 00000000000000..3b8a1c83b2f697
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/srem-vec-crash.ll
@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+;RUN: llc -mtriple=aarch64-unknown-unknown < %s | FileCheck %s
+
+define i32 @f(i1 %0) {
+; CHECK-LABEL: f:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w0, #1 // =0x1
+; CHECK-NEXT:    ret
+  %new0 = srem i1 %0, true
+  %last = zext i1 %new0 to i32
+  %2 = icmp ne i32 %last, 0
+  %3 = select i1 %2, i32 0, i32 1
+  ret i32 %3
+}

>From e50661404c1d4cf14af1920d1685bccd570f85dd Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor at mac.com>
Date: Sat, 23 Mar 2024 21:46:52 -0500
Subject: [PATCH 2/2] Update test based on review feedback.

---
 llvm/test/CodeGen/AArch64/srem-vec-crash.ll | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/llvm/test/CodeGen/AArch64/srem-vec-crash.ll b/llvm/test/CodeGen/AArch64/srem-vec-crash.ll
index 3b8a1c83b2f697..0fce8de30d4d4b 100644
--- a/llvm/test/CodeGen/AArch64/srem-vec-crash.ll
+++ b/llvm/test/CodeGen/AArch64/srem-vec-crash.ll
@@ -1,14 +1,15 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
-;RUN: llc -mtriple=aarch64-unknown-unknown < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-unknown-unknown < %s | FileCheck %s
 
-define i32 @f(i1 %0) {
-; CHECK-LABEL: f:
-; CHECK:       // %bb.0:
+define i32 @pr84830(i1 %arg) {
+; CHECK-LABEL: pr84830:
+; CHECK:       // %bb.0: // %bb
 ; CHECK-NEXT:    mov w0, #1 // =0x1
 ; CHECK-NEXT:    ret
-  %new0 = srem i1 %0, true
+bb:
+  %new0 = srem i1 %arg, true
   %last = zext i1 %new0 to i32
-  %2 = icmp ne i32 %last, 0
-  %3 = select i1 %2, i32 0, i32 1
-  ret i32 %3
+  %i = icmp ne i32 %last, 0
+  %i1 = select i1 %i, i32 0, i32 1
+  ret i32 %i1
 }