[Mlir-commits] [mlir] [mlir] Populate properties before parser-time inferReturnTypes (PR #194658)

[Zmicier Prybysh]

https://github.com/dimp-pl updated https://github.com/llvm/llvm-project/pull/194658

>From a33e691e0141d4685d0176e2a7352fa0223828fe Mon Sep 17 00:00:00 2001
From: Zmicier Prybysh <zprybysh at baylibre.com>
Date: Wed, 6 May 2026 17:13:09 +0200
Subject: [PATCH] [MLIR] Populate properties from attr-dict at parse time

Fixes #193284.

It's possible to segfault inferReturnTypes when accessing an
attribute that's stored as property under specific conditions:

- Attribute must be inherent.
- `DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>`.
- Attribute is read in `inferReturnTypes` via `Adaptor`.
- Op omits `type($result)` from assembly format.

The fix works by emitting an additional setInherentAttr per each
attribute in `attr-dict` parser after erification step.
---
 mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 13 +++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td     |  9 +++++++++
 mlir/test/mlir-tblgen/op-format.mlir      |  4 ++++
 mlir/tools/mlir-tblgen/OpFormatGen.cpp    | 12 ++++++++++++
 4 files changed, 38 insertions(+)

diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 340b44b14dd96..38177d22c8bd0 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1295,6 +1295,19 @@ LogicalResult TestOpWithPropertiesAndInferredType::inferReturnTypes(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TestOpWithAttrInferredType
+//===----------------------------------------------------------------------===//
+
+LogicalResult TestOpWithAttrInferredType::inferReturnTypes(
+    MLIRContext *context, std::optional<Location>, ValueRange operands,
+    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  Adaptor adaptor(operands, attributes, properties, regions);
+  inferredReturnTypes.push_back(IntegerType::get(context, adaptor.getLhs()));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // LoopBlockOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 348ff5d7f4ea0..5ac64ada612f7 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3527,6 +3527,15 @@ def TestOpWithPropertiesAndInferredType
   let results = (outs AnyType:$result);
 }
 
+def TestOpWithAttrInferredType
+  : TEST_Op<"with_attr_inferred_type", [
+    DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>
+  ]> {
+  let assemblyFormat = "$input attr-dict `:` type($input)";
+  let arguments = (ins I32:$input, I32Attr:$lhs);
+  let results = (outs AnyType:$output);
+}
+
 // Demonstrate how to wrap an existing C++ class named MyPropStruct.
 def MyStructProperty : Property<"MyPropStruct"> {
   let convertToAttribute = "return $_storage.asAttribute($_ctxt);";
diff --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index 7ff9091d5500d..3e8ee9ed330b4 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -534,6 +534,10 @@ test.with_properties_and_attr 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.with_attr_inferred_type %[[I32]] {lhs = 32 : i32} : i32
+%attr_inferred_i32 = test.with_attr_inferred_type %i32 {lhs = 32 : i32} : i32
+test.format_all_types_match_var %attr_inferred_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}>
 
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index cbcbc8e9bc102..6f70ce0a42489 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -676,6 +676,17 @@ const char *const inferReturnTypesParserCode = R"(
   result.addTypes(inferredReturnTypes);
 )";
 
+/// The code snippet used to copy inherent attributes from
+/// `result.attributes` into the inline properties storage.
+///
+/// {0}: The operation class name.
+const char *const attrDictPopulatePropertiesCode = R"(
+auto &odsAttrDictProps = result.getOrAddProperties<{0}::Properties>();
+for (const ::mlir::NamedAttribute &namedAttr : result.attributes)
+  {0}::setInherentAttr(odsAttrDictProps, namedAttr.getName().getValue(),
+      namedAttr.getValue());
+)";
+
 /// The code snippet used to generate a parser call for a region list.
 ///
 /// {0}: The name for the region list.
@@ -1640,6 +1651,7 @@ void OperationFormat::genElementParser(FormatElement *element, MethodBody &body,
               "result.name.getStringRef() << \"' op \";\n"
            << "  })))\n"
            << "  return ::mlir::failure();\n";
+      body << formatv(attrDictPopulatePropertiesCode, opCppClassName);
     }
     body.unindent() << "}\n";
     body.unindent();