diff --git a/mlir/test/IR/properties.mlir b/mlir/test/IR/properties.mlir index 9a1c49cb7dabf..b339a03812bad 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, resultSegmentSizes = array}> : (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 cfe19a2fd5c08..6752113cab8d4 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:$a1, Variadic:$a2); + let results = (outs Variadic:$b1, Variadic:$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($first) `+` custom($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 7e2b0694a860a..8d2e15a941370 100644 --- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp +++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp @@ -2008,10 +2008,26 @@ static void genNonDefaultValueCheck(MethodBody &body, const Operator &op, << "() != " << propElement.getVar()->prop.getDefaultValue(); } +/// Elide the variadic segment size attributes if necessary. +/// This pushes elided attribute names in `elidedStorage`. +static void genVariadicSegmentElision(OperationFormat &fmt, Operator &op, + MethodBody &body, + const char *elidedStorage) { + if (!fmt.allOperands && + op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) + body << " " << elidedStorage << ".push_back(\"operandSegmentSizes\");\n"; + if (!fmt.allResultTypes && + op.getTrait("::mlir::OpTrait::AttrSizedResultSegments")) + body << " " << elidedStorage << ".push_back(\"resultSegmentSizes\");\n"; +} + /// Generate the printer for the 'prop-dict' directive. static void genPropDictPrinter(OperationFormat &fmt, Operator &op, MethodBody &body) { body << " ::llvm::SmallVector<::llvm::StringRef, 2> elidedProps;\n"; + + genVariadicSegmentElision(fmt, op, body, "elidedProps"); + for (const NamedProperty *namedProperty : fmt.usedProperties) body << " elidedProps.push_back(\"" << namedProperty->name << "\");\n"; for (const NamedAttribute *namedAttr : fmt.usedAttributes) @@ -2057,13 +2073,9 @@ static void genPropDictPrinter(OperationFormat &fmt, Operator &op, static void genAttrDictPrinter(OperationFormat &fmt, Operator &op, MethodBody &body, bool withKeyword) { body << " ::llvm::SmallVector<::llvm::StringRef, 2> elidedAttrs;\n"; - // Elide the variadic segment size attributes if necessary. - if (!fmt.allOperands && - op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) - body << " elidedAttrs.push_back(\"operandSegmentSizes\");\n"; - if (!fmt.allResultTypes && - op.getTrait("::mlir::OpTrait::AttrSizedResultSegments")) - body << " elidedAttrs.push_back(\"resultSegmentSizes\");\n"; + + genVariadicSegmentElision(fmt, op, body, "elidedAttrs"); + for (const StringRef key : fmt.inferredAttributes.keys()) body << " elidedAttrs.push_back(\"" << key << "\");\n"; for (const NamedAttribute *attr : fmt.usedAttributes)