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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Zmicier Prybysh (dimp-pl)

<details>
<summary>Changes</summary>

Fixes #<!-- -->193284.

In inferReturnTypes it's possible to segfault 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 before the `inferReturnTypes` call.

This kind of repeats what `Operation::setAttrs` is doing in runtime, but only in parser-time so that `inferReturnTypes` sees a populated property.

Assisted-By: Claude Code

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


4 Files Affected:

- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+13) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+9) 
- (modified) mlir/test/mlir-tblgen/op-format.mlir (+4) 
- (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+16-1) 


``````````diff
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..30a9b50986224 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 attribute-form inherent attributes from
+/// `result.attributes` into the inline properties storage.
+///
+/// {0}: The operation class name
+const char *const inferReturnTypesPopulatePropertiesCode = R"(
+  auto &inferRetTypesProps = result.getOrAddProperties<{0}::Properties>();
+  for (const ::mlir::NamedAttribute &namedAttr : result.attributes)
+    {0}::setInherentAttr(inferRetTypesProps, 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.
@@ -1782,8 +1793,12 @@ void OperationFormat::genParserTypeResolution(Operator &op, MethodBody &body) {
   genParserOperandTypeResolution(op, body, emitTypeResolver);
 
   // Handle return type inference once all operands have been resolved
-  if (infersResultTypes)
+  if (infersResultTypes) {
+    if (useProperties)
+      body << formatv(inferReturnTypesPopulatePropertiesCode,
+                      op.getCppClassName());
     body << formatv(inferReturnTypesParserCode, op.getCppClassName());
+  }
 }
 
 void OperationFormat::genParserOperandTypeResolution(

``````````

</details>


https://github.com/llvm/llvm-project/pull/194658