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

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 for 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`.

Differential Revision: https://reviews.llvm.org/D155072

diff --git a/mlir/test/IR/properties.mlir b/mlir/test/IR/properties.mlir
index b073698b40aff44d10593a22dd37f8f19b0b8fa7..715d9351211f1892f770b11aa90ee950daa7c8ea 100644 (file)
--- a/mlir/test/IR/properties.mlir
+++ b/mlir/test/IR/properties.mlir
@@ -18,3 +18,9 @@ test.with_nice_properties "foo bar" is -3
 // 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]

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index a12a5a302672b1a1b1d4f107ff6f41a43556c033..034d1480776475c29ba9e89e4dcba2aa3a12ea96 100644 (file)
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -1933,6 +1933,18 @@ static ParseResult customParseProperties(OpAsmParser &parser,
   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"

diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 8c8243af19ec06572228225fdd7803890e80076c..7fbf751d5756a5b0bdd9f01126ac09cf0cad0386 100644 (file)
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3344,6 +3344,11 @@ def TestOpWithWrappedProperties : TEST_Op<"with_wrapped_properties"> {
   );
 }
 
+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.
 

diff --git a/mlir/test/mlir-tblgen/op-format-invalid.td b/mlir/test/mlir-tblgen/op-format-invalid.td
index a44a2e6a0c5c2d3bf8082d118cf091f70b8389fb..2723eab678ed336d12cd7385964112809d2f0991 100644 (file)
--- a/mlir/test/mlir-tblgen/op-format-invalid.td
+++ b/mlir/test/mlir-tblgen/op-format-invalid.td
@@ -478,6 +478,11 @@ def VariableInvalidN : TestFormat_Op<[{
   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
 //===----------------------------------------------------------------------===//

diff --git a/mlir/tools/mlir-tblgen/FormatGen.h b/mlir/tools/mlir-tblgen/FormatGen.h
index da30e35f135323f309d21f28521901cc80385f48..18a410277fc108aec941de9ca7484dff40502290 100644 (file)
--- a/mlir/tools/mlir-tblgen/FormatGen.h
+++ b/mlir/tools/mlir-tblgen/FormatGen.h
@@ -235,7 +235,15 @@ private:
 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; }

diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 3e6db51b59758eae486fca042f193ac11b1d4435..0a1b4e6e31551efa247f5a8be7a0cbb5e8f70bfa 100644 (file)
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -94,6 +94,10 @@ using RegionVariable = OpVariableElement<NamedRegion, VariableElement::Region>;
 /// This class represents a variable that refers to a successor.
 using SuccessorVariable =
     OpVariableElement<NamedSuccessor, VariableElement::Successor>;
+
+/// This class represents a variable that refers to a property argument.
+using PropertyVariable =
+    OpVariableElement<NamedProperty, VariableElement::Property>;
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -939,6 +943,9 @@ static void genCustomParameterParser(FormatElement *param, MethodBody &body) {
     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");
   }
@@ -1862,6 +1869,12 @@ static void genCustomDirectiveParameterPrinter(FormatElement *element,
     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");
   }
@@ -2457,6 +2470,7 @@ private:
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
+  llvm::DenseSet<const NamedProperty *> seenProperties;
 };
 } // namespace
 
@@ -2941,6 +2955,18 @@ OpFormatParser::parseVariableImpl(SMLoc loc, StringRef name, Context ctx) {
 
     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) {
@@ -3068,9 +3094,9 @@ OpFormatParser::parsePropDictDirective(SMLoc loc, Context context) {
 LogicalResult OpFormatParser::verifyCustomDirectiveArguments(
     SMLoc loc, ArrayRef<FormatElement *> arguments) {
   for (FormatElement *argument : arguments) {
-    if (!isa<StringElement, RefDirective, TypeDirective, AttrDictDirective,
-             AttributeVariable, OperandVariable, RegionVariable,
-             SuccessorVariable>(argument)) {
+    if (!isa<AttrDictDirective, AttributeVariable, OperandVariable,
+             PropertyVariable, RefDirective, RegionVariable, SuccessorVariable,
+             StringElement, TypeDirective>(argument)) {
       // TODO: FormatElement should have location info attached.
       return emitError(loc, "only variables and types may be used as "
                             "parameters to a custom directive");