[Mlir-commits] [mlir] f072942 - [mlir] ODS: support operations with resizable operand lists

[Alex Zinenko]

Author: Alex Zinenko
Date: 2020-04-16T23:30:34+02:00
New Revision: f072942fe2f94ba19482e62427e89864fb875782

URL: https://github.com/llvm/llvm-project/commit/f072942fe2f94ba19482e62427e89864fb875782
DIFF: https://github.com/llvm/llvm-project/commit/f072942fe2f94ba19482e62427e89864fb875782.diff

LOG: [mlir] ODS: support operations with resizable operand lists

MLIR supports operations with resizable operand lists, but this property must
be indicated during the construction of such operations. It can be done
programmatically by calling a function on OperationState. Introduce an
ODS-internal trait `ResizableOperandList` to indicate such operations are use
it when generating the bodies of various `build` functions as well as the
`parse` function when the declarative assembly format is used.

Differential Revision: https://reviews.llvm.org/D78292

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/include/mlir/TableGen/Operator.h
    mlir/lib/TableGen/Operator.cpp
    mlir/test/mlir-tblgen/op-operand.td
    mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
    mlir/tools/mlir-tblgen/OpFormatGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 14ef45ed4c7f..282267daf339 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1653,6 +1653,10 @@ class HasParent<string op>
 def FirstAttrDerivedResultType :
   GenInternalOpTrait<"FirstAttrDerivedResultType">;
 
+// Op has a resizable operand list. Auto-generated build and parse functions
+// should construct it as such.
+def ResizableOperandList : GenInternalOpTrait<"ResizableOperandList">;
+
 // TODO(antiagainst): Turn the following into normal traits and generate
 // verification for them.
 

diff  --git a/mlir/include/mlir/TableGen/Operator.h b/mlir/include/mlir/TableGen/Operator.h
index e65bc55a84f5..91f10bd8e29c 100644
--- a/mlir/include/mlir/TableGen/Operator.h
+++ b/mlir/include/mlir/TableGen/Operator.h
@@ -165,6 +165,9 @@ class Operator {
   // requiring the raw MLIR trait here.
   const OpTrait *getTrait(llvm::StringRef trait) const;
 
+  // Returns "true" if Op has a ResizableOperandList trait.
+  bool hasResizableOperandList() const;
+
   // Regions.
   using const_region_iterator = const NamedRegion *;
   const_region_iterator region_begin() const;

diff  --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index 808ba7aabc76..f967b76d074f 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -169,6 +169,10 @@ const tblgen::OpTrait *tblgen::Operator::getTrait(StringRef trait) const {
   return nullptr;
 }
 
+bool tblgen::Operator::hasResizableOperandList() const {
+  return getTrait("OpTrait::ResizableOperandList") != nullptr;
+}
+
 auto tblgen::Operator::region_begin() const -> const_region_iterator {
   return regions.begin();
 }

diff  --git a/mlir/test/mlir-tblgen/op-operand.td b/mlir/test/mlir-tblgen/op-operand.td
index 2ffde33c5331..42aadfbdc119 100644
--- a/mlir/test/mlir-tblgen/op-operand.td
+++ b/mlir/test/mlir-tblgen/op-operand.td
@@ -58,3 +58,23 @@ def OpD : NS_Op<"mix_variadic_and_normal_inputs_op", [SameVariadicOperandSize]>
 // CHECK-NEXT: odsState.addOperands(input1);
 // CHECK-NEXT: odsState.addOperands(input2);
 // CHECK-NEXT: odsState.addOperands(input3);
+// CHECK-NOT: odsState.setOperandListToResizable
+
+// Check that resizable operand list flag is set up correctly in all generated
+// builders and in the parser.
+def OpE : NS_Op<"resizable_operand_list", [ResizableOperandList]> {
+  let arguments = (ins Variadic<AnyType>:$input);
+  let assemblyFormat = "$input attr-dict `:` type($input)";
+}
+
+// CHECK-LABEL: OpE::build(Builder *odsBuilder, OperationState &odsState, ValueRange
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::build(Builder *odsBuilder, OperationState &odsState, ArrayRef<Type>
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::build(Builder *, OperationState
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::parse
+// CHECK: result.setOperandListToResizable()

diff  --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 75edc1b7d256..41f392e67f62 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -685,6 +685,7 @@ void OpEmitter::genSeparateArgParamBuilder() {
     auto &m =
         opClass.newMethod("void", "build", paramList, OpMethod::MP_Static);
     auto &body = m.body();
+
     genCodeForAddingArgAndRegionForBuilder(
         body, /*isRawValueAttr=*/attrType == AttrParamKind::UnwrappedValue);
 
@@ -762,7 +763,9 @@ void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
   auto &body = m.body();
 
   // Operands
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
 
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
@@ -843,7 +846,10 @@ void OpEmitter::genUseAttrAsResultTypeBuilder() {
   }
 
   // Operands
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
+
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
 
@@ -929,7 +935,9 @@ void OpEmitter::genCollectiveParamBuilder() {
          << (numVariadicOperands != 0 ? " >= " : " == ")
          << numNonVariadicOperands
          << "u && \"mismatched number of parameters\");\n";
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
 
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
@@ -1099,6 +1107,8 @@ void OpEmitter::genCodeForAddingArgAndRegionForBuilder(OpMethodBody &body,
       body << "  if (" << argName << ")\n  ";
     body << "  " << builderOpState << ".addOperands(" << argName << ");\n";
   }
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n", builderOpState);
 
   // If the operation has the operand segment size attribute, add it here.
   if (op.getTrait("OpTrait::AttrSizedOperandSegments")) {

diff  --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index a8116e4290b4..f76a2d3af9e8 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -706,6 +706,10 @@ void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   genParserSuccessorResolution(op, body);
   genParserVariadicSegmentResolution(op, body);
 
+  // Mark the operation as having resizable operand list if required.
+  if (op.hasResizableOperandList())
+    body << "  result.setOperandListToResizable();\n";
+
   body << "  return success();\n";
 }