[SCEV] Match 'zext (trunc A to iB) to iY' as URem.

[fhahn]

Summary

URem operations with constant power-of-2 second operands are modeled as
such. This patch on its own has very little impact (e.g. no changes in
CodeGen for MultiSource/SPEC2000/SPEC2006 on X86 -O3 -flto), but I'll
soon post follow-up patches that make use of it to more accurately
determine the trip multiple.

[lebedev.ri]

I would expect that you should be able to showcase the effect with opt -analyze -scalar-evolution-based test:
https://godbolt.org/z/xddz8d (because there's zext(A % B) --> zext(A) % zext(B) matchURem()-driven fold)

llvm/lib/Analysis/ScalarEvolution.cpp

12857

What if Expr isn't an zext itself, but it is used by zext (which is the case in the single user of this method)?

12860–12861

What if we had zext(zext(trunc)), so we end up with a wider type than A was?

[mkazantsev]

llvm/lib/Analysis/ScalarEvolution.cpp

12867

I think getPowerOf2 would be a useful utility function in SCEV. Consider factoring out (maybe as separate patch).

12868

1u << X will overflowan become zero if X > 32. Consider using APInt.

llvm/unittests/Analysis/ScalarEvolutionTest.cpp

67

static?

[mkazantsev]

llvm/unittests/Analysis/ScalarEvolutionTest.cpp

1334

Could you please make a 64-bit version of this test to catch bug with 1u << size if it happens?

[fhahn]

Address comments: use APInt instead unsigned for shift, handle different zext sizes, add test with divisor > 1 << 32.

Thanks!

[fhahn]

In D89821#2342783, @lebedev.ri wrote:

I would expect that you should be able to showcase the effect with opt -analyze -scalar-evolution-based test:
https://godbolt.org/z/xddz8d (because there's zext(A % B) --> zext(A) % zext(B) matchURem()-driven fold)

Not sure if that's possible, because the expression is already in the form zext (trunc ... to iX) to iY and I think if there was an outer zext, it would be already folded into the inner zext.

llvm/lib/Analysis/ScalarEvolution.cpp

12857

Hm, I not sure what to do about that case. I think we need to zext (trunc) combo to match that.

12860–12861

Originally the code had a check to ensure the type of the starting expression matched the type of A, but I now updated things to do a ZExt on demand to make A match the type of Expr.

12868

Thanks, updated to use APInt and added a test where the LHS is > 1 << 32.

[fhahn]

make matchURem static in test class

[mkazantsev]

Looks good, thanks!

llvm/lib/Analysis/ScalarEvolution.cpp

12862

{} not needed

12863

if (auto *Trunc = ...

[fhahn]

Thanks! Adjusted the comments in the committed version

llvm/lib/Analysis/ScalarEvolution.cpp

for (auto &I : *BB) {
 
       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.
+// arbitrary expressions. Also match zext (trunc A to iB) to iY, which is used
+// for URem with constant power-of-2 second operands.
 // 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) {
+  // Try to match 'zext (trunc A to iB) to iY', which is used
+  // for URem with constant power-of-2 second operands. Make sure the size of
+  // the operand A matches the size of the whole expressions.
+  if (const auto *ZExt = dyn_cast<SCEVZeroExtendExpr>(Expr))
+    if (const auto *Trunc = dyn_cast<SCEVTruncateExpr>(ZExt->getOperand(0))) {
+      LHS = Trunc->getOperand();
+      if (LHS->getType() != Expr->getType())
+        LHS = getZeroExtendExpr(LHS, Expr->getType());
+      RHS = getConstant(APInt(getTypeSizeInBits(Expr->getType()), 1)
+                        << getTypeSizeInBits(Trunc->getType()));
+      return true;
+    }
   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)

llvm/unittests/Analysis/ScalarEvolutionTest.cpp

void runWithSE(
     Test(*F, *LI, SE);
   }
 
   static Optional<APInt> computeConstantDifference(ScalarEvolution &SE,
                                                    const SCEV *LHS,
                                                    const SCEV *RHS) {
     return SE.computeConstantDifference(LHS, RHS);
   }
+
+  static bool matchURem(ScalarEvolution &SE, const SCEV *Expr, const SCEV *&LHS,
+                        const SCEV *&RHS) {
+    return SE.matchURem(Expr, LHS, RHS);
+  }
 };
 
 TEST_F(ScalarEvolutionsTest, SCEVUnknownRAUW) {
   FunctionType *FTy = FunctionType::get(Type::getVoidTy(Context),
                                               std::vector<Type *>(), false);
   Function *F = Function::Create(FTy, Function::ExternalLinkage, "f", M);
   BasicBlock *BB = BasicBlock::Create(Context, "entry", F);
   ReturnInst::Create(Context, nullptr, BB);
TEST_F(ScalarEvolutionsTest, ProveImplicationViaNarrowing) {
   SMDiagnostic Err;
   std::unique_ptr<Module> M = parseAssemblyString(
       "define i32 @foo(i32 %start, i32* %q) { "
       "entry: "
       "  %wide.start = zext i32 %start to i64 "
       "  br label %loop "
       "loop: "
       "  %wide.iv = phi i64 [%wide.start, %entry], [%wide.iv.next, %backedge] "
       "  %iv = phi i32 [%start, %entry], [%iv.next, %backedge] "
       "  %cond = icmp eq i64 %wide.iv, 0 "
       "  br i1 %cond, label %exit, label %backedge "
       "backedge: "
       "  %iv.next = add i32 %iv, -1 "
       "  %index = zext i32 %iv.next to i64 "
       "  %load.addr = getelementptr i32, i32* %q, i64 %index "
       "  %stop = load i32, i32* %load.addr "
       "  %loop.cond = icmp eq i32 %stop, 0 "
runWithSE(*M, "foo", [](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
     // scalar-evolution-use-expensive-range-sharpening which is currently off.
     // Enable the check once it's switched true by default.
     // EXPECT_TRUE(SE.isBasicBlockEntryGuardedByCond(Backedge,
     //                                               ICmpInst::ICMP_UGT,
     //                                               IV, Zero));
   });
 }
 
+TEST_F(ScalarEvolutionsTest, MatchURem) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "target datalayout = \"e-m:e-p:32:32-f64:32:64-f80:32-n8:16:32-S128\" "
+      " "
+      "define void @test(i32 %a, i32 %b, i16 %c, i64 %d) {"
+      "entry: "
+      "  %rem1 = urem i32 %a, 2"
+      "  %rem2 = urem i32 %a, 5"
+      "  %rem3 = urem i32 %a, %b"
+      "  %c.ext = zext i16 %c to i32"
+      "  %rem4 = urem i32 %c.ext, 2"
+      "  %ext = zext i32 %rem4 to i64"
+      "  %rem5 = urem i64 %d, 17179869184"
+      "  ret void "
+      "} ",
+      Err, C);
+
+  assert(M && "Could not parse module?");
+  assert(!verifyModule(*M) && "Must have been well formed!");
+
+  runWithSE(*M, "test", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    for (auto *N : {"rem1", "rem2", "rem3", "rem5"}) {
+      auto *URemI = getInstructionByName(F, N);
+      auto *S = SE.getSCEV(URemI);
+      const SCEV *LHS, *RHS;
+      EXPECT_TRUE(matchURem(SE, S, LHS, RHS));
+      EXPECT_EQ(LHS, SE.getSCEV(URemI->getOperand(0)));
+      EXPECT_EQ(RHS, SE.getSCEV(URemI->getOperand(1)));
+      EXPECT_EQ(LHS->getType(), S->getType());
+      EXPECT_EQ(RHS->getType(), S->getType());
+    }
+
+    // Check the case where the urem operand is zero-extended. Make sure the
+    // match results are extended to the size of the input expression.
+    auto *Ext = getInstructionByName(F, "ext");
+    auto *URem1 = getInstructionByName(F, "rem4");
+    auto *S = SE.getSCEV(Ext);
+    const SCEV *LHS, *RHS;
+    EXPECT_TRUE(matchURem(SE, S, LHS, RHS));
+    EXPECT_NE(LHS, SE.getSCEV(URem1->getOperand(0)));
+    // RHS and URem1->getOperand(1) have different widths, so compare the
+    // integer values.
+    EXPECT_EQ(cast<SCEVConstant>(RHS)->getValue()->getZExtValue(),
+              cast<SCEVConstant>(SE.getSCEV(URem1->getOperand(1)))
+                  ->getValue()
+                  ->getZExtValue());
+    EXPECT_EQ(LHS->getType(), S->getType());
+    EXPECT_EQ(RHS->getType(), S->getType());
+  });
+}
+
 }  // end namespace llvm