Enable using constant attribute as matchers.

Update to allow constant attribute values to be used to match or as result in rewrite rule. Define variable ctx in the matcher to allow matchers to refer to the context of the operation being matched.

PiperOrigin-RevId: 231322019

2 files changed, 15 insertions, 2 deletions

diff --git a/mlir/include/mlir/IR/op_base.td b/mlir/include/mlir/IR/op_base.td
index 63a055f1469..7d85be29384 100644
--- a/mlir/include/mlir/IR/op_base.td
+++ b/mlir/include/mlir/IR/op_base.td
@@ -157,7 +157,7 @@ class ContainerType<Type etype, Pred containerPred, code elementTypeCall,
     Type<AllOf<[containerPred,
                 SubstLeaves<"{0}", !cast<string>(elementTypeCall),
 		            etype.predicate>]>,
-	 descr # "<" # etype.description # ">" > {
+	  descr # "<" # etype.description # ">" > {
   // The type of elements in the container.
   Type elementType = etype;
 
@@ -317,7 +317,13 @@ class DerivedTypeAttr<code body> : DerivedAttr<"Type", body>;
 // Represents a constant attribute of specific Attr type. A constant
 // attribute can be specified only of attributes that have a constant
 // builder call defined. The constant value is specified as a string.
-class ConstantAttr<Attr attribute, string val> {
+//
+// If used as a constraint, it generates a matcher on a constant attribute by
+// using the constant value builder of the attribute and the value.
+class ConstantAttr<Attr attribute, string val> : AttrConstraint<
+    CPred<"{0} == " #
+      !subst("{0}", "mlir::Builder(ctx)", !subst("{1}", val,
+        !cast<string>(attribute.constBuilderCall)))>> {
   Attr attr = attribute;
   string value = val;
 }
@@ -445,6 +451,12 @@ class Pat<dag pattern, dag result> : Pattern<pattern, [result]>;
 // that has to match. Used on the input attributes of a rewrite rule.
 class mAttr<Pred pred> : AttrConstraint<pred>;
 
+// Combine a list of attribute matchers into an attribute matcher that holds if
+// any of the original matchers does.
+class mAttrAnyOf<list<AttrConstraint> attrs> :
+  mAttr<AnyOf<!foldl([]<Pred>, attrs, prev, attr,
+                       !listconcat(prev, [attr.predicate]))>>;
+
 // Attribute transforms. This is the base class to specify a
 // transformation of a matched attribute. Used on the output of a rewrite
 // rule.
diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index d954db86178..0315bf03d83 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -204,6 +204,7 @@ void PatternEmitter::emitMatchMethod(DagNode tree) {
   PatternMatchResult match(OperationInst *op0) const override {
     // TODO: This just handle 1 result
     if (op0->getNumResults() != 1) return matchFailure();
+    auto ctx = op0->getContext(); (void)ctx;
     auto state = std::make_unique<MatchedState>();)"
      << "\n";
   emitOpMatch(tree, 0);