[mlir][linalg] adding operation to access the iteration index of enclosing linalg ops.

[gysit]

Summary

The linalg.index operation provides access to the iteration indexes of immediately enclosing linalg operations. It takes a dimension dim attribute and returns the iteration index in the given dimension. Having linalg.index allows us to unify linalg.generic and linalg.indexed_generic and also enables index access in named operations.

[nicolasvasilache]

Nice, thank you for starting this unification effort!

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

622

I see in MemRefOps.td we have something like:

Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment

I am not sure whether you can use with a C++ value instead of 0, but it should at least work for >= 0 (which isn't currently check IIUC).

[gysit]

Confine the dim attribute to values larger than zero.

[gysit]

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

622

Good point! I added IntMinValue to ensure the dim value needs to be larger than or equal to zero. I left the check for the upper bound of dim in the CPP since that check if very much operation specific.

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

InterfaceMethod<
       /*retTy=*/"std::string",
       /*methodName=*/"getLibraryCallName",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getLibraryCallName();
       }]
     >,
-
+    InterfaceMethod<
+      /*desc=*/[{
+         Return whether the op accesses the iteration indices.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"hasIndexSemantics",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/""
+    >,
     //===------------------------------------------------------------------===//
     // Linalg generalization hooks.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
         Hook to provide a custom AffineMap used to compute all the operand
         subshapes given loop bounds. This is used to answer the question: "given
         an iteration space over the codomain, what are the subshapes of the

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

OpResult getTiedOpResult(OpOperand& opOperand) {
 
     unsigned getNumLoops() { return step().size(); }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
+def Linalg_IndexOp : Linalg_Op<"index", [NoSideEffect]>,
+    Arguments<(ins I64Attr:$dim)>,
+    Results<(outs Index:$result)> {
+  let summary = "linalg index operation";
+  let description = [{
+    The `linalg.index` operation returns the iteration index of the immediately
+    enclosing linalg structured operation for the iteration dimension `dim`. The
+    `dim` attribute specifies the position of the accessed dimension in the
+    indexing map domain.
+
+    Example:
+
+    ```mlir
+    #map = affine_map<(i, j) -> (i, j)>
+    linalg.generic {indexing_maps = [#map, #map],
+                    iterator_types = ["parallel", "parallel"]}
+      outs(%I, %J : memref<?x?xindex>, memref<?x?xindex>) {
+      ^bb0(%arg0 : index, %arg1 : index):
+      // Access the outer iteration dimension i
+      %i = linalg.index 0 : index
+      // Access the inner iteration dimension j
+      %j = linalg.index 1 : index
+      linalg.yield %i, %j : index, index
+    }
+    ```
+
+    This may lower to IR resembling:
+
+    ```mlir
+    %0 = dim %I, %c0 : memref<?x?xindex>
+    %1 = dim %I, %c1 : memref<?x?xindex>
+    scf.for %i = %c0 to %0 step %c1 {
+      scf.for %j = %c0 to %1 step %c1 {
+        store %i, %I[%i, %j] : memref<?x?xindex>
+        store %j, %J[%i, %j] : memref<?x?xindex>
+      }
+    }
+    ```
+  }];
+  let builders = [
+    OpBuilder<(ins "int64_t":$dim),
+      [{ build($_builder, $_state, $_builder.getIndexType(), dim); }]>
+  ];
+
+  let assemblyFormat = [{ $dim attr-dict `:` type($result) }];
+}
 
 #endif // LINALG_OPS

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

class LinalgStructuredBase_Op<string mnemonic, list<OpTrait> props>
   code structuredOpsBaseDecls = [{
     // Return the number of induction variables in the basic block. This should
     // always be 0 for index-free linalg ops. For IndexedGeneric, this must be
     // equal to numLoops.
     unsigned getNumPayloadInductionVariables() {
       return isa<IndexedGenericOp>(this->getOperation()) ? getNumLoops() : 0;
     }
 
+    // Return whether the op accesses the iteration indices.
+    bool hasIndexSemantics() {
+      Operation *op = this->getOperation();
+      if(op->getNumRegions() == 0 || op->getRegion(0).empty())
+        return false;
+      return !op->getRegion(0).front().getOps<IndexOp>().empty();
+    }
+
     LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
         SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
       return cast<LinalgOp>(getOperation()).reifyReturnTypeShapesPerResultDim(b,
           reifiedReturnShapes);
     }
   }];
 }
 

mlir/lib/Conversion/LinalgToStandard/LinalgToStandard.cpp

 }
 
 LogicalResult mlir::linalg::LinalgOpToLibraryCallRewrite::matchAndRewrite(
     LinalgOp op, PatternRewriter &rewriter) const {
   // Only LinalgOp for which there is no specialized pattern go through this.
   if (isa<CopyOp>(op) || isa<IndexedGenericOp>(op))
     return failure();
 
+  // TODO: remove once index ops are supported.
+  if (op.hasIndexSemantics())
+    return failure();
+
   auto libraryCallName = getLibraryCallSymbolRef(op, rewriter);
   if (!libraryCallName)
     return failure();
 
   rewriter.replaceOpWithNewOp<mlir::CallOp>(
       op, libraryCallName.getValue(), TypeRange(),
       createTypeCanonicalizedMemRefOperands(rewriter, op->getLoc(),
                                             op->getOperands()));

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

void TiledLoopOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
   results.insert<TiledLoopResultsFolder>(context);
 }
 
 LogicalResult TiledLoopOp::fold(ArrayRef<Attribute>,
                                 SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
+//===----------------------------------------------------------------------===//
+// IndexOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(IndexOp op) {
+  auto linalgOp = dyn_cast<LinalgOp>(op->getParentOp());
+  if (!linalgOp)
+    return op.emitOpError("expected parent op with LinalgOp interface");
+  if (linalgOp.getNumLoops() <= op.dim())
+    return op.emitOpError("expected dim (")
+           << op.dim() << ") to be lower than the number of loops ("
+           << linalgOp.getNumLoops() << ") of the enclosing LinalgOp";
+  return success();
+}
+
 /////// Operations corresponding to library calls defined with Tablegen ////////
 
 template <typename LinalgPoolingOp>
 static LogicalResult verifyStrideOrDilation(LinalgPoolingOp op,
                                             ArrayRef<Attribute> attrs,
                                             bool isStride) {
   auto strideOrDilation = isStride ? "stride" : "dilation";
   if (attrs.size() != op.getNumWindowLoops())

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

 
 namespace {
 /// Pattern to fold unit-trip count loops in GenericOps.
 template <typename GenericOpTy>
 struct FoldUnitDimLoops : public OpRewritePattern<GenericOpTy> {
   using OpRewritePattern<GenericOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOpTy op,
                                 PatternRewriter &rewriter) const override {
+    // TODO: remove once index ops are supported.
+    if (op.hasIndexSemantics())
+      return failure();
+
     SmallVector<AffineMap, 4> indexingMaps = op.getIndexingMaps();
     if (indexingMaps.empty())
       return failure();
 
     // Check if any of the iteration dimensions are unit-trip count. They will
     // end up being unit-trip count if they are used to index into a unit-dim
     // tensor/memref.
     AffineMap invertedMap = inversePermutation(concatAffineMaps(indexingMaps));

 namespace {
 
 /// Pattern to replace tensors operands/results that are unit extents.
 template <typename GenericOpTy>
 struct ReplaceUnitExtentTensors : public OpRewritePattern<GenericOpTy> {
   using OpRewritePattern<GenericOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOpTy op,
                                 PatternRewriter &rewriter) const override {
+    // TODO: remove once index ops are supported.
+    if (op.hasIndexSemantics())
+      return failure();
+
     if (!op.hasTensorSemantics())
       return failure();
 
     MLIRContext *context = rewriter.getContext();
     Location loc = op.getLoc();
 
     SmallVector<AffineMap, 4> newIndexingMaps;
     SmallVector<ArrayAttr, 4> reassociationMaps;

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

 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 /// Implementation of fusion of generic ops and indexed_generic ops.
 static bool areElementwiseOpsFusable(LinalgOp producer, LinalgOp consumer,
                                      unsigned consumerIdx) {
+  // TODO: remove once index ops are supported.
+  if (producer.hasIndexSemantics() || consumer.hasIndexSemantics())
+    return false;
+
   // Producer and consumer must have tensor semantics.
   if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics())
     return false;
 
   // Verify that
   // - the producer has all "parallel" iterator type.
   if (producer.getNumParallelLoops() != producer.getNumLoops())
     return false;

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

 public:
   LinalgRewritePattern(MLIRContext *context,
                        ArrayRef<unsigned> interchangeVector)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
         interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    if (!isa<LinalgOp>(op))
+    auto linalgOp = dyn_cast<LinalgOp>(op);
+    // TODO: remove hasIndexSemantics check once index ops are supported.
+    if (!linalgOp || linalgOp.hasIndexSemantics())
       return failure();
     if (!linalgOpToLoopsImpl<LoopType>(op, rewriter, interchangeVector))
       return failure();
     rewriter.eraseOp(op);
     return success();
   }
 
 private:

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

mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
     Operation *op, PatternRewriter &rewriter, TiledLinalgOp &result) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
+  // TODO: remove hasIndexSemantics check once index ops are supported.
+  if (!linalgOp || linalgOp.hasIndexSemantics())
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
     : RewritePattern(opName, benefit, context, {}),
       dependenceGraph(dependenceGraph), tilingOptions(tilingOptions),
       fusionOptions(fusionOptions), filter(filter),
       fusedOpMarker(fusedOpMarker), originalOpMarker(originalOpMarker) {}
 
 LogicalResult mlir::linalg::LinalgBaseTileAndFusePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
+  // TODO: remove hasIndexSemantics check once index ops are supported.
+  if (!linalgOp || linalgOp.hasIndexSemantics())
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   DenseSet<Operation *> producers;
   producers.insert(linalgOp);
   for (auto dependence : dependenceGraph.getDependentOperationsInto(linalgOp)) {
     Optional<unsigned> operandNumber = dependence.getIndexingOpViewOperandNum();
mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
     ArrayRef<unsigned> interchangeVector, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(filter),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}
 
 LogicalResult mlir::linalg::LinalgBaseInterchangePattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
+  // TODO: remove hasIndexSemantics check once index ops are supported.
+  if (!linalgOp || linalgOp.hasIndexSemantics())
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
   if (failed(interchangeGenericLinalgOpPrecondition(op, interchangeVector)))
     return failure();
 
   // TODO: figure out how this interplays with named ops. In particular this
   // should break the named op property.

 mlir::linalg::LinalgBaseVectorizationPattern::LinalgBaseVectorizationPattern(
     StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
     PatternBenefit benefit)
     : RewritePattern(opName, benefit, context, {}), filter(filter) {}
 
 LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-  if (!linalgOp)
+  // TODO: remove hasIndexSemantics check once index ops are supported.
+  if (!linalgOp || linalgOp.hasIndexSemantics())
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
   SmallVector<Value> newResults;
   if (failed(vectorizeLinalgOp(rewriter, op, newResults)))
     return failure();
   if (!newResults.empty())
     rewriter.replaceOp(op, newResults);

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

LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
   auto linalgOp = cast<linalg::LinalgOp>(op);
   // All types must be static shape to go to vector.
   for (Value operand : linalgOp.getShapedOperands())
     if (!operand.getType().cast<ShapedType>().hasStaticShape())
       return failure();
   for (Type outputTensorType : linalgOp.getOutputTensorTypes())
     if (!outputTensorType.cast<ShapedType>().hasStaticShape())
       return failure();
+  // TODO: remove once index ops are supported.
+  if (linalgOp.hasIndexSemantics())
+    return failure();
   if (isElementwise(op))
     return success();
   return success(isaContractionOpInterface(linalgOp));
 }
 
 LogicalResult
 mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op,
                                 SmallVectorImpl<Value> &newResults) {

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

 //      CHECK:   %[[FILL:.+]] = linalg.fill(%[[INIT]], %{{.+}})
 //      CHECK:   %[[RESULT:.+]] = linalg.generic
 // CHECK-SAME:     indexing_maps = [#[[MAP2]], #[[MAP3]]]
 // CHECK-SAME:     iterator_types = ["parallel", "reduction"]
 // CHECK-SAME:     ins(%[[RESHAPE]] : tensor<?x?xf32>)
 // CHECK-SAME:     outs(%[[FILL]] : tensor<?xf32>)
 //      CHECK:   %[[RESULT_RESHAPE:.+]] = linalg.tensor_reshape %[[RESULT]] [#[[MAP2]]]
 //      CHECK:   return %[[RESULT_RESHAPE]]
+
+// -----
+
+//  CHECK: #{{.+}} = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: @index_op
+func @index_op(%arg0: memref<1x8xindex>) {
+  linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  outs(%arg0 : memref<1x8xindex>) {
+  ^bb0(%arg1: index):   // no predecessors
+    %0 = linalg.index 1 : index
+    linalg.yield %0 : index
+  }
+  return
+}

mlir/test/Dialect/Linalg/fusion-indexed-generic.mlir

 // CHECK-NOT:  scf.parallel
 // CHECK:      linalg.indexed_generic
 // CHECK:        ^bb0([[i:%.*]]: index, [[j:%.*]]: index
 // CHECK:          [[j_new:%.*]] = addi [[j]], [[J]] : index
 // CHECK:          {{.*}} = index_cast [[i]] : index to i32
 // CHECK:          {{.*}} = index_cast [[j_new]] : index to i32
 // CHECK:      linalg.generic
 // CHECK:          addf
+
+// -----
+
+#map = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+#id_2d = affine_map<(d0, d1) -> (d0, d1)>
+#pointwise_2d_trait = {
+  indexing_maps = [#id_2d],
+  iterator_types = ["parallel", "parallel"]
+}
+func @index_op(%A: memref<?x?xindex>,
+               %B: memref<?x?xindex>) {
+  linalg.generic #pointwise_2d_trait
+   outs(%B : memref<?x?xindex>) {
+  ^bb0(%arg6: index):   // no predecessors
+    %2 = constant 0 : index
+    linalg.yield %2 : index
+  }
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c25 = constant 25 : index
+  %c10 = constant 10 : index
+  %0 = memref.dim %A, %c0 : memref<?x?xindex>
+  %1 = memref.dim %A, %c1 : memref<?x?xindex>
+  %2 = memref.dim %B, %c0 : memref<?x?xindex>
+  %3 = memref.dim %B, %c1 : memref<?x?xindex>
+  scf.for %arg2 = %c0 to %0 step %c10 {
+    scf.for %arg3 = %c0 to %1 step %c25 {
+      %4 = memref.subview %A[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+          memref<?x?xindex> to memref<?x?xindex, #map>
+      %5 = memref.subview %B[%arg2, %arg3][%c10, %c25][%c1, %c1] :
+          memref<?x?xindex> to memref<?x?xindex, #map>
+      linalg.generic {
+        indexing_maps = [#id_2d, #id_2d],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%4 : memref<?x?xindex, #map>)
+       outs(%5 : memref<?x?xindex, #map>) {
+      ^bb0(%arg6: index, %arg7: index):
+        %6 = linalg.index 0 : index
+        linalg.yield %6 : index
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @index_op
+// CHECK:  linalg.generic
+// CHECK:  scf.for
+// CHECK:    scf.for
+// CHECK-NOT:  scf.for
+// CHECK:      linalg.generic

mlir/test/Dialect/Linalg/fusion-tensor.mlir

func @generic_index_op2(%arg0: tensor<1x8xf64>, %arg1: tensor<1x8xi32>) -> tensor<1x8xi32> {
   ^bb0(%a: f64, %b: i32):
     %r = call @compute2(%a, %b) : (f64, i32) -> i32
     linalg.yield %r : i32
   } -> tensor<1x8xi32>
 
   // CHECK-NEXT:   return %[[R]] : tensor<1x8xi32>
   return %1 : tensor<1x8xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @index_op(
+// CHECK-COUNT-2: linalg.generic
+func @index_op(%arg0: tensor<1x8xindex>, %arg1: tensor<1x8xindex>) -> tensor<1x8xindex> {
+  %0 = linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  outs(%arg0 : tensor<1x8xindex>) {
+  ^bb0(%a: index):   // no predecessors
+    %2 = linalg.index 1 : index
+    linalg.yield %2 : index
+  } -> tensor<1x8xindex>
+  %1 = linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>, affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  ins(%0 : tensor<1x8xindex>)
+  outs(%arg1 : tensor<1x8xindex>) {
+  ^bb0(%a: index, %b: index):   // no predecessors
+    %2 = linalg.index 0 : index
+    %3 = addi %2, %a : index
+    linalg.yield %3 : index
+  } -> tensor<1x8xindex>
+  return %1 : tensor<1x8xindex>
+}

mlir/test/Dialect/Linalg/invalid.mlir

 
 func @yield_parent(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
   // expected-error @+1 {{op expected parent op with LinalgOp interface}}
   linalg.yield %arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>
 }
 
 // -----
 
+func @index_parent() {
+  // expected-error @+1 {{op expected parent op with LinalgOp interface}}
+  linalg.index 0 : index
+}
+
+// -----
+
+func @index_dim_lower_than_number_of_loops(%arg0: memref<f32>) {
+  // expected-error @+6 {{op expected dim (2) to be lower than the number of loops (0) of the enclosing LinalgOp}}
+  linalg.generic {
+      indexing_maps =  [ affine_map<() -> ()> ],
+      iterator_types = []}
+      outs(%arg0 : memref<f32>) {
+    ^bb(%0: f32):
+      linalg.index 2 : index
+      linalg.yield %0 : f32
+  }
+}
+
+// -----
+
 func @generic_no_region(%arg0: memref<f32>) {
   // expected-error @+5 {{expected '{' to begin a region}}
   linalg.generic {
     indexing_maps =  [ affine_map<() -> (0)> ],
     iterator_types = []
   } ins(%arg0 : memref<f32>)
 }
 

mlir/test/Dialect/Linalg/loop-order.mlir

-// RUN: mlir-opt %s -convert-linalg-to-loops="interchange-vector=4,0,3,1,2" | FileCheck --check-prefix=LOOP %s
-// RUN: mlir-opt %s -convert-linalg-to-parallel-loops="interchange-vector=4,0,3,1,2" | FileCheck --check-prefix=PARALLEL %s
-// RUN: mlir-opt %s -convert-linalg-to-affine-loops="interchange-vector=4,0,3,1,2" | FileCheck --check-prefix=AFFINE %s
+// RUN: mlir-opt %s -convert-linalg-to-loops="interchange-vector=4,0,3,1,2" -split-input-file | FileCheck --check-prefix=LOOP %s
+// RUN: mlir-opt %s -convert-linalg-to-parallel-loops="interchange-vector=4,0,3,1,2" -split-input-file | FileCheck --check-prefix=PARALLEL %s
+// RUN: mlir-opt %s -convert-linalg-to-affine-loops="interchange-vector=4,0,3,1,2" -split-input-file | FileCheck --check-prefix=AFFINE %s
 
 func @copy(%input: memref<1x2x3x4x5xf32>, %output: memref<1x2x3x4x5xf32>) {
   linalg.copy(%input, %output): memref<1x2x3x4x5xf32>, memref<1x2x3x4x5xf32>
   return
 }
 
 // LOOP: scf.for %{{.*}} = %c0 to %c5 step %c1
 // LOOP:   scf.for %{{.*}} = %c0 to %c1 step %c1
 // LOOP:     scf.for %{{.*}} = %c0 to %c4 step %c1
 // LOOP:       scf.for %{{.*}} = %c0 to %c2 step %c1
 // LOOP:         scf.for %{{.*}} = %c0 to %c3 step %c1
 
 // PARALLEL:      scf.parallel
 // PARALLEL-SAME:   to (%c5, %c1, %c4, %c2, %c3)
 
 // AFFINE: affine.for %{{.*}} = 0 to 5
 // AFFINE:   affine.for %{{.*}} = 0 to 1
 // AFFINE:     affine.for %{{.*}} = 0 to 4
 // AFFINE:       affine.for %{{.*}} = 0 to 2
 // AFFINE:         affine.for %{{.*}} = 0 to 3
 
+// -----
+
+func @index_op(%arg0: memref<4x8xindex>) {
+  linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  outs(%arg0 : memref<4x8xindex>) {
+  ^bb0(%arg1: index):   // no predecessors
+    %0 = linalg.index 1 : index
+    linalg.yield %0 : index
+  }
+  return
+}
+// LOOP-LABEL: @index_op
+//      LOOP:   linalg.generic
+
+// PARALLEL-LABEL: @index_op
+//      PARALLEL:   linalg.generic
+
+// AFFINE-LABEL: @index_op
+//      AFFINE:   linalg.generic

mlir/test/Dialect/Linalg/roundtrip.mlir

 
 func @generic_region(%arg0: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>,
                      %arg1: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>) {
   linalg.generic #trait_3
        ins(%arg0 : memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>)
       outs(%arg1 : memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>)
       attrs = {foo = 1} {
     ^bb(%a: vector<3x4xi4>, %b: f32) :
+      %0 = linalg.index 0 : index
+      %1 = linalg.index 1 : index
+      %2 = linalg.index 2 : index
       linalg.yield %b : f32
   }
   return
 }
 // CHECK-LABEL: func @generic_region
 //       CHECK:   linalg.generic {
 //  CHECK-SAME:     indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}],
 //  CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel"],
 //  CHECK-SAME:     library_call = "some_external_function_name_2"
 //  CHECK-SAME:      ins({{.*}} : memref<?x?xvector<3x4xi4>, #[[$strided2D]]>)
 //  CHECK-SAME:     outs({{.*}} : memref<?x?x?xf32, #[[$strided3D]]>)
 //  CHECK-SAME:     attrs = {foo = 1 : i64} {
 //       CHECK:  ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32):
+//       CHECK:    %{{.*}} = linalg.index 0 : index
+//       CHECK:    %{{.*}} = linalg.index 1 : index
+//       CHECK:    %{{.*}} = linalg.index 2 : index
 //       CHECK:    linalg.yield %{{.*}} : f32
 
 func @indexed_generic(%arg0: memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>,
                       %arg1: memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>) {
   linalg.indexed_generic #trait_3
        ins(%arg0 : memref<?x?xvector<3x4xi4>, offset: ?, strides: [?, 1]>)
       outs(%arg1 : memref<?x?x?xf32, offset: ?, strides: [?, ?, 1]>)
       attrs = {foo = 1} {

mlir/test/Dialect/Linalg/tile.mlir

 //   TILE-002-NOT:   for
 //       TILE-002:     linalg.generic
 
 // TILE-234-LABEL: func @pointwise
 //       TILE-234:   for
 //       TILE-234:     for
 //   TILE-234-NOT:   for
 //       TILE-234:       linalg.generic
+
+// TILE-2-LABEL: func @index_op
+//   TILE-2-NOT:   for
+//       TILE-2:   linalg.generic
+func @index_op(%arg0: memref<?x?xindex>) {
+  linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  outs(%arg0 : memref<?x?xindex>) {
+  ^bb0(%arg1: index):   // no predecessors
+    %0 = linalg.index 1 : index
+    linalg.yield %0 : index
+  }
+  return
+}

mlir/test/Dialect/Linalg/vectorization.mlir

 func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
                   %pad_value: f32) -> tensor<6x?x?x?xf32> {
   %0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
       linalg.yield %pad_value : f32
     } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
   return %0 : tensor<6x?x?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @index_op
+//       CHECK:   linalg.generic
+func @index_op(%arg0: memref<4x8xindex>) {
+  linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]}
+  outs(%arg0 : memref<4x8xindex>) {
+  ^bb0(%arg1: index):   // no predecessors
+    %0 = linalg.index 1 : index
+    linalg.yield %0 : index
+  }
+  return
+}

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

 
 static LogicalResult fuseLinalgOpsGreedily(FuncOp f) {
   OpBuilder b(f);
   DenseSet<Operation *> eraseSet;
 
   // Save original Linalg ops, we only want to make a pass over those.
   SmallVector<LinalgOp, 8> linalgOps;
   f.walk([&](LinalgOp op) {
+    // TODO: remove hasIndexSemantics check once index ops are supported.
+    if (op.hasIndexSemantics())
+      return;
+
     // TODO: support multi-results.
     if (op->getNumResults() <= 1)
       linalgOps.push_back(op);
   });
 
   // Tile and Fuse for tensors inputs (TODO: all tensor operands).
   bool changed = false;
   for (LinalgOp linalgOp : llvm::reverse(linalgOps)) {