[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

297–299

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

301–304

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

 };
 
 /// 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 name_op_matcher {
-  StringRef opName;
-  name_op_matcher(StringRef opN) : opName(opN) {}
+  StringRef name;
+  name_op_matcher(StringRef name) : name(name) {}
 
-  bool match(Operation *op) { return op->getName().getStringRef() == opName; }
+  bool match(Operation *op) { return op->getName().getStringRef() == name; }
 };
 
 /// The matcher that matches operations that have the specified attribute name.
 struct attr_op_matcher {
-  StringRef opAttrName;
-  attr_op_matcher(StringRef attrN) : opAttrName(attrN) {}
+  StringRef attr_name;
+  attr_op_matcher(StringRef attr_name) : attr_name(attr_name) {}
 
-  bool match(Operation *op) { return op->hasAttr(opAttrName); }
+  bool match(Operation *op) { return op->hasAttr(attr_name); }
 };
 
 /// 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;
 
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 have the specified attribute
+/// name, and binds the attribute value.
+template <typename AttrT>
+struct attr_op_binder {
+  StringRef attr_name;
+  AttrT *bind_value;
+
+  /// Creates a matcher instance that binds the attribute value to
+  /// bind_value if match succeeds.
+  attr_op_binder(StringRef attr_name, AttrT *bind_value)
+      : attr_name(attr_name), bind_value(bind_value) {}
+  /// Creates a matcher instance that doesn't bind if match succeeds.
+  attr_op_binder(StringRef attr_name)
+      : attr_name(attr_name), bind_value(nullptr) {}
+
+  bool match(Operation *op) {
+    if (!op->hasAttr(attr_name))
+      return false;
+
+    if (auto attr = op->getAttr(attr_name).dyn_cast<AttrT>()) {
+      if (bind_value)
+        *bind_value = attr;
+      return true;
+    }
+    return false;
+  }
+};
+
 /// 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::attr_op_matcher m_AttrName(StringRef attrN) {
-  return detail::attr_op_matcher(attrN);
+inline detail::attr_op_matcher m_Attr(StringRef attr_name) {
+  return detail::attr_op_matcher(attr_name);
 }
 
 /// Matches a named operation.
-inline detail::name_op_matcher m_Name(StringRef opN) {
-  return detail::name_op_matcher(opN);
+inline detail::name_op_matcher m_Op(StringRef op_name) {
+  return detail::name_op_matcher(op_name);
 }
 
 /// 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::attr_op_binder<AttrT> m_Attr(StringRef attr_name,
+                                            AttrT *bind_value) {
+  return detail::attr_op_binder<AttrT>(attr_name, bind_value);
+}
+
 /// 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 @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_Name("test.name"))) matched
-//       CHECK:   Pattern m_AttrName("fastmath") matched
+//       CHECK:   Pattern mul(*, add(*, m_Op("test.name"))) matched
+//       CHECK:   Pattern m_Attr("fastmath") matched and bound constant to: fast

mlir/test/lib/IR/TestMatchers.cpp

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) {
-  auto p = m_Op<arith::MulFOp>(
-      m_Any(), m_Op<arith::AddFOp>(m_Any(), m_Name("test.name")));
-  auto p1 = m_AttrName("fastmath");
+  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_Name(\"test.name\"))) matched\n";
+    llvm::outs() << "Pattern mul(*, add(*, m_Op(\"test.name\"))) matched\n";
   if (p1.match(lastOp))
-    llvm::outs() << "Pattern m_AttrName(\"fastmath\") matched\n";
+    llvm::outs()
+        << "Pattern m_Attr(\"fastmath\") matched and bound constant 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