[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

269–271

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

273–276

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 name_op_matcher {
+  StringRef opName;
+  name_op_matcher(StringRef opN) : opName(opN) {}
+
+  bool match(Operation *op) { return op->getName().getStringRef() == opName; }
+};
+
+/// The matcher that matches operations that have the specified attribute name.
+struct attr_op_matcher {
+  StringRef opAttrName;
+  attr_op_matcher(StringRef attrN) : opAttrName(attrN) {}
+
+  bool match(Operation *op) { return op->hasAttr(opAttrName); }
+};
+
 /// 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.
bool match(Operation *op) {
     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) {}
 
   bool match(Operation *op) {
     Attribute attr;
     if (!constant_op_binder<Attribute>(&attr).match(op))
       return false;
     auto type = op->getResult(0).getType();
 
     if (type.isa<FloatType>())
       return attr_value_binder<FloatAttr>(bind_value).match(attr);
     if (type.isa<VectorType, RankedTensorType>()) {
       if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
         return attr_value_binder<FloatAttr>(bind_value)
             .match(splatAttr.getSplatValue<Attribute>());
       }
     }
     return false;
   }
 };
 

 
 } // 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);
+}
+
+/// Matches a named operation.
+inline detail::name_op_matcher m_Name(StringRef opN) {
+  return detail::name_op_matcher(opN);
+}
+
 /// 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 scalar / vector splat / tensor splat float (both positive

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_Name("test.name"))) matched
+//       CHECK:   Pattern m_AttrName("fastmath") matched

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) {
+  auto p = m_Op<arith::MulFOp>(
+      m_Any(), m_Op<arith::AddFOp>(m_Any(), m_Name("test.name")));
+  auto p1 = m_AttrName("fastmath");
+  // 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";
+  if (p1.match(lastOp))
+    llvm::outs() << "Pattern m_AttrName(\"fastmath\") matched\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