[llvm] [SCEV] BECount to zero if `((-C + (C smax %x)) /u %x), C > 0` holds (PR #104580)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto updated https://github.com/llvm/llvm-project/pull/104580

>From a5689c3c1bc90f7caa5c45f58f3ded35d9b4ddc7 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Fri, 16 Aug 2024 11:51:05 +0200
Subject: [PATCH 1/2] [SCEV] Introduce test for PR104580 (NFC)

---
 .../udiv-of-x-xsmaxone-fold.ll                | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)
 create mode 100644 llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll

diff --git a/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll b/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll
new file mode 100644
index 00000000000000..35b1550aba2c28
--- /dev/null
+++ b/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll
@@ -0,0 +1,45 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py
+; RUN: opt -disable-output -passes="print<scalar-evolution>" < %s 2>&1 | FileCheck %s
+
+ at g_var = external global i32, align 4
+
+define void @test(i32 %x) {
+; CHECK-LABEL: 'test'
+; CHECK-NEXT:  Classifying expressions for: @test
+; CHECK-NEXT:    %g_var.promoted = load i32, ptr @g_var, align 4
+; CHECK-NEXT:    --> %g_var.promoted U: full-set S: full-set
+; CHECK-NEXT:    %smax = tail call i32 @llvm.smax.i32(i32 %x, i32 1)
+; CHECK-NEXT:    --> (1 smax %x) U: [1,-2147483648) S: [1,-2147483648)
+; CHECK-NEXT:    %add = add nsw i32 %smax, -1
+; CHECK-NEXT:    --> (-1 + (1 smax %x))<nsw> U: [0,2147483647) S: [0,2147483647)
+; CHECK-NEXT:    %udiv = udiv i32 %add, %x
+; CHECK-NEXT:    --> ((-1 + (1 smax %x))<nsw> /u %x) U: [0,2147483647) S: [0,2147483647)
+;
+; CHECK-NEXT:  Determining loop execution counts for: @test
+; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Unpredictable constant max backedge-taken count.
+; CHECK-NEXT:  Loop %for.body: Unpredictable symbolic max backedge-taken count.
+;
+entry:
+  %g_var.promoted = load i32, ptr @g_var, align 4
+  %smax = tail call i32 @llvm.smax.i32(i32 %x, i32 1)
+  %add = add nsw i32 %smax, -1
+  %udiv = udiv i32 %add, %x
+  %cond = icmp ult i32 %udiv, 7
+  br i1 %cond, label %for.body, label %exit
+
+for.body:                                         ; preds = %for.body, %entry
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %for.body ]
+  %val = phi i32 [ %g_var.promoted, %entry ], [ %add.1, %for.body ]
+  %add.1 = add nsw i32 %val, %x
+  %iv.next = add nsw i32 %iv, %x
+  %cond.1 = icmp slt i32 %iv.next, 1
+  br i1 %cond.1, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  store i32 %add.1, ptr @g_var, align 4
+  br label %exit
+
+exit:                                             ; preds = %entry, %for.cond.cleanup
+  ret void
+}

>From 2024380bb4a0fdf288f16c833c8506fe36543782 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Fri, 16 Aug 2024 12:44:20 +0200
Subject: [PATCH 2/2] [SCEV] BECount to zero if `((-C + (C smax %x)) /u %x), C
 > 0` holds

The SCEV expression `((-C + (C smax %x)) /u %x)` can be folded
to zero for positive %x and for any positive constant C.

Proof: https://alive2.llvm.org/ce/z/_dLm8C.
---
 llvm/lib/Analysis/ScalarEvolution.cpp         | 19 +++++++++++++++++++
 .../udiv-of-x-xsmaxone-fold.ll                |  2 +-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 21a1c74eefc071..a93ce15bd27b99 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -3547,6 +3547,25 @@ const SCEV *ScalarEvolution::getUDivExpr(const SCEV *LHS,
     }
   }
 
+  // ((-C + (C smax %x)) /u %x) evaluates to zero, for any positive %x and any
+  // positive constant C.
+  if (const auto *AE = dyn_cast<SCEVAddExpr>(LHS);
+      AE && AE->getNumOperands() == 2) {
+    if (const auto *VC = dyn_cast<SCEVConstant>(AE->getOperand(0))) {
+      APInt NegC = VC->getAPInt();
+      if (NegC.isNegative() &&
+          NegC != APInt::getSignedMinValue(NegC.getBitWidth())) {
+        const auto *MME = dyn_cast<SCEVSMaxExpr>(AE->getOperand(1));
+        if (MME && MME->getNumOperands() == 2) {
+          if (isa<SCEVConstant>(MME->getOperand(0)) &&
+              cast<SCEVConstant>(MME->getOperand(0))->getAPInt() == -NegC)
+            if (MME->getOperand(1) == RHS)
+              return getZero(LHS->getType());
+        }
+      }
+    }
+  }
+
   // The Insertion Point (IP) might be invalid by now (due to UniqueSCEVs
   // changes). Make sure we get a new one.
   IP = nullptr;
diff --git a/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll b/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll
index 35b1550aba2c28..a316d5164c85b1 100644
--- a/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll
+++ b/llvm/test/Analysis/ScalarEvolution/udiv-of-x-xsmaxone-fold.ll
@@ -13,7 +13,7 @@ define void @test(i32 %x) {
 ; CHECK-NEXT:    %add = add nsw i32 %smax, -1
 ; CHECK-NEXT:    --> (-1 + (1 smax %x))<nsw> U: [0,2147483647) S: [0,2147483647)
 ; CHECK-NEXT:    %udiv = udiv i32 %add, %x
-; CHECK-NEXT:    --> ((-1 + (1 smax %x))<nsw> /u %x) U: [0,2147483647) S: [0,2147483647)
+; CHECK-NEXT:    --> 0 U: [0,1) S: [0,1)
 ;
 ; CHECK-NEXT:  Determining loop execution counts for: @test
 ; CHECK-NEXT:  Loop %for.body: Unpredictable backedge-taken count.