[mlir][ods] Fix missing property elision for variadic segment properties

[Moxinilian]

This fixes a bug where variadic segment properties would not be elided when printing prop-dict.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Théo Degioanni (Moxinilian)

Changes

This fixes a bug where variadic segment properties would not be elided when printing prop-dict.

---

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

3 Files Affected:

• (modified) mlir/test/IR/properties.mlir (+9-1)
• (modified) mlir/test/lib/Dialect/Test/TestOps.td (+9)
• (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+9)

diff --git a/mlir/test/IR/properties.mlir b/mlir/test/IR/properties.mlir
index 9a1c49cb7dabf3..b339a03812badb 100644
--- a/mlir/test/IR/properties.mlir
+++ b/mlir/test/IR/properties.mlir
@@ -1,4 +1,4 @@
-// # RUN: mlir-opt %s -split-input-file | mlir-opt  |FileCheck %s
+// # RUN: mlir-opt %s -split-input-file | mlir-opt | FileCheck %s
 // # RUN: mlir-opt %s -mlir-print-op-generic -split-input-file  | mlir-opt -mlir-print-op-generic | FileCheck %s --check-prefix=GENERIC
 
 // CHECK:   test.with_properties
@@ -38,6 +38,14 @@ test.using_property_in_custom [1, 4, 20]
 // GENERIC-SAME: }>
 test.using_property_ref_in_custom 1 + 4 = 5
 
+// Tests that the variadic segment size properties are elided.
+// CHECK: %[[CI64:.*]] = arith.constant
+// CHECK-NEXT: test.variadic_segment_prop %[[CI64]], %[[CI64]] : %[[CI64]] : i64, i64 : i64 end
+// GENERIC: %[[CI64:.*]] = "arith.constant"()
+// GENERIC-NEXT: "test.variadic_segment_prop"(%[[CI64]], %[[CI64]], %[[CI64]]) <{operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>}> : (i64, i64, i64) -> (i64, i64, i64)
+%ci64 = arith.constant 0 : i64
+test.variadic_segment_prop %ci64, %ci64 : %ci64 : i64, i64 : i64 end
+
 // 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 cfe19a2fd5c08b..6752113cab8d41 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3047,6 +3047,15 @@ def TestOpUsingPropertyInCustomAndOther
   );
 }
 
+def TestOpWithVariadicSegmentProperties : TEST_Op<"variadic_segment_prop",
+    [AttrSizedOperandSegments, AttrSizedResultSegments]> {
+  let arguments = (ins Variadic<I64>:$a1, Variadic<I64>:$a2);
+  let results = (outs Variadic<I64>:$b1, Variadic<I64>:$b2);
+  let assemblyFormat = [{
+    $a1 `:` $a2 `:` type($b1) `:` type($b2) prop-dict attr-dict `end`
+  }];
+}
+
 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/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 7e2b0694a860a3..121e7fb4309805 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -2012,6 +2012,15 @@ static void genNonDefaultValueCheck(MethodBody &body, const Operator &op,
 static void genPropDictPrinter(OperationFormat &fmt, Operator &op,
                                MethodBody &body) {
   body << "  ::llvm::SmallVector<::llvm::StringRef, 2> elidedProps;\n";
+
+  // Elide the variadic segment size properties if necessary.
+  if (!fmt.allOperands &&
+      op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments"))
+    body << "  elidedProps.push_back(\"operandSegmentSizes\");\n";
+  if (!fmt.allResultTypes &&
+      op.getTrait("::mlir::OpTrait::AttrSizedResultSegments"))
+    body << "  elidedProps.push_back(\"resultSegmentSizes\");\n";
+
   for (const NamedProperty *namedProperty : fmt.usedProperties)
     body << "  elidedProps.push_back(\"" << namedProperty->name << "\");\n";
   for (const NamedAttribute *namedAttr : fmt.usedAttributes)

[zero9178]

LGTM!