[llvm] [SCEV] Fix BinomialCoefficient Iteration to fit in W bits (PR #88010)

[via llvm-commits]

https://github.com/annamthomas updated https://github.com/llvm/llvm-project/pull/88010

>From 14b0015271a766df6f44f40f9cfa8e5005fb3b43 Mon Sep 17 00:00:00 2001
From: Anna Thomas <anna at azul.com>
Date: Mon, 8 Apr 2024 12:05:06 -0400
Subject: [PATCH] [SCEV] Fix BinomialCoefficient Iteration to fit in W bits

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
---
 llvm/lib/Analysis/ScalarEvolution.cpp         |  6 ++--
 .../test/Transforms/IndVarSimplify/pr87798.ll | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/IndVarSimplify/pr87798.ll

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index e030b9fc7dac4f..9fcce797f55976 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -928,11 +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);
-    OddFactorial *= Mult;
+    OddFactorial *= (i >> TwoFactors);
   }
 
   // We need at least W + T bits for the multiplication step
diff --git a/llvm/test/Transforms/IndVarSimplify/pr87798.ll b/llvm/test/Transforms/IndVarSimplify/pr87798.ll
new file mode 100644
index 00000000000000..6e5c6a609877f6
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/pr87798.ll
@@ -0,0 +1,36 @@
+; RUN: opt -S -passes='print<scalar-evolution>,indvars' -verify-scev < %s | FileCheck %s
+
+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"
+
+; print<scalar-evolution> is used to compute SCEVs for all values in the
+; function.
+; 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
+}