[mlir][tblgen] Generate builders with inferred return types and unwrapped attributes

[jfurtek]

Summary

This diff causes mlir-tblgen to generate code for an additional builder for an
operation argument with a return type that can be inferred *AND* an attribute in
the argument list can be "unwrapped." (Previously, the unwrapped build function
was only generated for builders with explicit return types in separate or
aggregate form.) As an example, this builder might be used by code that creates
operations that implement the SameOperandsAndResultType interface. A test case
was created.

[jfurtek]

Adding @mehdi_amini and @rriddle to reviewers to either review or suggest a reviewer...

[jpienaar]

Nice, thanks (I need to remove the hard coding on SameElement... there but that's not related here directly)

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1656

if any attribute ? (just to avoid a reading that all needs to be unwrappable)

1690

Is the lambda even needed? Seems like we could just have it inlined here with returns changed to continue and horizontal whitespace reduced. Alternatively I'd avoid the array and do

emit(AttrParamKind::WrappedAttr);
if (canGenerateUnwrappedBuilder(op)))

emit(AttrParamKind::UnwrappedValue);

to make it more locally obvious

[jfurtek]

• Improve code flow and rebase

[jfurtek]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1690

I was mimicking the code in genSeparateArgParamBuilder() - but that has a more elaborate emit() scenario. I agree that your suggestion is better - changed in updated commit.

[jfurtek]

• Rebase

mlir/docs/OpDefinitions.md

 }
 ```
 
 The following builders are generated:
 
 ```c++
 // All result-types/operands/attributes have one aggregate parameter.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
-                  ArrayRef<Type> resultTypes,
+                  TypeRange resultTypes,
                   ValueRange operands,
                   ArrayRef<NamedAttribute> attributes);
 
 // Each result-type/operand/attribute has a separate parameter. The parameters
 // for attributes are of mlir::Attribute types.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
                   Type i32_result, Type f32_result, ...,
                   Value i32_operand, Value f32_operand, ...,
                   IntegerAttr i32_attr, FloatAttr f32_attr, ...);
 
 // Each result-type/operand/attribute has a separate parameter. The parameters
 // for attributes are raw values unwrapped with mlir::Attribute instances.
 // (Note that this builder will not always be generated. See the following
 // explanation for more details.)
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
                   Type i32_result, Type f32_result, ...,
                   Value i32_operand, Value f32_operand, ...,
                   APInt i32_attr, StringRef f32_attr, ...);
 
 // Each operand/attribute has a separate parameter but result type is aggregate.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
-                  ArrayRef<Type> resultTypes,
+                  TypeRange resultTypes,
                   Value i32_operand, Value f32_operand, ...,
                   IntegerAttr i32_attr, FloatAttr f32_attr, ...);
 
 // All operands/attributes have aggregate parameters.
 // Generated if return type can be inferred.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
                   ValueRange operands, ArrayRef<NamedAttribute> attributes);
 

mlir/test/mlir-tblgen/op-decl-and-defs.td

 // CHECK:   ::mlir::IntegerAttr getAttr1Attr()
 // CHECK:   uint32_t getAttr1();
 // CHECK:   ::mlir::FloatAttr getAttr2Attr()
 // CHECK:   ::llvm::Optional< ::llvm::APFloat > getAttr2();
 // CHECK:   ::mlir::Attribute removeAttr2Attr();
 // CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value val);
 // CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, int integer = 0);
 // CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::TypeRange s, ::mlir::Value a, ::mlir::ValueRange b, ::mlir::IntegerAttr attr1, /*optional*/::mlir::FloatAttr attr2, unsigned someRegionsCount)
+// CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::ValueRange b, ::mlir::IntegerAttr attr1, /*optional*/::mlir::FloatAttr attr2, unsigned someRegionsCount);
 // CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::TypeRange s, ::mlir::Value a, ::mlir::ValueRange b, uint32_t attr1, /*optional*/::mlir::FloatAttr attr2, unsigned someRegionsCount)
+// CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::ValueRange b, uint32_t attr1, /*optional*/::mlir::FloatAttr attr2, unsigned someRegionsCount);
 // CHECK:   static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes, unsigned numRegions)
 // CHECK:   static ::mlir::ParseResult parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result);
 // CHECK:   void print(::mlir::OpAsmPrinter &p);
 // CHECK:   ::mlir::LogicalResult verifyInvariants();
 // CHECK:   static void getCanonicalizationPatterns(::mlir::RewritePatternSet &results, ::mlir::MLIRContext *context);
 // CHECK:   ::mlir::LogicalResult fold(::llvm::ArrayRef<::mlir::Attribute> operands, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);
 // CHECK:   // Display a graph for debugging purposes.
 // CHECK:   void displayGraph();

 }
 
 def NS_KWithTraitOp : NS_Op<"KWithTrait", [NS_TestTrait]>;
 
 // CHECK-LABEL: NS::KWithTraitOp declarations
 // CHECK:       class KWithTraitOp : public ::mlir::Op<KWithTraitOp
 // CHECK-SAME:      SomeNamespace::TestTrait
 
+def NS_LOp : NS_Op<"op_with_same_operands_and_result_types_unwrapped_attr", [SameOperandsAndResultType]> {
+  let arguments = (ins AnyType:$a, AnyType:$b, I32Attr:$attr1);
+  let results = (outs AnyType:$r);
+}
+// CHECK_LABEL: class NS_LOp :
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
+// CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, ::mlir::IntegerAttr attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b, uint32_t attr1);
+// CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
+
+
 // Test that type defs have the proper namespaces when used as a constraint.
 // ---
 
 def Test_Dialect2 : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::dialect2";
 }
 def TestDialect2Type : TypeDef<Test_Dialect2, "Dialect2Type">;

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

void OpEmitter::genInferredTypeCollectiveParamBuilder() {
   body << formatv(R"(
   } else {{
     ::llvm::report_fatal_error("Failed to infer result type(s).");
   })",
                   opClass.getClassName(), builderOpState);
 }
 
 void OpEmitter::genUseOperandAsResultTypeSeparateParamBuilder() {
+  SmallVector<AttrParamKind, 2> attrBuilderType;
+  attrBuilderType.push_back(AttrParamKind::WrappedAttr);
+  // Generate additional builder(s) if attributes can be "unwrapped"
+  if (canGenerateUnwrappedBuilder(op))
+    attrBuilderType.push_back(AttrParamKind::UnwrappedValue);
+
+  auto emit = [&](AttrParamKind attrType) {
     SmallVector<MethodParameter> paramList;
     SmallVector<std::string, 4> resultNames;
     llvm::StringSet<> inferredAttributes;
     buildParamList(paramList, inferredAttributes, resultNames,
-                 TypeParamKind::None);
+                   TypeParamKind::None, attrType);
 
     auto *m = opClass.addStaticMethod("void", "build", std::move(paramList));
     // If the builder is redundant, skip generating the method
     if (!m)
       return;
     auto &body = m->body();
-  genCodeForAddingArgAndRegionForBuilder(body, inferredAttributes);
+    genCodeForAddingArgAndRegionForBuilder(body, inferredAttributes,
+                                           /*isRawValueAttr=*/attrType ==
+                                               AttrParamKind::UnwrappedValue);
 
     auto numResults = op.getNumResults();
     if (numResults == 0)
       return;
 
     // Push all result types to the operation state
     const char *index = op.getOperand(0).isVariadic() ? ".front()" : "";
     std::string resultType =
         formatv("{0}{1}.getType()", getArgumentName(op, 0), index).str();
     body << "  " << builderOpState << ".addTypes({" << resultType;
     for (int i = 1; i != numResults; ++i)
       body << ", " << resultType;
     body << "});\n\n";
+  };
+
+  for (auto attrType : attrBuilderType)
+    emit(attrType);
 }
 
 void OpEmitter::genUseAttrAsResultTypeBuilder() {
   SmallVector<MethodParameter> paramList;
   paramList.emplace_back("::mlir::OpBuilder &", "odsBuilder");
   paramList.emplace_back("::mlir::OperationState &", builderOpState);
   paramList.emplace_back("::mlir::ValueRange", "operands");
   paramList.emplace_back("::llvm::ArrayRef<::mlir::NamedAttribute>",