[mlir][Linalg] Add support to lower named ops to loops.

[nicolasvasilache]

Summary

This revision adds support to allow named ops to lower to loops.
Linalg.batch_matmul successfully lowers to loops and to LLVM.

In the process, this test also activates linalg to affine loops.
However padded convolutions to not lower to affine.load atm so this revision overrides the type of underlying load / store operation.

[rriddle]

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

1119

Can we use std::array here instead?

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

125–154

nit: Please use /// for comments.

185

Depends on what you need, but you could use mlir::inlineRegion. You just need to provide an InlinerInterface.

[nicolasvasilache]

Update affine loops.

[nicolasvasilache]

Address review.

[nicolasvasilache]

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

185

Ack, keeping that for a global change in a separate revision.
Thanks!

[bondhugula]

Looks good... some minor comments.

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

357–358

A couple of typos here: access and StdIndexedValue

358

I might have missed this but a test case exercising this path is probably missing.

695

Nit: typo operands -> operand.

mlir/test/Dialect/Linalg/loops.mlir

3

Nit: I think you have a separate file affine.mlir for -convert-linalg-to-affine-loops?

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

1574

A comment here and perhaps above will help.

[ftynse]

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

126

Does "1." actually emit stores?

133

Nit: // -> ///

mlir/test/Dialect/Linalg/loops.mlir

3

Why do you need to disable pass threading here?

6

I'm surprised this ever worked on non-unix system

892

The comment in the code says loads are emitted in the same order as operands, so why "DAG" ?

895

DAG here looks like it would allow breaking use-def chains...

[nicolasvasilache]

Address review comments.

mlir/test/Dialect/Linalg/loops.mlir

6

Hmm, I've picked this up from here IIRC

[mravishankar]

Overall looks fine to me. Thanks

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

127

This is not necessary for this CL, but you only need to emit loads for output views if there is a use of the corresponding arguments. Something I have been wanting to address for a while now.

458

Ack :)

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

public:
   }
 };
 
 /// 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();
+  static SmallVector<StringRef, 8> referenceIterators(TypeRange inputTypes,
+                                                      TypeRange outputTypes);
+
+  static SmallVector<AffineMap, 8> referenceIndexingMaps(TypeRange inputTypes,
+                                                         TypeRange outputTypes);
 };
 
 } // namespace linalg
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_LINALGTRAITS_H_

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

 #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);
+static void buildNamedStructuredOpRegionAndAttributes(
+    Builder &builder, OperationState &result, TypeRange operandTypes,
 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);
 
OpFoldResult TransposeOp::fold(ArrayRef<Attribute>) {
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // Auto-generated Linalg named ops.
 //===----------------------------------------------------------------------===//
 
 template <typename NamedStructuredOpType>
-void buildNamedStructuredOpRegion(Builder &builder, OperationState &result,
+void buildNamedStructuredOpRegionAndAttributes(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);
+
+  auto indexingMaps = builder.getAffineMapArrayAttr(
+      NamedStructuredOpType::referenceIndexingMaps(operandTypes,
+                                                   tensorResultTypes));
+  result.addAttribute(getIndexingMapsAttrName(), indexingMaps);
+
+  auto iterators =
+      builder.getStrArrayAttr(NamedStructuredOpType::referenceIterators(
+          operandTypes, tensorResultTypes));
+  result.addAttribute(getIteratorTypesAttrName(), iterators);
 }
 
 template <typename NamedStructuredOpType>
 static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op) {
+  SmallVector<StringRef, 8> silentAttrNames{getIndexingMapsAttrName(),
+                                            getIteratorTypesAttrName()};
   p << op.getOperationName() << ' ';
-  p.printOptionalAttrDict(op.getAttrs());
+  p.printOptionalAttrDict(op.getAttrs(), silentAttrNames);
   p << ' ' << op.getOperands();
   p << ": (" << op.getOperandTypes() << ")";
   auto outputTensorTypes = op.getResultTypes();
   if (!outputTensorTypes.empty())
     p << " -> (" << outputTensorTypes << ")";
 }
 
 template <typename NamedStructuredOpType>
static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
   // 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>(
+  buildNamedStructuredOpRegionAndAttributes<NamedStructuredOpType>(
       parser.getBuilder(), result, operandTypes, tensorResultTypes);
 
   return parser.resolveOperands(operandsInfo, operandTypes,
                                 parser.getCurrentLocation(), result.operands);
 }
 
 template <typename NamedStructuredOpType>
 static LogicalResult verifyNamedStructuredOp(NamedStructuredOpType op) {

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

SmallVector<Value, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
   SmallVector<Value, 4> res;
   for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
     res.push_back(
         linalg_range(std_constant_index(0), sizes[idx], std_constant_index(1)));
   }
   return res;
 }
 
-template <typename OpType>
-static void
-inlineRegionAndEmitStdStore(OpType op, ArrayRef<Value> indexedValues,
+template <typename IndexedValueType, typename OpType>
+static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
                                      ArrayRef<SmallVector<Value, 8>> indexing,
                                      ArrayRef<Value> outputBuffers) {
   auto &b = ScopedContext::getBuilder();
   auto &block = op.region().front();
   BlockAndValueMapping map;
   map.map(block.getArguments(), indexedValues);
   for (auto &op : block.without_terminator()) {
     assert(op.getNumRegions() == 0 && "expected a non-nested region");
     auto *newOp = b.clone(op, map);
     map.map(op.getResults(), newOp->getResults());
   }
 
   Operation &terminator = block.back();
   assert(isa<YieldOp>(terminator) &&
-         "expected an yield op in the end of the region");
+         "expected a yield op in the end of the region");
   for (unsigned i = 0, e = terminator.getNumOperands(); i < e; ++i) {
-    std_store(map.lookupOrDefault(terminator.getOperand(i)), outputBuffers[i],
-              ArrayRef<Value>{indexing[i].begin(), indexing[i].end()});
+    IndexedValueType O(outputBuffers[i]);
+    O(indexing[i]) = map.lookupOrDefault(terminator.getOperand(i));
   }
 }
 
 // Returns a pair that contains input indices and output indices of a
 // SingleInputPoolingOp `op`.
 struct InputAndOutputIndices {
   SmallVector<Value, 8> inputs;
   SmallVector<Value, 8> outputs;
 };
 template <typename SingleInputPoolingOp>
 static InputAndOutputIndices getInputAndOutputIndices(ArrayRef<Value> allIvs,
                                                       SingleInputPoolingOp op) {
   auto &b = ScopedContext::getBuilder();
   auto loc = ScopedContext::getLocation();
   auto mapsRange = op.indexing_maps().template getAsRange<AffineMapAttr>();
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   return InputAndOutputIndices{
       makeCanonicalAffineApplies(b, loc, maps[0], allIvs),
       makeCanonicalAffineApplies(b, loc, maps[2], allIvs)};
 }
 
 namespace {
 
-// Generic loop emitter, to be specialized on an op-per op basis.
-// TODO: Hook up to named ops interface and, later, retire when all named ops
-// are auto-generated.
+// Emits the MLIR for the scalar part of the generic op by:
+//   1. Emitting std_load and std_store ops for each input and output
+//      view in order. This is achieved by applying the appropriate input or
+//      output map to the enclosing induction variables.
+//   2. Emitting a call to `op.fun()` that takes as arguments the scalars
+//      from point 1. above.
+//   3. Emitting std_store to store the results of 2. to the output
+//      views.
+//
+// An example output may resemble:
+//
+// ```
+//    loop.for %i = %c0 to %0 step %c1 {
+//      loop.for %j = %c0 to %1 step %c1 {
+//        loop.for %k = %c0 to %4 step %c1 {
+//          %11 = load %arg0[%i, %j] :
+//            memref<?x?xf32, stride_specification>
+//          %12 = load %arg1[%i, %j, %k] :
+//            memref<?x?x?xf32, stride_specification>
+//          %13 = load %arg2[%i, %k, %j] :
+//            memref<?x?x?xf32, stride_specification>
+//          %14:2 = call @foo(%11, %12, %13) : (f32, f32, f32) -> (f32, f32)
+//          store %14#0, %arg1[%i, %j, %k] :
+//            memref<?x?x?Xf32, stride_specification>
+//          store %14#1, %arg2[%i, %k, %j] :
+//            memref<?x?x?Xf32, stride_specification>
+//       }
+//      }
+//    }
+// ```
 template <typename IndexedValueType, typename LinalgOpType>
 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);
+    auto b = ScopedContext::getBuilder();
+    auto loc = ScopedContext::getLocation();
+    unsigned nInputs = linalgOp.getNumInputs();
+    unsigned nOutputs = linalgOp.getNumOutputs();
+    SmallVector<Value, 4> indexedValues(nInputs + nOutputs);
+
+    // 1.a. Emit std_load from input views.
+    for (unsigned i = 0; i < nInputs; ++i) {
+      auto indexing = makeCanonicalAffineApplies(
+          b, loc, linalgOp.getInputIndexingMap(i), allIvs);
+      // Pass input[i] through IndexedValueType emits the proper load operation.
+      indexedValues[i] = IndexedValueType(linalgOp.getInput(i))(indexing);
+    }
+
+    // 1.b. Emit std_load from output views.
+    // TODO(mravishankar): Avoid the loads if the corresponding argument of the
+    // region has no uses.
+    for (unsigned i = 0; i < nOutputs; ++i) {
+      auto indexing = makeCanonicalAffineApplies(
+          b, loc, linalgOp.getOutputIndexingMap(i), allIvs);
+      // Pass input[i] through IndexedValueType emits the proper load operation.
+      indexedValues[nInputs + i] =
+          IndexedValueType(linalgOp.getOutputBuffer(i))(indexing);
+    }
+
+    // TODO(ntv): When a region inliner exists, use it.
+    // 2. Inline region, currently only works for a single basic block.
+    // 3. Emit std_store.
+    SmallVector<SmallVector<Value, 8>, 8> indexing;
+    SmallVector<Value, 8> outputBuffers;
+    for (unsigned i = 0; i < nOutputs; ++i) {
+      indexing.push_back(makeCanonicalAffineApplies(
+          b, loc, linalgOp.getOutputIndexingMap(i), allIvs));
+      outputBuffers.push_back(linalgOp.getOutputBuffer(i));
+    }
+    inlineRegionAndEmitStore<IndexedValueType>(linalgOp, indexedValues,
+                                               indexing, outputBuffers);
   }
 };
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, CopyOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
     assert(copyOp.hasBufferSemantics() &&
public:
   }
 };
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, ConvOp> {
 public:
   /// Returns the input value of convOp. If the indices in `imIdx` is out of
   /// boundary, returns 0 instead.
-  static Value getConvOpInput(ConvOp convOp, IndexedValueType im,
+  static Value getConvOpInput(ConvOp convOp, StdIndexedValue im,
                               MutableArrayRef<Value> imIdx) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (!convOp.padding())
       return im(imIdx);
 
     auto *context = ScopedContext::getContext();
     Value zeroIndex = std_constant_index(0);
     SmallVector<Value, 8> conds;
static void emitScalarImplementation(ArrayRef<Value> allIvs, ConvOp convOp) {
     auto maps = llvm::to_vector<8>(llvm::map_range(
         mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
     SmallVector<Value, 8> fIdx(
         makeCanonicalAffineApplies(b, loc, maps[0], allIvs));
     SmallVector<Value, 8> imIdx(
         makeCanonicalAffineApplies(b, loc, maps[1], allIvs));
     SmallVector<Value, 8> oIdx(
         makeCanonicalAffineApplies(b, loc, maps[2], allIvs));
-    IndexedValueType F(convOp.filter()), I(convOp.input()), O(convOp.output());
+
+    // Padded conv involves a affine.max in the memory acces which is not
+    // allowed by affine.load. Override to always use an StdIndecexValue.
+    StdIndexedValue F(convOp.filter()), I(convOp.input()), O(convOp.output());
 
     // Emit scalar form.
     Value paddedInput = getConvOpInput(convOp, I, imIdx);
     O(oIdx) += F(fIdx) * paddedInput;
   }
 };
 
 template <typename IndexedValueType>
static void emitScalarImplementation(ArrayRef<Value> allIvs,
     auto indices = getInputAndOutputIndices(allIvs, op);
     IndexedValueType input(op.input()), output(op.output());
 
     // Emit scalar form.
     output(indices.outputs) += input(indices.inputs);
   }
 };
 
-// Emits the MLIR for the scalar part of the generic op by:
-//   1. Emitting std_load and std_store ops for each input and output
-//      view in order. This is achieved by applying the appropriate input or
-//      output map to the enclosing induction variables.
-//   2. Emitting a call to `op.fun()` that takes as arguments the scalars
-//      from point 1. above.
-//   3. Emitting std_store to store the results of 2. to the output
-//      views.
-//
-// An example output may resemble:
-//
-// ```
-//    loop.for %i = %c0 to %0 step %c1 {
-//      loop.for %j = %c0 to %1 step %c1 {
-//        loop.for %k = %c0 to %4 step %c1 {
-//          %11 = load %arg0[%i, %j] :
-//            memref<?x?xf32, stride_specification>
-//          %12 = load %arg1[%i, %j, %k] :
-//            memref<?x?x?xf32, stride_specification>
-//          %13 = load %arg2[%i, %k, %j] :
-//            memref<?x?x?xf32, stride_specification>
-//          %14:2 = call @foo(%11, %12, %13) : (f32, f32, f32) -> (f32, f32)
-//          store %14#0, %arg1[%i, %j, %k] :
-//            memref<?x?x?Xf32, stride_specification>
-//          store %14#1, %arg2[%i, %k, %j] :
-//            memref<?x?x?Xf32, stride_specification>
-//       }
-//      }
-//    }
-// ```
-template <typename IndexedValueType>
-class LinalgScopedEmitter<IndexedValueType, GenericOp> {
-public:
-  static void emitScalarImplementation(ArrayRef<Value> allIvs,
-                                       GenericOp genericOp) {
-    assert(genericOp.hasBufferSemantics() &&
-           "expected linalg op with buffer semantics");
-    auto b = ScopedContext::getBuilder();
-    auto loc = ScopedContext::getLocation();
-    unsigned nInputs = genericOp.getNumInputs();
-    unsigned nOutputs = genericOp.getNumOutputs();
-    SmallVector<Value, 4> indexedValues(nInputs + nOutputs);
-
-    // 1.a. Emit std_load from input views.
-    for (unsigned i = 0; i < nInputs; ++i) {
-      auto indexing = makeCanonicalAffineApplies(
-          b, loc, genericOp.getInputIndexingMap(i), allIvs);
-      indexedValues[i] = std_load(genericOp.getInput(i), indexing);
-    }
-
-    // 1.b. Emit std_load from output views.
-    // TODO(mravishankar): Avoid the loads if the corresponding argument of the
-    // region has no uses.
-    for (unsigned i = 0; i < nOutputs; ++i) {
-      Value output = genericOp.getOutputBuffer(i);
-      auto indexing = makeCanonicalAffineApplies(
-          b, loc, genericOp.getOutputIndexingMap(i), allIvs);
-      indexedValues[nInputs + i] = std_load(output, indexing);
-    }
-
-    // TODO(ntv): When a region inliner exists, use it.
-    // 2. Inline region, currently only works for a single basic block.
-    // 3. Emit std_store.
-    SmallVector<SmallVector<Value, 8>, 8> indexing;
-    SmallVector<Value, 8> outputBuffers;
-    for (unsigned i = 0; i < nOutputs; ++i) {
-      indexing.push_back(makeCanonicalAffineApplies(
-          b, loc, genericOp.getOutputIndexingMap(i), allIvs));
-      outputBuffers.push_back(genericOp.getOutputBuffer(i));
-    }
-    inlineRegionAndEmitStdStore(genericOp, indexedValues, indexing,
-                                outputBuffers);
-  }
-};
-
 // Emits the MLIR for the scalar part of the indexed generic op by:
 //   1. Emitting std_load and std_store ops for each input and output view in
 //      order. This is achieved by applying the appropriate input or output map
 //      to the enclosing induction variables.
 //   2. Emitting a call to `op.fun()` that takes as arguments the induction
 //      variables and the scalars from point 1. above.
 //   3. Emitting std_store to store the results of 2. to the output views.
 //
static void emitScalarImplementation(ArrayRef<Value> allIvs,
     unsigned nOutputs = indexedGenericOp.getNumOutputs();
     unsigned nLoops = allIvs.size();
     SmallVector<Value, 4> indexedValues(nLoops + nInputs + nOutputs);
 
     for (unsigned i = 0; i < nLoops; ++i) {
       indexedValues[i] = allIvs[i];
     }
 
     // 1.a. Emit std_load from input views.
     for (unsigned i = 0; i < nInputs; ++i) {
       Value input = indexedGenericOp.getInput(i);
       auto indexing = makeCanonicalAffineApplies(
           b, loc, indexedGenericOp.getInputIndexingMap(i), allIvs);
       indexedValues[nLoops + i] = std_load(input, indexing);
     }
 
     // 1.b. Emit std_load from output views.
static void emitScalarImplementation(ArrayRef<Value> allIvs,
     // 3. Emit std_store.
     SmallVector<SmallVector<Value, 8>, 8> indexing;
     SmallVector<Value, 8> outputBuffers;
     for (unsigned i = 0; i < nOutputs; ++i) {
       indexing.push_back(makeCanonicalAffineApplies(
           b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));
       outputBuffers.push_back(indexedGenericOp.getOutputBuffer(i));
     }
-    inlineRegionAndEmitStdStore(indexedGenericOp, indexedValues, indexing,
-                                outputBuffers);
+    inlineRegionAndEmitStore<IndexedValueType>(indexedGenericOp, indexedValues,
+                                               indexing, outputBuffers);
   }
 };
 
 // This struct is for factoring out the implementation and support template
 // instantiations in the following 2 cases:
 //   1. Appending to a list of patterns via RewritePatternList.
 //   2. Direct invocation via `linalgOpToLoops` and `linalgOpToAffineLoops`.
 // The implementation must work both in DRR and inside a RewritePattern. As a

 //   1. Zero operand + returns a single AffineConstantExpr
 //   2. One operand + returns a single AffineDimExpr
 //   3. One operands + returns a single AffineSymbolExpr
 //
 // In the first case, the AffineApplyOp is replaced by a new constant. In the
 // other cases, it is replaced by its unique operand.
 struct FoldAffineOp : public RewritePattern {
   FoldAffineOp(MLIRContext *context)
       : RewritePattern(AffineApplyOp::getOperationName(), 0, context) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     AffineApplyOp affineApplyOp = cast<AffineApplyOp>(op);
     auto map = affineApplyOp.getAffineMap();
     if (map.getNumResults() != 1 || map.getNumInputs() > 1)
       return failure();
 

mlir/test/Dialect/Linalg/loops.mlir

 // RUN: mlir-opt %s -convert-linalg-to-loops | FileCheck --check-prefix=CHECKLOOP %s
 // RUN: mlir-opt %s -convert-linalg-to-parallel-loops | FileCheck --check-prefix=CHECKPARALLEL %s
+// RUN: mlir-opt %s -convert-linalg-to-affine-loops --disable-pass-threading | FileCheck --check-prefix=CHECKAFFINE %s
 
 // Test that we can lower all the way to LLVM without crashing, don't check results here.
-// RUN: mlir-opt %s --convert-linalg-to-llvm -o=/dev/null 2>&1
+// RUN: mlir-opt %s -convert-linalg-to-loops --convert-linalg-to-llvm -o=/dev/null 2>&1
 
 // CHECKLOOP-DAG: #[[strided1D:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECKLOOP-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 // CHECKLOOP-DAG: #[[strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
 // CHECKLOOP-DAG: #[[strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
 // CHECKLOOP-DAG: #[[clampMinMap:.*]] = affine_map<(d0) -> (d0, 0)>
 
 // CHECKLOOP-DAG: #[[Stride1Dilation1:.*]] = affine_map<(d0, d1) -> (d0  + d1)>

 //  CHECKPARALLEL-SAME: %[[ARG0]]: memref<f32>
 //  CHECKPARALLEL-SAME: %[[ARG1]]: memref<f32>
 //  CHECKPARALLEL-SAME: %[[ARG2]]: memref<f32>
 //   CHECKPARALLEL-NOT: loop.for
 //   CHECKPARALLEL-DAG: load %[[ARG0]][]
 //   CHECKPARALLEL-DAG: load %[[ARG1]][]
 //       CHECKPARALLEL: addf
 //       CHECKPARALLEL: store %{{.*}}, %[[ARG2]][]
+
+//----------------------------------------------------------------------------//
+// Named ops to loops.
+//----------------------------------------------------------------------------//
+func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) {
+  linalg.batch_matmul %A, %B, %C : (memref<?x?x?xf32>, memref<?x?x?xf32>, memref<?x?x?xf32>) -> ()
+  return
+}
+// CHECKLOOP-LABEL: @named_batch_matmul
+//  CHECKLOOP-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKLOOP-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKLOOP-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//       CHECKLOOP: %[[B:.*]] = dim %[[mA]], 0 : memref<?x?x?xf32>
+//       CHECKLOOP: %[[M:.*]] = dim %[[mA]], 1 : memref<?x?x?xf32>
+//       CHECKLOOP: %[[K:.*]] = dim %[[mA]], 2 : memref<?x?x?xf32>
+//       CHECKLOOP: %[[N:.*]] = dim %[[mB]], 2 : memref<?x?x?xf32>
+//       CHECKLOOP: loop.for %[[b:.*]] = %{{.*}} to %[[B]] step %{{.*}} {
+//       CHECKLOOP:   loop.for %[[m:.*]] = %{{.*}} to %[[M]] step %{{.*}} {
+//       CHECKLOOP:     loop.for %[[n:.*]] = %{{.*}} to %[[N]] step %{{.*}} {
+//       CHECKLOOP:       loop.for %[[k:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
+//   CHECKLOOP-DAG:       %[[va:.*]] = load %[[mA]][%[[b]], %[[m]], %[[k]]] : memref<?x?x?xf32>
+//   CHECKLOOP-DAG:       %[[vb:.*]] = load %[[mB]][%[[b]], %[[k]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKLOOP-DAG:       %[[vc:.*]] = load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKLOOP-DAG:       %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
+//   CHECKLOOP-DAG:       %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
+//       CHECKLOOP:       store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
+
+// CHECKAFFINE-LABEL: @named_batch_matmul
+//  CHECKAFFINE-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKAFFINE-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKAFFINE-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//       CHECKAFFINE: %[[B:.*]] = dim %[[mA]], 0 : memref<?x?x?xf32>
+//       CHECKAFFINE: %[[M:.*]] = dim %[[mA]], 1 : memref<?x?x?xf32>
+//       CHECKAFFINE: %[[K:.*]] = dim %[[mA]], 2 : memref<?x?x?xf32>
+//       CHECKAFFINE: %[[N:.*]] = dim %[[mB]], 2 : memref<?x?x?xf32>
+//       CHECKAFFINE: affine.for %[[b:.*]] = 0 to %[[B]] {
+//       CHECKAFFINE:   affine.for %[[m:.*]] = 0 to %[[M]] {
+//       CHECKAFFINE:     affine.for %[[n:.*]] = 0 to %[[N]] {
+//       CHECKAFFINE:       affine.for %[[k:.*]] = 0 to %[[K]] {
+//   CHECKAFFINE-DAG:       %[[va:.*]] = affine.load %[[mA]][%[[b]], %[[m]], %[[k]]] : memref<?x?x?xf32>
+//   CHECKAFFINE-DAG:       %[[vb:.*]] = affine.load %[[mB]][%[[b]], %[[k]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKAFFINE-DAG:       %[[vc:.*]] = affine.load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKAFFINE-DAG:       %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
+//   CHECKAFFINE-DAG:       %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
+//       CHECKAFFINE:       affine.store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
+
+// CHECKPARALLEL-LABEL: @named_batch_matmul
+//  CHECKPARALLEL-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKPARALLEL-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//  CHECKPARALLEL-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[B:.*]] = dim %[[mA]], 0 : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[M:.*]] = dim %[[mA]], 1 : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[K:.*]] = dim %[[mA]], 2 : memref<?x?x?xf32>
+//       CHECKPARALLEL: %[[N:.*]] = dim %[[mB]], 2 : memref<?x?x?xf32>
+//       CHECKPARALLEL: loop.parallel (%[[b:.*]], %[[m:.*]], %[[n:.*]]) = ({{.*}}) to (%[[B]], %[[M]], %[[N]]) step ({{.*}}) {
+//       CHECKPARALLEL:   loop.for %[[k:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
+//   CHECKPARALLEL-DAG:       %[[va:.*]] = load %[[mA]][%[[b]], %[[m]], %[[k]]] : memref<?x?x?xf32>
+//   CHECKPARALLEL-DAG:       %[[vb:.*]] = load %[[mB]][%[[b]], %[[k]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKPARALLEL-DAG:       %[[vc:.*]] = load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
+//   CHECKPARALLEL-DAG:       %[[inc:.*]] = mulf %[[va]], %[[vb]] : f32
+//   CHECKPARALLEL-DAG:       %[[res:.*]] = addf %[[vc]], %[[inc]] : f32
+//       CHECKPARALLEL:       store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>

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
 
 // ODS-LABEL: def Test1Op : LinalgNamedStructured_Op<"test1", [
 //  ODS-NEXT:   NInputs<2>
 //  ODS-NEXT:   NOutputs<1>
 //  ODS-NEXT:   NamedStructuredOpTraits
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  Test1Op::referenceIterators() {
-//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+// IMPL-LABEL:  SmallVector<StringRef, 8> Test1Op::referenceIterators
+//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  Test1Op::referenceIndexingMaps() {
+//       IMPL:  SmallVector<AffineMap, 8> 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:  Test1Op::regionBuilder(Block &block) {
+//       IMPL:  void Test1Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 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 Test2Op : LinalgNamedStructured_Op<"test2", [
 //  ODS-NEXT:   NInputs<2>
 //  ODS-NEXT:   NOutputs<1>
 //  ODS-NEXT:   NamedStructuredOpTraits
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  Test2Op::referenceIterators() {
-//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+// IMPL-LABEL:  SmallVector<StringRef, 8> Test2Op::referenceIterators
+//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  Test2Op::referenceIndexingMaps() {
+//       IMPL:  SmallVector<AffineMap, 8> 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:  Test2Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 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 Test3Op : LinalgNamedStructured_Op<"test3", [
 //  ODS-NEXT:   NInputs<2>
 //  ODS-NEXT:   NOutputs<1>
 //  ODS-NEXT:   NamedStructuredOpTraits
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
-// IMPL-LABEL:  Test3Op::referenceIterators() {
-//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+// IMPL-LABEL:  SmallVector<StringRef, 8> Test3Op::referenceIterators
+//       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
-//       IMPL:  Test3Op::referenceIndexingMaps() {
+//       IMPL:  SmallVector<AffineMap, 8> 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:  Test3Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 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

SingleBlockImplicitTerminator<"YieldOp">]> {
       let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
       let regions = (region SizedRegion<1>:$region);
       let builders = [OpBuilder<
         "OpBuilder &b, OperationState &result, TypeRange outputTypes, "
         # "ValueRange views",
         [{{
           result.addOperands(views);
           result.addTypes(outputTypes);
-          buildNamedStructuredOpRegion<{0}>(
+          buildNamedStructuredOpRegionAndAttributes<{0}>(
             b, result, TypeRange(views), outputTypes);
         }]>
       ];
       let parser = [{
         return ::parseNamedStructuredOp<{0}>(parser, result);
       }];
       let extraClassDeclaration = [{{
         llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
+        static SmallVector<StringRef, 8> referenceIterators(
+          TypeRange inputTypes, TypeRange outputTypes);
+
         llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
+        static SmallVector<AffineMap, 8> referenceIndexingMaps(
+          TypeRange inputTypes, TypeRange outputTypes);
+
         static void regionBuilder(Block &block);
       }];
   })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 cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(
-    llvm::Optional<SmallVector<StringRef, 8>> {0}::referenceIterators() {
+    // This is temporary until we transition out of manually specified ops
+    // that should be auto-generated with linalg-ods-gen.
+    llvm::Optional<SmallVector<StringRef, 8>> {0}::referenceIterators() {{
+      llvm_unreachable("Unexpected missing `iterator_types` attribute.");
+    }
+    SmallVector<StringRef, 8> {0}::referenceIterators(
+      TypeRange inputTypes, TypeRange outputTypes) {
       return SmallVector<StringRef, 8>{{ {1} };
     })FMT";
 
   std::string iteratorsStr;
   llvm::raw_string_ostream ss(iteratorsStr);
   unsigned pos = 0;
   llvm::interleaveComma(
       state.dims, ss, [&](std::pair<StringRef, AffineExpr> p) {

 }
 
 /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
 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();
+  // This is temporary until we transition out of manually specified ops that
+  // should be auto-generated with linalg-ods-gen.
+  llvm::Optional<SmallVector<AffineMap, 8>> {0}::referenceIndexingMaps() {{
+    llvm_unreachable("Unexpected missing `indexing_maps` attribute.");
+  }
+  SmallVector<AffineMap, 8> {0}::referenceIndexingMaps(
+    TypeRange inputTypes, TypeRange outputTypes) {
+    assert(!inputTypes.empty() && "At least one input expected");
+    MLIRContext *context = (*inputTypes.begin()).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)
+  for (const 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}, context)";
     if (u.indexingMap.isEmpty()) {
       mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), "context");
       return;
     }