[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

2346

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.

2365

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

2351

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

2351

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

2350

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]);
-    Type *DstType = Trunc->getType();
-    Type *SrcType = Trunc->getOperand()->getType();
+  auto FindTruncSrcType = [&]() -> Type * {
+    // We're ultimately looking to fold an addrec of truncs and muls of only
+    // constants and truncs, so if we find any other types of SCEV
+    // as operands of the addrec then we bail and return nullptr here.
+    // Otherwise, we return the type of the operand of a trunc that we find.
+    if (auto *T = dyn_cast<SCEVTruncateExpr>(Ops[Idx]))
+      return T->getOperand()->getType();
+    if (const auto *Mul = dyn_cast<SCEVMulExpr>(Ops[Idx])) {
+      const auto *LastOp = Mul->getOperand(Mul->getNumOperands() - 1);
+      if (const auto *T = dyn_cast<SCEVTruncateExpr>(LastOp))
+        return T->getOperand()->getType();
+    }
+    return nullptr;
+  };
+  if (auto *SrcType = FindTruncSrcType()) {
     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) {
for (unsigned i = 0, e = Ops.size(); i != e; ++i) {
         break;
       }
     }
     if (Ok) {
       // Evaluate the expression in the larger type.
       const SCEV *Fold = getAddExpr(LargeOps, Flags, Depth + 1);
       // If it folds to something simple, use it. Otherwise, don't.
       if (isa<SCEVConstant>(Fold) || isa<SCEVUnknown>(Fold))
-        return getTruncateExpr(Fold, DstType);
+        return getTruncateExpr(Fold, Ty);
     }
   }
 
   // Skip past any other cast SCEVs.
   while (Idx < Ops.size() && Ops[Idx]->getSCEVType() < scAddExpr)
     ++Idx;
 
   // If there are add operands they would be next.

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