[llvm] [SCEVExpander] Don't expand a UDiv with a possibly-poison divisor (PR #202378)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-analysis

Author: lijinpei-amd

<details>
<summary>Changes</summary>

SCEVExpander::isSafeToExpand only check divisor isKnownNonZero, which
ignore the possibility of poison. For the following divisor:
```
%ct = call i32 @<!-- -->llvm.cttz.i32(i32 %x, i1 true)
%divisor = add i32 %ct, 1
...
%rem = urem i32 1, %divisor
```
The urem may be hoisted unsafely.

Fix by also check divisor isGuaranteedNotToBePoison.

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

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


3 Files Affected:

- (modified) llvm/include/llvm/Analysis/ScalarEvolution.h (+3-3) 
- (modified) llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp (+2-1) 
- (added) llvm/test/Transforms/IndVarSimplify/X86/exit-value-safe-udiv.ll (+44) 


``````````diff
diff --git a/llvm/include/llvm/Analysis/ScalarEvolution.h b/llvm/include/llvm/Analysis/ScalarEvolution.h
index 15b891b5b4da9..50af763614a31 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolution.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolution.h
@@ -1235,6 +1235,9 @@ class ScalarEvolution {
   /// Test if the given expression is known to be non-zero.
   LLVM_ABI bool isKnownNonZero(const SCEV *S);
 
+  /// Returns true if \p Op is guaranteed to not be poison.
+  LLVM_ABI static bool isGuaranteedNotToBePoison(const SCEV *Op);
+
   /// Test if the given expression is known to be a power of 2.  OrNegative
   /// allows matching negative power of 2s, and OrZero allows matching 0.
   LLVM_ABI bool isKnownToBeAPowerOfTwo(const SCEV *S, bool OrZero = false,
@@ -2427,9 +2430,6 @@ class ScalarEvolution {
   /// Returns true if \p Op is guaranteed not to cause immediate UB.
   bool isGuaranteedNotToCauseUB(const SCEV *Op);
 
-  /// Returns true if \p Op is guaranteed to not be poison.
-  static bool isGuaranteedNotToBePoison(const SCEV *Op);
-
   /// Return true if the SCEV corresponding to \p I is never poison.  Proving
   /// this is more complex than proving that just \p I is never poison, since
   /// SCEV commons expressions across control flow, and you can have cases
diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index d747c7fb2d162..b2e45c3340b9f 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -2449,7 +2449,8 @@ struct SCEVFindUnsafe {
 
   bool follow(const SCEV *S) {
     if (const SCEVUDivExpr *D = dyn_cast<SCEVUDivExpr>(S)) {
-      if (!SE.isKnownNonZero(D->getRHS())) {
+      if (!SE.isKnownNonZero(D->getRHS()) ||
+          !SE.isGuaranteedNotToBePoison(D->getRHS())) {
         IsUnsafe = true;
         return false;
       }
diff --git a/llvm/test/Transforms/IndVarSimplify/X86/exit-value-safe-udiv.ll b/llvm/test/Transforms/IndVarSimplify/X86/exit-value-safe-udiv.ll
new file mode 100644
index 0000000000000..453d93cea2708
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/X86/exit-value-safe-udiv.ll
@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -passes=indvars < %s | FileCheck %s
+
+; Reduced from https://github.com/llvm/llvm-project/issues/202028.
+
+target triple = "x86_64-unknown-linux-gnu"
+
+declare i32 @llvm.cttz.i32(i32, i1)
+
+define i32 @speculated_exit_value_udiv(i32 %x, i1 %guard, i1 %backedge) {
+; CHECK-LABEL: @speculated_exit_value_udiv(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CT:%.*]] = call i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 true)
+; CHECK-NEXT:    [[D:%.*]] = add i32 [[CT]], 1
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br i1 [[GUARD:%.*]], label [[EXIT_EARLY:%.*]], label [[BODY:%.*]]
+; CHECK:       body:
+; CHECK-NEXT:    [[REM:%.*]] = urem i32 1, [[D]]
+; CHECK-NEXT:    br i1 [[BACKEDGE:%.*]], label [[LOOP]], label [[EXIT_LATE:%.*]]
+; CHECK:       exit.early:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       exit.late:
+; CHECK-NEXT:    [[REM_LCSSA:%.*]] = phi i32 [ [[REM]], [[BODY]] ]
+; CHECK-NEXT:    ret i32 [[REM_LCSSA]]
+;
+entry:
+  %ct = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %d = add i32 %ct, 1
+  br label %loop
+
+loop:
+  br i1 %guard, label %exit.early, label %body
+
+body:
+  %rem = urem i32 1, %d
+  br i1 %backedge, label %loop, label %exit.late
+
+exit.early:
+  ret i32 0
+
+exit.late:
+  ret i32 %rem
+}

``````````

</details>


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