[mlir][scf] Match any constants instead of arith.constant

[Mogball]

Summary

By matching arith.constant specifically, SCF canonicalizers/folders
are incompatible with other kinds of constants. Use the generic
matchers instead.

[rriddle]

nit on the description, these matchers are not EDSCs.

mlir/lib/Dialect/SCF/IR/SCF.cpp

1887

Can you sink this to the use?

[Mogball]

comments

mlir/lib/Dialect/SCF/IR/SCF.cpp

if (iterArgs.empty() && !bodyBuilder) {
     OpBuilder::InsertionGuard guard(builder);
     builder.setInsertionPointToStart(&bodyBlock);
     bodyBuilder(builder, result.location, bodyBlock.getArgument(0),
                 bodyBlock.getArguments().drop_front());
   }
 }
 
 LogicalResult ForOp::verify() {
-  if (auto cst = getStep().getDefiningOp<arith::ConstantIndexOp>())
-    if (cst.value() <= 0)
-      return emitOpError("constant step operand must be positive");
+  IntegerAttr step;
+  if (matchPattern(getStep(), m_Constant(&step)) && step.getInt() <= 0)
 
   auto opNumResults = getNumResults();
   if (opNumResults == 0)
     return success();
   // If ForOp defines values, check that the number and types of
   // the defined values match ForOp initial iter operands and backedge
   // basic block arguments.
   if (getNumIterOperands() != opNumResults)
LogicalResult matchAndRewrite(scf::ForOp forOp,
     return success();
   }
 };
 
 /// Util function that tries to compute a constant diff between u and l.
 /// Returns llvm::None when the difference between two AffineValueMap is
 /// dynamic.
 static Optional<int64_t> computeConstDiff(Value l, Value u) {
-  auto clb = l.getDefiningOp<arith::ConstantOp>();
-  auto cub = u.getDefiningOp<arith::ConstantOp>();
-  if (cub && clb) {
-    llvm::APInt lbValue = clb.getValue().cast<IntegerAttr>().getValue();
-    llvm::APInt ubValue = cub.getValue().cast<IntegerAttr>().getValue();
+  IntegerAttr clb, cub;
+  if (matchPattern(l, m_Constant(&clb)) && matchPattern(u, m_Constant(&cub))) {
+    llvm::APInt lbValue = clb.getValue();
+    llvm::APInt ubValue = cub.getValue();
     return (ubValue - lbValue).getSExtValue();
   }
 
   // Else a simple pattern match for x + c or c + x
   llvm::APInt diff;
   if (matchPattern(
           u, m_Op<arith::AddIOp>(matchers::m_Val(l), m_ConstantInt(&diff))) ||
       matchPattern(
if (!diff)
       return failure();
 
     // If the loop is known to have 0 iterations, remove it.
     if (*diff <= 0) {
       rewriter.replaceOp(op, op.getIterOperands());
       return success();
     }
 
-    auto step = op.getStep().getDefiningOp<arith::ConstantOp>();
-    if (!step)
+    IntegerAttr step;
+    if (!matchPattern(op.getStep(), m_Constant(&step)))
       return failure();
 
     // If the loop is known to have 1 iteration, inline its body and remove the
     // loop.
-    llvm::APInt stepValue = step.getValue().cast<IntegerAttr>().getValue();
+    llvm::APInt stepValue = step.getValue();
     if (stepValue.sge(*diff)) {
       SmallVector<Value, 4> blockArgs;
       blockArgs.reserve(op.getNumIterOperands() + 1);
       blockArgs.push_back(op.getLowerBound());
       llvm::append_range(blockArgs, op.getIterOperands());
       replaceOpWithRegion(rewriter, op, op.getLoopBody(), blockArgs);
       return success();
     }
struct RemoveUnusedResults : public OpRewritePattern<IfOp> {
   }
 };
 
 struct RemoveStaticCondition : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(IfOp op,
                                 PatternRewriter &rewriter) const override {
-    auto constant = op.getCondition().getDefiningOp<arith::ConstantOp>();
-    if (!constant)
+    BoolAttr condition;
+    if (!matchPattern(op.getCondition(), m_Constant(&condition)))
       return failure();
 
-    if (constant.getValue().cast<BoolAttr>().getValue())
+    if (condition.getValue())
       replaceOpWithRegion(rewriter, op, op.getThenRegion());
     else if (!op.getElseRegion().empty())
       replaceOpWithRegion(rewriter, op, op.getElseRegion());
     else
       rewriter.eraseOp(op);
 
     return success();
   }

 ///
 struct ConditionPropagation : public OpRewritePattern<IfOp> {
   using OpRewritePattern<IfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(IfOp op,
                                 PatternRewriter &rewriter) const override {
     // Early exit if the condition is constant since replacing a constant
     // in the body with another constant isn't a simplification.
-    if (op.getCondition().getDefiningOp<arith::ConstantOp>())
+    if (matchPattern(op.getCondition(), m_Constant()))
       return failure();
 
     bool changed = false;
     mlir::Type i1Ty = rewriter.getI1Type();
 
     // These variables serve to prevent creating duplicate constants
     // and hold constant true or false values.
     Value constantTrue = nullptr;
for (auto [trueResult, falseResult, opResult] :
       if (trueResult == falseResult) {
         if (!opResult.use_empty()) {
           opResult.replaceAllUsesWith(trueResult);
           changed = true;
         }
         continue;
       }
 
-      auto trueYield = trueResult.getDefiningOp<arith::ConstantOp>();
-      if (!trueYield)
-        continue;
-
-      if (!trueYield.getType().isInteger(1))
-        continue;
-
-      auto falseYield = falseResult.getDefiningOp<arith::ConstantOp>();
-      if (!falseYield)
+      BoolAttr trueYield, falseYield;
+      if (!matchPattern(trueResult, m_Constant(&trueYield)) ||
+          !matchPattern(falseResult, m_Constant(&falseYield)))
         continue;
 
-      bool trueVal = trueYield.getValue().cast<BoolAttr>().getValue();
-      bool falseVal = falseYield.getValue().cast<BoolAttr>().getValue();
+      bool trueVal = trueYield.getValue();
+      bool falseVal = falseYield.getValue();
       if (!trueVal && falseVal) {
         if (!opResult.use_empty()) {
+          Dialect *constDialect = trueResult.getDefiningOp()->getDialect();
           Value notCond = rewriter.create<arith::XOrIOp>(
               op.getLoc(), op.getCondition(),
-              rewriter.create<arith::ConstantOp>(
-                  op.getLoc(), i1Ty, rewriter.getIntegerAttr(i1Ty, 1)));
+              constDialect
+                  ->materializeConstant(rewriter,
+                                        rewriter.getIntegerAttr(i1Ty, 1), i1Ty,
+                                        op.getLoc())
+                  ->getResult(0));
           opResult.replaceAllUsesWith(notCond);
           changed = true;
         }
       }
       if (trueVal && !falseVal) {
         if (!opResult.use_empty()) {
           opResult.replaceAllUsesWith(op.getCondition());
           changed = true;