[Mlir-commits] [mlir] [mlir]Add resultSegmentSizes/operandSegmentSizes to prop-dict. (PR #211222)

Wed Jul 22 03:00:06 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir

Author: Natanael Cintean (natanael-cintean)

<details>
<summary>Changes</summary>

Fixes #211220

## Summary

`setPropertiesFromParsedAttr` (generated by `OpFormatGen.cpp` for ops using a
custom `assemblyFormat`) rejects the trait-injected `operandSegmentSizes` /
`resultSegmentSizes` properties when they appear in a `prop-dict`, even though
the printer emits them there for ops whose format uses a bulk operand/result
type directive (`type(operands)`, `type(results)`,
`functional-type(operands, results)`).

This breaks round-tripping: any op combining `AttrSizedOperandSegments` /
`AttrSizedResultSegments` with such a bulk directive fails to re-parse its own
printed output, because printer-side elision of these keys
(introduced in #115930) is intentionally skipped in that case — the sizes
can't be reconstructed from individually-typed operand/result groups, so they
must survive in the text, but the custom parser was never taught to read them
back.

`setPropertiesFromAttr` (used by the generic-form parser, bytecode, and C++
construction) already special-cases these two keys, so this change makes
`setPropertiesFromParsedAttr` mirror that behavior: keys are accepted from the
dictionary attr when the corresponding trait is present, and required exactly
when the format can't reconstruct the sizes itself (i.e. when
`fmt.allOperands` / `fmt.allResultTypes` is true). Otherwise the key remains
optional, since `genParserVariadicSegmentResolution` overwrites it with the
sizes inferred from the parsed operand/result groups.

## Changes

- `mlir/tools/mlir-tblgen/OpFormatGen.cpp`: `genParsedAttrPropertiesSetter` now
  emits a handler for `operandSegmentSizes` / `resultSegmentSizes` right after
  `usedKeys` is declared, for ops with `AttrSizedOperandSegments` /
  `AttrSizedResultSegments`.
- `mlir/test/lib/Dialect/Test/TestOps.td`: added
  `test.variadic_segment_prop_bulk_type`
  (`TestOpWithBulkTypesAndSegmentProperties`), which combines both
  `AttrSized*Segments` traits with a bulk `functional-type(operands, results)`
  directive — the pattern that triggers the bug.
- `mlir/test/IR/properties.mlir`: round-trip test (custom form and generic
  form) for the new op, exercising both the previously-broken parse path and
  confirming the printer output stays stable.

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


3 Files Affected:

- (modified) mlir/test/IR/properties.mlir (+11) 
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+16) 
- (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+32) 


``````````diff

diff --git a/mlir/test/IR/properties.mlir b/mlir/test/IR/properties.mlir
index 4dbb8ebfa5db3..4d83038f31cdd 100644
--- a/mlir/test/IR/properties.mlir
+++ b/mlir/test/IR/properties.mlir
@@ -46,6 +46,17 @@ test.using_property_ref_in_custom 1 + 4 = 5
 %ci64 = arith.constant 0 : i64
 test.variadic_segment_prop %ci64, %ci64 : %ci64 : i64, i64 : i64 end
 
+// Tests that the variadic segment size properties survive a round-trip
+// through the *custom* (non-generic) parser/printer when the assembly format
+// uses a bulk `functional-type(operands, results)` directive, which prevents
+// the printer from eliding `operandSegmentSizes` / `resultSegmentSizes` from
+// `<{...}>`. Without the parser-side fix, re-parsing the CHECK line below
+// (which is exactly what the printer emits) fails with "duplicate or unknown
+// key 'operandSegmentSizes' in dictionary attribute".
+// CHECK: test.variadic_segment_prop_bulk_type(%[[CI64]], %[[CI64]], %[[CI64]]) : (i64, i64, i64) -> (i64, i64, i64) <{operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>}>
+// GENERIC: "test.variadic_segment_prop_bulk_type"(%[[CI64]], %[[CI64]], %[[CI64]]) <{operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>}> : (i64, i64, i64) -> (i64, i64, i64)
+test.variadic_segment_prop_bulk_type(%ci64, %ci64, %ci64) : (i64, i64, i64) -> (i64, i64, i64) <{operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>}>
+
 // CHECK:   test.with_default_valued_properties na{{$}}
 // GENERIC: "test.with_default_valued_properties"()
 // GENERIC-SAME: <{a = 0 : i32, b = "", c = -1 : i32, unit = false}> : () -> ()
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index c12e2abb3ab75..db55b90c294e1 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3686,6 +3686,22 @@ def TestOpWithVariadicSegmentProperties : TEST_Op<"variadic_segment_prop",
   }];
 }
 
+// Same as `TestOpWithVariadicSegmentProperties`, but the assembly format uses
+// the bulk `type(operands)` / `type(results)` directives instead of spelling
+// out the type of each individual variadic group. This prevents the printer
+// from eliding `operandSegmentSizes` / `resultSegmentSizes` from the
+// `<{...}>` properties dictionary (since there is no other way to recover the
+// segment sizes on re-parse), so this op exercises the parser's ability to
+// read those trait-injected properties back out of `prop-dict`.
+def TestOpWithBulkTypesAndSegmentProperties : TEST_Op<"variadic_segment_prop_bulk_type",
+    [AttrSizedOperandSegments, AttrSizedResultSegments]> {
+  let arguments = (ins Variadic<I64>:$a1, Variadic<I64>:$a2);
+  let results = (outs Variadic<I64>:$b1, Variadic<I64>:$b2);
+  let assemblyFormat = [{
+    `(` operands `)` `:` functional-type(operands, results) prop-dict attr-dict
+  }];
+}
+
 def TestOpUsingPropertyRefInCustom : TEST_Op<"using_property_ref_in_custom"> {
   let assemblyFormat = "custom<IntProperty>($first) `+` custom<SumProperty>($second, ref($first)) attr-dict";
   let arguments = (ins IntProp<"int64_t">:$first, IntProp<"int64_t">:$second);
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index cbcbc8e9bc102..f9d07f6b3ba0f 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -1313,6 +1313,38 @@ ::mlir::MLIRContext *ctx = dict.getContext();
 (void)ctx;
 )decl";
 
+  // `operandSegmentSizes`/`resultSegmentSizes` are trait-injected properties
+  // not enumerated by `op.getProperties()`, so they need to be special-cased
+  // here, mirroring the handling in `setPropertiesFromAttr` in
+  // OpDefinitionsGen.cpp.
+  //
+  // {0}: segment sizes property name
+  // {1}: isRequired
+  const char *segmentSizesFromAttrFmt = R"decl(
+auto {0}AttrName = ::mlir::StringAttr::get(ctx, "{0}");
+usedKeys.insert({0}AttrName);
+auto attr = dict.get({0}AttrName);
+if (!attr && {1}) {{
+  emitError() << "expected key entry for {0} in DictionaryAttr to set "
+             "Properties.";
+  return ::mlir::failure();
+}
+if (attr && ::mlir::failed(::mlir::convertFromAttribute(prop.{0}, attr, [&]() {{
+      return emitError() << "for `{0}`: ";
+    })))
+  return ::mlir::failure();
+)decl";
+  if (op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) {
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+    body << formatv(segmentSizesFromAttrFmt, "operandSegmentSizes",
+                    fmt.allOperands);
+  }
+  if (op.getTrait("::mlir::OpTrait::AttrSizedResultSegments")) {
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+    body << formatv(segmentSizesFromAttrFmt, "resultSegmentSizes",
+                    fmt.allResultTypes);
+  }
+
   // {0}: fromAttribute call
   // {1}: property name
   // {2}: isRequired

``````````

</details>


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