[mlir][linalg] Start a named ops to generic ops pass

[antiagainst]

Summary

This commit starts a new pass and patterns for converting Linalg
named ops to generic ops. This enables us to leverage the flexbility
from generic ops during transformations. Right now only linalg.conv
is supported; others will be added when useful.

[nicolasvasilache]

Great, thanks for starting this Lei!

It is easy to support any LinalgOp (i.e. that implements LinalgStructuredOpInterface) with one of the MatchAnyTag patterns.
See e.g. https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp#L200.

If you just added that you could have linalg.matmul generate automatically and then all the named ops defined in the .tc file.
The remaining ops in LinalgStructuredOps.td would be for future work but ideally the .tc lang + parser evolves to support all cases automatically :)

Since you already wrote 90% of the logic in LinalgGeneralizationPattern, could I ask you to also write the pattern that matches any LinalgOp?

Thanks much!

[mravishankar]

Same comment as Nicolas. Seems like it is easily generalizable to all Named ops.

mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp

88

Each NamedOp seems to have a region builder that should give you this body?

[antiagainst]

Thanks for the comment. I looked into it a bit and feel this might not be as trivial as expected. Or maybe it's just that I didn't find the proper API for it.

We have two categories of named ops right now: 1) tc generated vs 2) manually written. For 1) they have regionBuilder but not for 2). So in order to handle all 1) cases, I need to have a way to differentiate between 1) and 2). (Long-term we want to convert all 2) into 1) but I think that's a separate task on itself.) But AFAICT, there isn't an op interface or something like that to allow me to get all 1) ops in a generic way so I can access the regionBuilder. So then I'm not sure how the MatchAnyTag and then filter for 1) works. The way I can think of is actually use template for each of the 1) ops. But there is another issue. I don't have a TableGen'generated list that I can just #include to get the list of ops in 1): the TC generated TD file is included in the main TD file. This probably is just a small issue as I can separate them out or manually instantiate for all 1) ops. But I wanted to know whether this is the approach we want to choose: using template patterns for all the tc generated ops. Please let me know.

[nicolasvasilache]

In D91357#2400928, @antiagainst wrote:

Thanks for the comment. I looked into it a bit and feel this might not be as trivial as expected. Or maybe it's just that I didn't find the proper API for it.

We have two categories of named ops right now: 1) tc generated vs 2) manually written. For 1) they have regionBuilder but not for 2). So in order to handle all 1) cases, I need to have a way to differentiate between 1) and 2). (Long-term we want to convert all 2) into 1) but I think that's a separate task on itself.) But AFAICT, there isn't an op interface or something like that to allow me to get all 1) ops in a generic way so I can access the regionBuilder. So then I'm not sure how the MatchAnyTag and then filter for 1) works. The way I can think of is actually use template for each of the 1) ops. But there is another issue. I don't have a TableGen'generated list that I can just #include to get the list of ops in 1): the TC generated TD file is included in the main TD file. This probably is just a small issue as I can separate them out or manually instantiate for all 1) ops. But I wanted to know whether this is the approach we want to choose: using template patterns for all the tc generated ops. Please let me know.

Yes I was just thinking of tackling 1) and not worry about 2).
The regionBuilder is currently auto-generated here https://github.com/llvm/llvm-project/blob/ecca7852d9d75aba859a3d8d001bfb2dda1345db/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp#L1524 and does not yet use interfaces.
The idea would be to expose that (maybe to the NamedStructuredOpTrait (https://github.com/llvm/llvm-project/blob/73ca690df88ad5df77afe29b984720963ee37183/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h#L95) or otherwise create a new interface if it is better).

I actually spent some effort removing all places where we used to rely on list of ops so I'd love it if we did not reintroduce them :)

[antiagainst]

• [mlir][linalg] Start a named ops to generic ops pass
• Add another pattern to handle all tc-generated named ops

[antiagainst]

Okay I created another pattern for tc-generated named ops. PTAL.

[nicolasvasilache]

Looks good, thanks Lei!
Approving conditioned on an extra tensor test.

mlir/test/Dialect/Linalg/generalize-named-ops.mlir

55

Add one test for matmul with tensors please, your code seems to handle the case.

[antiagainst]

• Add a test for tensor matmul

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

ArrayAttr indexing_maps() {
       unsigned outputRank = getOutputShapedType(0).getRank();
       return Builder(getContext()).getAffineMapArrayAttr({
           extractOrIdentityMap(maybeInputMap, inputRank, context),
           extractOrIdentityMap(maybeOutputMap, outputRank, context)});
     }
 
     Value getSource() { return input();}
     Value getTarget() { return output(); }
+
+    static std::function<void(Block &)> getRegionBuilder() {
+      return nullptr;
+    }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", [
let extraClassDeclaration = libraryCallName # [{
     }
 
     ArrayAttr indexing_maps() {
       MLIRContext *context = getContext();
       // filling_value -> O(ivs)
       return Builder(getContext()).getAffineMapArrayAttr({
           extractOrIdentityMap(llvm::None, getNumParallelLoops(), context)});
     }
+
+    static std::function<void(Block &)> getRegionBuilder() {
+      return nullptr;
+    }
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
int64_t getLowPad(unsigned i) {
       return padding().getValue().getValue<int64_t>({i, 0});
     }
 
     int64_t getHighPad(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!padding().hasValue()) return 0;
       return padding().getValue().getValue<int64_t>({i, 1});
     }
+
+    static std::function<void(Block &)> getRegionBuilder() {
+      return nullptr;
+    }
   }];
 }
 
 def ConvOp : PoolingBase_Op<"conv", [
     NInputs<2>,
     // Despite having reductions, this manually defined ConvOp may only take
     // memref operands and can never have init tensors.
     ZeroInitTensors,
std::string getLibraryCallName() {
       return library_call().hasValue() ?
         library_call()->str() : "op_has_no_registered_library_name";
     }
     llvm::Optional<unsigned> getSymbolSource() {
       auto ss = symbol_source();
       return ss.hasValue() ?
         llvm::Optional<unsigned>(ss.getValue()) : llvm::None;
     }
+
+    static std::function<void(Block &)> getRegionBuilder() {
+      return nullptr;
+    }
   }];
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parseGenericOp(parser, result); }];
 }
 
 /// Index-free GenericOp.
 def GenericOp : GenericOpBase<"generic"> {
   let description = [{

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

InterfaceMethod<
         unsigned numRegions = $_op.getOperation()->getNumRegions();
         Operation *res = create(b, loc, resultTypes, operands, $_op.getAttrs());
         assert(res->getNumRegions() == numRegions && "inconsistent # regions");
         for (unsigned ridx = 0; ridx < numRegions; ++ridx)
           $_op.getOperation()->getRegion(ridx).cloneInto(
             &res->getRegion(ridx), map);
         return res;
       }]
+    >,
+    StaticInterfaceMethod<
+      /*desc=*/[{
+        Returns the region builder for constructing the body for linalg.generic.
+        Returns a null function if this named op does not define a region
+        builder.
+      }],
+      /*retTy=*/"std::function<void(Block &)>",
+      /*methodName=*/"getRegionBuilder",
+      (ins),
+      [{ return ConcreteOp::getRegionBuilder(); }]
     >
   ];
 
   let extraClassDeclaration = [{
     /// Returns all the operands past the inputs, output_buffers and
     /// init_tensors operands. Asserts that these operands are value types to
     /// allow transformations like tiling to just use the values when cloning
     /// `linalgOp`.

mlir/include/mlir/Dialect/Linalg/Passes.h

 /// buffers instead.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass();
 
 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
 /// parallel loops.
 void populateElementwiseToLinalgConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx);
 
+/// Create a pass to conver named Linalg operations to Linalg generic
+/// operations.
+std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();
+
 /// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
 /// producer (consumer) generic operation by expanding the dimensionality of the
 /// loop in the generic op.
 void populateFoldReshapeOpsByExpansionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns);
 
 /// Patterns to fold a collapsing (expanding) tensor_reshape operation with its
 /// producer (consumer) generic/indexed_generic operation by linearizing the

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgTilingToParallelLoops
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t",
                "Test generation of dynamic promoted buffers",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
   let dependentDialects = ["AffineDialect", "linalg::LinalgDialect", "scf::SCFDialect"];
 }
 
+def LinalgGeneralization : FunctionPass<"linalg-generalize-named-ops"> {
+  let summary = "Convert named ops into generic ops";
+  let constructor = "mlir::createLinalgGeneralizationPass()";
+  let dependentDialects = ["linalg::LinalgDialect"];
+}
+
 #endif // MLIR_DIALECT_LINALG_PASSES

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgMarker marker;
   /// Controls whether the pattern lowers to library calls, scf.for, affine.for
   /// or scf.parallel.
   LinalgLoweringType loweringType;
 };
 
+/// Linalg generalization patterns
+
+/// Populates `patterns` with patterns to convert spec-generated named ops to
+/// linalg.generic ops.
+void populateLinalgNamedOpsGeneralizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    LinalgMarker marker = LinalgMarker());
+
+/// Populates `patterns` with patterns to convert linalg.conv ops to
+/// linalg.generic ops.
+void populateLinalgConvGeneralizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    LinalgMarker marker = LinalgMarker());
+
 //===----------------------------------------------------------------------===//
 // Op-specific patterns.
 //===----------------------------------------------------------------------===//
 /// Match and rewrite for the pattern:
 /// ```
 ///    %alloc = ...
 ///    [optional] %view = std.view %alloc ...
 ///    %subView = subview %allocOrView ...

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
   Bufferize.cpp
   CodegenStrategy.cpp
   DropUnitDims.cpp
   ElementwiseToLinalg.cpp
   Fusion.cpp
   FusionOnTensors.cpp
+  Generalization.cpp
   Hoisting.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
   Sparsification.cpp
   Tiling.cpp
   Transforms.cpp
   Vectorization.cpp

mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp

+//===- Generalization.cpp - linalg named ops to generic ops  --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the Linalg generalization pass. It converts named
+// Linalg ops to linalg.generic ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/EDSC/Builders.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "linalg-generalization"
+
+using namespace mlir;
+
+// Creates a linalg.generic op from the given `namedOp`. Returns a null op if
+// the given `namedOp` does not have a region builder.
+static linalg::GenericOp createGenericOpFromNamedOp(linalg::LinalgOp namedOp,
+                                                    OpBuilder &builder) {
+  auto regionBuilder = namedOp.getRegionBuilder();
+  if (!regionBuilder) {
+    LLVM_DEBUG(llvm::dbgs() << "no region builder for op: " << namedOp << "\n");
+    return nullptr;
+  }
+
+  SmallVector<AffineMap, 4> indexingMaps = namedOp.getIndexingMaps();
+  auto iterators = llvm::to_vector<4>(
+      namedOp.iterator_types().getAsValueRange<StringAttr>());
+  auto resultTypes = namedOp.getOutputTensorTypes();
+  SmallVector<Type, 4> types(resultTypes.begin(), resultTypes.end());
+
+  return builder.create<linalg::GenericOp>(
+      namedOp.getLoc(), types, namedOp.getInputs(), namedOp.getOutputBuffers(),
+      namedOp.getInitTensors(), indexingMaps, iterators,
+      [&regionBuilder](OpBuilder &bodyBuilder, Location loc, ValueRange) {
+        edsc::ScopedContext scope(bodyBuilder, loc);
+        regionBuilder(*bodyBuilder.getBlock());
+      });
+}
+
+namespace {
+
+/// Base class for all linalg generalization patterns. A subclass must provide
+/// the following method:
+///   linalg::GenericOp createGenericOp(RootOp, PatternRewriter &)
+/// for creating the generic op.
+// TODO: remove this pattern after migrating all manually-written named ops
+// into auto-generated ones.
+template <typename ConcretePattern, typename RootOp>
+struct LinalgGeneralizationPattern : OpRewritePattern<RootOp> {
+  LinalgGeneralizationPattern(MLIRContext *context, linalg::LinalgMarker marker,
+                              PatternBenefit benefit = 1)
+      : OpRewritePattern<RootOp>(context, benefit), marker(std::move(marker)) {}
+
+  LogicalResult matchAndRewrite(RootOp rootOp,
+                                PatternRewriter &rewriter) const override {
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp.getOperation());
+    if (!linalgOp)
+      return failure();
+    if (failed(marker.checkAndNotify(rewriter, linalgOp)))
+      return failure();
+
+    auto *pattern = static_cast<const ConcretePattern *>(this);
+    linalg::GenericOp genericOp = pattern->createGenericOp(rootOp, rewriter);
+    if (!genericOp)
+      return failure();
+
+    rewriter.replaceOp(rootOp, genericOp.getResults());
+    marker.replaceLinalgMarker(rewriter, genericOp.getOperation());
+    return success();
+  }
+
+private:
+  linalg::LinalgMarker marker;
+};
+
+struct GeneralizeConvOp
+    : public LinalgGeneralizationPattern<GeneralizeConvOp, linalg::ConvOp> {
+  using LinalgGeneralizationPattern::LinalgGeneralizationPattern;
+
+  linalg::GenericOp createGenericOp(linalg::ConvOp, OpBuilder &rewriter) const;
+};
+
+/// Catch-all pattern for converting all named ops with a region builder into
+/// linalg.generic.
+struct LinalgNamedOpGeneralizationPattern : RewritePattern {
+  LinalgNamedOpGeneralizationPattern(MLIRContext *context,
+                                     linalg::LinalgMarker marker,
+                                     PatternBenefit benefit = 1)
+      : RewritePattern(benefit, MatchAnyOpTypeTag()),
+        marker(std::move(marker)) {}
+
+  LogicalResult matchAndRewrite(Operation *rootOp,
+                                PatternRewriter &rewriter) const override {
+    auto linalgOp = dyn_cast<linalg::LinalgOp>(rootOp);
+    if (!linalgOp)
+      return failure();
+    if (failed(marker.checkAndNotify(rewriter, linalgOp)))
+      return failure();
+
+    // No nothing to do for linalg.generic and linalg.indexed_generic.
+    if (isa<linalg::GenericOp, linalg::IndexedGenericOp>(rootOp))
+      return failure();
+
+    linalg::GenericOp genericOp =
+        createGenericOpFromNamedOp(linalgOp, rewriter);
+    if (!genericOp)
+      return failure();
+
+    rewriter.replaceOp(rootOp, genericOp.getResults());
+    marker.replaceLinalgMarker(rewriter, genericOp.getOperation());
+    return success();
+  }
+
+private:
+  linalg::LinalgMarker marker;
+};
+
+struct LinalgGeneralizationPass
+    : public LinalgGeneralizationBase<LinalgGeneralizationPass> {
+  void runOnFunction() override;
+};
+
+} // namespace
+
+void LinalgGeneralizationPass::runOnFunction() {
+  FuncOp func = getFunction();
+  OwningRewritePatternList patterns;
+  linalg::populateLinalgConvGeneralizationPatterns(&getContext(), patterns);
+  linalg::populateLinalgNamedOpsGeneralizationPatterns(&getContext(), patterns);
+  applyPatternsAndFoldGreedily(func.getBody(), std::move(patterns));
+}
+
+linalg::GenericOp GeneralizeConvOp::createGenericOp(linalg::ConvOp convOp,
+                                                    OpBuilder &builder) const {
+  SmallVector<AffineMap, 4> indexingMaps = convOp.getIndexingMaps();
+  auto iterators =
+      llvm::to_vector<4>(convOp.iterator_types().getAsValueRange<StringAttr>());
+  return builder.create<linalg::GenericOp>(
+      convOp.getLoc(), /*resultTensorTypes=*/ArrayRef<Type>(),
+      convOp.getInputBuffers(), convOp.getOutputBuffers(),
+      /*initTensors=*/ValueRange(), indexingMaps, iterators,
+      [](OpBuilder &bodyBuilder, Location bodyLoc, ValueRange bodyArgs) {
+        Value mul =
+            bodyBuilder.create<MulFOp>(bodyLoc, bodyArgs[0], bodyArgs[1]);
+        Value add = bodyBuilder.create<AddFOp>(bodyLoc, mul, bodyArgs[2]);
+        bodyBuilder.create<linalg::YieldOp>(bodyLoc, add);
+      });
+}
+
+void mlir::linalg::populateLinalgConvGeneralizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    linalg::LinalgMarker marker) {
+  patterns.insert<GeneralizeConvOp>(context, marker);
+}
+
+void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(
+    MLIRContext *context, OwningRewritePatternList &patterns,
+    linalg::LinalgMarker marker) {
+  patterns.insert<LinalgNamedOpGeneralizationPattern>(context, marker);
+}
+
+std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgGeneralizationPass() {
+  return std::make_unique<LinalgGeneralizationPass>();
+}

mlir/test/Dialect/Linalg/generalize-named-ops.mlir

+// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
+
+func @generalize_conv(%input : memref<1x225x225x3xf32>, %filter: memref<3x3x3x32xf32>, %output: memref<1x112x112x32xf32>) {
+  linalg.conv(%filter, %input, %output) {dilations = [2, 3], strides = [4, 5]} : memref<3x3x3x32xf32>, memref<1x225x225x3xf32>, memref<1x112x112x32xf32>
+  return
+}
+
+// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>
+// CHECK:  #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 4 + d3 * 2, d2 * 5 + d4 * 3, d5)>
+// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d6)>
+
+// CHECK: func @generalize_conv
+// CHECK-SAME:  %[[INPUT:.+]]: memref<1x225x225x3xf32>
+// CHECK-SAME: %[[FILTER:.+]]: memref<3x3x3x32xf32>
+// CHECK-SAME: %[[OUTPUT:.+]]: memref<1x112x112x32xf32>
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[FILTER_MAP]], #[[INPUT_MAP]], #[[OUTPUT_MAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "window", "window", "reduction", "parallel"]
+// CHECK-SAME:  ins(%[[FILTER]], %[[INPUT]]
+// CHECK-SAME: outs(%[[OUTPUT]]
+
+// CHECK: ^{{.*}}(%[[FILTER_ARG:.+]]: f32, %[[INPUT_ARG:.+]]: f32, %[[OUTPUT_ARG:.+]]: f32)
+// CHECK:   %[[MUL:.+]] = mulf %[[FILTER_ARG]], %[[INPUT_ARG]]
+// CHECK:   %[[ADD:.+]] = addf %[[MUL]], %[[OUTPUT_ARG]]
+// CHECK:   linalg.yield %[[ADD]]
+
+// -----
+
+func @generalize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
+  linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>) outs(%C: memref<16x32xf32>)
+  return
+}
+
+
+// CHECK: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
+// CHECK: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK: func @generalize_matmul_buffer
+// CHECK-SAME: %[[A:.+]]: memref<16x8xf32>
+// CHECK-SAME: %[[B:.+]]: memref<8x32xf32>
+// CHECK-SAME: %[[C:.+]]: memref<16x32xf32>
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[A]], %[[B]]
+// CHECK-SAME: outs(%[[C]]
+
+// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
+// CHECK:   %[[MUL:.+]] = mulf %[[A_ARG]], %[[B_ARG]] : f32
+// CHECK:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
+// CHECK:   linalg.yield %[[ADD]] : f32
+
+// -----
+
+func @generalize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
+  %0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>) init(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
+  return %0: tensor<16x32xf32>
+}
+
+// CHECK: func @generalize_matmul_tensor
+
+// CHECK: linalg.generic
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
+// CHECK-SAME: init(%{{.+}} : tensor<16x32xf32>)
+
+// CHECK:      ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
+// CHECK-NEXT:   %[[MUL:.+]] = mulf %[[A_ARG]], %[[B_ARG]] : f32
+// CHECK-NEXT:   %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
+// CHECK-NEXT:   linalg.yield %[[ADD]] : f32
+// CHECK-NEXT: -> tensor<16x32xf32>

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

SingleBlockImplicitTerminator<"YieldOp">]> {
       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";