[mlir] Allow using non-attribute properties in declarative rewrite patterns (#143071)

This commit adds support for non-attribute properties (such as
StringProp and I64Prop) in declarative rewrite patterns. The handling
for properties follows the handling for attributes in most cases,
including in the generation of static matchers.

Constraints that are shared between multiple types are supported by
making the constraint matcher a templated function, which is the
equivalent to passing ::mlir::Attribute for an arbitrary C++ type.

Author: krzysz00

mlir/docs/DeclarativeRewrites.md

@@ -380,6 +380,11 @@ template. The string can be an arbitrary C++ expression that evaluates into some
 C++ object expected at the `NativeCodeCall` site (here it would be expecting an
 array attribute). Typically the string should be a function call.
 
+In the case of properties, the return value of the `NativeCodeCall` should
+be in terms of the _interface_ type of a property. For example, the `NativeCodeCall`
+for a `StringProp` should return a `StringRef`, which will copied into the underlying
+`std::string`, just as if it were an argument to the operation's builder.
+
 ##### `NativeCodeCall` placeholders
 
 In `NativeCodeCall`, we can use placeholders like `$_builder`, `$N` and `$N...`.
@@ -416,14 +421,20 @@ must be either passed by reference or pointer to the variable used as argument
 so that the matched value can be returned. In the same example, `$val` will be
 bound to a variable with `Attribute` type (as `I32Attr`) and the type of the
 second argument in `Foo()` could be `Attribute&` or `Attribute*`. Names with
-attribute constraints will be captured as `Attribute`s while everything else
-will be treated as `Value`s.
+attribute constraints will be captured as `Attribute`s, names with
+property constraints (which must have a concrete interface type) will be treated
+as that type, and everything else will be treated as `Value`s.
 
 Positional placeholders will be substituted by the `dag` object parameters at
 the `NativeCodeCall` use site. For example, if we define `SomeCall :
 NativeCodeCall<"someFn($1, $2, $0)">` and use it like `(SomeCall $in0, $in1,
 $in2)`, then this will be translated into C++ call `someFn($in1, $in2, $in0)`.
 
+In the case of properties, the placeholder will be bound to a value of the _interface_
+type of the property. For example, passing in a `StringProp` as an argument to a `NativeCodeCall` will pass a `StringRef` (as if the getter of the matched
+operation were called) and not a `std::string`. See
+`mlir/include/mlir/IR/Properties.td` for details on interface vs. storage type.
+
 Positional range placeholders will be substituted by multiple `dag` object
 parameters at the `NativeCodeCall` use site. For example, if we define
 `SomeCall : NativeCodeCall<"someFn($1...)">` and use it like `(SomeCall $in0,

mlir/include/mlir/IR/Properties.td

@@ -401,6 +401,21 @@ class ConfinedProperty<Property p, Pred pred, string newSummary = "">
   : ConfinedProp<p, pred, newSummary>,
     Deprecated<"moved to shorter name ConfinedProp">;
 
+/// Defines a constant value of type `prop` to be used in pattern matching.
+/// When used as a constraint, forms a matcher that tests that the property is
+/// equal to the given value (and matches any other constraints on the property).
+/// The constant value is given as a string and should be of the _interface_ type
+/// of the attribute.
+///
+/// This requires that the given property's inference type be comparable to the
+/// given value with `==`, and does require specify a concrete property type.
+class ConstantProp<Property prop, string val>
+    : ConfinedProp<prop,
+        CPred<"$_self == " # val>,
+        "constant '" # prop.summary # "': " # val> {
+  string value = val;
+}
+
 //===----------------------------------------------------------------------===//
 // Primitive property combinators
 //===----------------------------------------------------------------------===//

mlir/include/mlir/TableGen/Pattern.h

@@ -73,12 +73,23 @@ class DagLeaf {
   // specifies an attribute constraint.
   bool isAttrMatcher() const;
 
+  // Returns true if this DAG leaf is matching a property. That is, it
+  // specifies a property constraint.
+  bool isPropMatcher() const;
+
+  // Returns true if this DAG leaf is describing a property. That is, it
+  // is a subclass of `Property` in tablegen.
+  bool isPropDefinition() const;
+
   // Returns true if this DAG leaf is wrapping native code call.
   bool isNativeCodeCall() const;
 
   // Returns true if this DAG leaf is specifying a constant attribute.
   bool isConstantAttr() const;
 
+  // Returns true if this DAG leaf is specifying a constant property.
+  bool isConstantProp() const;
+
   // Returns true if this DAG leaf is specifying an enum case.
   bool isEnumCase() const;
 
@@ -88,9 +99,19 @@ class DagLeaf {
   // Returns this DAG leaf as a constraint. Asserts if fails.
   Constraint getAsConstraint() const;
 
+  // Returns this DAG leaf as a property constraint. Asserts if fails. This
+  // allows access to the interface type.
+  PropConstraint getAsPropConstraint() const;
+
+  // Returns this DAG leaf as a property definition. Asserts if fails.
+  Property getAsProperty() const;
+
   // Returns this DAG leaf as an constant attribute. Asserts if fails.
   ConstantAttr getAsConstantAttr() const;
 
+  // Returns this DAG leaf as an constant property. Asserts if fails.
+  ConstantProp getAsConstantProp() const;
+
   // Returns this DAG leaf as an enum case.
   // Precondition: isEnumCase()
   EnumCase getAsEnumCase() const;
@@ -279,6 +300,10 @@ class SymbolInfoMap {
     //   the DAG of the operation, `operandIndexOrNumValues` specifies the
     //   operand index, and `variadicSubIndex` must be set to `std::nullopt`.
     //
+    // * Properties not associated with an operation (e.g. as arguments to
+    //   native code) have their corresponding PropConstraint stored in the
+    //   `dag` field. This constraint is only used when
+    //
     // * If a symbol is defined in a `variadic` DAG, `dag` specifies the DAG
     //   of the parent operation, `operandIndexOrNumValues` specifies the
     //   declared operand index of the variadic operand in the parent
@@ -364,12 +389,20 @@ class SymbolInfoMap {
 
     // What kind of entity this symbol represents:
     // * Attr: op attribute
+    // * Prop: op property
     // * Operand: op operand
     // * Result: op result
     // * Value: a value not attached to an op (e.g., from NativeCodeCall)
     // * MultipleValues: a pack of values not attached to an op (e.g., from
     //   NativeCodeCall). This kind supports indexing.
-    enum class Kind : uint8_t { Attr, Operand, Result, Value, MultipleValues };
+    enum class Kind : uint8_t {
+      Attr,
+      Prop,
+      Operand,
+      Result,
+      Value,
+      MultipleValues
+    };
 
     // Creates a SymbolInfo instance. `dagAndConstant` is only used for `Attr`
     // and `Operand` so should be std::nullopt for `Result` and `Value` kind.
@@ -384,6 +417,15 @@ class SymbolInfoMap {
     static SymbolInfo getAttr() {
       return SymbolInfo(nullptr, Kind::Attr, std::nullopt);
     }
+    static SymbolInfo getProp(const Operator *op, int index) {
+      return SymbolInfo(op, Kind::Prop,
+                        DagAndConstant(nullptr, index, std::nullopt));
+    }
+    static SymbolInfo getProp(const PropConstraint *constraint) {
+      // -1 for anthe `operandIndexOrNumValues` is a sentinel value.
+      return SymbolInfo(nullptr, Kind::Prop,
+                        DagAndConstant(constraint, -1, std::nullopt));
+    }
     static SymbolInfo
     getOperand(DagNode node, const Operator *op, int operandIndex,
                std::optional<int> variadicSubIndex = std::nullopt) {
@@ -488,6 +530,10 @@ class SymbolInfoMap {
   // is already bound.
   bool bindAttr(StringRef symbol);
 
+  // Registers the given `symbol` as bound to a property that satisfies the
+  // given `constraint`. `constraint` must name a concrete interface type.
+  bool bindProp(StringRef symbol, const PropConstraint &constraint);
+
   // Returns true if the given `symbol` is bound.
   bool contains(StringRef symbol) const;
 

mlir/include/mlir/TableGen/Property.h

@@ -32,9 +32,9 @@ class Pred;
 // Wrapper class providing helper methods for accesing property constraint
 // values.
 class PropConstraint : public Constraint {
+public:
   using Constraint::Constraint;
 
-public:
   static bool classof(const Constraint *c) { return c->getKind() == CK_Prop; }
 
   StringRef getInterfaceType() const;
@@ -143,6 +143,10 @@ class Property : public PropConstraint {
   // property constraints, this function is added for future-proofing)
   Property getBaseProperty() const;
 
+  // Returns true if this property is backed by a TableGen definition and that
+  // definition is a subclass of `className`.
+  bool isSubClassOf(StringRef className) const;
+
 private:
   // Elements describing a Property, in general fetched from the record.
   StringRef summary;
@@ -169,6 +173,21 @@ struct NamedProperty {
   Property prop;
 };
 
+// Wrapper class providing helper methods for processing constant property
+// values defined using the `ConstantProp` subclass of `Property`
+// in TableGen.
+class ConstantProp : public Property {
+public:
+  explicit ConstantProp(const llvm::DefInit *def) : Property(def) {
+    assert(isSubClassOf("ConstantProp"));
+  }
+
+  static bool classof(Property *p) { return p->isSubClassOf("ConstantProp"); }
+
+  // Return the constant value of the property as an expression
+  // that produces an interface-type constant.
+  StringRef getValue() const;
+};
 } // namespace tblgen
 } // namespace mlir
 

mlir/lib/TableGen/CodeGenHelpers.cpp

@@ -205,10 +205,14 @@ static ::llvm::LogicalResult {0}(
 
 /// Code for a pattern type or attribute constraint.
 ///
-/// {3}: "Type type" or "Attribute attr".
-static const char *const patternAttrOrTypeConstraintCode = R"(
+/// {0}: name of function
+/// {1}: Condition template
+/// {2}: Constraint summary
+/// {3}: "::mlir::Type type" or "::mlirAttribute attr" or "propType prop".
+/// Can be "T prop" for generic property constraints.
+static const char *const patternConstraintCode = R"(
 static ::llvm::LogicalResult {0}(
-    ::mlir::PatternRewriter &rewriter, ::mlir::Operation *op, ::mlir::{3},
+    ::mlir::PatternRewriter &rewriter, ::mlir::Operation *op, {3},
     ::llvm::StringRef failureStr) {
   if (!({1})) {
     return rewriter.notifyMatchFailure(op, [&](::mlir::Diagnostic &diag) {
@@ -265,15 +269,31 @@ void StaticVerifierFunctionEmitter::emitPatternConstraints() {
   FmtContext ctx;
   ctx.addSubst("_op", "*op").withBuilder("rewriter").withSelf("type");
   for (auto &it : typeConstraints) {
-    os << formatv(patternAttrOrTypeConstraintCode, it.second,
+    os << formatv(patternConstraintCode, it.second,
                   tgfmt(it.first.getConditionTemplate(), &ctx),
-                  escapeString(it.first.getSummary()), "Type type");
+                  escapeString(it.first.getSummary()), "::mlir::Type type");
   }
   ctx.withSelf("attr");
   for (auto &it : attrConstraints) {
-    os << formatv(patternAttrOrTypeConstraintCode, it.second,
+    os << formatv(patternConstraintCode, it.second,
                   tgfmt(it.first.getConditionTemplate(), &ctx),
-                  escapeString(it.first.getSummary()), "Attribute attr");
+                  escapeString(it.first.getSummary()),
+                  "::mlir::Attribute attr");
+  }
+  ctx.withSelf("prop");
+  for (auto &it : propConstraints) {
+    PropConstraint propConstraint = cast<PropConstraint>(it.first);
+    StringRef interfaceType = propConstraint.getInterfaceType();
+    // Constraints that are generic over multiple interface types are
+    // templatized under the assumption that they'll be used correctly.
+    if (interfaceType.empty()) {
+      interfaceType = "T";
+      os << "template <typename T>";
+    }
+    os << formatv(patternConstraintCode, it.second,
+                  tgfmt(propConstraint.getConditionTemplate(), &ctx),
+                  escapeString(propConstraint.getSummary()),
+                  Twine(interfaceType) + " prop");
   }
 }
 
@@ -367,10 +387,15 @@ void StaticVerifierFunctionEmitter::collectOpConstraints(
 void StaticVerifierFunctionEmitter::collectPatternConstraints(
     const ArrayRef<DagLeaf> constraints) {
   for (auto &leaf : constraints) {
-    assert(leaf.isOperandMatcher() || leaf.isAttrMatcher());
-    collectConstraint(
-        leaf.isOperandMatcher() ? typeConstraints : attrConstraints,
-        leaf.isOperandMatcher() ? "type" : "attr", leaf.getAsConstraint());
+    assert(leaf.isOperandMatcher() || leaf.isAttrMatcher() ||
+           leaf.isPropMatcher());
+    Constraint constraint = leaf.getAsConstraint();
+    if (leaf.isOperandMatcher())
+      collectConstraint(typeConstraints, "type", constraint);
+    else if (leaf.isAttrMatcher())
+      collectConstraint(attrConstraints, "attr", constraint);
+    else if (leaf.isPropMatcher())
+      collectConstraint(propConstraints, "prop", constraint);
   }
 }