[mlir][Linalg] Create a named batch_matmul op and pipe it through.

[nicolasvasilache]

Summary

This revision is the first in a set of improvements that aim at allowing
more generalized named Linalg op generation from a mathematical
specification.

This revision allows creating a new op and checks that the parser,
printer and verifier are hooked up properly.

This opened up a few design points that will be addressed in the future:

1. A named linalg op has a static region builder instead of an explicitly parsed region. This is not currently compatible with assemblyFormat so a custom parser / printer are needed.
2. The convention for structured ops and tensor return values needs to evolve to allow tensor-land and buffer land specifications to agree
3. ReferenceIndexingMaps and referenceIterators will need to become static to allow building attributes at parse time.
4. Error messages will be improved once we have 3. and we pretty print in custom form.

[nicolasvasilache]

Rebase

[bondhugula]

Building region from an OperationState at parse time is not straightforward because the region is not attached to an op, which is required to fill it with a builder. A temporary fake op is
created for this purpose and the resulting region is taken from it.

Is this comment relevant to this revision? What region are you trying to "build" at parse time here? Why can't the typical addRegion() followed parseRegionBody be used?

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

123–124

of -> or
buffer -> memref

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

273

You don't need getBlocks() - region.front() will work.

mlir/test/Dialect/Linalg/roundtrip.mlir

628

For the custom format, better to drop the parentheses
linalg.batchmatmul %a3, %b3, %c3 :

FWIW, it'd be better with an underscore batch_matmul.

[bondhugula]

Building region from an OperationState at parse time is not straightforward because the region is not attached to an op, which is required to fill it with a builder. A temporary fake op is
created for this purpose and the resulting region is taken from it.

I now see the relevant method. But to use the builder, you only need a block or a region, and not an op(?). You could do an addRegion(), followed by pushing a new block into it, and then create the builder to insert. Am I missing something?

[nicolasvasilache]

I now see the relevant method. But to use the builder, you only need a block or a region, and not an op(?). You could do an addRegion(), followed by pushing a new block into it, and then create the builder to insert. Am I missing something?

This is what I started with but along the line I realized that region.getContext() requires region to be attached to a Operation*.
See: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/IR/Region.h#L31

However this was a red herring only due to the fact that I was constructing OpBuilder from region.
I.e. this fails:

Region *region = result.addRegion();
OpBuilder opBuilder(region);

This works fine:

Region *region = result.addRegion();
OpBuilder opBuilder(context);
opBuilder.setInsertionPoint(&region.front(), region.front().begin());

Thanks for noting!

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

123–124

First typo is corrected.
For the second, Linalg uses the terminology tensors and buffers consistently.
Memref is a particular implementation of buffers, the only one supported for now.

mlir/test/Dialect/Linalg/roundtrip.mlir

628

Fair enough.

[bondhugula]

In D78327#1988920, @nicolasvasilache wrote:

I now see the relevant method. But to use the builder, you only need a block or a region, and not an op(?). You could do an addRegion(), followed by pushing a new block into it, and then create the builder to insert. Am I missing something?

This is what I started with but along the line I realized that region.getContext() requires region to be attached to a Operation*.
See: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/IR/Region.h#L31

However this was a red herring only due to the fact that I was constructing OpBuilder from region.
I.e. this fails:

Region *region = result.addRegion();
OpBuilder opBuilder(region);

This works fine:

Region *region = result.addRegion();
OpBuilder opBuilder(context);
opBuilder.setInsertionPoint(&region.front(), region.front().begin());

Thanks for noting!

That's great. In fact, the following is less verbose and should work:

auto opBuilder = OpBuilder::atBlockBegin(&region.front());

[nicolasvasilache]

That's great. In fact, the following is less verbose and should work:

auto opBuilder = OpBuilder::atBlockBegin(&region.front());

It doesn't unfortunately because of exactly the same problem: there is no context in the block or the region.
The OpBuilder must be constructed with a context and then step into the region.

[nicolasvasilache]

Address comments.

[bondhugula]

In D78327#1989405, @nicolasvasilache wrote:

That's great. In fact, the following is less verbose and should work:

auto opBuilder = OpBuilder::atBlockBegin(&region.front());

It doesn't unfortunately because of exactly the same problem: there is no context in the block or the region.
The OpBuilder must be constructed with a context and then step into the region.

I see - it won't have the context with the block based ctors as well. Nit: then you could just use setInsertionPointToStart(block) after the context arg ctor.

[bondhugula]

Mostly superficial comments...

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

94

of -> or

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

763–764

Doc comments please.

791–794

All of this looks problematic. Do you need this on YieldOp's verifier instead of LinalgOp's verifier? (Check for its terminator there and verify?)

1086

Nit: bodyRegion is weird. Just region is good or opRegion.

1097

setInsertionPointToStart(&bodyRegion.front())

1104

Micronit: " " -> ' '; likewise below.

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

124

Doc comment please.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

What is the reference for this specification ? ONNX/TF both seem to have a batch dimension for B as well. Without that this is effectively broadcasting B

[silvas]

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

This isn't enough to legalize e.g. tf.BatchMatMul or torch.matmul, which allow leading batch dimensions on both sides.

https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/batch-mat-mul
https://pytorch.org/docs/stable/torch.html#torch.matmul

In IREE we have a batch matmul op that handles batch on both sides:
https://github.com/google/iree/blob/f80f39c7e96c2af15741e9c774eb8b54bf38df28/iree/compiler/Dialect/VMLA/IR/VMLAOps.td#L323

I expect that in a typical lowering flow, we will legalize tf.BatchMatMul or torch.matmul by reshaping all the batch dimensions into a single dimension on both sides (possibly a dummy "1" dimension in case of no batch on one side). Then we can expand this op into generic form and fuse/cleanup those reshapes which will eliminate batch dimensions on either side.

I don't see a situation where we would create this op.

My intuition is that batch matmul with a batch dimension only on one side is not that interesting, because fundamentally it is the same as a regular matmul, because you just fold the batch dimension into the free dimension of the respective operand (e.g. in the case you have here, you can just reshape the two dimensions Batch,M in the LHS into a single dimension of extent Batch*M). Batch matmul is only interesting from a lowering perspective when you have a batch dimension on both sides, which introduces a distinct data-reuse behavior as compared to a normal matmul.

So in terms of defining a set of "primitives" or lowering to library calls (e.g. https://devblogs.nvidia.com/cublas-strided-batched-matrix-multiply/), having a batch on both sides seems to be the only relevant case. So I would recommend defining this as:

def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {

[nicolasvasilache]

Address review.
Hacking tablgen.

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

It's just something to get started, the semantics are variadic and will require extensions.

Once these are implemented it will be easy to update.
If you have a strong preference for another op, let me know what you'd prefer (an op that also exercises reduction).
It can't be dot/matvec/matmul for now because that's already taken and more work is needed to replace them.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

791–794

This is consistent with LoopOps::YieldOp and ReturnOp, what would justify diverging from that in Linalg specifically?
If relevant this seems like it would warrant a global change.
Note that LinalgOp is an interface though, not an Op per se.
I'll have a followup NFC to rename globally.

Still, I was missing the SingleBlockImplicitTerminator<"YieldOp"> so I added it where relevant + updated some tests.

[nicolasvasilache]

The tablegen story is awful atm, suggestions most welcome.

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

I went with @silvas ' suggestion, we can iterate on the semantics later once we have variadic support.

[nicolasvasilache]

Another tablgen hack at a distance.

[nicolasvasilache]

Using @mehdi_amini's version for CMake dependencies.

[nicolasvasilache]

Ok, the cmake story now seems as good as it is going to get in the near future.

[nicolasvasilache]

Add missing comment.

[bondhugula]

Minor thing: please update 'batch_matmul' in revision title.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

Thanks for the update. I am not sure I follow what the semantics is variadic implies, i.e. I dont see anything variadic about the op as defined here, but I might be misreading the terms.
My concern was merely if the named ops are supposed to have implicit broadcast semantics (in thoery it can, but that seems to lead to complications when it comes to things like dynamic broadcasting, etc. based on discussion on discourse). As it was defined previously, I read B as having broadcast semantics. Anyway, its OK now so thanks for taking care of it.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

3

Actually strike the last comment. The spec has nothing to do with broadcasting. But, the current spec is indeed more preferable.

[ftynse]

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

351

Nit: why plural "Traits" ?

356

Nit: consider adding a short doc and/or vertical whitespace around declarations

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

6 ↗

(On Diff #258478)

Nit typo: depend

[bondhugula]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

791–794

I actually also meant to highlight the lines above as well (788-789). This change will go against unifying the loop dialect and linalg dialect yield ops (including any other) into a single std yield - but we don't have to worry about it all now. It can be easily adjusted when/if an std.yield is added later.

mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt

+# Declare a function to generate ODS with mlir-linalg-ods-gen
+function(add_linalg_ods_gen tc_filename output_file)
+  set(TC_SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/${tc_filename})
+  set(GEN_ODS_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.td)
+  set(GEN_CPP_FILE ${CMAKE_CURRENT_BINARY_DIR}/${output_file}.cpp.inc)
+  set_source_files_properties(
+    ${GEN_ODS_FILE}
+    PROPERTIES GENERATED TRUE)
+  set_source_files_properties(
+    ${GEN_CPP_FILE}
+    PROPERTIES GENERATED TRUE)
+  add_custom_command(
+    OUTPUT ${GEN_ODS_FILE} ${GEN_CPP_FILE}
+    COMMAND mlir-linalg-ods-gen -gen-ods-decl ${TC_SOURCE} > ${GEN_ODS_FILE}
+    COMMAND mlir-linalg-ods-gen -gen-impl ${TC_SOURCE} > ${GEN_CPP_FILE}
+    MAIN_DEPENDENCY
+    ${TC_SOURCE}
+    DEPENDS
+    mlir-linalg-ods-gen
+    VERBATIM)
+  add_custom_target(
+    MLIR${output_file}IncGen
+    DEPENDS
+    mlir-linalg-ods-gen
+    ${GEN_ODS_FILE} ${GEN_CPP_FILE})
+endfunction()
+
 add_mlir_dialect(LinalgOps linalg)
 add_mlir_doc(LinalgDoc -gen-op-doc LinalgOps Dialects/)
 
+add_linalg_ods_gen(LinalgNamedStructuredOpsSpec.tc LinalgNamedStructuredOps)
 set(LLVM_TARGET_DEFINITIONS LinalgStructuredOps.td)
+set(TABLEGEN_OUTPUT
+    ${TABLEGEN_OUTPUT}
+    MLIRLinalgNamedStructuredOpsIncGen)
 mlir_tablegen(LinalgStructuredOps.h.inc -gen-op-decls)
 mlir_tablegen(LinalgStructuredOps.cpp.inc -gen-op-defs)
 add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LinalgStructuredOpsInterface.td)
 mlir_tablegen(LinalgStructuredOpsInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(LinalgStructuredOpsInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRLinalgStructuredOpsInterfaceIncGen)

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

+ods_def<BatchMatmulOp>:
+def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {
+  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(b, k, n)));
+}

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

def LinalgOperand: Type<
   Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>;
 
 class LinalgOperandOfRank<int rank>: Type<
   And<[
     LinalgOperand.predicate,
     CPred<"$_self.cast<ShapedType>().getRank() == " # rank>]
   >>;
 
-class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, []> {
+class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic,
+    [SingleBlockImplicitTerminator<"YieldOp">]> {
   let arguments = (ins Variadic<LinalgOperand>:$views,
                    I64Attr:$args_in,
                    I64Attr:$args_out,
                    AffineMapArrayAttr:$indexing_maps,
                    ArrayAttr:$iterator_types,
                    OptionalAttr<StrAttr>:$doc,
                    OptionalAttr<StrAttr>:$library_call);
   let results = (outs Variadic<AnyRankedTensor>:$output_tensors);

 
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as a declarative configurations of generic ops.
 //===----------------------------------------------------------------------===//
 
 def NamedStructuredOpTraits : NativeOpTrait<"linalg::NamedStructuredOpTraits">;
 
 class LinalgNamedStructured_Op<string mnemonic, list<OpTrait> props>
-  : Op<Linalg_Dialect, mnemonic,
-       !listconcat(props, [StructuredOpTraits, LinalgStructuredInterface])> {
+    : LinalgStructuredBase_Op<mnemonic, props> {
   string spec = ?;
-  let assemblyFormat = "`(` operands `)` attr-dict `:` "
-    "functional-type(operands, results)";
+  // We cannot use an assemblyFormat atm because we need to hook in a custom-
+  // built implicit region from a static OpClass method.
+  // TODO: Revisit in the future if/when appropriate.
+  // let assemblyFormat = "`(` operands `)` attr-dict `:` "
+  //   "functional-type(operands, results)";
+
+  // The parser needs to specialize on the OpType so it has to be auto-generated
+  // in the linalg-ods tool.
+  let printer = [{ return ::printNamedStructuredOp(p, *this); }];
+  let verifier = [{ return ::verifyNamedStructuredOp(*this); }];
+  let hasFolder = 1;
 }
 
+// This file is auto-generated from a tc specification.
+include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.td"
+
 #endif // LINALG_STRUCTURED_OPS

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

InterfaceMethod<[{
       }],
       "llvm::Optional<unsigned>", "getIndexOfInput", (ins "Value":$v)
     >,
     InterfaceMethod<
       "Return the input operands from the current operation.",
       "Operation::operand_range", "getInputs"
     >,
     InterfaceMethod<[{
-        Return the type of the input shape at the given index.
+        Return the `i`-th input shaped type, irrespective of buffer or tensor
+        type.
       }], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>,
     InterfaceMethod<[{
         Return the subset of input operands that are of ranked tensor type.
       }], "SmallVector<RankedTensorType, 4>", "getInputTensorTypes">,
 
     //===------------------------------------------------------------------===//
     // Output arguments handling.
     //===------------------------------------------------------------------===//
InterfaceMethod<[{
         not part of the output buffers.
       }],
       "llvm::Optional<unsigned>", "getIndexOfOutputBuffer", (ins "Value":$view)
     >,
     InterfaceMethod<[{
         Return the type of the output buffer at the given index.
       }], "MemRefType", "getOutputBufferType", (ins "unsigned":$i)>,
     InterfaceMethod<[{
+        Return the `i`-th output shaped type, irrespective of buffer or tensor
+        type.
+      }], "ShapedType", "getOutputShapedType", (ins "unsigned":$i)>,
+    InterfaceMethod<[{
         Return the results that are of ranked tensor type.
       }], "SmallVector<RankedTensorType, 4>", "getOutputTensorTypes">,
     InterfaceMethod<
       "Return the output buffers (operands) from the current operation.",
       "Operation::operand_range", "getOutputBuffers"
     >,
 
     //===------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

 //===- LinalgTraits.h - Linalg Traits ---------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_LINALGTRAITS_H_
 #define MLIR_DIALECT_LINALG_LINALGTRAITS_H_
 
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Function.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace OpTrait {
 namespace linalg {
 
public:
   /// Return the index of `value` in the list of inputs if found, llvm::None
   /// otherwise.
   Optional<unsigned> getIndexOfInput(Value value) {
     auto it = llvm::find(getInputs(), value);
     if (it != getInputs().end())
       return it - getInputs().begin();
     return llvm::None;
   }
-  /// Return the `i`-th input buffer type.
+  /// Return the `i`-th input shaped type, irrespective of buffer or tensor
+  /// type.
   ShapedType getInputShapedType(unsigned i) {
     return getInput(i).getType().template cast<ShapedType>();
   }
   /// Return the range over inputs.
   Operation::operand_range getInputs() {
     auto range = this->getOperation()->getOperands();
     return {range.begin(), range.begin() + nInputs()};
   }
public:
   static LogicalResult verifyTrait(Operation *op) {
     auto nOperands = cast<ConcreteType>(op).getNumInputsAndOutputBuffers();
     if (failed(OpTrait::impl::verifyAtLeastNOperands(op, nOperands)))
       return failure();
     return success();
   }
 };
 
+/// This class provides the API for named Linalg StructuredOps.
+template <typename ConcreteType>
+class NamedStructuredOpTraits
+    : public OpTrait::TraitBase<ConcreteType, NamedStructuredOpTraits> {
+public:
+  llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
+  llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
+  std::function<void(OpBuilder &, Location, ArrayRef<Value>)>
+  emitScalarImplementation();
+};
+
 } // namespace linalg
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_LINALGTRAITS_H_

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 //===- LinalgOps.cpp - Implementation of the linalg operations ------------===//
 //
 // 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 operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
+/// Forward declarations.
+template <typename NamedStructuredOpType>
+static void buildNamedStructuredOpRegion(Builder &builder,
+                                         OperationState &result,
+                                         TypeRange operandTypes,
+                                         TypeRange tensorResultTypes);
+template <typename NamedStructuredOpType>
+static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op);
+template <typename NamedStructuredOpType>
+static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
+                                          OperationState &result);
+template <typename NamedStructuredOpType>
+static LogicalResult verifyNamedStructuredOp(NamedStructuredOpType op);
+
 /// Determines whether it is possible to fold it away in the parent Linalg op:
 ///
 /// ```mlir
 ///   %1 = memref_cast %0 : memref<8x16xf32> to memref<?x?xf32>
 ///   %2 = linalg.slice %1 ... : memref<?x?xf32> ...
 ///   // or
 ///   %1 = memref_cast %0 : memref<8x16xf32, affine_map<(i, j)->(16 * i + j)>>
 ///          to memref<?x?xf32>
static ParseResult parseGenericOp(OpAsmParser &parser, OperationState &result) {
   if (parser.parseOptionalArrowTypeList(tensorResultTypes))
     return failure();
   if (!tensorResultTypes.empty())
     result.addTypes(tensorResultTypes);
   return parser.resolveOperands(operandsInfo, operandTypes,
                                 parser.getCurrentLocation(), result.operands);
 }
 
-LogicalResult verifyBlockArgs(GenericOp op, Block &block) {
+template <typename GenericOpType>
+struct BlockArgsVerifier {
+  static LogicalResult verify(GenericOpType op, Block &block);
+};
+
+template <typename GenericOpType>
+LogicalResult BlockArgsVerifier<GenericOpType>::verify(GenericOpType op,
+                                                       Block &block) {
   auto nOperands = op.getNumOperands();
   if (block.getNumArguments() != nOperands)
     return op.emitOpError("expected number of block arguments to match number "
                           "of operands");
 
   // Note: the number and type of yield values are checked in the YieldOp.
   auto nInputViews = op.getNumInputs();
   for (unsigned i = 0; i < nOperands; ++i) {
     auto viewType = op.getShapedType(i);
     if (viewType.getElementType() != block.getArgument(i).getType())
       return op.emitOpError("expected block argument ")
              << (i + 1) << " of the same type as elemental type of "
              << ((i < nInputViews) ? "input " : "output ")
              << "operand: " << viewType;
   }
   return success();
 }
 
-LogicalResult verifyBlockArgs(IndexedGenericOp op, Block &block) {
+template <>
+LogicalResult BlockArgsVerifier<IndexedGenericOp>::verify(IndexedGenericOp op,
+                                                          Block &block) {
   auto nInputViews = op.getNumInputs();
   auto nLoops = op.getNumLoops();
   auto nOperands = op.getNumOperands();
   if (block.getNumArguments() != nOperands + nLoops)
     return op.emitOpError(
         "expected number of block arguments to match number of operands + "
         "number of loops");
 
static LogicalResult verifyGenericOp(GenericOpType op) {
   if (nInputsAndOutputBuffers != llvm::size(op.views()))
     return op.emitOpError("expected exactly ")
            << nInputsAndOutputBuffers
            << " inputs (tensor or buffer) and output buffer operands";
 
   auto &region = op.region();
   if (region.getBlocks().size() != 1)
     return op.emitOpError("expected region with 1 block");
-  if (failed(verifyBlockArgs(op, region.getBlocks().front())))
+  if (failed(BlockArgsVerifier<GenericOpType>::verify(op, region.front())))
     return failure();
 
   SmallVector<AffineMap, 4> indexingMaps;
   indexingMaps.reserve(op.indexing_maps().size());
   for (auto en : llvm::enumerate(op.indexing_maps())) {
     auto idx = en.index();
     auto m = en.value().template cast<AffineMapAttr>().getValue();
     indexingMaps.push_back(m); // Save reference to map for further checks.
     auto view = (idx < nInputViews) ? op.getInputShapedType(idx)
static ParseResult parseYieldOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 2> opInfo;
   SmallVector<Type, 2> types;
   llvm::SMLoc loc = parser.getCurrentLocation();
   return failure(parser.parseOperandList(opInfo) ||
                  parser.parseOptionalAttrDict(result.attributes) ||
                  (!opInfo.empty() && parser.parseColonTypeList(types)) ||
                  parser.resolveOperands(opInfo, types, loc, result.operands));
 }
 
-template <typename GenericOpType>
-static LogicalResult verifyYield(YieldOp op, GenericOpType genericOp) {
-  // The operand number and types must match the view element types.
-  auto nOutputs = genericOp.getNumOutputs();
+// Check the operand number and types must match the element types of the
+// LinalgOp interface's shaped operands.
+static LogicalResult verifyYield(YieldOp op, LinalgOp linalgOpInterface) {
+  auto nOutputs = linalgOpInterface.getNumOutputs();
   if (op.getNumOperands() != nOutputs)
     return op.emitOpError("expected number of yield values (")
            << nOutputs << ") to match the number of operands of the enclosing "
-           << "linalg.generic op (" << op.getNumOperands() << ")";
+           << "LinalgOp (" << op.getNumOperands() << ")";
 
   for (unsigned i = 0; i != nOutputs; ++i) {
-    auto elementType = genericOp.getOutputShapedType(i).getElementType();
+    auto elementType =
+        linalgOpInterface.getOutputShapedType(i).getElementType();
     if (op.getOperand(i).getType() != elementType)
       return op.emitOpError("type of yield operand ")
              << (i + 1) << " (" << op.getOperand(i).getType()
              << ") doesn't match "
              << "the element type of the enclosing linalg.generic op ("
              << elementType << ")";
   }
   return success();
 }
 
 static LogicalResult verify(YieldOp op) {
   auto *parentOp = op.getParentOp();
   if (parentOp->getNumRegions() != 1 || parentOp->getRegion(0).empty())
     return op.emitOpError("expected single non-empty parent region");
 
-  auto genericOp = dyn_cast<GenericOp>(parentOp);
-  if (genericOp)
-    return verifyYield(op, genericOp);
-
-  auto indexedGenericOp = dyn_cast<IndexedGenericOp>(parentOp);
-  if (indexedGenericOp)
-    return verifyYield(op, indexedGenericOp);
-
-  return op.emitOpError("expected '")
-         << GenericOp::getOperationName() << "' or '"
-         << IndexedGenericOp::getOperationName() << "' parent op";
+  if (auto linalgOp = dyn_cast<LinalgOp>(parentOp))
+    return verifyYield(op, cast<LinalgOp>(parentOp));
+
+  return op.emitOpError("expected parent op with LinalgOp interface");
 }
 
 /////// Operations corresponding to library calls defined with Tablegen ////////
 
 static LogicalResult verify(FillOp op) {
   auto viewType = op.getOutputShapedType(0);
   auto fillType = op.value().getType();
   if (viewType.getElementType() != fillType)
if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return {};
 }
 OpFoldResult TransposeOp::fold(ArrayRef<Attribute>) {
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return {};
 }
+
+//===----------------------------------------------------------------------===//
+// Auto-generated Linalg named ops.
+//===----------------------------------------------------------------------===//
+
+template <typename NamedStructuredOpType>
+void buildNamedStructuredOpRegion(Builder &builder, OperationState &result,
+                                  TypeRange operandTypes,
+                                  TypeRange tensorResultTypes) {
+  Region &region = *result.addRegion();
+  Block *body = new Block();
+  // TODO: atm all operands go through getElementTypeOrSelf,
+  // reconsider when we have evidence we need to.
+  for (auto t : operandTypes)
+    body->addArgument(getElementTypeOrSelf(t));
+  for (auto t : tensorResultTypes)
+    body->addArgument(getElementTypeOrSelf(t));
+  region.push_back(body);
+
+  OpBuilder opBuilder(builder.getContext());
+  opBuilder.setInsertionPointToStart(&region.front());
+  mlir::edsc::ScopedContext scope(opBuilder, builder.getUnknownLoc());
+  NamedStructuredOpType::regionBuilder(*body);
+}
+
+template <typename NamedStructuredOpType>
+static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op) {
+  p << op.getOperationName() << ' ';
+  p.printOptionalAttrDict(op.getAttrs());
+  p << ' ' << op.getOperands();
+  p << ": (" << op.getOperandTypes() << ")";
+  auto outputTensorTypes = op.getResultTypes();
+  if (!outputTensorTypes.empty())
+    p << " -> (" << outputTensorTypes << ")";
+}
+
+template <typename NamedStructuredOpType>
+static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
+                                          OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 8> operandsInfo;
+
+  // Optional attributes may be added.
+  if (parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseOperandList(operandsInfo))
+    return failure();
+
+  SmallVector<Type, 8> operandTypes;
+  if (parser.parseColon() || parser.parseLParen() ||
+      parser.parseTypeList(operandTypes) || parser.parseRParen())
+    return failure();
+
+  // Generic ops may specify that a subset of its outputs are tensors. Such
+  // outputs are specified in the result type.
+  SmallVector<Type, 8> tensorResultTypes;
+  if (parser.parseOptionalArrowTypeList(tensorResultTypes))
+    return failure();
+
+  if (!tensorResultTypes.empty())
+    result.addTypes(tensorResultTypes);
+
+  buildNamedStructuredOpRegion<NamedStructuredOpType>(
+      parser.getBuilder(), result, operandTypes, tensorResultTypes);
+
+  return parser.resolveOperands(operandsInfo, operandTypes,
+                                parser.getCurrentLocation(), result.operands);
+}
+
+template <typename NamedStructuredOpType>
+static LogicalResult verifyNamedStructuredOp(NamedStructuredOpType op) {
+  return verifyGenericOp<NamedStructuredOpType>(op);
+}
+
+#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.cpp.inc"
+
+// TODO: Determine whether we can generate the folders and verifiers.
+LogicalResult BatchMatmulOp::fold(ArrayRef<Attribute>,
+                                  SmallVectorImpl<OpFoldResult> &) {
+  return foldMemRefCast(*this);
+}

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

getInputAndOutputIndices(ArrayRef<Value> allIvs, SingleInputPoolingOp op) {
   SmallVector<ValueHandle, 8> iIdx(
       makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
   SmallVector<ValueHandle, 8> oIdx(
       makeCanonicalAffineApplies(b, loc, maps[2], allIvs));
   return {iIdx, oIdx};
 }
 
 namespace {
+
+// Generic loop emitter, to be specialized on an op-per op basis.
+// TODO:
 template <typename IndexedValueType, typename LinalgOpType>
-class LinalgScopedEmitter {};
+class LinalgScopedEmitter {
+public:
+  static void emitScalarImplementation(ArrayRef<Value> allIvs,
+                                       LinalgOpType linalgOp) {
+    assert(linalgOp.hasBufferSemantics() &&
+           "expected linalg op with buffer semantics");
+    llvm_unreachable("NYI");
+    linalgOp.emitScalarImplementation()(ScopedContext::getBuilder(),
+                                        ScopedContext::getLocation(), allIvs);
+  }
+};
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, CopyOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
     assert(copyOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
     auto nPar = copyOp.getNumParallelLoops();

mlir/test/Dialect/Linalg/invalid.mlir

 func @transpose_bad_rank(%v : memref<?x?xf32, affine_map<(i, j)[off, M]->(off + M * i + j)>>) {
   // expected-error @+1 {{expected a permutation map of same rank as the view}}
   linalg.transpose %v (i) -> (i) : memref<?x?xf32, affine_map<(i, j)[off, M]->(off + M * i + j)>>
 }
 
 // -----
 
 func @yield_parent(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
-  // expected-error @+1 {{op expected 'linalg.generic' or 'linalg.indexed_generic' parent op}}
+  // expected-error @+1 {{op expected parent op with LinalgOp interface}}
   linalg.yield %arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>
 }
 
 // -----
 
 func @generic_no_region(%arg0: memref<f32>) {
   // expected-error @+6 {{expected '{' to begin a region}}
   linalg.generic {
linalg.generic {
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = []
   } %arg0, %arg0, %arg0 {}: memref<f32>, memref<f32>, memref<f32>
 }
 
 // -----
 
 func @generic_mismatched_num_returns(%arg0: memref<f32>) {
-  // expected-error @+8 {{op expected number of yield values (1) to match the number of operands of the enclosing linalg.generic op (0)}}
+  // expected-error @+8 {{op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (0)}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<() -> ()> ],
     iterator_types = []
   } %arg0 {
     ^bb(%0: f32):
       linalg.yield
   }: memref<f32>
 }
 
 // -----
 
 func @generic_symbol_in_map(%arg0: memref<i32>) {
   // expected-error @+1 {{op expected indexing_map #0 to have no symbols}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<()[N] -> (0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%i : i32):
+    linalg.yield %i : i32
   }: memref<i32>
 }
 
 // -----
 
 func @generic_wrong_dim_in_map(%arg0: memref<1xi32>) {
   // expected-error @+1 {{op expected indexing_map #0 to have 1 dim(s) to match the number of loops}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%i : i32):
+    linalg.yield %i : i32
   }: memref<1xi32>
 }
 
 // -----
 
 func @generic_one_d_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
   // expected-error @+1 {{op expected indexing_map #0 results to match view rank: 'memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>>'}}
   linalg.generic {

 
 ////////////////////////////////////////////////////////////////////////////////
 ///////////////////////////// Region tests /////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
 // -----
 
 func @generic_empty_region(%arg0: memref<f32>) {
-  // expected-error @+1 {{op expected region with 1 block}}
+  %f0 = constant 0.0: f32
+  // expected-error @+1 {{op expects region #0 to have 0 or 1 blocks}}
   linalg.generic {
     args_in = 1,
     args_out = 1,
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = []
   } %arg0, %arg0 {
     ^bb1:
+      linalg.yield %f0: f32
     ^bb2:
+      linalg.yield %f0: f32
+  }: memref<f32>, memref<f32>
+}
+
+// -----
+
+func @generic_empty_region(%arg0: memref<f32>) {
+  %f0 = constant 0.0: f32
+  // expected-error @+1 {{linalg.generic' op expected region with 1 block}}
+  linalg.generic {
+    args_in = 1,
+    args_out = 1,
+    indexing_maps =  [ affine_map<() -> (0)> ],
+    iterator_types = []
+  } %arg0, %arg0 {
   }: memref<f32>, memref<f32>
 }
 
 // -----
 
 func @generic_mismatched_num_arguments(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected number of block arguments to match number of operands}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = []
   } %arg0 {
-    ^bb:
+    ^bb(%f: f32, %g: f32):
+      linalg.yield %f: f32
   }: memref<f32>
 }
 
 // -----
 
 func @generic_block_arg_type(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected block argument 1 of the same type as elemental type of output operand: 'memref<f32>'}}
   linalg.generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = []
   } %arg0 {
     ^bb(%i: i1):
+    linalg.yield %i : i1
   }: memref<f32>
 }
 
 // -----
 
 func @indexed_generic_block_arg_count(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected number of block arguments to match number of operands + number of loops}}
   linalg.indexed_generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<(d0) -> (d0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%f: f32):
+      linalg.yield %f : f32
   }: memref<f32>
 }
 
 // -----
 
 func @indexed_generic_block_induction_var_arg_type(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected block argument 1 to be an index}}
   linalg.indexed_generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<(d0) -> (d0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%i: f64, %f: f32):
+    linalg.yield %f: f32
   }: memref<f32>
 }
 
 // -----
 
 func @indexed_generic_block_arg_type(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected block argument 2 of the same type as elemental type of output operand: 'memref<f32>'}}
   linalg.indexed_generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<(d0) -> (d0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%i: index, %f: i1):
+    linalg.yield %i: index
   }: memref<f32>
 }
 
 // -----
 
 func @indexed_generic_arg_count(%arg0: memref<f32>) {
   // expected-error @+1 {{op expected number of block arguments to match number of operands + number of loops}}
   linalg.indexed_generic {
linalg.indexed_generic {
     ^bb(%0: i32, %1: f32):
       linalg.yield %1: f32
   } : memref<f32>
 }
 
 // -----
 
 func @indexed_generic_result_count(%arg0: memref<?xf32>) {
-  // expected-error @+8 {{op expected number of yield values (1) to match the number of operands of the enclosing linalg.generic op (2)}}
+  // expected-error @+8 {{op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (2)}}
   linalg.indexed_generic {
     args_in = 0,
     args_out = 1,
     indexing_maps =  [ affine_map<(d0) -> (d0)> ],
     iterator_types = ["parallel"]
   } %arg0 {
     ^bb(%i: index, %val: f32):
       linalg.yield %val, %val: f32, f32
func @generic_result_tensor_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
   } %arg0 {
     ^bb(%i: f32):
       linalg.yield %i: f32
   }: memref<?xf32, affine_map<(i)[off]->(off + i)>> -> f32
 }
 
 // -----
 
+func @generic_result_tensor_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
+  // expected-error @+1 {{op result #0 must be ranked tensor of any type values, but got 'f32'}}
+  %0 = linalg.generic {
+    args_in = 0,
+    args_out = 1,
+    indexing_maps = [ affine_map<(i) -> (i)> ],
+    iterator_types = ["parallel"]
+  } %arg0 {
+    ^bb(%i: f32):
+      linalg.yield %i: f32
+  }: memref<?xf32, affine_map<(i)[off]->(off + i)>> -> f32
+}
+
+// -----
+
+func @generic(%arg0: memref<?x?xi4>) {
+  // expected-error @+2 {{op expects regions to end with 'linalg.yield', found 'std.addf'}}
+  // expected-note @+1 {{in custom textual format, the absence of terminator implies 'linalg.yield'}}
+  linalg.generic  {
+    args_in = 0,
+    args_out = 1,
+    indexing_maps = [ affine_map<(i) -> (i)> ],
+    iterator_types = ["parallel"]
+  } %arg0 {
+    ^bb(%0: i4) :
+      %1 = std.addf %0, %0: i4
+  } : memref<?x?xi4>
+  return
+}
+
+// -----
+
 func @generic_result_0_element_type(%arg0: memref<?xf32>) {
   // expected-error @+1 {{'linalg.dot' op expected 3 operands, but found 2}}
   linalg.dot(%arg0, %arg0): memref<?xf32>, memref<?xf32>
 }
 
 // -----
 
 // expected-error @+1 {{unknown Linalg type}}

 func @pooling_rank_mismatch(%arg0: memref<?x?x?xf32>,
                             %arg1: memref<2x3xf32>,
                             %arg2: memref<?x?x?xf32>) {
   // expected-error @+1 {{expects memref ranks to match}}
   linalg.pooling_max(%arg0, %arg1, %arg2) {strides = [2, 1, 2]}:
     memref<?x?x?xf32>, memref<2x3xf32>, memref<?x?x?xf32>
   return
 }
+
+// -----
+
+func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?xf32>, %c3: memref<?x?x?xf32>) {
+  // expected-error @+1 {{op expected indexing_map #1 results to match view rank: 'memref<?x?xf32>'}}
+  linalg.batch_matmul %a3, %b3, %c3 : (memref<?x?x?xf32>, memref<?x?xf32>, memref<?x?x?xf32>) -> ()
+  return
+}

mlir/test/Dialect/Linalg/roundtrip.mlir

 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<?x?x?xf32, #[[strided3DOFF0]]> into memref<?x?xf32, #[[strided2DOFF0]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<?x?xf32, #[[strided2DOFF0]]> into memref<?x?x?xf32, #[[strided3DOFF0]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<?x?x?xf32, #[[strided3D]]> into memref<?x?xf32, #[[strided2D]]>
 //       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
 //  CHECK-SAME:     memref<?x?xf32, #[[strided2D]]> into memref<?x?x?xf32, #[[strided3D]]>
+
+
+// TODO: Return tensors need a semantics convention update.
+func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,
+                %ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>) {
+  linalg.batch_matmul %a3, %b3, %c3 : (memref<?x?x?xf32>, memref<?x?x?xf32>, memref<?x?x?xf32>) -> ()
+  linalg.batch_matmul %ta3, %tb3, %c3 : (tensor<?x?x?xf32>, tensor<?x?x?xf32>, memref<?x?x?xf32>) -> ()
+  return
+}
+// CHECK-LABEL: func @named_ops
+//       CHECK:   linalg.batch_matmul
+//       CHECK:   linalg.batch_matmul
+

mlir/test/lib/DeclarativeTransforms/CMakeLists.txt

 set(LLVM_TARGET_DEFINITIONS TestLinalgTransformPatterns.td)
+# Silent dependency at a distance .. LinalgStructuredOps.td includes
+# an auto generated .td tracked by target MLIRLinalgNamedStructuredOpsIncGen
+set(TABLEGEN_OUTPUT
+    ${TABLEGEN_OUTPUT}
+        MLIRLinalgNamedStructuredOpsIncGen)
 mlir_tablegen(TestLinalgTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestLinalgTransformPatternsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestVectorTransformPatterns.td)
 mlir_tablegen(TestVectorTransformPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestVectorTransformPatternsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestLinalgMatmulToVectorPatterns.td)
+# Silent dependency at a distance .. LinalgStructuredOps.td includes
+# an auto generated .td tracked by target MLIRLinalgNamedStructuredOpsIncGen
+set(TABLEGEN_OUTPUT
+    ${TABLEGEN_OUTPUT}
+        MLIRLinalgNamedStructuredOpsIncGen)
 mlir_tablegen(TestLinalgMatmulToVectorPatterns.h.inc -gen-rewriters)
 add_public_tablegen_target(MLIRTestLinalgMatmulToVectorPatternsIncGen)

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

 // RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 | FileCheck %s --check-prefix=ODS
 // RUN: mlir-linalg-ods-gen %s -gen-impl=1 | FileCheck %s --check-prefix=IMPL
 
-// RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 -test-emit-include-td-header \
-// RUN: | mlir-tblgen -gen-op-decls -I %S/../../include
-
-// ODS-LABEL: def matvecOp : LinalgNamedStructured_Op<"matvec", [
-//  ODS-NEXT:   NInputs<2>,
-//  ODS-NEXT:   NOutputs<1>,
-//  ODS-NEXT:   NamedStructuredOpTraits]>
+// ODS-LABEL: def Test1Op : LinalgNamedStructured_Op<"test1", [
+//  ODS-NEXT:   NInputs<2>
+//  ODS-NEXT:   NOutputs<1>
+//  ODS-NEXT:   NamedStructuredOpTraits
+//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  matvec::referenceIterators() {
+// IMPL-LABEL:  Test1Op::referenceIterators() {
 //  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  matvec::referenceIndexingMaps() {
-//       IMPL:  AffineMap::get(2, 0, {d0, d1}),
-//  IMPL-NEXT:  AffineMap::get(2, 0, {d1}),
-//  IMPL-NEXT:  AffineMap::get(2, 0, {d0}) };
+//       IMPL:  Test1Op::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(2, 0, {d0, d1}, context),
+//  IMPL-NEXT:  AffineMap::get(2, 0, {d1}, context),
+//  IMPL-NEXT:  AffineMap::get(2, 0, {d0}, context) };
 //
-//       IMPL:  matvec::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  Test1Op::regionBuilder(Block &block) {
 //       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
-def matvec(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
+ods_def<Test1Op> :
+def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
   C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
 }
 
-// ODS-LABEL: def matmulOp : LinalgNamedStructured_Op<"matmul", [
-//  ODS-NEXT:   NInputs<2>,
-//  ODS-NEXT:   NOutputs<1>,
-//  ODS-NEXT:   NamedStructuredOpTraits]>
+// ODS-LABEL: def Test2Op : LinalgNamedStructured_Op<"test2", [
+//  ODS-NEXT:   NInputs<2>
+//  ODS-NEXT:   NOutputs<1>
+//  ODS-NEXT:   NamedStructuredOpTraits
+//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  matmul::referenceIterators() {
+// IMPL-LABEL:  Test2Op::referenceIterators() {
 //  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  matmul::referenceIndexingMaps() {
-//       IMPL:  AffineMap::get(3, 0, {d0, d2}),
-//  IMPL-NEXT:  AffineMap::get(3, 0, {d2, d1}),
-//  IMPL-NEXT:  AffineMap::get(3, 0, {d0, d1}) };
+//       IMPL:  Test2Op::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(3, 0, {d0, d2}, context),
+//  IMPL-NEXT:  AffineMap::get(3, 0, {d2, d1}, context),
+//  IMPL-NEXT:  AffineMap::get(3, 0, {d0, d1}, context) };
 //
-//       IMPL:  matmul::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  Test2Op::regionBuilder(Block &block) {
 //       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
-def matmul(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
+ods_def<Test2Op> :
+def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
   C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
 }
 
-// ODS-LABEL: def batchmatmulOp : LinalgNamedStructured_Op<"batchmatmul", [
-//  ODS-NEXT:   NInputs<2>,
-//  ODS-NEXT:   NOutputs<1>,
-//  ODS-NEXT:   NamedStructuredOpTraits]>
+// ODS-LABEL: def Test3Op : LinalgNamedStructured_Op<"test3", [
+//  ODS-NEXT:   NInputs<2>
+//  ODS-NEXT:   NOutputs<1>
+//  ODS-NEXT:   NamedStructuredOpTraits
+//  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  batchmatmul::referenceIterators() {
+// IMPL-LABEL:  Test3Op::referenceIterators() {
 //  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  batchmatmul::referenceIndexingMaps() {
-//       IMPL:  AffineMap::get(4, 0, {d0, d1, d3}),
-//  IMPL-NEXT:  AffineMap::get(4, 0, {d3, d2}),
-//  IMPL-NEXT:  AffineMap::get(4, 0, {d0, d1, d2}) };
+//       IMPL:  Test3Op::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(4, 0, {d0, d1, d3}, context),
+//  IMPL-NEXT:  AffineMap::get(4, 0, {d3, d2}, context),
+//  IMPL-NEXT:  AffineMap::get(4, 0, {d0, d1, d2}, context) };
 //
-//       IMPL:  batchmatmul::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  Test3Op::regionBuilder(Block &block) {
 //       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
-//       TBLGEN: batchmatmulOp
-def batchmatmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
+ods_def<Test3Op> :
+def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
 }

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

enum class Kind {
     r_brace,
     r_paren,
     semicolon,
     star,
 
     // Keywords.
     kw_def,
     FIRST_KEYWORD = kw_def,
+    kw_ods_def,
     kw_floordiv,
     kw_ceildiv,
     kw_mod,
     LAST_KEYWORD = kw_mod,
 
     // String valued tokens.
     id,
     integer,
Token Lexer::lexIdentifier(const char *tokStart) {
   // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
   while (isalnum(*curPtr) || *curPtr == '_' || *curPtr == '-')
     ++curPtr;
 
   // Check to see if this identifier is a keyword.
   StringRef str(tokStart, curPtr - tokStart);
   Token::Kind kind = llvm::StringSwitch<Token::Kind>(str)
                          .Case("def", Token::Kind::kw_def)
+                         .Case("ods_def", Token::Kind::kw_ods_def)
                          .Case("floordiv", Token::Kind::kw_floordiv)
                          .Case("ceildiv", Token::Kind::kw_ceildiv)
                          .Case("mod", Token::Kind::kw_mod)
                          .Default(Token::Kind::id);
 
   return Token(kind, str);
 }
 

 ///   op-arg-list ::= op-arg (`,` op-arg)*
 ///   tensor-expr ::= op-spec `(` op-arg-list `)`
 ///
 /// Underlying op-arg are stored by unique_ptr to base class.
 struct TensorExpr : public Expression {
   TensorExpr(StringRef name,
              SmallVectorImpl<std::unique_ptr<Expression>> &&exprs,
              ArrayRef<unsigned> reductionDims)
-      : Expression(Kind::TensorExpr), opId(name), expressions(std::move(exprs)),
+      : Expression(Kind::TensorExpr), operationName(name),
+        expressions(std::move(exprs)),
         reductionDimensions(reductionDims.begin(), reductionDims.end()) {}
 
   static bool classof(const Expression *e) {
     return e->kind == Kind::TensorExpr;
   }
 
   bool operator==(const TensorExpr &other) const {
-    if (opId != other.opId)
+    if (operationName != other.operationName)
       return false;
     if (expressions.size() != other.expressions.size())
       return false;
     for (unsigned i = 0, e = expressions.size(); i < e; ++i)
       if (*expressions[i] != *other.expressions[i])
         return false;
     for (unsigned i = 0, e = reductionDimensions.size(); i < e; ++i)
       if (reductionDimensions[i] != other.reductionDimensions[i])
         return false;
     return true;
   }
 
   /// Visitation function. Performs preorder or postorder traversal depending on
   /// `PreOrder` and applies `callback` on each node.
   template <typename Lambda, bool PreOrder>
   void visit(Lambda callback) const;
 
-  StringRef opId;
+  StringRef operationName;
   SmallVector<std::unique_ptr<Expression>, 4> expressions;
   SetVector<unsigned> reductionDimensions;
 };
 
 Expression::~Expression() {}
 
 bool Expression::operator==(const Expression &e) const {
   if (this->kind != e.kind)
LogicalResult parseOneComprehension(StringRef cppOpName,
                                       StringRef linalgOpName,
                                       ComprehensionParsingState &state);
 
   /// Parse and print the information for a TC def.
   /// When `gen-ods-decl` is used, this prints the ODS declaration for the TC.
   /// When `gen-impl` is used, this prints the C++ implementation for the extra
   /// methods defined in ODS (referenceIterators, referenceIndexingMaps and
   /// regionBuilder).
-  LogicalResult parseAndEmitTCDef(llvm::raw_ostream &os);
+  LogicalResult parseAndEmitODSDef(llvm::raw_ostream &os);
 
   /// Print the ODS class that defines a new `cppOpName` for a `linalgOpName`.
   void printODS(llvm::raw_ostream &os, StringRef cppOpName,
                 StringRef linalgOpName);
 
   /// Print the C++ StructuredOpsInterface impl of `referenceIterators`.
-  void printReferenceIterators(llvm::raw_ostream &os, StringRef opId,
+  void printReferenceIterators(llvm::raw_ostream &os, StringRef cppOpName,
                                ComprehensionParsingState &state);
 
   /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
-  void printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef opId,
+  void printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef cppOpName,
                                   ComprehensionParsingState &state);
 
   /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
-  void printRegionBuilder(llvm::raw_ostream &os, StringRef opId,
+  void printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                           ComprehensionParsingState &state);
 
 private:
   //===--------------------------------------------------------------------===//
   // Internal bookkeeping of tensors.
   //===--------------------------------------------------------------------===//
   struct RegisteredTensor {
     StringRef type;
visitPostorder(*pExpr, [&](const Expression &e) {
       state.orderedTensorArgs[use] = tensor.index;
     });
   if (failed)
     return failure();
 
   return success();
 }
 
-/// Parse and print the information for a TC def.
+/// Parse and print the information for a ODS def.
 ///
 ///   tensor-def-list ::= tensor-def (`,` tensor-def )*
 ///
 ///   comprehension-list ::= comprehension comprehension*
 ///
 ///   tc-def ::= `def` bare-id `(`tensor-def-list`)` `->` `(` tensor-def-list`)`
 ///     `{` comprehension-list `}`
 ///
+///   ods-def ::= `ods_def` `<` bare-id `>` `:` tc-def
+///
 /// All the affine-expr in a `tensor-typedef` must be dimensionless (i.e.
 /// contain only expressions involving symbols and constants), but can
 /// otherwise contain arbitrary affine expressions.
-LogicalResult TCParser::parseAndEmitTCDef(llvm::raw_ostream &os) {
+LogicalResult TCParser::parseAndEmitODSDef(llvm::raw_ostream &os) {
+  if (failed(parser.parseToken(Token::Kind::kw_ods_def,
+                               "expected 'ods_def' to define a TC ODS")) ||
+      failed(parser.parseToken(Token::Kind::lt, "expected '<'")))
+    return failure();
+  StringRef cppOpName = parser.curToken.getSpelling();
+  LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing ODS: " << cppOpName << "\n");
+
+  if (failed(parser.parseToken(Token::Kind::id, "expected id")) ||
+      failed(parser.parseToken(Token::Kind::gt, "expected '>'")) ||
+      failed(parser.parseToken(Token::Kind::colon, "expected ':'")))
+    return failure();
   if (failed(parser.parseToken(Token::Kind::kw_def,
                                "expected 'def' to define a TC")))
     return failure();
 
   StringRef tcName = parser.curToken.getSpelling();
-  LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing tc: " << tcName << "\n");
+  LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing TC: " << tcName << "\n");
   if (failed(parser.parseToken(Token::Kind::id, "expected id")) ||
       failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
     return failure();
 
   auto parseInputDef = [&]() -> LogicalResult {
     return parseTensorDef(/*isOutput=*/false);
   };
   if (failed(parser.parseCommaSeparatedListUntil(
LogicalResult TCParser::parseAndEmitODSDef(llvm::raw_ostream &os) {
   }
 
   if (failed(parser.parseToken(Token::Kind::l_brace, "expected '{'")))
     return failure();
 
   SmallVector<ComprehensionParsingState, 4> perComprehensionStates;
   while (parser.curToken.isNot(Token::Kind::r_brace)) {
     perComprehensionStates.push_back(ComprehensionParsingState());
-    if (failed(parseOneComprehension(tcName, tcName,
+    if (failed(parseOneComprehension(cppOpName, tcName,
                                      perComprehensionStates.back())))
       return failure();
   };
   parser.parseToken(Token::Kind::r_brace, "expected '}'");
 
   // Print.
   auto nComprehensions = perComprehensionStates.size();
   if (nComprehensions != 1) {
     parser.emitError("only 1 comprehension supported for now, got: " +
                      llvm::Twine(nComprehensions));
     return failure();
   }
   if (genODSDecl) {
-    printODS(os, tcName, tcName);
+    printODS(os, cppOpName, tcName);
     os << "\n";
   }
   if (genODSImpl) {
     auto &state = perComprehensionStates.back();
     std::string extraMethods;
     llvm::raw_string_ostream ss(extraMethods);
-    printReferenceIterators(ss, tcName, state);
-    printReferenceIndexingMaps(ss, tcName, state);
-    printRegionBuilder(ss, tcName, state);
+    printReferenceIterators(ss, cppOpName, state);
+    printReferenceIndexingMaps(ss, cppOpName, state);
+    printRegionBuilder(ss, cppOpName, state);
     ss.flush();
     os << extraMethods << "\n";
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Printing functions
 //===----------------------------------------------------------------------===//
 
 /// Print the ODS class that defines a new `cppOpName` for a `linalgOpName`.
 void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName,
                         StringRef linalgOpName) {
-  const char *header = R"FMT(  def {0}Op : LinalgNamedStructured_Op<"{1}", [
+  const char *header = R"FMT(  def {0} : LinalgNamedStructured_Op<"{1}", [
     NInputs<{2}>,
     NOutputs<{3}>,
-    NamedStructuredOpTraits]> {
+    NamedStructuredOpTraits,
+    SingleBlockImplicitTerminator<"YieldOp">]> {
       let arguments = (ins Variadic<LinalgOperand>:$views);
       let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
+      let regions = (region SizedRegion<1>:$region);
+      let builders = [OpBuilder<
+        "Builder *b, OperationState &result, TypeRange outputTypes, "
+        # "ValueRange views",
+        [{{
+          result.addOperands(views);
+          result.addTypes(outputTypes);
+          buildNamedStructuredOpRegion<{0}>(
+            *b, result, TypeRange(views), outputTypes);
+        }]>
+      ];
+      let parser = [{
+        return ::parseNamedStructuredOp<{0}>(parser, result);
+      }];
       let extraClassDeclaration = [{{
         llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
         llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
-        void regionBuilder(ArrayRef<BlockArgument> args);
+        static void regionBuilder(Block &block);
       }];
-      let hasFolder = 1;
   })FMT";
 
   unsigned nInputs = 0, nOutputs = 0;
   for (auto &t : registeredTensors) {
     if (t.getValue().isOutput)
       nOutputs++;
     else
       nInputs++;
   }
 
   os << llvm::formatv(header, cppOpName, linalgOpName, nInputs, nOutputs);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `referenceIterators`.
-void TCParser::printReferenceIterators(llvm::raw_ostream &os, StringRef opId,
+void TCParser::printReferenceIterators(llvm::raw_ostream &os,
+                                       StringRef cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(
     llvm::Optional<SmallVector<StringRef, 8>> {0}::referenceIterators() {
       return SmallVector<StringRef, 8>{{ {1} };
     })FMT";
 
   std::string iteratorsStr;
llvm::interleaveComma(
             break;
         }
         ss << (reduction ? "getReductionIteratorTypeName()"
                          : "getParallelIteratorTypeName()");
         pos++;
       });
   ss.flush();
 
-  os << llvm::formatv(referenceReferenceIteratorsFmt, opId, iteratorsStr);
+  os << llvm::formatv(referenceReferenceIteratorsFmt, cppOpName, iteratorsStr);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
-void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef opId,
+void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
+                                          StringRef cppOpName,
                                           ComprehensionParsingState &state) {
   const char *referenceIndexingMapsFmt =
       R"FMT(
   llvm::Optional<SmallVector<AffineMap, 8>> {0}::referenceIndexingMaps() {
     MLIRContext *context = getContext();
     AffineExpr {1};
     bindDims(context, {1});
     return SmallVector<AffineMap, 8>{{ {2} };
   })FMT";
 
   std::string dimsStr;
   llvm::raw_string_ostream ss(dimsStr);
   llvm::interleaveComma(
       state.dims, ss,
       [&](std::pair<StringRef, AffineExpr> p) { ss << p.second; });
   ss.flush();
 
   std::string mapsStr;
   llvm::raw_string_ostream mapsStringStream(mapsStr);
   SmallVector<TensorUse, 4> orderedUses(state.orderedTensorArgs.size());
   for (auto it : state.orderedTensorArgs)
     orderedUses[it.second] = it.first;
   llvm::interleaveComma(orderedUses, mapsStringStream, [&](TensorUse u) {
     assert(u.indexingMap);
-    const char *mapFmt = "\n\tAffineMap::get({0}, 0, {1})";
+    const char *mapFmt = "\n\tAffineMap::get({0}, 0, {1}, context)";
     if (u.indexingMap.isEmpty()) {
       mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), "context");
       return;
     }
 
     std::string exprsStr;
     llvm::raw_string_ostream exprsStringStream(exprsStr);
     exprsStringStream << "{";
     llvm::interleaveComma(u.indexingMap.getResults(), exprsStringStream);
     exprsStringStream << "}";
     exprsStringStream.flush();
 
     mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), exprsStr);
   });
   mapsStringStream.flush();
 
-  os << llvm::formatv(referenceIndexingMapsFmt, opId, dimsStr, mapsStr);
+  os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr);
 }
 
 /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
-void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef opId,
+void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                                   ComprehensionParsingState &state) {
   unsigned count = state.orderedTensorArgs.size();
   llvm::DenseMap<const TensorExpr *, unsigned> subExprsMap;
   std::function<void(llvm::raw_ostream & os, const Expression &)> printExpr;
   printExpr = [&](llvm::raw_ostream &os, const Expression &e) -> void {
     if (auto *pUse = dyn_cast<TensorUse>(&e)) {
       os << "_" << state.orderedTensorArgs.find(*pUse)->second;
       return;
     }
     auto *pTensorExpr = cast<TensorExpr>(&e);
     if (subExprsMap.count(pTensorExpr) > 0) {
       os << "_" << subExprsMap[pTensorExpr];
     } else {
       std::string subExprs;
       llvm::raw_string_ostream subExprsStringStream(subExprs);
       llvm::interleaveComma(pTensorExpr->expressions, subExprsStringStream,
                             [&](const std::unique_ptr<Expression> &e) {
                               printExpr(subExprsStringStream, *e);
                             });
       subExprsStringStream.flush();
       const char *tensorExprFmt = "\n    ValueHandle _{0} = {1}({2});";
-      os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->opId, subExprs);
+      os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->operationName,
+                          subExprs);
       subExprsMap[pTensorExpr] = count;
     }
   };
 
   const char *regionBuilderFmt = R"FMT(
-  void {0}::regionBuilder(ArrayRef<BlockArgument> args) {
+  void {0}::regionBuilder(Block &block) {
     using namespace edsc;
     using namespace intrinsics;
+    auto args = block.getArguments();
     ValueHandle {1};
     {2}
     (linalg_yield(ValueRange{ {3} }));
   })FMT";
 
   unsigned idx = 0;
   std::string valueHandleStr;
   llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);
llvm::interleaveComma(state.expressions, yieldStringStream,
                         [&](const std::unique_ptr<Expression> &e) {
                           printExpr(yieldStringStream, *e);
                         });
 
   valueHandleStringStream.flush();
   expressionStringStream.flush();
   yieldStringStream.flush();
 
-  os << llvm::formatv(regionBuilderFmt, opId, valueHandleStr, expressionsStr,
-                      yieldStr);
+  os << llvm::formatv(regionBuilderFmt, cppOpName, valueHandleStr,
+                      expressionsStr, yieldStr);
 }
 
 /// Iterate over each Tensor Comprehension def.
 LogicalResult parseAndEmitAllTensorComprehensions(llvm::raw_ostream &os,
                                                   Parser &parser) {
   while (parser.curToken.getKind() != Token::Kind::eof) {
     TCParser tcParser(parser);
-    if (failed(tcParser.parseAndEmitTCDef(os)))
+    if (failed(tcParser.parseAndEmitODSDef(os)))
       return failure();
   }
   return success();
 }
 
 int main(int argc, char **argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "Linalg ODS Gen");