add folder

Signed-off-by: James Newling <[email protected]>

Author: newling

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

@@ -2999,6 +2999,7 @@ def Vector_StepOp : Vector_Op<"step", [
   }];
   let results = (outs VectorOfRankAndType<[1], [Index]>:$result);
   let assemblyFormat = "attr-dict `:` type($result)";
+  let hasCanonicalizer = 1;
 }
 
 def Vector_YieldOp : Vector_Op<"yield", [

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -7602,6 +7602,120 @@ void StepOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
   setResultRanges(getResult(), result);
 }
 
+namespace {
+
+/// Constant fold vector.step when it is compared to constant with arith.cmpi
+/// and the result is the same at all indices. For example, rewrite:
+///
+/// %cst = arith.constant dense<7> : vector<3xindex>
+/// %0 = vector.step : vector<3xindex>
+/// %1 = arith.cmpi ugt, %0, %cst : vector<3xindex>
+///
+/// as
+///
+/// %out = arith.constant dense<false> : vector<3xi1>
+///
+/// Above [0, 1, 2] > [7, 7, 7] => [false, false, false]. Because the result is
+/// false at ALL indices we fold to the constant. false. If the constant was 1,
+/// then [0, 1, 2] > [1, 1, 1] => [false, false, true] and we do not constant
+/// fold, preferring the more 'compact' vector.step representation.
+struct StepCompareFolder : public OpRewritePattern<StepOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(StepOp stepOp,
+                                PatternRewriter &rewriter) const override {
+
+    int64_t stepSize = stepOp.getResult().getType().getNumElements();
+
+    for (auto &use : stepOp.getResult().getUses()) {
+      if (auto cmpiOp = dyn_cast<arith::CmpIOp>(use.getOwner())) {
+        unsigned stepOperandNumber = use.getOperandNumber();
+
+        // arith.cmpi has a canonicalizer to put constants on operand 1. Let it
+        // run first.
+        if (stepOperandNumber != 0) {
+          continue;
+        }
+
+        // Check that operand 1 is a constant.
+        unsigned otherOperandNumber = 1;
+        Value otherOperand = cmpiOp.getOperand(otherOperandNumber);
+        auto maybeConstValue = getConstantIntValue(otherOperand);
+        if (!maybeConstValue.has_value())
+          continue;
+        int64_t constValue = maybeConstValue.value();
+
+        arith::CmpIPredicate pred = cmpiOp.getPredicate();
+
+        auto maybeSplat = [&]() -> std::optional<bool> {
+          // Handle ult (unsigned less than) and uge (unsigned greater equal).
+          // Examples where stepSize = constValue = 3, for the 4
+          // cases of [ult, uge] x [stepOperandNumber = 0, 1]:
+          //
+          // pred stepOperandNumber
+          // ==== =================
+          // ult  0                 [0, 1, 2] < 3  ==> true.
+          // ult  1                 3 < [0, 1, 2]  ==> false.
+          // uge  0                 [0, 1, 2] >= 3 ==> true.
+          // uge  1                 3 >= [0, 1, 2] ==> false.
+          //
+          // If constValue is any smaller, the comparison is not constant.
+          if (pred == arith::CmpIPredicate::ult ||
+              pred == arith::CmpIPredicate::uge) {
+            if (stepSize <= constValue) {
+              return pred == arith::CmpIPredicate::ult;
+            }
+          }
+
+          // Handle ule and ugt.
+          //
+          // pred stepOperandNumber
+          // ==== =================
+          // ule  0                 [0, 1, 2] <= 2  ==> true
+          // (stepSize = 3, constValue = 2).
+          if (pred == arith::CmpIPredicate::ule ||
+              pred == arith::CmpIPredicate::ugt) {
+            if (stepSize <= constValue + 1) {
+              return pred == arith::CmpIPredicate::ule;
+            }
+          }
+
+          // Handle eq and ne
+          if (pred == arith::CmpIPredicate::eq ||
+              pred == arith::CmpIPredicate::ne) {
+            if (stepSize <= constValue) {
+              return pred == arith::CmpIPredicate::ne;
+            }
+          }
+
+          return std::optional<bool>();
+        }();
+
+        if (!maybeSplat.has_value())
+          continue;
+
+        rewriter.setInsertionPointAfter(cmpiOp);
+        auto boolConst = mlir::arith::ConstantOp::create(
+            rewriter, cmpiOp.getLoc(),
+            rewriter.getBoolAttr(maybeSplat.value()));
+        auto splat = vector::BroadcastOp::create(
+            rewriter, cmpiOp.getLoc(), cmpiOp.getResult().getType(), boolConst);
+
+        rewriter.replaceOp(cmpiOp, splat.getResult());
+        return success();
+      }
+    }
+
+    return failure();
+  }
+};
+} // namespace
+
+void StepOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                         MLIRContext *context) {
+  results.add<StepCompareFolder>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // Vector Masking Utilities
 //===----------------------------------------------------------------------===//