[llvm] [ConstraintElim] Add missing check to make sure the bound is loop-invariant (PR #70555)

[Yingwei Zheng via llvm-commits]

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

This PR adds a missing check for https://github.com/llvm/llvm-project/commit/e6a1657fa30c747f4412fc47f567660ebe861a9e to make sure the bound is loop-invariant.

Example:
```
define i32 @main() {
entry:
  br label %for.header

for.header:
  %ind1 = phi i64 [ %ind1.i, %for.latch ], [ 0, %entry ]
  %ind2 = phi i64 [ %ind2.i, %for.latch ], [ 5, %entry ]
  %rem = srem i64 %ind2, 4
  %cmp6 = icmp eq i64 %ind1, %rem
  br i1 %cmp6, label %exit1, label %for.latch

for.latch:
  %ind2.i = add i64 %ind2, 1
  %ind1.i = add i64 %ind1, 1
  %cond = icmp eq i64 %ind1.i, 8
  br i1 %cond, label %exit2, label %for.header

exit2:
  %cmp = icmp eq i64 %rem, 0
  %ret = zext i1 %cmp to i32
  ret i32 %ret

exit1:
  ret i32 0
}
```
We cannot infer `%ind1 <u %rem` from `icmp eq i64 %ind1, %rem` in the loop since `%rem` is not a loop-invariant.

Fixes #70510.

>From dc58f7e2ca3b67f0ac4a105f580ccf7ac41e3f8e Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 28 Oct 2023 21:45:03 +0800
Subject: [PATCH 1/2] [ConstraintElim] Add pre-commit tests for PR70510. NFC.

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

diff --git a/llvm/test/Transforms/ConstraintElimination/pr70510.ll b/llvm/test/Transforms/ConstraintElimination/pr70510.ll
new file mode 100644
index 000000000000000..f211146dd38e601
--- /dev/null
+++ b/llvm/test/Transforms/ConstraintElimination/pr70510.ll
@@ -0,0 +1,49 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
+
+; FIXME: This is a miscompilation. %rem is not loop-invariant.
+define i32 @main() {
+; CHECK-LABEL: define i32 @main() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_HEADER:%.*]]
+; CHECK:       for.header:
+; CHECK-NEXT:    [[IND1:%.*]] = phi i64 [ [[IND1_I:%.*]], [[FOR_LATCH:%.*]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[IND2:%.*]] = phi i64 [ [[IND2_I:%.*]], [[FOR_LATCH]] ], [ 5, [[ENTRY]] ]
+; CHECK-NEXT:    [[REM:%.*]] = srem i64 [[IND2]], 4
+; CHECK-NEXT:    [[CMP6:%.*]] = icmp eq i64 [[IND1]], [[REM]]
+; CHECK-NEXT:    br i1 [[CMP6]], label [[EXIT1:%.*]], label [[FOR_LATCH]]
+; CHECK:       for.latch:
+; CHECK-NEXT:    [[IND2_I]] = add i64 [[IND2]], 1
+; CHECK-NEXT:    [[IND1_I]] = add i64 [[IND1]], 1
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[IND1_I]], 8
+; CHECK-NEXT:    br i1 [[COND]], label [[EXIT2:%.*]], label [[FOR_HEADER]]
+; CHECK:       exit2:
+; CHECK-NEXT:    [[RET:%.*]] = zext i1 false to i32
+; CHECK-NEXT:    ret i32 [[RET]]
+; CHECK:       exit1:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  br label %for.header
+
+for.header:
+  %ind1 = phi i64 [ %ind1.i, %for.latch ], [ 0, %entry ]
+  %ind2 = phi i64 [ %ind2.i, %for.latch ], [ 5, %entry ]
+  %rem = srem i64 %ind2, 4
+  %cmp6 = icmp eq i64 %ind1, %rem
+  br i1 %cmp6, label %exit1, label %for.latch
+
+for.latch:
+  %ind2.i = add i64 %ind2, 1
+  %ind1.i = add i64 %ind1, 1
+  %cond = icmp eq i64 %ind1.i, 8
+  br i1 %cond, label %exit2, label %for.header
+
+exit2:
+  %cmp = icmp eq i64 %rem, 0
+  %ret = zext i1 %cmp to i32
+  ret i32 %ret
+
+exit1:
+  ret i32 0
+}

>From fd56196ed46c3562531fad45bf94a94cc406b7aa Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 28 Oct 2023 21:47:21 +0800
Subject: [PATCH 2/2] [ConstraintElim] Add missing check to make sure the bound
 is loop-invariant

---
 llvm/lib/Transforms/Scalar/ConstraintElimination.cpp  | 4 ++++
 llvm/test/Transforms/ConstraintElimination/pr70510.ll | 4 ++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
index 37f720ec40f4e54..2d829ee017be6e4 100644
--- a/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/ConstraintElimination.cpp
@@ -899,6 +899,10 @@ void State::addInfoForInductions(BasicBlock &BB) {
   else
     return;
 
+  // Make sure the bound B is loop-invariant.
+  if (!L->isLoopInvariant(B))
+    return;
+
   // Handle negative steps.
   if (StepOffset.isNegative()) {
     // TODO: Extend to allow steps > -1.
diff --git a/llvm/test/Transforms/ConstraintElimination/pr70510.ll b/llvm/test/Transforms/ConstraintElimination/pr70510.ll
index f211146dd38e601..888fb9c0bd57ada 100644
--- a/llvm/test/Transforms/ConstraintElimination/pr70510.ll
+++ b/llvm/test/Transforms/ConstraintElimination/pr70510.ll
@@ -1,7 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
 ; RUN: opt -passes=constraint-elimination -S %s | FileCheck %s
 
-; FIXME: This is a miscompilation. %rem is not loop-invariant.
 define i32 @main() {
 ; CHECK-LABEL: define i32 @main() {
 ; CHECK-NEXT:  entry:
@@ -18,7 +17,8 @@ define i32 @main() {
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[IND1_I]], 8
 ; CHECK-NEXT:    br i1 [[COND]], label [[EXIT2:%.*]], label [[FOR_HEADER]]
 ; CHECK:       exit2:
-; CHECK-NEXT:    [[RET:%.*]] = zext i1 false to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i64 [[REM]], 0
+; CHECK-NEXT:    [[RET:%.*]] = zext i1 [[CMP]] to i32
 ; CHECK-NEXT:    ret i32 [[RET]]
 ; CHECK:       exit1:
 ; CHECK-NEXT:    ret i32 0