[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

Author: annamthomas

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

llvm/test/Analysis/ScalarEvolution/pr87798.ll

@@ -0,0 +1,59 @@
+; RUN: opt -disable-output -passes='print<scalar-evolution>' -verify-scev < %s 2>&1 | 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.
+
+; CHECK: Classifying expressions for: @pr87798
+; CHECK:  %phi = phi i32 [ 0, %bb ], [ %add4, %bb1 ]
+; CHECK-NEXT:  -->  {0,+,0,+,0,+,2,+,3}<%bb1> U: full-set S: full-set		Exits: 0		LoopDispositions: { %bb1: Computable }
+
+; CHECK-NEXT:  %phi2 = phi i32 [ 0, %bb ], [ %add, %bb1 ]
+; CHECK-NEXT:  -->  {0,+,0,+,1}<%bb1> U: full-set S: full-set		Exits: 0		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %phi3 = phi i32 [ 0, %bb ], [ %add5, %bb1 ]
+; CHECK-NEXT:  -->  {0,+,1}<nuw><nsw><%bb1> U: [0,1) S: [0,1)		Exits: 0		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %add = add i32 %phi2, %phi3
+; CHECK-NEXT:  -->  {0,+,1,+,1}<%bb1> U: full-set S: full-set		Exits: 0		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %mul = mul i32 %phi2, %phi3
+; CHECK-NEXT:  -->  {0,+,0,+,2,+,3}<%bb1> U: full-set S: full-set		Exits: 0		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %add4 = add i32 %mul, %phi
+; CHECK-NEXT:  -->  {0,+,0,+,2,+,5,+,3}<%bb1> U: full-set S: full-set		Exits: 0		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %and = and i32 %phi, 1
+; CHECK-NEXT:  -->  (zext i1 {false,+,false,+,false,+,false,+,true}<%bb1> to i32) U: [0,2) S: [0,2)		Exits: 0	LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %add5 = add i32 %phi3, 1
+; CHECK-NEXT:  -->  {1,+,1}<nuw><nsw><%bb1> U: [1,2) S: [1,2)		Exits: 1		LoopDispositions: { %bb1: Computable }
+; CHECK-NEXT:  %phi9 = phi i32 [ %and, %bb1 ]
+; CHECK-NEXT:  -->  (zext i1 {false,+,false,+,false,+,false,+,true}<%bb1> to i32) U: [0,2) S: [0,2)  -->  0 U: [0,1) S: [0,1)
+; CHECK-NEXT:  %zext = zext i32 %phi9 to i64
+; CHECK-NEXT:  -->  poison U: full-set S: full-set
+
+define i32 @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
+}