[mlir][linalg] Support generating builders for named op attributes

[antiagainst]

Summary

This commit adds support to generate an additional builder for
each named op that has attributes. This gives better experience
when creating the named ops.

Along the way adds support for i64.

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 }
 
 // Test attribute definitions
 // ODS-LABEL: def Test4Op
 // ODS: F32ArrayAttr:$array_attr,
 // ODS: F32:$f32_attr,
 // ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
 // ODS: I32:$i32_attr,
+// ODS: I64:$i64_attr,
 // ODS: RankedI32ElementsAttr<[5, 6]>:$ivec_attr,
 // ODS: OptionalAttr<F32>:$optional_attr
 //
 ods_def<Test4Op> :
 def test4(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
 attr(
   f32_attr: f32,
   i32_attr: i32,
+  i64_attr: i64,
   fvec_attr: 4xf32,
   ivec_attr: 5x6xi32,
   array_attr : f32[],
   optional_attr? : f32
 ) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
 }
 

 //
 ods_def<Test5Op>:
 def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
      attr(strides: 2xi32) {
   O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
     I(n, h * strides[0] + kh, w * strides[1] + kw, c), K(f, kh, kw, c)));
 }
 
+// Test documentation
 // ODS-LABEL: def Test6Op
 // ODS:       let summary = [{ My magic op. }];
 // ODS-NEXT:  let description = [{
 // ODS-NEXT:    It has two inputs.
 // ODS-NEXT:    It has one output.
 // ODS-NEXT:  }];
 //
 ods_def<Test6Op>:
 def test6(A: f32(M, K), B: f32(K)) -> (C: f32(M))
 """
 My magic op.
 
 It has two inputs.
 It has one output.
 """
 {
   C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
 }
+
+// Test attribute builder
+// ODS-LABEL: def Test7Op
+// ODS:         OpBuilderDAG<
+// ODS:           (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+// ODS:            "ValueRange":$outputs, "Attribute":$attr_a, "Attribute":$attr_b)
+// ODS:           $_state.addAttribute("attr_a", attr_a);
+// ODS:           $_state.addAttribute("attr_b", attr_b);
+//
+ods_def<Test7Op>:
+def test7(A: f32(M, K), B: f32(K)) -> (C: f32(M))
+     attr(attr_a: f32, attr_b: 4xi32)
+{
+  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
+}

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName,
   SmallVector<std::string, 4> attributes;
   for (const auto &attr : registeredAttrs) {
     llvm::StringRef name = attr.first;
 
     llvm::StringRef elementType = attr.second.elementType;
     std::string odsType = llvm::StringSwitch<std::string>(elementType)
                               .Case("f32", "F32")
                               .Case("i32", "I32")
+                              .Case("i64", "I64")
                               .Default("");
     if (odsType.empty()) {
       parser.emitError("unimplemented support for attribute element type: " +
                        elementType);
       return;
     }
 
     const auto &dims = attr.second.vectorDims;
SingleBlockImplicitTerminator<"YieldOp">]> {
       let arguments = (ins
         Variadic<AnyShaped>:$inputs,
         Variadic<AnyShaped>:$outputs{3}
       );
       let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
       let regions = (region AnyRegion:$region);
 
       let skipDefaultBuilders = 1;
-      let builders = [ OpBuilderDAG<
+      let builders = [
+        OpBuilderDAG<
         (ins "ValueRange":$inputs, "ValueRange":$outputs),
         [{{
           $_state.addOperands(inputs);
           $_state.addOperands(outputs);
           $_state.addAttribute(
             "operand_segment_sizes",
             $_builder.getI32VectorAttr({{
               static_cast<int32_t>(inputs.size()),
               static_cast<int32_t>(outputs.size())}));
           buildNamedStructuredOpRegionAndAttributes<{0}>(
             $_builder,
             $_state,
             TypeRange(inputs),
             TypeRange(outputs));
-        }]>, OpBuilderDAG<
+        }]>,
+        OpBuilderDAG<
         (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
              "ValueRange":$outputs),
         [{{
           $_state.addOperands(inputs);
           $_state.addOperands(outputs);
           $_state.addTypes(resultTensorTypes);
           $_state.addAttribute(
             "operand_segment_sizes",
             $_builder.getI32VectorAttr({{
               static_cast<int32_t>(inputs.size()),
               static_cast<int32_t>(outputs.size())}));
           buildNamedStructuredOpRegionAndAttributes<{0}>(
             $_builder,
             $_state,
             TypeRange(inputs),
             TypeRange(outputs));
-        }]>, OpBuilderDAG<
+        }]>,
+        OpBuilderDAG<
         (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
              CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
         [{{
           $_state.addOperands(operands);
           $_state.addAttributes(attributes);
           $_state.addTypes(resultTensorTypes);
           (void)$_state.addRegion();
         }]>
+        {5}
       ];
       let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
       let parser = [{{ return ::parseNamedStructuredOp<{0}>(parser, result); }];
       let hasFolder = 1;
       let hasCanonicalizer = 1;
 
       let extraClassDeclaration = [{{
         // Auto-generated.
         ArrayAttr iterator_types();
         ArrayAttr indexing_maps();
         static void regionBuilder(Block &block);
         static std::function<void(Block &)> getRegionBuilder() {{ return regionBuilder; }
 
         // Generic methods.
         static unsigned getNumRegionArgs() {{ return {4}; }
         std::string getLibraryCallName() {{
           return generateLibraryCallName(getOperation());
         }
       }];
   })FMT";
 
+  // Generate documentation.
   std::string doc;
-
   if (!docString.empty()) {
     const char *docFmt = R"FMT(
       let summary = [{ {0} }];
       let description = [{
         {1}
       }];
     )FMT";
 
     StringRef summary, description;
     std::tie(summary, description) = docString.trim().split('\n');
     doc = llvm::formatv(docFmt, summary.trim(), description.trim());
   }
 
+  // Generate an additional builder that has parameters for attributes.
+  std::string attrBuilder;
+  if (!registeredAttrs.empty()) {
+    SmallVector<std::string, 4> attrParams, attrStmts;
+    for (const auto &attr : registeredAttrs) {
+      llvm::StringRef name = attr.first;
+      attrParams.push_back(llvm::formatv("\"Attribute\":${0}", name));
+      attrStmts.push_back(
+          llvm::formatv("$_state.addAttribute(\"{0}\", {0});", name));
+    }
+    std::string attrParamsList = llvm::join(attrParams, ", ");
+    std::string attrStmtsList = llvm::join(attrStmts, "\n");
+
+    const char *builderFmt = R"FMT(
+      , OpBuilderDAG<
+      (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+           "ValueRange":$outputs, {1}),
+      [{{
+        $_state.addOperands(inputs);
+        $_state.addOperands(outputs);
+        $_state.addTypes(resultTensorTypes);
+        $_state.addAttribute(
+          "operand_segment_sizes",
+          $_builder.getI32VectorAttr({{
+            static_cast<int32_t>(inputs.size()),
+            static_cast<int32_t>(outputs.size())}));
+        buildNamedStructuredOpRegionAndAttributes<{0}>(
+          $_builder,
+          $_state,
+          TypeRange(inputs),
+          TypeRange(outputs));
+        {2}
+      }]>
+    )FMT";
+    attrBuilder =
+        llvm::formatv(builderFmt, cppOpName, attrParamsList, attrStmtsList);
+  }
+
+  // Finally put everything together.
   os << llvm::formatv(header, cppOpName, linalgOpName, doc, attrList,
-                      state.orderedTensorArgs.size());
+                      state.orderedTensorArgs.size(), attrBuilder);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
 void TCParser::printReferenceIterators(llvm::raw_ostream &os,
                                        StringRef cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(
if (!vectorDims.empty()) {
     SmallVector<std::string, 4> indexStrs;
     for (uint64_t index : indices)
       indexStrs.push_back(std::to_string(index));
     std::string indexList = llvm::join(indexStrs, ", ");
     if (elementType == "f32")
       return llvm::formatv("getValue<float>({ {0} })", indexList);
     if (elementType == "i32")
       return llvm::formatv("getValue<int>({ {0} })", indexList);
+    if (elementType == "i64")
+      return llvm::formatv("getValue<int64_t>({ {0} })", indexList);
 
     return "";
   }
 
   if (elementType == "f32")
     return "getValue().convertToFloat()";
-  if (elementType == "i32")
+  if (elementType == "i32" || elementType == "i64")
     return "getInt()";
   return "";
 }
 
 /// Iterate over each Tensor Comprehension def.
 LogicalResult parseAndEmitAllTensorComprehensions(llvm::raw_ostream &os,
                                                   Parser &parser) {
   while (parser.curToken.getKind() != Token::Kind::eof) {