[mlir][ODS] Add support for passing properties to `custom`

[zero9178]

Summary

Printing and parsing properties of ops is currently only possible through the prop-dict attribute. This forces a specific place that the property is printed at and is generally not very flexible.

This patch adds support of passing properties to the custom directive, making it possible to incorporate them with more complex syntax. This makes it possible to parse and print them without generic syntax and without the use of prop-dict.

[zero9178]

clang-format

[Mogball]

Does ref($prop) work in a custom directive?

mlir/test/lib/Dialect/Test/TestOps.td

3349–3351

mlir/tools/mlir-tblgen/OpFormatGen.cpp

98

This is missing the typical comment

[zero9178]

In D155072#4493942, @Mogball wrote:

Does ref($prop) work in a custom directive?

Not right now no. It currently errors out as they are currently only valid in the context of a custom directive. Wasn't planning to implement it this revision but can do so if you like.

[Mogball]

Doesn't have to be in this one. Thanks for the patch!

[jpienaar]

Thanks!

mlir/tools/mlir-tblgen/OpFormatGen.cpp

3098

Could you sort these while here?

[mehdi_amini]

Thanks! (if ref can be added it'd be even better :))

mlir/test/IR/properties.mlir

 // GENERIC-SAME:  <{prop = {label = "foo bar", value = -3 : i32}}> : () -> ()
 test.with_nice_properties "foo bar" is -3
 
 // CHECK:   test.with_wrapped_properties
 // CHECK-SAME:    "content for properties"
 // GENERIC: "test.with_wrapped_properties"()
 // GENERIC-SAME:  <{prop = "content for properties"}> : () -> ()
 test.with_wrapped_properties <{prop = "content for properties"}>
+
+// CHECK: test.using_property_in_custom
+// CHECK-SAME: [1, 4, 20]
+// GENERIC: "test.using_property_in_custom"()
+// GENERIC-SAME: prop = array<i64: 1, 4, 20>
+test.using_property_in_custom [1, 4, 20]

mlir/test/lib/Dialect/Test/TestDialect.cpp

static ParseResult customParseProperties(OpAsmParser &parser,
   std::string label;
   if (parser.parseKeywordOrString(&label) || parser.parseKeyword("is") ||
       parser.parseInteger(prop.value))
     return failure();
   prop.label = std::make_shared<std::string>(std::move(label));
   return success();
 }
 
+static bool parseUsingPropertyInCustom(OpAsmParser &parser, int64_t value[3]) {
+  return parser.parseLSquare()
+         || parser.parseInteger(value[0])
+         || parser.parseComma()
+         || parser.parseInteger(value[1])
+         || parser.parseComma()
+         || parser.parseInteger(value[2])
+         || parser.parseRSquare();
+}
+
+static void printUsingPropertyInCustom(OpAsmPrinter &printer, Operation *op,
+                                       ArrayRef<int64_t> value) {
+  printer << '[' << value << ']';
+}
+
 #include "TestOpEnums.cpp.inc"
 #include "TestOpInterfaces.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/Dialect/Test/TestOps.td

 
 def TestOpWithWrappedProperties : TEST_Op<"with_wrapped_properties"> {
   let assemblyFormat = "prop-dict attr-dict";
   let arguments = (ins
     MyStructProperty:$prop
   );
 }
 
+def TestOpUsingPropertyInCustom : TEST_Op<"using_property_in_custom"> {
+  let assemblyFormat = "custom<UsingPropertyInCustom>($prop) attr-dict";
+  let arguments = (ins
+    ArrayProperty<"int64_t", 3>:$prop
+  );
+}
+
 // Op with a properties struct defined out-of-line. The struct has custom
 // printer/parser.
 
 def PropertiesWithCustomPrint : Property<"PropertiesWithCustomPrint"> {
   let convertToAttribute = [{
     getPropertiesAsAttribute($_ctxt, $_storage)
   }];
   let convertFromAttribute = [{

mlir/test/mlir-tblgen/op-format-invalid.td

 }
 // CHECK: error: region #0, named 'region', not found
 def VariableInvalidN : TestFormat_Op<[{
   attr-dict
 }]> {
   let regions = (region AnyRegion:$region);
 }
 
+// CHECK: error: property 'prop' is already bound
+def VariableInvalidO : TestFormat_Op<[{
+  custom<Test>($prop, $prop) attr-dict
+}]>, Arguments<(ins IntProperty<"int64_t">:$prop)>;
+
 //===----------------------------------------------------------------------===//
 // Coverage Checks
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: type of result #0, named 'result', is not buildable and a buildable type cannot be inferred
 // CHECK: note: suggest adding a type constraint to the operation or adding a 'type($result)' directive to the custom assembly format
 def ZCoverageInvalidA : TestFormat_Op<[{
   attr-dict

mlir/tools/mlir-tblgen/FormatGen.h

 };
 
 /// This class represents a variable element. A variable refers to some part of
 /// the object being parsed, e.g. an attribute or operand on an operation or a
 /// parameter on an attribute.
 class VariableElement : public FormatElementBase<FormatElement::Variable> {
 public:
   /// These are the kinds of variables.
-  enum Kind { Attribute, Operand, Region, Result, Successor, Parameter };
+  enum Kind { Attribute, Operand, Region, Result, Successor, Parameter, Property };
 
   /// Get the kind of variable.
   Kind getKind() const { return kind; }
 
 protected:
   /// Create a variable with a kind.
   VariableElement(Kind kind) : kind(kind) {}
 

mlir/tools/mlir-tblgen/OpFormatGen.cpp

using ResultVariable =
     OpVariableElement<NamedTypeConstraint, VariableElement::Result>;
 
 /// This class represents a variable that refers to a region.
 using RegionVariable = OpVariableElement<NamedRegion, VariableElement::Region>;
 
 /// This class represents a variable that refers to a successor.
 using SuccessorVariable =
     OpVariableElement<NamedSuccessor, VariableElement::Successor>;
+
+using PropertyVariable =
+    OpVariableElement<NamedProperty, VariableElement::Property>;
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // DirectiveElement
 
 namespace {
 /// This class represents the `operands` directive. This directive represents
 /// all of the operands of an operation.
else
       body << formatv("{0}RawTypes[0]", listName);
 
   } else if (auto *string = dyn_cast<StringElement>(param)) {
     FmtContext ctx;
     ctx.withBuilder("parser.getBuilder()");
     ctx.addSubst("_ctxt", "parser.getContext()");
     body << tgfmt(string->getValue(), &ctx);
 
+  } else if (auto *property = dyn_cast<PropertyVariable>(param)) {
+    body << llvm::formatv("result.getOrAddProperties<Properties>().{0}",
+                          property->getVar()->name);
   } else {
     llvm_unreachable("unknown custom directive parameter");
   }
 }
 
 /// Generate the parser for a custom directive.
 static void genCustomDirectiveParser(CustomDirective *dir, MethodBody &body,
                                      bool useProperties,
else
       body << name << "().getType()";
 
   } else if (auto *string = dyn_cast<StringElement>(element)) {
     FmtContext ctx;
     ctx.withBuilder("::mlir::Builder(getContext())");
     ctx.addSubst("_ctxt", "getContext()");
     body << tgfmt(string->getValue(), &ctx);
 
+  } else if (auto *property = dyn_cast<PropertyVariable>(element)) {
+    FmtContext ctx;
+    ctx.addSubst("_ctxt", "getContext()");
+    const NamedProperty *namedProperty = property->getVar();
+    ctx.addSubst("_storage", "getProperties()." + namedProperty->name);
+    body << tgfmt(namedProperty->prop.getConvertFromStorageCall(), &ctx);
   } else {
     llvm_unreachable("unknown custom directive parameter");
   }
 }
 
 /// Generate the printer for a custom directive.
 static void genCustomDirectivePrinter(CustomDirective *customDir,
                                       const Operator &op, MethodBody &body) {
private:
   bool hasPropDict = false;
   bool hasAllRegions = false, hasAllSuccessors = false;
   bool canInferResultTypes = false;
   llvm::SmallBitVector seenOperandTypes, seenResultTypes;
   llvm::SmallSetVector<const NamedAttribute *, 8> seenAttrs;
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
+  llvm::DenseSet<const NamedProperty *> seenProperties;
 };
 } // namespace
 
 LogicalResult OpFormatParser::verify(SMLoc loc,
                                      ArrayRef<FormatElement *> elements) {
   // Check that the attribute dictionary is in the format.
   if (!hasAttrDict)
     return emitError(loc, "'attr-dict' directive not found in "
if (ctx == RefDirectiveContext) {
         return emitError(loc, "attribute '" + name +
                                   "' must be bound before it is referenced");
     } else if (!seenAttrs.insert(attr)) {
       return emitError(loc, "attribute '" + name + "' is already bound");
     }
 
     return create<AttributeVariable>(attr);
   }
+
+  if (const NamedProperty *property = findArg(op.getProperties(), name)) {
+    if (ctx != CustomDirectiveContext)
+      return emitError(
+          loc, "properties currently only supported in `custom` directive");
+
+    if (!seenProperties.insert(property).second)
+      return emitError(loc, "property '" + name + "' is already bound");
+
+    return create<PropertyVariable>(property);
+  }
+
   // Operands
   if (const NamedTypeConstraint *operand = findArg(op.getOperands(), name)) {
     if (ctx == TopLevelContext || ctx == CustomDirectiveContext) {
       if (fmt.allOperands || !seenOperands.insert(operand).second)
         return emitError(loc, "operand '" + name + "' is already bound");
     } else if (ctx == RefDirectiveContext && !seenOperands.count(operand)) {
       return emitError(loc, "operand '" + name +
                                 "' must be bound before it is referenced");
OpFormatParser::parsePropDictDirective(SMLoc loc, Context context) {
   return create<PropDictDirective>();
 }
 
 LogicalResult OpFormatParser::verifyCustomDirectiveArguments(
     SMLoc loc, ArrayRef<FormatElement *> arguments) {
   for (FormatElement *argument : arguments) {
     if (!isa<StringElement, RefDirective, TypeDirective, AttrDictDirective,
              AttributeVariable, OperandVariable, RegionVariable,
-             SuccessorVariable>(argument)) {
+             SuccessorVariable, PropertyVariable>(argument)) {
       // TODO: FormatElement should have location info attached.
       return emitError(loc, "only variables and types may be used as "
                             "parameters to a custom directive");
     }
     if (auto *type = dyn_cast<TypeDirective>(argument)) {
       if (!isa<OperandVariable, ResultVariable>(type->getArg())) {
         return emitError(loc, "type directives within a custom directive may "
                               "only refer to variables");