[mlir] Generalize m_ConstantIndexOp

[chelini]

Summary

Generalize m_ContantIndexOp by using the m_Constant matcher and checking the operation for an index type.

[merge_guards_bot]

Unit tests: pass. 61914 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

mlir/include/mlir/IR/Matchers.h

88

Please add explicit to avoid spurious implicit conversions

[chelini]

Add explicit keyword for the single-parameter constructor.

[rriddle]

mlir/include/mlir/IR/Matchers.h

83

This is a lot of duplication for what should be one or two lines of code.

[rriddle]

mlir/include/mlir/IR/Matchers.h

83

I would imagine that the constant_op_binder could optionally constrain the type(using isa<>) as well (via templates).

mlir/test/lib/IR/TestMatchers.cpp

141

Drop the assert

143

nit: merge these on one line.

mlir/include/mlir/IR/Matchers.h

if (succeeded(op->fold(/*operands=*/llvm::None, foldedOp))) {
           return true;
         }
       }
     }
     return false;
   }
 };
 
+/// The matcher that matches a constant index operation.
+struct constant_index_op_binder {
+  IntegerAttr *bind_value;
+
+  /// Creates a matcher instance that binds the index
+  /// attribute value if match succeeds.
+  explicit constant_index_op_binder(IntegerAttr *bv) : bind_value(bv) {}
+  /// Create a matcher instance that doesn't bind if match succeeds.
+  constant_index_op_binder() : bind_value(nullptr) {}
+
+  bool match(Operation *op) {
+    IntegerAttr attr;
+    if (!constant_op_binder<IntegerAttr>(&attr).match(op))
+      return false;
+    if (!op->getResult(0).getType().isIndex())
+      return false;
+    if (bind_value)
+      *bind_value = attr;
+    return true;
+  }
+};
+
 /// The matcher that matches a constant scalar / vector splat / tensor splat
 /// integer operation and binds the constant integer value.
 struct constant_int_op_binder {
   IntegerAttr::ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
   constant_int_op_binder(IntegerAttr::ValueType *bv) : bind_value(bv) {}
 

 
 /// Matches a value from a constant foldable operation and writes the value to
 /// bind_value.
 template <typename AttrT>
 inline detail::constant_op_binder<AttrT> m_Constant(AttrT *bind_value) {
   return detail::constant_op_binder<AttrT>(bind_value);
 }
 
+/// Matches a constant index operation.
+inline detail::constant_index_op_binder m_ConstantIndex() {
+  return detail::constant_index_op_binder();
+}
+
+/// Matches a constant index operation and
+/// writes the integer value to bind_value.
+inline detail::constant_index_op_binder
+m_ConstantIndex(IntegerAttr *bind_value) {
+  return detail::constant_index_op_binder(bind_value);
+}
+
 /// Matches a constant scalar / vector splat / tensor splat integer one.
 inline detail::constant_int_value_matcher<1> m_One() {
   return detail::constant_int_value_matcher<1>();
 }
 
 /// Matches the given OpClass.
 template <typename OpClass> inline detail::op_matcher<OpClass> m_Op() {
   return detail::op_matcher<OpClass>();

mlir/lib/Dialect/StandardOps/Ops.cpp

ParseResult mlir::parseDimAndSymbolList(OpAsmParser &parser,
   auto indexTy = parser.getBuilder().getIndexType();
   if (parser.parseOperandList(opInfos,
                               OpAsmParser::Delimiter::OptionalSquare) ||
       parser.resolveOperands(opInfos, indexTy, operands))
     return failure();
   return success();
 }
 
-/// Matches a ConstantIndexOp.
-/// TODO: This should probably just be a general matcher that uses m_Constant
-/// and checks the operation for an index type.
-static detail::op_matcher<ConstantIndexOp> m_ConstantIndex() {
-  return detail::op_matcher<ConstantIndexOp>();
-}
-
 //===----------------------------------------------------------------------===//
 // Common canonicalization pattern support logic
 //===----------------------------------------------------------------------===//
 
 /// This is a common class used for patterns of the form
 /// "someop(memrefcast) -> someop".  It folds the source of any memref_cast
 /// into the root operation directly.
 static LogicalResult foldMemRefCast(Operation *op) {

mlir/test/IR/test-matchers.mlir

func @test2(%a: f32) -> f32 {
   %1 = addf %a, %0: f32
   %2 = mulf %a, %1: f32
   return %2: f32
 }
 
 // CHECK-LABEL: test2
 //       CHECK:   Pattern add(add(a, constant), a) matched and bound constant to: 1.000000e+00
 //       CHECK:   Pattern add(add(a, constant), a) matched
+
+func @test3(%a: index) -> index {
+  %0 = constant 23: index
+  %1 = muli %0, %a: index
+  return %1: index
+}
+
+// CHECK-LABEL: test3
+//       CHECK:   Pattern mul(a, constantIndex) matched and bound constant to: 23
+//       CHECK:   Pattern mul(a, constant) matched and bound constant to: 23

mlir/test/lib/IR/TestMatchers.cpp

void test2(FuncOp f) {
   if (p.match(lastOp))
     llvm::outs()
         << "Pattern add(add(a, constant), a) matched and bound constant to: "
         << floatAttr.getValueAsDouble() << "\n";
   if (p1.match(lastOp))
     llvm::outs() << "Pattern add(add(a, constant), a) matched\n";
 }
 
+void test3(FuncOp f) {
+  assert(f.getNumArguments() == 1 && "matcher test funcs must have 1 args");
+  auto a = m_Val(f.getArgument(0));
+  IntegerAttr integerAttrConstantIndex;
+  IntegerAttr integerAttrConstant;
+  auto mConstantIndex =
+      m_Op<MulIOp>(m_ConstantIndex(&integerAttrConstantIndex), a);
+  Operation *lastOp = f.getBody().front().back().getPrevNode();
+  if (mConstantIndex.match(lastOp))
+    llvm::outs()
+        << "Pattern mul(a, constantIndex) matched and bound constant to: "
+        << integerAttrConstantIndex.getInt() << "\n";
+  auto mConstant = m_Op<MulIOp>(m_Constant(&integerAttrConstant), a);
+  if (mConstant.match(lastOp))
+    llvm::outs() << "Pattern mul(a, constant) matched and bound constant to: "
+                 << integerAttrConstant.getInt() << "\n";
+}
+
 void TestMatchers::runOnFunction() {
   auto f = getFunction();
   llvm::outs() << f.getName() << "\n";
   if (f.getName() == "test1")
     test1(f);
   if (f.getName() == "test2")
     test2(f);
+  if (f.getName() == "test3")
+    test3(f);
 }
 
 static PassRegistration<TestMatchers> pass("test-matchers",
                                            "Test C++ pattern matchers.");