[mlir][ODS] Make prop-dict behave closer to attr-dict

[zero9178]

attr-dict currently prints any attribute (inherent or discardable) that does not occur elsewhere in the assembly format. prop-dict on the other hand, always prints and parses all properties (including inherent attributes stored as properties) as part of the property dictionary, regardless of whether the properties or attributes are used elsewhere. (with the exception of default-valued attributes implemented recently in #87970).

This PR changes the behavior of prop-dict to only print and parse attributes and properties that do not occur elsewhere in the assembly format. This is achieved by 1) adding used attributes and properties to the elision list when printing and 2) using a custom version of setPropertiesFromAttr called setPropertiesFromParsedAttr that is sensitive to the assembly format and auto-generated by ODS.

The current and new behavior of prop-dict and attr-dict were also documented.

Happens to also fix #88506

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Markus Böck (zero9178)

Changes

attr-dict currently prints any attribute (inherent or discardable) that does not occur elsewhere in the assembly format. prop-dict on the other hand, always prints and parses all properties (including inherent attributes stored as properties) as part of the property dictionary, regardless of whether the properties or attributes are used elsewhere. (with the exception of default-valued attributes implemented recently in #87970).

This PR changes the behavior of prop-dict to only print and parse attributes and properties that do not occur elsewhere in the assembly format. This is achieved by 1) adding used attributes and properties to the elision list when printing and 2) using a custom version of setPropertiesFromAttr called setPropertiesFromParsedAttr that is sensitive to the assembly format and auto-generated by ODS.

The current and new behavior of prop-dict and attr-dict were also documented.

Happens to also fix #88506

---

Full diff: https://github.com/llvm/llvm-project/pull/88659.diff

6 Files Affected:

• (modified) mlir/docs/DefiningDialects/Operations.md (+11-1)
• (modified) mlir/include/mlir/IR/OpDefinition.h (+29-3)
• (modified) mlir/test/lib/Dialect/Test/TestDialect.cpp (+11)
• (modified) mlir/test/lib/Dialect/Test/TestOps.td (+26)
• (modified) mlir/test/mlir-tblgen/op-format.mlir (+11)
• (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+114-4)

diff --git a/mlir/docs/DefiningDialects/Operations.md b/mlir/docs/DefiningDialects/Operations.md
index b27330319f6594..729393d5362673 100644
--- a/mlir/docs/DefiningDialects/Operations.md
+++ b/mlir/docs/DefiningDialects/Operations.md
@@ -640,13 +640,23 @@ The available directives are as follows:
 
 *   `attr-dict`
 
-    -   Represents the attribute dictionary of the operation.
+    -   Represents the attribute dictionary of the operation. Any inherent 
+    -   attributes that are not used elsewhere in the format are printed as
+    -   part of the attribute dictionary unless a `prop-dict` is present.
+    -   Discardable attributes are always part of the `attr-dict`.  
 
 *   `attr-dict-with-keyword`
 
     -   Represents the attribute dictionary of the operation, but prefixes the
         dictionary with an `attributes` keyword.
 
+*   `prop-dict`
+
+    -   Represents the properties of the operation converted to a dictionary.
+    -   Any property or inherent attribute that are not used elsewhere in the
+    -   format are parsed and printed as part of this dictionary.
+    -   If present, the `attr-dict` will not contain any inherent attributes.
+
 *   `custom` < UserDirective > ( Params )
 
     -   Represents a custom directive implemented by the user in C++.
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index f4355c8ce26ace..59f094d6690991 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1993,8 +1993,10 @@ class Op : public OpState, public Traits<ConcreteType>... {
         p, ConcreteType::getPropertiesAsAttr(ctx, properties), elidedProps);
   }
 
-  /// Parser the properties. Unless overridden, this method will print by
-  /// converting the properties to an Attribute.
+  /// Parses 'prop-dict' for the operation. Unless overridden, the method will
+  /// parse the properties using the generic property dictionary using the
+  /// '<{ ... }>' syntax. The resulting properties are stored within the
+  /// property structure of 'result', accessible via 'getOrAddProperties'.
   template <typename T = ConcreteType>
   static ParseResult parseProperties(OpAsmParser &parser,
                                      OperationState &result) {
@@ -2002,7 +2004,31 @@ class Op : public OpState, public Traits<ConcreteType>... {
       return parseProperties(
           parser, result.getOrAddProperties<InferredProperties<T>>());
     }
-    return genericParseProperties(parser, result.propertiesAttr);
+
+    Attribute propertyDictionary;
+    if (genericParseProperties(parser, propertyDictionary))
+      return failure();
+
+    // The generated 'setPropertiesFromParsedAttr', like
+    // 'setPropertiesFromAttr', expects a 'DictionaryAttr' that is not null.
+    // Use an empty dictionary in the case that the whole dictionary is
+    // optional.
+    if (!propertyDictionary)
+      propertyDictionary = DictionaryAttr::get(result.getContext());
+
+    auto emitError = [&]() {
+      return mlir::emitError(result.location, "invalid properties ")
+             << propertyDictionary << " for op " << result.name.getStringRef()
+             << ": ";
+    };
+
+    // Copy the data from the dictionary attribute into the property struct of
+    // the operation. This method is generated by ODS by default if there are
+    // any occurrences of 'prop-dict' in the assembly format and should set
+    // any properties that aren't parsed elsewhere.
+    return ConcreteOpType::setPropertiesFromParsedAttr(
+        result.getOrAddProperties<InferredProperties<T>>(), propertyDictionary,
+        emitError);
   }
 
 private:
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 380c74a47e509a..74378633032ce7 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -740,6 +740,17 @@ LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
   return success();
 }
 
+LogicalResult TestOpWithPropertiesAndInferredType::inferReturnTypes(
+    MLIRContext *context, std::optional<Location>, ValueRange operands,
+    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+
+  Adaptor adaptor(operands, attributes, properties, regions);
+  inferredReturnTypes.push_back(IntegerType::get(
+      context, adaptor.getLhs() + adaptor.getProperties().rhs));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffect interfaces
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index edca05fde5a524..c88f85b8b6b361 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2852,6 +2852,23 @@ def TestOpWithProperties : TEST_Op<"with_properties"> {
   );
 }
 
+def TestOpWithPropertiesAndAttr
+  : TEST_Op<"with_properties_and_attr"> {
+  let assemblyFormat = "$lhs prop-dict attr-dict";
+
+  let arguments = (ins I32Attr:$lhs, IntProperty<"int64_t">:$rhs);
+}
+
+def TestOpWithPropertiesAndInferredType
+  : TEST_Op<"with_properties_and_inferred_type", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface>
+  ]> {
+  let assemblyFormat = "$lhs prop-dict attr-dict";
+
+  let arguments = (ins I32Attr:$lhs, IntProperty<"int64_t">:$rhs);
+  let results = (outs AnyType:$result);
+}
+
 // Demonstrate how to wrap an existing C++ class named MyPropStruct.
 def MyStructProperty : Property<"MyPropStruct"> {
   let convertToAttribute = "$_storage.asAttribute($_ctxt)";
@@ -2871,6 +2888,15 @@ def TestOpUsingPropertyInCustom : TEST_Op<"using_property_in_custom"> {
   let arguments = (ins ArrayProperty<"int64_t", 3>:$prop);
 }
 
+def TestOpUsingPropertyInCustomAndOther
+  : TEST_Op<"using_property_in_custom_and_other"> {
+  let assemblyFormat = "custom<UsingPropertyInCustom>($prop) prop-dict attr-dict";
+  let arguments = (ins
+    ArrayProperty<"int64_t", 3>:$prop,
+    IntProperty<"int64_t">:$other
+  );
+}
+
 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);
diff --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index 14e1cdb07db39d..46d272649caedc 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -480,6 +480,17 @@ test.format_infer_variadic_type_from_non_variadic %i64, %i64 : i64
 // CHECK: test.format_infer_type_variadic_operands(%[[I32]], %[[I32]] : i32, i32) (%[[I64]], %[[I64]] : i64, i64)
 %ignored_res13:4 = test.format_infer_type_variadic_operands(%i32, %i32 : i32, i32) (%i64, %i64 : i64, i64)
 
+// CHECK: test.with_properties_and_attr 16 < {rhs = 16 : i64}>
+test.with_properties_and_attr 16 <{rhs = 16 : i64}>
+
+// CHECK: test.with_properties_and_inferred_type 16 < {rhs = 16 : i64}>
+%should_be_i32 = test.with_properties_and_inferred_type 16 <{rhs = 16 : i64}>
+// Assert through the verifier that its inferred as i32.
+test.format_all_types_match_var %should_be_i32, %i32 : i32
+
+// CHECK: test.using_property_in_custom_and_other [1, 4, 20] < {other = 16 : i64}>
+test.using_property_in_custom_and_other [1, 4, 20] <{other = 16 : i64}>
+
 //===----------------------------------------------------------------------===//
 // Check DefaultValuedStrAttr
 //===----------------------------------------------------------------------===//
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 5963b5e689da74..806991035e6685 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -379,8 +379,11 @@ struct OperationFormat {
   std::vector<TypeResolution> operandTypes, resultTypes;
 
   /// The set of attributes explicitly used within the format.
-  SmallVector<const NamedAttribute *, 8> usedAttributes;
+  llvm::SmallSetVector<const NamedAttribute *, 8> usedAttributes;
   llvm::StringSet<> inferredAttributes;
+
+  /// The set of properties explicitly used within the format.
+  llvm::SmallSetVector<const NamedProperty *, 8> usedProperties;
 };
 } // namespace
 
@@ -1183,6 +1186,105 @@ static void genAttrParser(AttributeVariable *attr, MethodBody &body,
   }
 }
 
+// Generates the 'setPropertiesFromParsedAttr' used to set properties from a
+// 'prop-dict' dictionary attr.
+static void genParsedAttrPropertiesSetter(OperationFormat &fmt, Operator &op,
+                                          OpClass &opClass) {
+  // Not required unless 'prop-dict' is present.
+  if (!fmt.hasPropDict)
+    return;
+
+  SmallVector<MethodParameter> paramList;
+  paramList.emplace_back("Properties &", "prop");
+  paramList.emplace_back("::mlir::Attribute", "attr");
+  paramList.emplace_back("::llvm::function_ref<::mlir::InFlightDiagnostic()>",
+                         "emitError");
+
+  Method *method = opClass.addStaticMethod("::mlir::LogicalResult",
+                                           "setPropertiesFromParsedAttr",
+                                           std::move(paramList));
+  MethodBody &body = method->body().indent();
+
+  body << R"decl(
+::mlir::DictionaryAttr dict = ::llvm::dyn_cast<::mlir::DictionaryAttr>(attr);
+if (!dict) {
+  emitError() << "expected DictionaryAttr to set properties";
+  return ::mlir::failure();
+}
+)decl";
+
+  // TODO: properties might be optional as well.
+  const char *propFromAttrFmt = R"decl(
+auto setFromAttr = [] (auto &propStorage, ::mlir::Attribute propAttr,
+         ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError) {{
+  {0};
+};
+auto attr = dict.get("{1}");
+if (!attr) {{
+  emitError() << "expected key entry for {1} in DictionaryAttr to set "
+             "Properties.";
+  return ::mlir::failure();
+}
+if (::mlir::failed(setFromAttr(prop.{1}, attr, emitError)))
+  return ::mlir::failure();
+)decl";
+
+  // Generate the setter for any property not parsed elsewhere.
+  for (const NamedProperty &namedProperty : op.getProperties()) {
+    if (fmt.usedProperties.contains(&namedProperty))
+      continue;
+
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+
+    StringRef name = namedProperty.name;
+    const Property &prop = namedProperty.prop;
+    FmtContext fctx;
+    body << formatv(propFromAttrFmt,
+                    tgfmt(prop.getConvertFromAttributeCall(),
+                          &fctx.addSubst("_attr", "propAttr")
+                               .addSubst("_storage", "propStorage")
+                               .addSubst("_diag", "emitError")),
+                    name);
+  }
+
+  // Generate the setter for any attribute not parsed elsewhere.
+  for (const NamedAttribute &namedAttr : op.getAttributes()) {
+    if (fmt.usedAttributes.contains(&namedAttr))
+      continue;
+
+    const Attribute &attr = namedAttr.attr;
+    // Derived attributes do not need to be parsed.
+    if (attr.isDerivedAttr())
+      continue;
+
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+
+    // If the attribute has a default value or is optional, it does not need to
+    // be present in the parsed dictionary attribute.
+    bool isRequired = !attr.isOptional() && !attr.hasDefaultValue();
+    body << formatv(R"decl(
+auto &propStorage = prop.{0};
+auto attr = dict.get("{0}");
+if (attr || /*isRequired=*/{1}) {{
+  if (!attr) {{
+    emitError() << "expected key entry for {0} in DictionaryAttr to set "
+               "Properties.";
+    return ::mlir::failure();
+  }
+  auto convertedAttr = ::llvm::dyn_cast<std::remove_reference_t<decltype(propStorage)>>(attr);
+  if (convertedAttr) {{
+    propStorage = convertedAttr;
+  } else {{
+    emitError() << "Invalid attribute `{0}` in property conversion: " << attr;
+    return ::mlir::failure();
+  }
+}
+)decl",
+                    namedAttr.name, isRequired);
+  }
+  body << "return ::mlir::success();\n";
+}
+
 void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   SmallVector<MethodParameter> paramList;
   paramList.emplace_back("::mlir::OpAsmParser &", "parser");
@@ -1214,6 +1316,8 @@ void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   genParserTypeResolution(op, body);
 
   body << "  return ::mlir::success();\n";
+
+  genParsedAttrPropertiesSetter(*this, op, opClass);
 }
 
 void OperationFormat::genElementParser(FormatElement *element, MethodBody &body,
@@ -1776,6 +1880,11 @@ const char *enumAttrBeginPrinterCode = R"(
 static void genPropDictPrinter(OperationFormat &fmt, Operator &op,
                                MethodBody &body) {
   body << "  ::llvm::SmallVector<::llvm::StringRef, 2> elidedProps;\n";
+  for (const NamedProperty *namedProperty : fmt.usedProperties)
+    body << "  elidedProps.push_back(\"" << namedProperty->name << "\");\n";
+  for (const NamedAttribute *namedAttr : fmt.usedAttributes)
+    body << "  elidedProps.push_back(\"" << namedAttr->name << "\");\n";
+
   // Add code to check attributes for equality with the default value
   // for attributes with the elidePrintingDefaultValue bit set.
   for (const NamedAttribute &namedAttr : op.getAttributes()) {
@@ -2543,7 +2652,7 @@ class OpFormatParser : public FormatParser {
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
-  llvm::DenseSet<const NamedProperty *> seenProperties;
+  llvm::SmallSetVector<const NamedProperty *, 8> seenProperties;
 };
 } // namespace
 
@@ -2589,7 +2698,8 @@ LogicalResult OpFormatParser::verify(SMLoc loc,
     return failure();
 
   // Collect the set of used attributes in the format.
-  fmt.usedAttributes = seenAttrs.takeVector();
+  fmt.usedAttributes = std::move(seenAttrs);
+  fmt.usedProperties = std::move(seenProperties);
 
   // Set whether prop-dict is used in the format
   fmt.hasPropDict = hasPropDict;
@@ -3042,7 +3152,7 @@ OpFormatParser::parseVariableImpl(SMLoc loc, StringRef name, Context ctx) {
         return emitError(loc, "property '" + name +
                                   "' must be bound before it is referenced");
     } else {
-      if (!seenProperties.insert(property).second)
+      if (!seenProperties.insert(property))
         return emitError(loc, "property '" + name + "' is already bound");
     }
 

[zero9178]

Small note: These are for the most part copied from the code that also generates setPropertiesFromAttr but I opted to not deduplicate these after trying such a version as 1) both formats are actually very context sensitive (make assumptions about code generated beforehand) and 2) it allows for diverging implementations in the future.

Feel free to tell me otherwise if you'd rather deduplicate these

[zero9178]

The odd space after the < is an odd bug/behaviour unrelated to this patch