[SCEV] Re-apply r335197 (with Polly fixes).

[timshen]

Summary

This initiates a discussion on changing Polly accordingly while re-applying r335197 (D48338).

I have never worked on Polly. The proposed change to param_div_div_div_2.ll is not educated, but just patterns that match the output.

All LLVM files are already reviewed in D48338.

[efriedma]

LGTM on the polly bits.

The polly output is changing because polly strips "zext" operations off SCEV expressions, but doesn't touch extensions inside a SCEV expression. The updated test still tests what it's supposed to test.

llvm/include/llvm/Analysis/ScalarEvolution.h

private:
   /// A loop is considered "used" by an expression if it contains
   /// an add rec on said loop.
   void getUsedLoops(const SCEV *S, SmallPtrSetImpl<const Loop *> &LoopsUsed);
 
   /// Find all of the loops transitively used in \p S, and update \c LoopUsers
   /// accordingly.
   void addToLoopUseLists(const SCEV *S);
 
+  /// Try to match the pattern generated by getURemExpr(A, B). If successful,
+  /// Assign A and B to LHS and RHS, respectively.
+  bool matchURem(const SCEV *Expr, const SCEV *&LHS, const SCEV *&RHS);
+
   FoldingSet<SCEV> UniqueSCEVs;
   FoldingSet<SCEVPredicate> UniquePreds;
   BumpPtrAllocator SCEVAllocator;
 
   /// This maps loops to a list of SCEV expressions that (transitively) use said
   /// loop.
   DenseMap<const Loop *, SmallVector<const SCEV *, 4>> LoopUsers;
 

llvm/lib/Analysis/ScalarEvolution.cpp

if (AR->isAffine()) {
       if (proveNoWrapByVaryingStart<SCEVZeroExtendExpr>(Start, Step, L)) {
         const_cast<SCEVAddRecExpr *>(AR)->setNoWrapFlags(SCEV::FlagNUW);
         return getAddRecExpr(
             getExtendAddRecStart<SCEVZeroExtendExpr>(AR, Ty, this, Depth + 1),
             getZeroExtendExpr(Step, Ty, Depth + 1), L, AR->getNoWrapFlags());
       }
     }
 
+  // zext(A % B) --> zext(A) % zext(B)
+  {
+    const SCEV *LHS;
+    const SCEV *RHS;
+    if (matchURem(Op, LHS, RHS))
+      return getURemExpr(getZeroExtendExpr(LHS, Ty, Depth + 1),
+                         getZeroExtendExpr(RHS, Ty, Depth + 1));
+  }
+
+  // zext(A / B) --> zext(A) / zext(B).
+  if (auto *Div = dyn_cast<SCEVUDivExpr>(Op))
+    return getUDivExpr(getZeroExtendExpr(Div->getLHS(), Ty, Depth + 1),
+                       getZeroExtendExpr(Div->getRHS(), Ty, Depth + 1));
+
   if (auto *SA = dyn_cast<SCEVAddExpr>(Op)) {
     // zext((A + B + ...)<nuw>) --> (zext(A) + zext(B) + ...)<nuw>
     if (SA->hasNoUnsignedWrap()) {
       // If the addition does not unsign overflow then we can, by definition,
       // commute the zero extension with the addition operation.
       SmallVector<const SCEV *, 4> Ops;
       for (const auto *Op : SA->operands())
         Ops.push_back(getZeroExtendExpr(Op, Ty, Depth + 1));
for (auto &I : *BB) {
       if (II->second.second == Expr)
         continue;
 
       OS.indent(Depth) << "[PSE]" << I << ":\n";
       OS.indent(Depth + 2) << *Expr << "\n";
       OS.indent(Depth + 2) << "--> " << *II->second.second << "\n";
     }
 }
+
+// Match the mathematical pattern A - (A / B) * B, where A and B can be
+// arbitrary expressions.
+// It's not always easy, as A and B can be folded (imagine A is X / 2, and B is
+// 4, A / B becomes X / 8).
+bool ScalarEvolution::matchURem(const SCEV *Expr, const SCEV *&LHS,
+                                const SCEV *&RHS) {
+  const auto *Add = dyn_cast<SCEVAddExpr>(Expr);
+  if (Add == nullptr || Add->getNumOperands() != 2)
+    return false;
+
+  const SCEV *A = Add->getOperand(1);
+  const auto *Mul = dyn_cast<SCEVMulExpr>(Add->getOperand(0));
+
+  if (Mul == nullptr)
+    return false;
+
+  const auto MatchURemWithDivisor = [&](const SCEV *B) {
+    // (SomeExpr + (-(SomeExpr / B) * B)).
+    if (Expr == getURemExpr(A, B)) {
+      LHS = A;
+      RHS = B;
+      return true;
+    }
+    return false;
+  };
+
+  // (SomeExpr + (-1 * (SomeExpr / B) * B)).
+  if (Mul->getNumOperands() == 3 && isa<SCEVConstant>(Mul->getOperand(0)))
+    return MatchURemWithDivisor(Mul->getOperand(1)) ||
+           MatchURemWithDivisor(Mul->getOperand(2));
+
+  // (SomeExpr + ((-SomeExpr / B) * B)) or (SomeExpr + ((SomeExpr / B) * -B)).
+  if (Mul->getNumOperands() == 2)
+    return MatchURemWithDivisor(Mul->getOperand(1)) ||
+           MatchURemWithDivisor(Mul->getOperand(0)) ||
+           MatchURemWithDivisor(getNegativeSCEV(Mul->getOperand(1))) ||
+           MatchURemWithDivisor(getNegativeSCEV(Mul->getOperand(0)));
+  return false;
+}

llvm/test/Analysis/ScalarEvolution/zext-divrem.ll

+; RUN: opt -analyze -scalar-evolution -S < %s | FileCheck %s
+
+define i64 @test1(i32 %a, i32 %b) {
+; CHECK-LABEL: @test1
+  %div = udiv i32 %a, %b
+  %zext = zext i32 %div to i64
+; CHECK: %zext
+; CHECK-NEXT: -->  ((zext i32 %a to i64) /u (zext i32 %b to i64))
+  ret i64 %zext
+}
+
+define i64 @test2(i32 %a, i32 %b) {
+; CHECK-LABEL: @test2
+  %rem = urem i32 %a, %b
+  %zext = zext i32 %rem to i64
+; CHECK: %zext
+; CHECK-NEXT: -->  ((zext i32 %a to i64) + (-1 * (zext i32 %b to i64) * ((zext i32 %a to i64) /u (zext i32 %b to i64))))
+  ret i64 %zext
+}
+
+define i64 @test3(i32 %a, i32 %b) {
+; CHECK-LABEL: @test3
+  %div = udiv i32 %a, %b
+  %mul = mul i32 %div, %b
+  %sub = sub i32 %a, %mul
+  %zext = zext i32 %sub to i64
+; CHECK: %zext
+; CHECK-NEXT: -->  ((zext i32 %a to i64) + (-1 * (zext i32 %b to i64) * ((zext i32 %a to i64) /u (zext i32 %b to i64))))
+  ret i64 %zext
+}
+
+define i64 @test4(i32 %t) {
+; CHECK-LABEL: @test4
+  %a = udiv i32 %t, 2
+  %div = udiv i32 %t, 112
+  %mul = mul i32 %div, 56
+  %sub = sub i32 %a, %mul
+  %zext = zext i32 %sub to i64
+; CHECK: %zext
+; CHECK-NEXT: -->  ((-56 * ((zext i32 %t to i64) /u 112)) + ((zext i32 %t to i64) /u 2))
+  ret i64 %zext
+}

polly/test/Isl/CodeGen/param_div_div_div_2.ll

 ; RUN: opt %loadPolly -polly-scops -analyze < %s | FileCheck %s
 ; RUN: opt %loadPolly -polly-codegen -S < %s | FileCheck %s --check-prefix=IR
 ;
 ; Check that we guard the divisions because we moved them and thereby increased
 ; their domain.
 ;
 ; CHECK:         Invalid Context:
-; CHECK-NEXT:    [p_0] -> {  : p_0 < 0 }
-; CHECK-NEXT:    p0: ((%a /u %b) /u (%c /u %d))
+; CHECK-NEXT:    [p_0] -> {  : false }
+; CHECK-NEXT:    p0: (((zext i32 %a to i64) /u (zext i32 %b to i64)) /u ((zext i32 %c to i64) /u (zext i32 %d to i64)))
 ;
 ;    void f(unsigned *A, unsigned a, unsigned b, unsigned c, unsigned d) {
 ;      for (unsigned i; i < 100; i++)
 ;        A[i] += A[(a / b) / (c / d)];
 ;    }
 ;
-; IR:       %[[R0:[.a-zA-Z0-9]*]] = icmp ugt i32 %b, 1
-; IR-NEXT:  %[[R1:[.a-zA-Z0-9]*]] = select i1 %[[R0]], i32 %b, i32 1
-; IR-NEXT:  %[[R2:[.a-zA-Z0-9]*]] = udiv i32 %a, %[[R1]]
-; IR-NEXT:  %[[R5:[.a-zA-Z0-9]*]] = icmp ugt i32 %d, 1
-; IR-NEXT:  %[[R6:[.a-zA-Z0-9]*]] = select i1 %[[R5]], i32 %d, i32 1
-; IR-NEXT:  %[[R7:[.a-zA-Z0-9]*]] = udiv i32 %c, %[[R6]]
-; IR-NEXT:  %[[R3:[.a-zA-Z0-9]*]] = icmp ugt i32 %[[R7]], 1
-; IR-NEXT:  %[[R4:[.a-zA-Z0-9]*]] = select i1 %[[R3]], i32 %[[R7]], i32 1
-; IR-NEXT:  %[[R8:[.a-zA-Z0-9]*]] = udiv i32 %[[R2]], %[[R4]]
+; IR:       %[[A:[.a-zA-Z0-9]*]] = zext i32 %a to i64
+; IR-NEXT:  %[[B:[.a-zA-Z0-9]*]] = zext i32 %b to i64
+; IR-NEXT:  %[[R0:[.a-zA-Z0-9]*]] = icmp ugt i64 %[[B]], 1
+; IR-NEXT:  %[[R1:[.a-zA-Z0-9]*]] = select i1 %[[R0]], i64 %[[B]], i64 1
+; IR-NEXT:  %[[R2:[.a-zA-Z0-9]*]] = udiv i64 %[[A]], %[[R1]]
+; IR-NEXT:  %[[C:[.a-zA-Z0-9]*]] = zext i32 %c to i64
+; IR-NEXT:  %[[D:[.a-zA-Z0-9]*]] = zext i32 %d to i64
+; IR-NEXT:  %[[R5:[.a-zA-Z0-9]*]] = icmp ugt i64 %[[D]], 1
+; IR-NEXT:  %[[R6:[.a-zA-Z0-9]*]] = select i1 %[[R5]], i64 %[[D]], i64 1
+; IR-NEXT:  %[[R7:[.a-zA-Z0-9]*]] = udiv i64 %[[C]], %[[R6]]
+; IR-NEXT:  %[[R3:[.a-zA-Z0-9]*]] = icmp ugt i64 %[[R7]], 1
+; IR-NEXT:  %[[R4:[.a-zA-Z0-9]*]] = select i1 %[[R3]], i64 %[[R7]], i64 1
+; IR-NEXT:  %[[R8:[.a-zA-Z0-9]*]] = udiv i64 %[[R2]], %[[R4]]
 ;
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 define void @f(i32* %A, i32 %a, i32 %b, i32 %c, i32 %d) {
 entry:
   br label %for.cond
 
 for.cond:                                         ; preds = %for.inc, %entry