[llvm] dc81e69 - [IndVars] Check expansion safety in makeIVComparisonInvariant() (PR62992)

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2023-05-31T11:21:35+02:00
New Revision: dc81e69eb126ea1dbab202733ae34b1a3d6eb0cf

URL: https://github.com/llvm/llvm-project/commit/dc81e69eb126ea1dbab202733ae34b1a3d6eb0cf
DIFF: https://github.com/llvm/llvm-project/commit/dc81e69eb126ea1dbab202733ae34b1a3d6eb0cf.diff

LOG: [IndVars] Check expansion safety in makeIVComparisonInvariant() (PR62992)

Make sure the invariant expressions are safe to expand. In
particular, we should not speculative a trapping division into
the preheader.

Fixes https://github.com/llvm/llvm-project/issues/62992.

Added: 
    llvm/test/Transforms/IndVarSimplify/pr62992.ll

Modified: 
    llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
index 1458e6184861c..62821f149f444 100644
--- a/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyIndVar.cpp
@@ -217,8 +217,10 @@ bool SimplifyIndvar::makeIVComparisonInvariant(ICmpInst *ICmp,
 
   // Do not generate something ridiculous.
   auto *PHTerm = Preheader->getTerminator();
-  if (Rewriter.isHighCostExpansion({ InvariantLHS, InvariantRHS }, L,
-                                   2 * SCEVCheapExpansionBudget, TTI, PHTerm))
+  if (Rewriter.isHighCostExpansion({InvariantLHS, InvariantRHS}, L,
+                                   2 * SCEVCheapExpansionBudget, TTI, PHTerm) ||
+      !Rewriter.isSafeToExpandAt(InvariantLHS, PHTerm) ||
+      !Rewriter.isSafeToExpandAt(InvariantRHS, PHTerm))
     return false;
   auto *NewLHS =
       Rewriter.expandCodeFor(InvariantLHS, IVOperand->getType(), PHTerm);

diff  --git a/llvm/test/Transforms/IndVarSimplify/pr62992.ll b/llvm/test/Transforms/IndVarSimplify/pr62992.ll
new file mode 100644
index 0000000000000..c8f47b57f1eda
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/pr62992.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
+; RUN: opt -S -passes=indvars < %s | FileCheck %s
+
+declare void @use(i1)
+
+; Make sure the division does not get expanded into the preheader.
+
+define i32 @test(i32 %arg) {
+; CHECK-LABEL: define i32 @test
+; CHECK-SAME: (i32 [[ARG:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 false, label [[IF:%.*]], label [[LOOP_LATCH:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    [[DIV:%.*]] = udiv i32 7, [[ARG]]
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp ult i32 1, [[DIV]]
+; CHECK-NEXT:    call void @use(i1 [[CMP2]])
+; CHECK-NEXT:    br label [[LOOP_LATCH]]
+; CHECK:       loop.latch:
+; CHECK-NEXT:    br i1 false, label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop.latch ]
+  %iv.next = add i32 %iv, 1
+  %cmp = icmp eq i32 %iv, 1
+  br i1 %cmp, label %if, label %loop.latch
+
+if:
+  %div = udiv i32 7, %arg
+  %cmp2 = icmp ult i32 %iv.next, %div
+  call void @use(i1 %cmp2)
+  br label %loop.latch
+
+loop.latch:
+  br i1 false, label %loop, label %exit
+
+exit:
+  %inc.lcssa = phi i32 [ %iv.next, %loop.latch ]
+  ret i32 %inc.lcssa
+}