[llvm] [ConstraintElim] Drop invalid rows instead of failing the elimination (PR #76299)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/76299

This patch tries to drop invalid rows when overflow occurs instead of failing the elimination immediately.
Fixes the regression in [velox/buffer/Buffer.h](https://github.com/facebookincubator/velox/blob/50187434e32bffcbebcd6501898763c56de40065/velox/buffer/Buffer.h#L347-L350), which was introduced by #76262.

See also https://github.com/dtcxzyw/llvm-opt-benchmark/issues/35#issuecomment-1868362725.


>From 9c21c21f9794af835594daf68c55bc0534ce6e90 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sun, 24 Dec 2023 03:46:39 +0800
Subject: [PATCH 1/2] [ConstraintElim] Add pre-commit tests. NFC.

---
 .../constraint-overflow.ll                    | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll

diff --git a/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll b/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll
new file mode 100644
index 00000000000000..07e98e7581b649
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll
@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
+
+define void @f(i64 %a3, i64 %numElements) {
+; CHECK-LABEL: define void @f(
+; CHECK-SAME: i64 [[A3:%.*]], i64 [[NUMELEMENTS:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i64 [[NUMELEMENTS]], 1152921504606846975
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[A1:%.*]] = shl nuw i64 [[NUMELEMENTS]], 4
+; CHECK-NEXT:    br label [[IF_END:%.*]]
+; CHECK:       if.end:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i64 [[A1]], [[A3]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_END_I:%.*]], label [[ABORT:%.*]]
+; CHECK:       if.end.i:
+; CHECK-NEXT:    [[CMP2_NOT_I:%.*]] = icmp ult i64 [[A1]], [[A3]]
+; CHECK-NEXT:    br i1 [[CMP2_NOT_I]], label [[ABORT]], label [[EXIT:%.*]]
+; CHECK:       abort:
+; CHECK-NEXT:    tail call void @llvm.trap()
+; CHECK-NEXT:    unreachable
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cond = icmp ule i64 %numElements, 1152921504606846975
+  call void @llvm.assume(i1 %cond)
+  %a1 = shl nuw i64 %numElements, 4
+  br label %if.end
+if.end:
+  %cmp = icmp ugt i64 %a1, %a3
+  br i1 %cmp, label %if.end.i, label %abort
+if.end.i:
+  %cmp2.not.i = icmp ult i64 %a1, %a3
+  br i1 %cmp2.not.i, label %abort, label %exit
+abort:
+  tail call void @llvm.trap()
+  unreachable
+exit:
+  ret void
+}
+
+declare void @llvm.trap()
+declare void @llvm.assume(i1)

>From 7eb9c5feb00f42b23acaf33dc973ea154f3fe9ac Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sun, 24 Dec 2023 04:01:38 +0800
Subject: [PATCH 2/2] [ConstraintElim] Drop invalid rows instead of failing the
 elimination

---
 llvm/lib/Analysis/ConstraintSystem.cpp        | 27 ++++++++++++-------
 .../constraint-overflow.ll                    |  3 +--
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Analysis/ConstraintSystem.cpp b/llvm/lib/Analysis/ConstraintSystem.cpp
index 8a802515b6f4fb..43bd5e5b0ccbdc 100644
--- a/llvm/lib/Analysis/ConstraintSystem.cpp
+++ b/llvm/lib/Analysis/ConstraintSystem.cpp
@@ -73,6 +73,8 @@ bool ConstraintSystem::eliminateUsingFM() {
       }
 
       SmallVector<Entry, 8> NR;
+      bool Valid = true;
+      uint32_t NextGCD = NewGCD;
       unsigned IdxUpper = 0;
       unsigned IdxLower = 0;
       auto &LowerRow = RemainingRows[LowerR];
@@ -96,27 +98,34 @@ bool ConstraintSystem::eliminateUsingFM() {
           IdxUpper++;
         }
 
-        if (MulOverflow(UpperV, ((-1) * LowerLast / GCD), M1))
-          return false;
+        if (MulOverflow(UpperV, ((-1) * LowerLast / GCD), M1)) {
+          Valid = false;
+          break;
+        }
         if (IdxLower < LowerRow.size() && LowerRow[IdxLower].Id == CurrentId) {
           LowerV = LowerRow[IdxLower].Coefficient;
           IdxLower++;
         }
 
-        if (MulOverflow(LowerV, (UpperLast / GCD), M2))
-          return false;
-        if (AddOverflow(M1, M2, N))
-          return false;
+        if (MulOverflow(LowerV, (UpperLast / GCD), M2)) {
+          Valid = false;
+          break;
+        }
+        if (AddOverflow(M1, M2, N)) {
+          Valid = false;
+          break;
+        }
         if (N == 0)
           continue;
         NR.emplace_back(N, CurrentId);
 
-        NewGCD =
-            APIntOps::GreatestCommonDivisor({32, (uint32_t)N}, {32, NewGCD})
+        NextGCD =
+            APIntOps::GreatestCommonDivisor({32, (uint32_t)N}, {32, NextGCD})
                 .getZExtValue();
       }
-      if (NR.empty())
+      if (!Valid || NR.empty())
         continue;
+      NewGCD = NextGCD;
       Constraints.push_back(std::move(NR));
       // Give up if the new system gets too big.
       if (Constraints.size() > 500)
diff --git a/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll b/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll
index 07e98e7581b649..6200f3adf0088d 100644
--- a/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll
+++ b/llvm/test/Transforms/ConstraintElimination/constraint-overflow.ll
@@ -13,8 +13,7 @@ define void @f(i64 %a3, i64 %numElements) {
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i64 [[A1]], [[A3]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF_END_I:%.*]], label [[ABORT:%.*]]
 ; CHECK:       if.end.i:
-; CHECK-NEXT:    [[CMP2_NOT_I:%.*]] = icmp ult i64 [[A1]], [[A3]]
-; CHECK-NEXT:    br i1 [[CMP2_NOT_I]], label [[ABORT]], label [[EXIT:%.*]]
+; CHECK-NEXT:    br i1 false, label [[ABORT]], label [[EXIT:%.*]]
 ; CHECK:       abort:
 ; CHECK-NEXT:    tail call void @llvm.trap()
 ; CHECK-NEXT:    unreachable