[mlir][ods] Make OpBuilder and OperationState optional

[jpienaar]

Summary

The OpBuilder is required to start with OpBuilder and OperationState, so remove
the need for the user to specify it. To make it simpler to update callers,
retain the legacy behavior for now and skip injecting OpBuilder/OperationState
when params start with OpBuilder.

Related to bug 47442.

[lattner]

Super awesome -- and a great way to stage this in. Thank you for taking this!

mlir/docs/OpDefinitions.md

 
 #### Custom builder methods
 
 However, if the above cases cannot satisfy all needs, you can define additional
 convenience build methods with `OpBuilder`.
 
 `OpBuilder` is a class that takes the parameter list and the optional `build()`
 method body. They are separated because we need to generate op declaration and
-definition into separate files. The parameter list should _include_ `Builder
-*builder, OperationState &state`. If the `body` is not provided, _only_ the
-builder declaration will be generated; this provides a way to define complicated
-builders entirely in C++ files.
+definition into separate files. The parameter list should not include `OpBuilder
+&builder, OperationState &state` as they will be inserted automatically and the
+placeholders `$_builder` and `$_state` used. For legacy/to be deprecated reason
+if the `OpBuilder` parameter starts with `OpBuilder` param, then the parameter
+is used. If the `body` is not provided, only the builder declaration will be
+generated; this provides a way to define complicated builders entirely in C++
+files.
 
 For example, for the following op:
 
 ```tablegen
 def MyOp : Op<"my_op", []> {
   let arguments = (ins F32Attr:$attr);
 
   let results = (outs);
 }
 ```
 
 If we want to define a builder with a default value for the only attribute, we
 can add into `MyOp`:
 
 ```tablegen
 def MyOp : ... {
   ...
 
   let builders = [
-    OpBuilder<"OpBuilder &builder, OperationState &state, float val = 0.5f", [{
-      state.addAttribute("attr", builder.getF32FloatAttr(val));
+    OpBuilder<"float val = 0.5f", [{
+      $_state.addAttribute("attr", $_builder.getF32FloatAttr(val));
     }]>
   ];
 }
 ```
 
 The generated builder will look like:
 
 ```c++

mlir/test/lib/Dialect/Test/TestOps.td

let results = (outs
     Variadic<I32>:$output2,
     Variadic<I32>:$output3
   );
 
   // We will use this op in a nested result pattern, where we cannot deduce the
   // result type. So need to provide a builder not requiring result types.
   let builders = [
     OpBuilder<
-      "OpBuilder &builder, OperationState &state, IntegerAttr count",
+      "IntegerAttr count",
       [{
-        auto i32Type = builder.getIntegerType(32);
-        state.addTypes(i32Type); // $output1
+        auto i32Type = $_builder.getIntegerType(32);
+        $_state.addTypes(i32Type); // $output1
         SmallVector<Type, 4> types(count.getInt(), i32Type);
-        state.addTypes(types); // $output2
-        state.addTypes(types); // $output3
-        state.addAttribute("count", count);
+        $_state.addTypes(types); // $output2
+        $_state.addTypes(types); // $output3
+        $_state.addAttribute("count", count);
       }]>
   ];
 }
 
 // Generates an op with variadic results using nested pattern.
 def : Pat<(OneI32ResultOp),
           (MixedVOperandOp3
               (MixedVResultOp3:$results__0 ConstantAttr<I32Attr, "2">),

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

 // CHECK:   ::mlir::Operation::result_range getODSResults(unsigned index);
 // CHECK:   ::mlir::Value r();
 // CHECK:   ::mlir::Region &someRegion();
 // CHECK:   ::mlir::MutableArrayRef<Region> someRegions();
 // CHECK:   ::mlir::IntegerAttr attr1Attr()
 // CHECK:   uint32_t attr1();
 // CHECK:   ::mlir::FloatAttr attr2Attr()
 // CHECK:   ::llvm::Optional< ::llvm::APFloat > attr2();
-// CHECK:   static void build(Value val);
+// CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value val);
 // CHECK:   static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::llvm::ArrayRef<::mlir::Type> 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::Type r, ::llvm::ArrayRef<::mlir::Type> s, ::mlir::Value a, ::mlir::ValueRange b, uint32_t attr1, /*optional*/::mlir::FloatAttr attr2, unsigned someRegionsCount)
 // CHECK:   static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::llvm::ArrayRef<::mlir::Type> 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 verify();
 // CHECK:   static void getCanonicalizationPatterns(::mlir::OwningRewritePatternList &results, ::mlir::MLIRContext *context);
 // CHECK:   ::mlir::LogicalResult fold(::llvm::ArrayRef<::mlir::Attribute> operands, ::llvm::SmallVectorImpl<::mlir::OpFoldResult> &results);

 def NS_SkipDefaultBuildersOp : NS_Op<"skip_default_builders", []> {
   let skipDefaultBuilders = 1;
   let builders = [OpBuilder<"Value val">];
 }
 
 // CHECK-LABEL: NS::SkipDefaultBuildersOp declarations
 // CHECK:     class SkipDefaultBuildersOp
 // CHECK-NOT:   static void build(::mlir::Builder
-// CHECK:       static void build(Value
+// CHECK:       static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, Value
 
 // Check leading underscore in op name
 // ---
 
 def NS__AOp : NS_Op<"_op_with_leading_underscore", []>;
 
 // CHECK-LABEL: NS::_AOp declarations
 // CHECK: class _AOp : public ::mlir::Op<_AOp

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

static cl::opt<std::string> opIncFilter(
     cl::cat(opDefGenCat));
 static cl::opt<std::string> opExcFilter(
     "op-exclude-regex",
     cl::desc("Regex of name of op's to exclude (no filter if empty)"),
     cl::cat(opDefGenCat));
 
 static const char *const tblgenNamePrefix = "tblgen_";
 static const char *const generatedArgName = "odsArg";
+static const char *const builder = "odsBuilder";
 static const char *const builderOpState = "odsState";
 
 // The logic to calculate the actual value range for a declared operand/result
 // of an op with variadic operands/results. Note that this logic is not for
 // general use; it assumes all variadic operands/results must have the same
 // number of values.
 //
 // {0}: The list of whether each declared operand/result is variadic.
void OpEmitter::genBuilder() {
   // Handle custom builders if provided.
   // TODO: Create wrapper class for OpBuilder to hide the native
   // TableGen API calls here.
   {
     auto *listInit = dyn_cast_or_null<ListInit>(def.getValueInit("builders"));
     if (listInit) {
       for (Init *init : listInit->getValues()) {
         Record *builderDef = cast<DefInit>(init)->getDef();
-        StringRef params = builderDef->getValueAsString("params");
+        StringRef params = builderDef->getValueAsString("params").trim();
+        // TODO: Remove this and just generate the builder/state always.
+        bool skipParamGen = params.startswith("OpBuilder") ||
+                            params.startswith("mlir::OpBuilder") ||
+                            params.startswith("::mlir::OpBuilder");
         StringRef body = builderDef->getValueAsString("body");
         bool hasBody = !body.empty();
 
         OpMethod::Property properties =
             hasBody ? OpMethod::MP_Static : OpMethod::MP_StaticDeclaration;
+        std::string paramStr =
+            skipParamGen ? params.str()
+                         : llvm::formatv("::mlir::OpBuilder &{0}, "
+                                         "::mlir::OperationState &{1}, {2}",
+                                         builder, builderOpState, params)
+                               .str();
         auto *method =
-            opClass.addMethodAndPrune("void", "build", properties, params);
-        if (hasBody)
+            opClass.addMethodAndPrune("void", "build", properties, paramStr);
+        if (hasBody) {
+          if (skipParamGen) {
             method->body() << body;
+          } else {
+            FmtContext fctx;
+            fctx.withBuilder(builder);
+            fctx.addSubst("_state", builderOpState);
+            method->body() << tgfmt(body, &fctx);
+          }
+        }
       }
     }
     if (op.skipDefaultBuilders()) {
       if (!listInit || listInit->empty())
         PrintFatalError(
             op.getLoc(),
             "default builders are skipped and no custom builders provided");
       return;