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

This is essentially a follow up to https://reviews.llvm.org/D155072

This adds support for also passing properties as `ref` parameter to `custom`. This requires the property to have been bound previously and will error otherwise. This makes it possible for an implementation of `custom` to take previously parsed data into account, creating nice context-dependent grammars :-)

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

5 files changed, 56 insertions, 3 deletions

diff --git a/mlir/test/IR/properties.mlir b/mlir/test/IR/properties.mlir
index 715d9351211f..2aef2abc81f7 100644
--- a/mlir/test/IR/properties.mlir
+++ b/mlir/test/IR/properties.mlir
@@ -24,3 +24,12 @@ test.with_wrapped_properties <{prop = "content for properties"}>
 // GENERIC: "test.using_property_in_custom"()
 // GENERIC-SAME: prop = array<i64: 1, 4, 20>
 test.using_property_in_custom [1, 4, 20]
+
+// CHECK: test.using_property_ref_in_custom
+// CHECK-SAME: 1 + 4 = 5
+// GENERIC: "test.using_property_ref_in_custom"()
+// GENERIC-SAME: <{
+// GENERIC-SAME: first = 1
+// GENERIC-SAME: second = 4
+// GENERIC-SAME: }>
+test.using_property_ref_in_custom 1 + 4 = 5
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 034d14807764..fba882d3ba1a 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -1945,6 +1945,34 @@ static void printUsingPropertyInCustom(OpAsmPrinter &printer, Operation *op,
   printer << '[' << value << ']';
 }
 
+static bool parseIntProperty(OpAsmParser &parser, int64_t &value) {
+  return failed(parser.parseInteger(value));
+}
+
+static void printIntProperty(OpAsmPrinter &printer, Operation *op,
+                             int64_t value) {
+  printer << value;
+}
+
+static bool parseSumProperty(OpAsmParser &parser, int64_t &second,
+                             int64_t first) {
+  int64_t sum;
+  auto loc = parser.getCurrentLocation();
+  if (parser.parseInteger(second) || parser.parseEqual() ||
+      parser.parseInteger(sum))
+    return true;
+  if (sum != second + first) {
+    parser.emitError(loc, "Expected sum to equal first + second");
+    return true;
+  }
+  return false;
+}
+
+static void printSumProperty(OpAsmPrinter &printer, Operation *op,
+                             int64_t second, int64_t first) {
+  printer << second << " = " << (second + first);
+}
+
 #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 7fbf751d5756..30f334c4404b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3349,6 +3349,11 @@ def TestOpUsingPropertyInCustom : TEST_Op<"using_property_in_custom"> {
   let arguments = (ins ArrayProperty<"int64_t", 3>:$prop);
 }
 
+def TestOpUsingPropertyRefInCustom : TEST_Op<"using_property_ref_in_custom"> {
+  let assemblyFormat = "custom<IntProperty>($first) `+` custom<SumProperty>($second, ref($first)) attr-dict";
+  let arguments = (ins IntProperty<"int64_t">:$first, IntProperty<"int64_t">:$second);
+}
+
 // 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 2723eab678ed..b038ddd5717b 100644
--- a/mlir/test/mlir-tblgen/op-format-invalid.td
+++ b/mlir/test/mlir-tblgen/op-format-invalid.td
@@ -483,6 +483,11 @@ def VariableInvalidO : TestFormat_Op<[{
   custom<Test>($prop, $prop) attr-dict
 }]>, Arguments<(ins IntProperty<"int64_t">:$prop)>;
 
+// CHECK: error: property 'prop' must be bound before it is referenced
+def VariableInvalidP : TestFormat_Op<[{
+  custom<Test>(ref($prop)) attr-dict
+}]>, Arguments<(ins IntProperty<"int64_t">:$prop)>;
+
 //===----------------------------------------------------------------------===//
 // Coverage Checks
 //===----------------------------------------------------------------------===//
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 0a1b4e6e3155..9a5e4d52c550 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -2957,12 +2957,18 @@ OpFormatParser::parseVariableImpl(SMLoc loc, StringRef name, Context ctx) {
   }
 
   if (const NamedProperty *property = findArg(op.getProperties(), name)) {
-    if (ctx != CustomDirectiveContext)
+    if (ctx != CustomDirectiveContext && ctx != RefDirectiveContext)
       return emitError(
           loc, "properties currently only supported in `custom` directive");
 
-    if (!seenProperties.insert(property).second)
-      return emitError(loc, "property '" + name + "' is already bound");
+    if (ctx == RefDirectiveContext) {
+      if (!seenProperties.count(property))
+        return emitError(loc, "property '" + name +
+                                  "' must be bound before it is referenced");
+    } else {
+      if (!seenProperties.insert(property).second)
+        return emitError(loc, "property '" + name + "' is already bound");
+    }
 
     return create<PropertyVariable>(property);
   }