From 95ddc8341ae2c27229ad3dcf1d55abebcec15d02 Mon Sep 17 00:00:00 2001 From: MaheshRavishankar Date: Thu, 19 Aug 2021 10:08:59 -0700 Subject: [PATCH] [mlir][Linalg] Allow all build methods of Structured ops to specify additional attributes. Differential Revision: https://reviews.llvm.org/D108338 --- .../Dialect/Linalg/IR/LinalgStructuredOps.td | 12 ++++++++---- mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 19 ++++++++++++------- .../test-linalg-ods-gen.tc | 3 ++- .../test-linalg-ods-yaml-gen.yaml | 1 + .../mlir-linalg-ods-gen.cpp | 12 +++++++++--- .../mlir-linalg-ods-yaml-gen.cpp | 8 ++++++-- 6 files changed, 38 insertions(+), 17 deletions(-) diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td index 33f4992e41f9..1d4e6d546067 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -620,18 +620,22 @@ def GenericOp : LinalgStructuredBase_Op<"generic", [ "ValueRange":$outputs, "ArrayRef":$indexingMaps, "ArrayRef":$iteratorTypes, "StringRef":$doc, "StringRef":$libraryCall, - CArg<"function_ref", "nullptr">)>, + CArg<"function_ref", "nullptr">, + CArg<"ArrayRef", "{}">:$attributes)>, OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers, "ArrayRef":$indexingMaps, "ArrayRef":$iteratorTypes, "StringRef":$doc, "StringRef":$libraryCall, - CArg<"function_ref", "nullptr">)>, + CArg<"function_ref", "nullptr">, + CArg<"ArrayRef", "{}">:$attributes)>, OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, "ValueRange":$outputs, "ArrayRef":$indexingMaps, "ArrayRef":$iteratorTypes, - CArg<"function_ref", "nullptr">)>, + CArg<"function_ref", "nullptr">, + CArg<"ArrayRef", "{}">:$attributes)>, OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers, "ArrayRef":$indexingMaps, "ArrayRef":$iteratorTypes, - CArg<"function_ref", "nullptr">)> + CArg<"function_ref", "nullptr">, + CArg<"ArrayRef", "{}">:$attributes)> ]; let extraClassDeclaration = structuredOpsBaseDecls # [{ diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp index 6b65a9ecd9e5..f4750ca390a8 100644 --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -502,13 +502,15 @@ void GenericOp::build( OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef indexingMaps, ArrayRef iteratorTypes, StringRef doc, StringRef libraryCall, - function_ref bodyBuild) { + function_ref bodyBuild, + ArrayRef attributes) { build(builder, result, resultTensorTypes, inputs, outputs, builder.getAffineMapArrayAttr(indexingMaps), builder.getStrArrayAttr(iteratorTypes), doc.empty() ? StringAttr() : builder.getStringAttr(doc), libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall)); + result.addAttributes(attributes); if (!bodyBuild) return; @@ -527,30 +529,33 @@ void GenericOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, ValueRange outputs, ArrayRef indexingMaps, ArrayRef iteratorTypes, StringRef doc, StringRef libraryCall, - function_ref bodyBuild) { + function_ref bodyBuild, + ArrayRef attributes) { build(builder, result, TypeRange{}, inputs, outputs, indexingMaps, - iteratorTypes, doc, libraryCall, bodyBuild); + iteratorTypes, doc, libraryCall, bodyBuild, attributes); } void GenericOp::build( OpBuilder &builder, OperationState &result, ValueRange inputs, ValueRange outputs, ArrayRef indexingMaps, ArrayRef iteratorTypes, - function_ref bodyBuild) { + function_ref bodyBuild, + ArrayRef attributes) { build(builder, result, inputs, outputs, indexingMaps, iteratorTypes, /*doc=*/"", - /*libraryCall=*/"", bodyBuild); + /*libraryCall=*/"", bodyBuild, attributes); } void GenericOp::build( OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, ValueRange inputs, ValueRange outputs, ArrayRef indexingMaps, ArrayRef iteratorTypes, - function_ref bodyBuild) { + function_ref bodyBuild, + ArrayRef attributes) { build(builder, result, resultTensorTypes, inputs, outputs, indexingMaps, iteratorTypes, /*doc=*/"", - /*libraryCall=*/"", bodyBuild); + /*libraryCall=*/"", bodyBuild, attributes); } static void print(OpAsmPrinter &p, GenericOp op) { diff --git a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc index 471961f837bf..743bdbdb12d6 100644 --- a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc +++ b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc @@ -169,7 +169,8 @@ It has one output. // ODS-LABEL: def Test7Op // ODS: OpBuilder< // ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, -// ODS: "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b) +// ODS: "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b, +// ODS: CArg<"ArrayRef", "{}">:$attributes) // ODS: $_state.addAttribute("attr_a", attr_a); // ODS: $_state.addAttribute("attr_b", attr_b); // diff --git a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml index 6613ab2a006f..f8393a232e43 100644 --- a/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml +++ b/mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml @@ -75,6 +75,7 @@ structured_op: !LinalgStructuredOpConfig # ODS-NEXT: $_builder.getI32VectorAttr({ # ODS-NEXT: static_cast(inputs.size()), # ODS-NEXT: static_cast(outputs.size())})); +# ODS-NEXT: $_state.addAttributes(attributes); # ODS-NEXT: createAndFillStructuredOpRegion( # ODS-NEXT: $_builder, # ODS-NEXT: $_state, diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp index 1bdb5b8806d0..590f17fdedfa 100644 --- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp +++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp @@ -1910,7 +1910,8 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, let skipDefaultBuilders = 1; let builders = [ OpBuilder< - (ins "ValueRange":$inputs, "ValueRange":$outputs), + (ins "ValueRange":$inputs, "ValueRange":$outputs, + CArg<"ArrayRef", "{{}">:$attributes), [{{ $_state.addOperands(inputs); $_state.addOperands(outputs); @@ -1919,6 +1920,7 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, $_builder.getI32VectorAttr({{ static_cast(inputs.size()), static_cast(outputs.size())})); + $_state.addAttributes(attributes); createAndFillStructuredOpRegion<{0}>( $_builder, $_state, @@ -1927,7 +1929,8 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, }]>, OpBuilder< (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, - "ValueRange":$outputs), + "ValueRange":$outputs, + CArg<"ArrayRef", "{{}">:$attributes), [{{ $_state.addOperands(inputs); $_state.addOperands(outputs); @@ -1937,6 +1940,7 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, $_builder.getI32VectorAttr({{ static_cast(inputs.size()), static_cast(outputs.size())})); + $_state.addAttributes(attributes); createAndFillStructuredOpRegion<{0}>( $_builder, $_state, @@ -2020,7 +2024,8 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, const char *builderFmt = R"FMT( , OpBuilder< (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, - "ValueRange":$outputs, {1}), + "ValueRange":$outputs, {1}, + CArg<"ArrayRef", "{{}">:$attributes), [{{ $_state.addOperands(inputs); $_state.addOperands(outputs); @@ -2030,6 +2035,7 @@ void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName, $_builder.getI32VectorAttr({{ static_cast(inputs.size()), static_cast(outputs.size())})); + $_state.addAttributes(attributes); createAndFillStructuredOpRegion<{0}>( $_builder, $_state, diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp index a0eb1dea8860..98e90b69d631 100644 --- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp +++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp @@ -457,7 +457,8 @@ def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([ let skipDefaultBuilders = 1; let builders = [ OpBuilder< - (ins "ValueRange":$inputs, "ValueRange":$outputs), + (ins "ValueRange":$inputs, "ValueRange":$outputs, + CArg<"ArrayRef", "{{}">:$attributes), [{{ $_state.addOperands(inputs); $_state.addOperands(outputs); @@ -471,6 +472,7 @@ def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([ $_builder.getI32VectorAttr({{ static_cast(inputs.size()), static_cast(outputs.size())})); + $_state.addAttributes(attributes); createAndFillStructuredOpRegion<{0}>( $_builder, $_state, @@ -539,7 +541,8 @@ def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([ static const char structuredOpBuilderFormat[] = R"FMT( , OpBuilder< (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, - "ValueRange":$outputs, {1}), + "ValueRange":$outputs, {1}, + CArg<"ArrayRef", "{{}">:$attributes), [{{ $_state.addOperands(inputs); $_state.addOperands(outputs); @@ -555,6 +558,7 @@ static const char structuredOpBuilderFormat[] = R"FMT( TypeRange(inputs), TypeRange(outputs)); {2} + $_state.addAttributes(attributes); }]> )FMT";