[mlir][Linalg] NFC - Refactor in preparation for automatic Linalg "named" ops.

[nicolasvasilache]

Summary

This revision prepares the ground for declaratively defining Linalg "named" ops. Such named ops form the backbone of operations that are ubiquitous in the ML application domain.

This revision closely related to the definition of a "Tensor Computation Primitives Dialect" and demonstrates that ops can be expressed as declarative configurations of the linalg.generic op.

[mravishankar]

Overall looks fine to me.

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

368

Curious why. Isnt it
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>?

392

Same as above. Is there some information missing that blocks this?

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

234

I think you dont need this. You already have an unreachable in referenceIteratorTypes

266

Same here. You probably dont need this. referenceIndexingMaps already has the unreachable where needed

[jpienaar]

Wanted to give some feedback before next meeting, but I didn't review the C++ code changes in depth.

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

137

You probably want to move these out into a different file. It is rare but if linalg ops get included along with other interfaces then it would result in generating the interfaces 2x and have some non-intuitive behavior :) Keeping op interfaces in separate file so that the separate mlir-tblgen invocations for op def and op interface def can be used keeps things simple.

138

This isn't very descriptive, e.g., what it the result being returned here?. Could we reformulate it something like Returns the reference iterator attribute ?

I'm also not sure what Type signify here.

141

What does a return of none signify?

157

Same here and below. "Query" makes me think of a layer of indirection.

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

256

Could we move this to a constant on the trait? Or would the trait generate an accessor for it?

[nicolasvasilache]

ddress review comments.

[nicolasvasilache]

ld breaking typo.

[nicolasvasilache]

Fix build breadking typo.

[nicolasvasilache]

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

137

yes I am doing that in the next commit along with redoing the tablegen part.
I'm considering this as followup work, please shout if you disagree.

138

I rephrased and renamed, thanks!

141

updated, thanks!

368

Right now it is still encoded in different pieces in C++.
The idea is I will give a declarative tablgen form to autogenerate.

If using the parser is acceptable for this it could work too but the problem it that:

1. it can't be a static method if we want rank polymorphism
2. it is recomputed at each call so calling the Parser seems overkill.

anyway, fun discussions that are bigger than just this NFC revision.

392

same as above too :)

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

234

This is setting the ground for the next revision, I agree it is redundant right now.

256

I am not sure what a constant on the trait would mean here.
The "referenceIndexingMaps" will be a constant in the trait after the next revision and tablegen additions.

This property is the one that is used for all transformations.
For named ops it will not be set but specified and autogenerated with Tablegen.

266

same as above re future evolution.

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

 //===- LinalgOps.h - Linalg Operations --------------------------*- 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_LINALGOPS_H_
 #define MLIR_DIALECT_LINALG_LINALGOPS_H_
 
 #include "mlir/Dialect/Linalg/IR/LinalgTraits.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"

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

let methods = [
     >,
     InterfaceMethod<[{
         Query the type of the input shape at the given index.
       }], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>,
     InterfaceMethod<[{
         Query the subset of input operands that are of ranked tensor type.
       }], "SmallVector<RankedTensorType, 4>", "getInputTensorTypes">,
 
-
     //========================================================================//
     // Output arguments handling.
     //========================================================================//
     InterfaceMethod<
       "Query the number of outputs from the current operation.",
       "unsigned", "getNumOutputs"
     >,
     InterfaceMethod<"Query the output buffer at the given index.",
let methods = [
     InterfaceMethod<
       "Return the range over inputs (irrespective of type) and output buffers.",
       "Operation::operand_range", "getInputsAndOutputBuffers"
     >,
 
     //========================================================================//
     // Other interface methods.
     //========================================================================//
     InterfaceMethod<
+      "Query the reference iterator types attribute for this named op. "
+      "Typically this would be a static method but in order to allow "
+      "rank-polymorphic ops, this needs to be per object instance.",
+      "llvm::Optional<SmallVector<StringRef, 8>>", "referenceIteratorTypes"
+    >,
+    InterfaceMethod<
+      "Query the reference indexing maps attribute for this named op. "
+      "Typically this would be a static method but in order to allow "
+      "rank-polymorphic ops, this needs to be per object instance.",
+      "llvm::Optional<SmallVector<AffineMap, 8>>", "referenceIndexingMaps"
+    >,
+    InterfaceMethod<
       "Query the iterator types attribute within the current operation.",
       "ArrayAttr", "iterator_types"
     >,
     InterfaceMethod<
       "Query the indexing maps attribute within the current operation.",
       "ArrayAttr", "indexing_maps"
     >,
+    InterfaceMethod<"Query the input or output indexing map at index `i`.",
+      "AffineMap", "getIndexingMap", (ins "unsigned":$i)
+    >,
+    InterfaceMethod<"Query the input indexing map at index `i`.",
+      "AffineMap", "getInputIndexingMap", (ins "unsigned":$i)
+    >,
+    InterfaceMethod<"Query the output indexing map at index `i`.",
+      "AffineMap", "getOutputIndexingMap", (ins "unsigned":$i)
+    >,
     InterfaceMethod<[{
         Query whether the op has only MemRef input and outputs.
       }], "bool", "hasBufferSemantics">,
     InterfaceMethod<[{
         Query whether the op has only RankedTensor input and outputs.
       }], "bool", "hasTensorSemantics">,
 
     //========================================================================//
code libraryCallName = [{
     std::string getLibraryCallName() {
       return generateLibraryCallName(getOperation());
     }
   }];
   let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
 }
 
 ////////////////////////////////////////////////////////////////////////////////
-// Concrete Linalg ops.
+// Named Linalg ops, implemented as special configurations of a generic op.
 ////////////////////////////////////////////////////////////////////////////////
 def CopyOp : LinalgStructured_Op<"copy", [NInputs<1>, NOutputs<1>]> {
   let description = [{
     Copies the data in the input view into the output view.
 
     Usage:
       ```mlir
       linalg.copy(%arg0, %arg1) : memref<?xf32, stride_specification>,
def CopyOp : LinalgStructured_Op<"copy", [NInputs<1>, NOutputs<1>]> {
   // TODO(ntv) this should go away once the usage of OptionalAttr triggers
   // emission of builders with default arguments left unspecified.
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value input, Value output", [{
     return build(
       builder, result, input, output, AffineMapAttr(), AffineMapAttr());
   }]>];
   let extraClassDeclaration = libraryCallName # [{
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
+    // Rank-polymorphic.
+    //   filling_value -> O(ivs) with parallel iterators.
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
       unsigned nPar = input().getType().cast<ShapedType>().getRank();
-      MLIRContext *ctx = getContext();
-      SmallVector<Attribute, 8> iters(
-        nPar, StringAttr::get(getParallelIteratorTypeName(), ctx));
-      return ArrayAttr::get(iters, ctx);
+      return SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName());
+    }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for CopyOp");
     }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
 def FillOp : LinalgStructured_Op<"fill", [NInputs<0>, NOutputs<1>]> {
-  let arguments = (ins AnyStridedMemRef:$input,
+  let arguments = (ins AnyStridedMemRef:$output,
                    AnyTypeOf<[AnyFloat, AnyInteger, AnyVector]>:$value);
   let extraClassDeclaration = libraryCallName # [{
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
-      unsigned nPar = input().getType().cast<ShapedType>().getRank();
-      MLIRContext *ctx = getContext();
-      SmallVector<Attribute, 8> iters(
-        nPar, StringAttr::get(getParallelIteratorTypeName(), ctx));
-      return ArrayAttr::get(iters, ctx);
+    // Rank-polymorphic.
+    //   filling_value -> O(ivs) with parallel iterators.
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
+      unsigned nPar = output().getType().cast<ShapedType>().getRank();
+      return SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName());
     }
-  }];
-  let verifier = [{ return ::verify(*this); }];
 
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for CopyOp");
+    }
+  }];
   let hasFolder = 1;
 }
 
 def DotOp : LinalgStructured_Op<"dot", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<0>);
   let extraClassDeclaration = libraryCallName # [{
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
-      MLIRContext *ctx = getContext();
-      return ArrayAttr::get(
-        StringAttr::get(getReductionIteratorTypeName(), ctx), ctx);
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
+      return SmallVector<StringRef, 8>{getReductionIteratorTypeName()};
+    }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for DotOp");
     }
   }];
 
   let hasFolder = 1;
 }
 
 def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<1>,
                        AnyStridedMemRefOfRank<1>);
   let extraClassDeclaration = libraryCallName # [{
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
-      MLIRContext *ctx = getContext();
-      Attribute iters[2]{
-        StringAttr::get(getParallelIteratorTypeName(), ctx),
-        StringAttr::get(getReductionIteratorTypeName(), ctx)};
-      return ArrayAttr::get(iters, ctx);
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
+      return SmallVector<StringRef, 8>{
+        getParallelIteratorTypeName(),
+        getReductionIteratorTypeName()};
+    }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for MatvecOp");
     }
   }];
 
   let hasFolder = 1;
 }
 
 def MatmulOp : LinalgStructured_Op<"matmul", [NInputs<2>, NOutputs<1>]> {
   let arguments = (ins AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<2>,
                        AnyStridedMemRefOfRank<2>);
   let extraClassDeclaration = libraryCallName # [{
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
-      MLIRContext *ctx = getContext();
-      Attribute iters[3]{
-        StringAttr::get(getParallelIteratorTypeName(), ctx),
-        StringAttr::get(getParallelIteratorTypeName(), ctx),
-        StringAttr::get(getReductionIteratorTypeName(), ctx)};
-      return ArrayAttr::get(iters, ctx);
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
+      return SmallVector<StringRef, 8>{
+        getParallelIteratorTypeName(),
+        getParallelIteratorTypeName(),
+        getReductionIteratorTypeName()};
+    }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for MatmulOp");
     }
   }];
 
   let hasFolder = 1;
 }
 
 def ConvOp : LinalgStructured_Op<"conv", [NInputs<2>, NOutputs<1>]> {
   let description = [{
let arguments = (ins AnyStridedMemRef:$filter, AnyStridedMemRef:$input,
                    OptionalAttr<I64ArrayAttr>:$dilations);
   let extraClassDeclaration = libraryCallName # [{
     // TODO(ntv) extend to support more than 1 dimensions and potentially
     // grouping too.
     unsigned getNumBatchDimensions() { return 1; }
     unsigned getNumInputFeatureDimensions() { return 1; }
     unsigned getNumOutputFeatureDimensions() { return 1; }
 
+    // Defined in C++ for now.
+    // TODO(ntv): auto-generate.
     ArrayAttr indexing_maps();
 
-    ArrayAttr iterator_types() {
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
       // Outer parallel loops are always the number of output dimensions; i.e.
-      // [ b, xs, q] in the TF notation above.
+      // [b, xs, q] in the TF notation above.
       unsigned nPar = getOutputShapedType(0).getRank();
       unsigned nRed = getNumInputFeatureDimensions();
       // Window loops are a special kind of reduction that is never tiled or
       // parallelized across; i.e. [zs] in the TF notation above whose number
       // match `xs` (i.e. 1 window loop per "image" dimension).
       // This may evolve in the future.
       unsigned nWin =
         nPar - getNumBatchDimensions() - getNumInputFeatureDimensions();
-      MLIRContext *ctx = getContext();
-      SmallVector<Attribute, 8> iters(
-        nPar, StringAttr::get(getParallelIteratorTypeName(), ctx));
+      SmallVector<StringRef, 8> iters(nPar, getParallelIteratorTypeName());
       iters.reserve(nPar + nRed + nWin);
-      iters.append(nRed, StringAttr::get(getReductionIteratorTypeName(), ctx));
-      iters.append(nWin, StringAttr::get(getWindowIteratorTypeName(), ctx));
-      return ArrayAttr::get(iters, ctx);
+      iters.append(nRed, getReductionIteratorTypeName());
+      iters.append(nWin, getWindowIteratorTypeName());
+      return iters;
     }
 
     int64_t getStride(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!strides().hasValue()) return 1;
       return strides()->getValue()[i]
         .cast<IntegerAttr>().getValue().getSExtValue();
     }
 
     int64_t getDilation(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!dilations().hasValue()) return 1;
       return dilations()->getValue()[i]
         .cast<IntegerAttr>().getValue().getSExtValue();
     }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable("NYI referenceIndexingMaps for MatmulOp");
+    }
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
 def LinalgOperand: Type<
   Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>;
 
 class LinalgOperandOfRank<int rank>: Type<
   And<[
     LinalgOperand.predicate,
     CPred<"$_self.cast<ShapedType>().getRank() == " # rank>]
   >>;
 
+////////////////////////////////////////////////////////////////////////////////
+// Generic Linalg ops.
+////////////////////////////////////////////////////////////////////////////////
 class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, []> {
   let arguments = (ins Variadic<LinalgOperand>:$views,
                    I64Attr:$args_in,
                    I64Attr:$args_out,
                    AffineMapArrayAttr:$indexing_maps,
                    ArrayAttr:$iterator_types,
                    OptionalAttr<StrAttr>:$doc,
                    OptionalAttr<FlatSymbolRefAttr>:$fun,
                    OptionalAttr<StrAttr>:$library_call);
   let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = [{
     SmallVector<StringRef, 8> linalgTraitAttrNames() {
       return SmallVector<StringRef, 8>{
         getArgsInAttrName(), getArgsOutAttrName(), getDocAttrName(),
         getFunAttrName(), getIndexingMapsAttrName(), getLibraryCallAttrName(),
         getIteratorTypesAttrName()
       };
     }
+
     unsigned getNumInputs() { return args_in().getSExtValue(); }
+
     unsigned getNumOutputs() { return args_out().getSExtValue(); }
+
     FuncOp getFunction() {
       auto moduleOp = getParentOfType<ModuleOp>();
       return fun().hasValue() ?
         moduleOp.lookupSymbol<FuncOp>(fun().getValue()) : FuncOp();
     }
+
     StringRef getLibraryCallName() {
       return library_call().hasValue() ? library_call().getValue() : "";
     }
-    AffineMap getIndexingMap(unsigned i) {
-      assert(i < getNumInputsAndOutputs());
-      return indexing_maps().getValue()[i].cast<AffineMapAttr>().getValue();
-    }
-    AffineMap getInputIndexingMap(unsigned i) {
-      assert(i < getNumInputs());
-      return indexing_maps().getValue()[i].cast<AffineMapAttr>().getValue();
-    }
-    AffineMap getOutputIndexingMap(unsigned i) {
-      assert(i < getNumOutputs());
-      return indexing_maps().getValue()[i + getNumInputs()]
-          .cast<AffineMapAttr>().getValue();
+
+    llvm::Optional<SmallVector<StringRef, 8>> referenceIteratorTypes() {
+      llvm_unreachable(
+        "No such thing as reference iterator types for a generic op.");
+     }
+
+    llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps() {
+      llvm_unreachable(
+        "No such thing as reference indexing maps for a generic op.");
      }
   }];
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parseGenericOp(parser, result); }];
 }
 
+/// Index-free GenericOp.
 def GenericOp : GenericOpBase<"generic"> {
   let description = [{
     Generic Linalg op form where the key properties of the computation are
     specified as attributes. In pretty form, a linalg.generic op is written as:
 
       ```mlir
         linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
           memref<?x?xf32, stride_specification>,
let description = [{
     future.
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
+/// GenericOp with Indexing (i.e. multi-for style in which the region is passed
+/// the enclosing loop induction variables)
 def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
   let description = [{
     Indexed Generic Linalg op form where the key properties of the computation
     are specified as attributes. In pretty form, a linalg.indexed_generic op is
     written as:
 
       ```mlir
         linalg.indexed_generic #trait_attribute %A, %B, %C {other-attributes} :

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/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace OpTrait {
 namespace linalg {
 
if (i < getNumInputsAndOutputBuffers())
       return getOutputBufferType(i - nInputs()).template cast<ShapedType>();
     return getOutputTensorTypes()[i - getNumInputsAndOutputBuffers()]
         .template cast<ShapedType>();
   }
 
   //==========================================================================//
   // Other interface methods.
   //==========================================================================//
+
+  // Get or build the indexing_maps ArrayAttr.
+  ArrayAttr iterator_types() {
+    // Return the attribute if it is present.
+    if (auto attr = this->getOperation()->getAttr("iterator_types"))
+      return attr.template cast<ArrayAttr>();
+
+    // If not, form the attribute using the reference iterator types for the
+    // ConcreteType.
+    auto maybeReferenceIteratorTypes =
+        cast<ConcreteType>(this->getOperation()).referenceIteratorTypes();
+
+    // If there is no reference, this must be a generic op.
+    // TODO(ntv): Traits are used to define ops. Split into cpp to avoid
+    // cyclic dependency.
+    auto name = this->getOperation()->getName().getStringRef();
+    if (!maybeReferenceIteratorTypes && name != "generic" &&
+        name != "indexed_generic") {
+      this->getOperation()->dump();
+      llvm_unreachable("Op missing referenceIteratorTypes");
+    }
+
+    // If we have a reference, build the reference attribute.
+    auto *ctx = this->getOperation()->getContext();
+    auto attrRange = llvm::map_range(*maybeReferenceIteratorTypes,
+                                     [ctx](StringRef str) -> Attribute {
+                                       return StringAttr::get(str, ctx);
+                                     });
+    auto attr = ArrayAttr::get(llvm::to_vector<4>(attrRange), ctx);
+    // TODO(ntv): Need to memoize this. Can't just store as an attribute atm as
+    // it will impact parser, printer and tests.
+    // this->getOperation()->setAttr("iterator_types", attr);
+    return attr;
+  }
+
+  // Get or build the indexing_maps ArrayAttr.
+  ArrayAttr indexing_maps() {
+    // Return the attribute if it is present.
+    if (auto attr = this->getOperation()->getAttr("indexing_maps"))
+      return attr.template cast<ArrayAttr>();
+
+    // If not, form the attribute using the reference indexing map for the
+    // ConcreteType.
+    auto maybeReferenceIndexingMaps =
+        cast<ConcreteType>(this->getOperation()).referenceIndexingMaps();
+
+    // If there is no reference, this must be a generic op.
+    auto name = this->getOperation()->getName().getStringRef();
+    if (!maybeReferenceIndexingMaps && name != "generic" &&
+        name != "indexed_generic") {
+      this->getOperation()->dump();
+      llvm_unreachable("Op missing referenceIndexingMaps");
+    }
+
+    // If we have a reference, build the reference attribute and set it in the
+    // op before returning.
+    auto *ctx = this->getOperation()->getContext();
+    auto attrRange =
+        llvm::map_range(*maybeReferenceIndexingMaps, [ctx](AffineMap map) {
+          // 0-D corner case because there is no such thing as a concrete empty
+          // map type.
+          if (!map)
+            map = AffineMap::get(0, 0, getAffineConstantExpr(0, ctx));
+          return AffineMapAttr::get(map);
+        });
+    SmallVector<Attribute, 4> attrs{attrRange.begin(), attrRange.end()};
+    auto attr = ArrayAttr::get(attrs, ctx);
+    // TODO(ntv): Need to memoize this. Can't just store as an attribute atm as
+    // it will impact parser, printer and tests.
+    // this->getOperation()->setAttr("indexing_maps", attr);
+    return attr;
+  }
+
+  AffineMap getIndexingMap(unsigned i) {
+    assert(i < getNumInputsAndOutputs());
+    return indexing_maps()
+        .getValue()[i]
+        .template cast<AffineMapAttr>()
+        .getValue();
+  }
+
+  AffineMap getInputIndexingMap(unsigned i) {
+    assert(i < nInputs());
+    return indexing_maps()
+        .getValue()[i]
+        .template cast<AffineMapAttr>()
+        .getValue();
+  }
+
+  AffineMap getOutputIndexingMap(unsigned i) {
+    assert(i < nOutputs());
+    return indexing_maps()
+        .getValue()[i + nInputs()]
+        .template cast<AffineMapAttr>()
+        .getValue();
+  }
+
   /// Query whether the op has only buffer inputs and no returns.
   bool hasBufferSemantics() {
     return this->getOperation()->getNumResults() == 0 &&
            llvm::all_of(getInputs(),
                         [](Value v) { return v.getType().isa<MemRefType>(); });
   }
 
   /// Query whether the op has only tensor inputs and outputs.

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

static LogicalResult verify(ConvOp op) {
   if (auto dilations = op.dilations()) {
     if (failed(verifyStrideOrDilation(op, dilations->getValue(),
                                       /*isStride=*/false)))
       return failure();
   }
   return success();
 }
 
+static AffineMap extractOrIdentityMap(Optional<AffineMap> maybeMap,
+                                      unsigned rank, MLIRContext *context) {
+  if (maybeMap)
+    return maybeMap.getValue();
+  if (rank == 0)
+    return AffineMap();
+  return AffineMap::getMultiDimIdentityMap(rank, context);
+}
+
 namespace mlir {
 namespace linalg {
 
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
 
 } // namespace linalg
 } // namespace mlir
 
-static AffineMap extractOrIdentityMap(Optional<AffineMap> maybeMap,
-                                      unsigned rank, MLIRContext *context) {
-  if (maybeMap)
-    return maybeMap.getValue();
-  if (rank == 0)
-    return AffineMap();
-  return AffineMap::getMultiDimIdentityMap(rank, context);
-}
-
 // Returns `num` AffineDimExpr dimensions at positions [curIdx, curIdx + num)
 // and increments `curIdx` to `curIdx + num`.
 static SmallVector<AffineExpr, 4>
 makeAffineDimExprs(unsigned num, unsigned &curIdx, MLIRContext *context) {
   SmallVector<AffineExpr, 4> res;
   res.reserve(num);
   for (unsigned i = 0; i < num; ++i)
     res.push_back(getAffineDimExpr(curIdx++, context));
return SmallVector<AffineMap, 4>{
         // filter[z[0], ..., z[N-1], q, k]
         AffineMap::get(idx, 0, concat(concat(zs, qs), ks)),
         // input[b,
         //       x[0]*s[0] + d[0]*z[0], ..., x[N-1]*s[N-1] + d[N-1]*z[N-1],
         //       q]
         AffineMap::get(idx, 0, concat(concat(bs, ws), qs)),
         // output[b, x[0], ..., x[N-1], k]
         AffineMap::get(idx, 0, concat(concat(bs, xs), ks))};
-  } else if (auto genericOp = dyn_cast<GenericOp>(op)) {
-    SmallVector<AffineMap, 4> res;
-    unsigned nViews = genericOp.getNumInputsAndOutputs();
-    res.reserve(nViews);
-    for (unsigned i = 0, e = nViews; i < e; ++i) {
-      res.push_back(genericOp.getIndexingMap(i));
-    }
-    return res;
-  } else if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(op)) {
-    SmallVector<AffineMap, 4> res;
-    unsigned nViews = indexedGenericOp.getNumInputsAndOutputs();
-    res.reserve(nViews);
-    for (unsigned i = 0, e = nViews; i < e; ++i)
-      res.push_back(indexedGenericOp.getIndexingMap(i));
-    return res;
-  }
-  llvm_unreachable("Missing loopToOperandRangesMaps for op");
-}
+  auto linalgOp = cast<LinalgOp>(op);
+  unsigned nViews = linalgOp.getNumInputsAndOutputs();
+    res.push_back(linalgOp.getIndexingMap(i));
+  assert(nViews == linalgOp.indexing_maps().size());
 
 static void appendMangledType(llvm::raw_string_ostream &ss, Type t) {
   if (auto memref = t.dyn_cast<MemRefType>()) {
     ss << "view";
     for (auto size : memref.getShape())
       if (size < 0)
         ss << "sx";
       else