[SCEV] Generalize folding of trunc(x)+n*trunc(y) into folding m*trunc(x)+n*trunc(y)

[dneilson]

Summary

A SCEV such as:
{%v2,+,((-1 * (trunc i64 (-1 * %v1) to i32)) + (-1 * (trunc i64 %v1 to i32)))}<%loop>

can be folded into, simply, {%v2,+,0}. However, the current code in ::getAddExpr()
will not try to apply the simplification m*trunc(x)+n*trunc(y) -> trunc(trunc(m)*x+trunc(n)*y)
because it only keys off having a non-multiplied trunc as the first term in the simplification.

This patch generalizes this code to try to do a more generic fold of these trunc
expressions.

[sanjoy]

lib/Analysis/ScalarEvolution.cpp

2347

This is doing a bit more work than it needs to -- if the first element after a constant isn't a trunc, then none of them are.

2366

Instead of searching from a trunc operation and going from there, how about phrasing this logic as:

1. First figure out the SrcType, if possible (by looking at truncates and multiplies)
2. If (1) was successful, re-traverse the operands and try to do the trunc extraction

This makes it clear that the second part is only using SrcType and not the whole truncate.

[dneilson]

Minor rearrangement of the patch to get srctype from the lambda instead of a trunc.
( Sorry for the delay; I was out of town & away from computers. )

[sanjoy]

lib/Analysis/ScalarEvolution.cpp

2376

I think we can rewrite this inner loop as:

auto *LastOp = Mul->getOperand(Mul->getNumOperands() - 1);
if (auto *T = dyn_cast<SCEVTruncateExpr>(LastOp))
  return T->getOperand()->getType();
if (!isa<SCEVConstant>(LastOp))
  break;

I think we can use this trick in the outer loop as well. If we just look at the last operand to the add then we'll miss (say) zero extend expressions etc. but that's fine because the loop that actually does the work will catch it. Here we just need to provide a quick "estimate".

What do you think?

[dneilson]

lib/Analysis/ScalarEvolution.cpp

2376

Sounds like a good idea to me. I wasn't aware that the operands of all SCEV expressions are sorted.

Should be able to apply this trick down below as well.

[dneilson]

Rewrite lambda to exploit SCEV operand ordering.

[sanjoy]

lib/Analysis/ScalarEvolution.cpp

2375

I think this can just be:

if (auto *Mul = dyn_cast<SCEVMulExpr>(Ops[Idx]))
  if (auto *Trunc = dyn_cast<SCEVTruncateExpr>(Mul->getOperand(Mul->getNumOperands() - 1)))
    return Trunc->getOperand()->getType();
return nullptr;

The reasoning being we don't need to worry about mul expressions with the last operand as a constant -- these should have been folded away.

[dneilson]

Remove the loop from the lambda.

[sanjoy]

LGTM!

lib/Analysis/ScalarEvolution.cpp

if (Ops[i] == Ops[i+1]) {      //  X + Y + Y  -->  X + Y*2
       --i; e -= Count - 1;
       FoundMatch = true;
     }
   if (FoundMatch)
     return getAddExpr(Ops, Flags);
 
   // Check for truncates. If all the operands are truncated from the same
   // type, see if factoring out the truncate would permit the result to be
-  // folded. eg., trunc(x) + m*trunc(n) --> trunc(x + trunc(m)*n)
+  // folded. eg., n*trunc(x) + m*trunc(y) --> trunc(trunc(m)*x + trunc(n)*y)
   // if the contents of the resulting outer trunc fold to something simple.
-  for (; Idx < Ops.size() && isa<SCEVTruncateExpr>(Ops[Idx]); ++Idx) {
-    const SCEVTruncateExpr *Trunc = cast<SCEVTruncateExpr>(Ops[Idx]);
+  auto FindTrunc = [&]() -> const SCEVTruncateExpr* {
+    // Go through the available Ops to see if we have a compatible trunc() to
+    // start processing.
+    for (unsigned i=Idx; i<Ops.size() && Ops[i]->getSCEVType() <= scMulExpr; ++i) {
+      if (isa<SCEVTruncateExpr>(Ops[i])) {
+        return cast<SCEVTruncateExpr>(Ops[i]);
+      } else if (const auto *Mul = dyn_cast<SCEVMulExpr>(Ops[i])) {
+        bool Ok = true;
+        for (unsigned j=0, e=Mul->getNumOperands(); Ok && j<e; ++j) {
+          const auto *Op = Mul->getOperand(j);
+          if (const auto *T = dyn_cast<SCEVTruncateExpr>(Op)) {
+            return T;
+          } else if (!isa<SCEVConstant>(Op)) {
+            Ok = false;
+          }
+        }
+        if (!Ok) break;
+      } else {
+        break;
+      }
+    }
+    return nullptr;
+  };
+  if (const auto *Trunc = FindTrunc()) {
     Type *DstType = Trunc->getType();
     Type *SrcType = Trunc->getOperand()->getType();
     SmallVector<const SCEV *, 8> LargeOps;
     bool Ok = true;
     // Check all the operands to see if they can be represented in the
     // source type of the truncate.
     for (unsigned i = 0, e = Ops.size(); i != e; ++i) {
       if (const SCEVTruncateExpr *T = dyn_cast<SCEVTruncateExpr>(Ops[i])) {
         if (T->getOperand()->getType() != SrcType) {
           Ok = false;
           break;
         }
         LargeOps.push_back(T->getOperand());
       } else if (const SCEVConstant *C = dyn_cast<SCEVConstant>(Ops[i])) {
         LargeOps.push_back(getAnyExtendExpr(C, SrcType));
       } else if (const SCEVMulExpr *M = dyn_cast<SCEVMulExpr>(Ops[i])) {
         SmallVector<const SCEV *, 8> LargeMulOps;
         for (unsigned j = 0, f = M->getNumOperands(); j != f && Ok; ++j) {

unittests/Analysis/ScalarEvolutionTest.cpp

TEST_F(ScalarEvolutionsTest, SCEVAddRecFromPHIwithLargeConstants) {
   ScalarEvolution SE = buildSE(*F);
   SCEVUnionPredicate Preds;
   const SCEV *Expr = SE.getSCEV(Phi);
   EXPECT_NE(nullptr, Expr);
   EXPECT_TRUE(isa<SCEVUnknown>(Expr));
   auto Result = SE.createAddRecFromPHIWithCasts(cast<SCEVUnknown>(Expr));
 }
 
+TEST_F(ScalarEvolutionsTest, SCEVFoldSumOfTruncs) {
+  // Verify that the following SCEV gets folded to a zero:
+  //  (-1 * (trunc i64 (-1 * %0) to i32)) + (-1 * (trunc i64 %0 to i32)
+  Type *ArgTy = Type::getInt64Ty(Context);
+  Type *Int32Ty = Type::getInt32Ty(Context);
+  SmallVector<Type *, 1> Types;
+  Types.push_back(ArgTy);
+  FunctionType *FTy = FunctionType::get(Type::getVoidTy(Context), Types, false);
+  Function *F = cast<Function>(M.getOrInsertFunction("f", FTy));
+  BasicBlock *BB = BasicBlock::Create(Context, "entry", F);
+  ReturnInst::Create(Context, nullptr, BB);
+
+  ScalarEvolution SE = buildSE(*F);
+
+  auto *Arg = &*(F->arg_begin());
+  const auto *ArgSCEV = SE.getSCEV(Arg);
+
+  // Build the SCEV
+  const auto *A0 = SE.getNegativeSCEV(ArgSCEV);
+  const auto *A1 = SE.getTruncateExpr(A0, Int32Ty);
+  const auto *A = SE.getNegativeSCEV(A1);
+
+  const auto *B0 = SE.getTruncateExpr(ArgSCEV, Int32Ty);
+  const auto *B = SE.getNegativeSCEV(B0);
+
+  const auto *Expr = SE.getAddExpr(A, B);
+  dbgs() << "DDN\nExpr: " << *Expr << "\n";
+  // Verify that the SCEV was folded to 0
+  const auto *ZeroConst = SE.getConstant(Int32Ty, 0);
+  EXPECT_EQ(Expr, ZeroConst);
+}
+
 }  // end anonymous namespace
 }  // end namespace llvm