[SCEV] Fix BinomialCoefficient Iteration to fit in W bits

[annamthomas]

BinomialCoefficient computes the value of W-bit IV at iteration It of a
loop. When W is 1, we can call multiplicative inverse on 0 which
triggers an assert since 1b76120.

Since the arithmetic is supposed to wrap if It or K does not fit in W
bits, do the truncation into W bits after we do the shift.

Fixes #87798

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: None (annamthomas)

Changes

BinomialCoefficient computes the value of W-bit IV at iteration It of a
loop. When W is 1, we can call multiplicative inverse on 0 which
triggers an assert since 1b76120.

Since the arithmetic is supposed to wrap if It or K does not fit in W
bits, do the truncation into W bits after we do the shift.

Fixes #87798

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/ScalarEvolution.cpp (+2-3)
• (added) llvm/test/Transforms/IndVarSimplify/pr87798.ll (+36)

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index e030b9fc7dac4f..4301dcfd9a9537 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -928,10 +928,9 @@ static const SCEV *BinomialCoefficient(const SCEV *It, unsigned K,
   APInt OddFactorial(W, 1);
   unsigned T = 1;
   for (unsigned i = 3; i <= K; ++i) {
-    APInt Mult(W, i);
-    unsigned TwoFactors = Mult.countr_zero();
+    unsigned TwoFactors = countr_zero(i);
     T += TwoFactors;
-    Mult.lshrInPlace(TwoFactors);
+    APInt Mult(W, i >> TwoFactors);
     OddFactorial *= Mult;
   }
 
diff --git a/llvm/test/Transforms/IndVarSimplify/pr87798.ll b/llvm/test/Transforms/IndVarSimplify/pr87798.ll
new file mode 100644
index 00000000000000..17a761e726b984
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/pr87798.ll
@@ -0,0 +1,36 @@
+; RUN: opt -S -passes='indvars' -verify-scev < %s | FileCheck %s
+
+; REQUIRES: asserts
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128-ni:1-p2:32:8:8:32-ni:2"
+target triple = "x86_64-unknown-linux-gnu"
+
+; We should not crash on multiplicative inverse called within SCEV's binomial
+; coefficient function.
+define i32 @pr87798() {
+; CHECK-LABEL: pr87798 
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb1, %bb
+  %phi = phi i32 [ 0, %bb ], [ %add4, %bb1 ]
+  %phi2 = phi i32 [ 0, %bb ], [ %add, %bb1 ]
+  %phi3 = phi i32 [ 0, %bb ], [ %add5, %bb1 ]
+  %add = add i32 %phi2, %phi3
+  %mul = mul i32 %phi2, %phi3
+  %add4 = add i32 %mul, %phi
+  %and = and i32 %phi, 1
+  %add5 = add i32 %phi3, 1
+  br i1 true, label %preheader, label %bb1
+
+preheader:                                              ; preds = %bb1
+  %phi9 = phi i32 [ %and, %bb1 ]
+  br label %loop
+
+loop:                                              ; preds = %preheader, %loop
+  br label %loop
+
+bb7:                                              ; No predecessors!
+  %zext = zext i32 %phi9 to i64
+  ret i32 0
+}

[efriedma-quic]

LGTM