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

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 llvm#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,68 @@
+; NOTE: Assertions have been autogenerated by utils/update_analyze_test_checks.py UTC_ARGS: --version 4
+; 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.
+
+define i32 @pr87798() {
+; CHECK-LABEL: 'pr87798'
+; CHECK-NEXT:  Classifying expressions for: @pr87798
+; CHECK-NEXT:    %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
+; CHECK-NEXT:  Determining loop execution counts for: @pr87798
+; CHECK-NEXT:  Loop %loop: <multiple exits> Unpredictable backedge-taken count.
+; CHECK-NEXT:  Loop %loop: Unpredictable constant max backedge-taken count.
+; CHECK-NEXT:  Loop %loop: Unpredictable symbolic max backedge-taken count.
+; CHECK-NEXT:  Loop %bb1: backedge-taken count is i1 false
+; CHECK-NEXT:  Loop %bb1: constant max backedge-taken count is i1 false
+; CHECK-NEXT:  Loop %bb1: symbolic max backedge-taken count is i1 false
+; CHECK-NEXT:  Loop %bb1: Trip multiple is 1
+;
+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
+}