[mlir][matchers] Add m_Op(StringRef) and m_Attr matchers

[devajithvs]

Summary

This patch introduces support for m_Op with a StringRef argument and m_Attr matchers. These matchers will be very useful for mlir-query that is being developed currently.

Submitting this patch separately to reduce the final patch size and make it easier to upstream mlir-query.

[rriddle]

mlir/include/mlir/IR/Matchers.h

292–294

I'd likely name this m_Attr, with an optional binding capture of the attribute itself.

296–299

This should likely be named m_Op, given it does the same thing (in a non-templated way).

[devajithvs]

This comment was removed by devajithvs.

[devajithvs]

• Renamed m_AttrName -> m_Attr, with an optional binding capture of the attribute.
• Renamed m_Name -> m_Op

[rriddle]

Looks good to me, nice!

mlir/include/mlir/IR/Matchers.h

56–61

Slight preference here (I know the style in this file is all over the place).

106–107

Can you shift these to the bottom of the class? Also, please use camelCase for variable names here and throughout the patch (I know this file is all over the place, some of this code is *old*)

118–119

Can you use getAttrOfType<AttrT> here instead? It does the same thing, and that would remove the double lookup.

mlir/test/lib/IR/TestMatchers.cpp

155–157

[devajithvs]

Thanks @rriddle, done!
Am I correct in assuming that struct names remain in the same format (snake_case)?

[rriddle]

In D147262#4235296, @devajithvs wrote:

Thanks @rriddle, done!
Am I correct in assuming that struct names remain in the same format (snake_case)?

I couldn't remember if they were all snake_case, but just checked and it's also inconsistent in this file, can you use the proper naming convention for those (i.e. CamelCase)?

[devajithvs]

Updated struct names to CamelCase

[devajithvs]

Fix clang-format issues

[devajithvs]

Not sure why build on windows fails.

[mehdi_amini]

If you click on the link and search for error: you can easily find it: ******************** TEST 'Flang :: Driver/omp-frontend-forwarding.f90' FAILED ********************. Seems like unrelated to your patch.

[devajithvs]

In D147262#4236070, @mehdi_amini wrote:

If you click on the link and search for error: you can easily find it: ******************** TEST 'Flang :: Driver/omp-frontend-forwarding.f90' FAILED ********************. Seems like unrelated to your patch.

Thanks. Yeah, it seems unrelated, I tried rebuilding and it still failed. Maybe, I'll try again.

[devajithvs]

Modify an inline comment in the patch to trigger proper rebuild

[devajithvs]

Can someone help me land this patch if it's ok?

mlir/include/mlir/IR/Matchers.h

struct attr_value_binder {
   }
 };
 
 /// The matcher that matches operations that have the `ConstantLike` trait.
 struct constant_op_matcher {
   bool match(Operation *op) { return op->hasTrait<OpTrait::ConstantLike>(); }
 };
 
+/// The matcher that matches operations that have the specified op name.
+struct NameOpMatcher {
+  NameOpMatcher(StringRef name) : name(name) {}
+  bool match(Operation *op) { return op->getName().getStringRef() == name; }
+
+  StringRef name;
+};
+
+/// The matcher that matches operations that have the specified attribute name.
+struct AttrOpMatcher {
+  AttrOpMatcher(StringRef attrName) : attrName(attrName) {}
+  bool match(Operation *op) { return op->hasAttr(attrName); }
+
+  StringRef attrName;
+};
+
 /// The matcher that matches operations that have the `ConstantLike` trait, and
 /// binds the folded attribute value.
 template <typename AttrT>
 struct constant_op_binder {
   AttrT *bind_value;
 
   /// Creates a matcher instance that binds the constant attribute value to
   /// bind_value if match succeeds.
if (auto attr = foldedOp.front().get<Attribute>().dyn_cast<AttrT>()) {
       if (bind_value)
         *bind_value = attr;
       return true;
     }
     return false;
   }
 };
 
+/// The matcher that matches operations that have the specified attribute
+/// name, and binds the attribute value.
+template <typename AttrT>
+struct AttrOpBinder {
+  /// Creates a matcher instance that binds the attribute value to
+  /// bind_value if match succeeds.
+  AttrOpBinder(StringRef attrName, AttrT *bindValue)
+      : attrName(attrName), bindValue(bindValue) {}
+  /// Creates a matcher instance that doesn't bind if match succeeds.
+  AttrOpBinder(StringRef attrName) : attrName(attrName), bindValue(nullptr) {}
+
+  bool match(Operation *op) {
+    if (auto attr = op->getAttrOfType<AttrT>(attrName)) {
+      if (bindValue)
+        *bindValue = attr;
+      return true;
+    }
+    return false;
+  }
+  StringRef attrName;
+  AttrT *bindValue;
+};
+
 /// The matcher that matches a constant scalar / vector splat / tensor splat
 /// float operation and binds the constant float value.
 struct constant_float_op_binder {
   FloatAttr::ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
   constant_float_op_binder(FloatAttr::ValueType *bv) : bind_value(bv) {}
 

 
 } // namespace detail
 
 /// Matches a constant foldable operation.
 inline detail::constant_op_matcher m_Constant() {
   return detail::constant_op_matcher();
 }
 
+/// Matches a named attribute operation.
+inline detail::AttrOpMatcher m_Attr(StringRef attrName) {
+  return detail::AttrOpMatcher(attrName);
+}
+
+/// Matches a named operation.
+inline detail::NameOpMatcher m_Op(StringRef opName) {
+  return detail::NameOpMatcher(opName);
+}
+
 /// 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 named attribute operation and writes the value to bind_value.
+template <typename AttrT>
+inline detail::AttrOpBinder<AttrT> m_Attr(StringRef attrName,
+                                          AttrT *bindValue) {
+  return detail::AttrOpBinder<AttrT>(attrName, bindValue);
+}
+
 /// Matches a constant scalar / vector splat / tensor splat float (both positive
 /// and negative) zero.
 inline detail::constant_float_predicate_matcher m_AnyZeroFloat() {
   return {[](const APFloat &value) { return value.isZero(); }};
 }
 
 /// Matches a constant scalar / vector splat / tensor splat float positive zero.
 inline detail::constant_float_predicate_matcher m_PosZeroFloat() {

mlir/test/IR/test-matchers.mlir

func.func @test2(%a: f32) -> f32 {
   %1 = arith.addf %a, %0: f32
   %2 = arith.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.func @test3(%a: f32) -> f32 {
+  %0 = "test.name"() {value = 1.0 : f32} : () -> f32
+  %1 = arith.addf %a, %0: f32
+  %2 = arith.mulf %a, %1 fastmath<fast>: f32
+  return %2: f32
+}
+
+// CHECK-LABEL: test3
+//       CHECK:   Pattern mul(*, add(*, m_Op("test.name"))) matched
+//       CHECK:   Pattern m_Attr("fastmath") matched and bound value to: fast

mlir/test/lib/IR/TestMatchers.cpp

void test2(FunctionOpInterface 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(FunctionOpInterface f) {
+  arith::FastMathFlagsAttr fastMathAttr;
+  auto p = m_Op<arith::MulFOp>(m_Any(),
+                               m_Op<arith::AddFOp>(m_Any(), m_Op("test.name")));
+  auto p1 = m_Attr("fastmath", &fastMathAttr);
+
+  // Last operation that is not the terminator.
+  Operation *lastOp = f.getFunctionBody().front().back().getPrevNode();
+  if (p.match(lastOp))
+    llvm::outs() << "Pattern mul(*, add(*, m_Op(\"test.name\"))) matched\n";
+  if (p1.match(lastOp))
+    llvm::outs() << "Pattern m_Attr(\"fastmath\") matched and bound value to: "
+                 << fastMathAttr.getValue() << "\n";
+}
+
 void TestMatchers::runOnOperation() {
   auto f = getOperation();
   llvm::outs() << f.getName() << "\n";
   if (f.getName() == "test1")
     test1(f);
   if (f.getName() == "test2")
     test2(f);
+  if (f.getName() == "test3")
+    test3(f);
 }
 
 namespace mlir {
 void registerTestMatchers() { PassRegistration<TestMatchers>(); }
 } // namespace mlir