[mlir][linalg] Add IndexedGenericOp to GenericOp canonicalization.

[gysit]

Summary

Replace all linalg.indexed_generic ops by linalg.generic ops that access the iteration indices using the linalg.index op.

[gysit]

Rebase.

[ThomasRaoux]

LGTM

[mravishankar]

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

742

Is the plan to eventually just remove linalg.indexed_generic operation? That would simplify some of the logic for operation fusion (hopefully)

[gysit]

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

742

Yes the goal is to remove the IndexedGenericOp. In a next step, after the patch landed, I will start to remove IndexedGenericOp support from all the passes and only in a last step I will remove it entirely.

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

OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
            "nullptr">)>,
     OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputBuffers,
       "ArrayRef<AffineMap>":$indexingMaps, "ArrayRef<StringRef>":$iteratorTypes,
       CArg<"function_ref<void(OpBuilder &, Location, ValueRange, ValueRange)>",
            "nullptr">)>
   ];
   let verifier = [{ return ::verify(*this); }];
 
-  let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as a declarative configurations of generic ops.
 //===----------------------------------------------------------------------===//
 
 // This file is auto-generated from a TC def specification.
 include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.tcgen.td"
 include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.td"
 
 #endif // LINALG_STRUCTURED_OPS

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

static LogicalResult verifyGenericOp(GenericOpType op) {
 
   return success();
 }
 
 static LogicalResult verify(GenericOp op) { return verifyGenericOp(op); }
 
 static LogicalResult verify(IndexedGenericOp op) { return verifyGenericOp(op); }
 
+namespace {
+
+/// Replace indexed_generic ops by generic ops that access the iteration indices
+/// using index operation calls.
+struct ConvertIndexedToGenericOp : OpRewritePattern<IndexedGenericOp> {
+  using OpRewritePattern<IndexedGenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(IndexedGenericOp indexedOp,
+                                PatternRewriter &rewriter) const override {
+    // Replace all uses of the index block arguments.
+    BlockAndValueMapping bvm;
+    if (Block *body = indexedOp.getBody()) {
+      rewriter.setInsertionPointToStart(body);
+      for (const auto &en : llvm::enumerate(
+               body->getArguments().take_front(indexedOp.getNumLoops()))) {
+        Value index = rewriter.create<IndexOp>(indexedOp.getLoc(), en.index());
+        bvm.map(en.value(), index);
+      }
+    }
+
+    // Create a generic replacement operation and clone the body.
+    rewriter.setInsertionPointAfter(indexedOp);
+    SmallVector<StringRef> iterators = llvm::to_vector<4>(
+        indexedOp.iterator_types().getAsValueRange<StringAttr>());
+    GenericOp genericOp = rewriter.create<GenericOp>(
+        indexedOp.getLoc(), indexedOp->getResultTypes(), indexedOp.getInputs(),
+        indexedOp.getOutputs(), indexedOp.getIndexingMaps(), iterators);
+    Region &genericRegion = genericOp.region();
+    Region &indexedRegion = indexedOp.region();
+    rewriter.cloneRegionBefore(indexedRegion, genericRegion,
+                               genericRegion.begin(), bvm);
+
+    rewriter.replaceOp(indexedOp, genericOp->getResults());
+    return success();
+  }
+};
+} // namespace
+
+void IndexedGenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                   MLIRContext *context) {
+  results.add<ConvertIndexedToGenericOp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // InitTensorOp
 //===----------------------------------------------------------------------===//
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {
   unsigned rank = sizes.size();
   SmallVector<Value, 4> dynamicSizes;

 
 CANONICALIZERS_AND_FOLDERS(ConvOp)
 CANONICALIZERS_AND_FOLDERS(PoolingMaxOp)
 CANONICALIZERS_AND_FOLDERS(PoolingMinOp)
 CANONICALIZERS_AND_FOLDERS(PoolingSumOp)
 CANONICALIZERS_AND_FOLDERS(CopyOp)
 CANONICALIZERS_AND_FOLDERS(FillOp)
 CANONICALIZERS_AND_FOLDERS(GenericOp)
-CANONICALIZERS_AND_FOLDERS(IndexedGenericOp)
 
 // All named ops canonicalizers and folders are auto-generated in the
 // .cpp.inc.

mlir/test/Dialect/Linalg/bufferize.mlir

 
 // -----
 
 #map0 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL:   func @multiple_results_indexed
 // CHECK:           %[[RESULT0:.*]] = memref.alloc() : memref<4xi32>
 // CHECK:           %[[RESULT1:.*]] = memref.alloc() : memref<4xi32>
-// CHECK:           linalg.indexed_generic
+// CHECK:           linalg.generic
 // CHECK-SAME:      ins(%{{.*}} : memref<4xi32>)
 // CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xi32>, memref<4xi32>)
-// CHECK-NEXT: ^bb0(%{{.*}}: index, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32):
+// CHECK-NEXT: ^bb0(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32):
 func @multiple_results_indexed(%arg0: tensor<4xi32>)
         -> (tensor<4xi32>, tensor<4xi32>) {
     %0, %1 = linalg.indexed_generic {
       indexing_maps = [#map0, #map0, #map0],
       iterator_types = ["parallel"]
     } ins(%arg0 : tensor<4xi32>)
       outs (%arg0, %arg0 : tensor<4xi32>, tensor<4xi32>) {
       ^bb0(%i: index, %gen_arg1: i32, %out1: i32, %out2: i32):

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

 // Test case: linalg.indexed_generic.
 // Other than the payload argument handling, everything else is the same.
 
 #map = affine_map<(d0) -> (d0)>
 
 // CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0)>
 // CHECK-LABEL: @indexed_generic
 func @indexed_generic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
-  // CHECK: linalg.indexed_generic
-  // CHECK:   ^bb0(%{{.*}}: index, %[[BBARG:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32):
+  // CHECK: linalg.generic
+  // CHECK:   ^bb0(%[[BBARG:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32):
   // CHECK:     addf %[[BBARG]], %[[BBARG]]
   %0 = linalg.indexed_generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]}
       ins(%arg0, %arg0 : tensor<?xf32>, tensor<?xf32>)
      outs(%arg0 : tensor<?xf32>) {
   ^bb0(%index: index, %arg1: f32, %arg2: f32, %arg3: f32):
     %1 = addf %arg1, %arg2 : f32
     linalg.yield %1 : f32
   } -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }

mlir/test/Dialect/Linalg/canonicalize.mlir

 //  CHECK-DAG:   %[[C1:.+]] = constant 1 : index
 //  CHECK-DAG:   %[[C2:.+]] = constant 2 : index
 //  CHECK-DAG:   %[[C12:.+]] = constant 12 : index
 //      CHECK:   %[[IN_DIM1:.+]] = memref.dim %[[ARG0]], %[[C1]]
 //      CHECK:   %[[OUT_DIM1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[IN_DIM1]]]
 //      CHECK:   %[[IN_DIM2:.+]] = memref.dim %[[ARG0]], %[[C2]]
 //      CHECK:   %[[OUT_DIM2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[IN_DIM2]]]
 //      CHECK:   return %[[C12]], %[[OUT_DIM1]], %[[OUT_DIM2]]
+
+// -----
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+
+func @indexed_generic(%arg0: memref<?x?xindex>, %arg1: memref<?x?xindex>) {
+  linalg.indexed_generic {
+      indexing_maps = [#map, #map],
+      iterator_types = ["parallel", "parallel"]}
+    ins(%arg0 : memref<?x?xindex>)
+   outs(%arg1 : memref<?x?xindex>) {
+  ^bb0(%arg4: index, %arg5: index, %arg6: index, %arg7: index):
+    %0 = addi %arg4, %arg5 : index
+    %1 = addi %0, %arg6 : index
+    %2 = addi %1, %arg7 : index
+    linalg.yield %2 : index
+  }
+  return
+}
+
+//      CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//      CHECK: func @indexed_generic
+// CHECK-NEXT:   linalg.generic {
+// CHECK-SAME:     indexing_maps = [#[[MAP]], #[[MAP]]], iterator_types = ["parallel", "parallel"]}
+// CHECK-SAME:      ins(%[[ARG0:[A-Za-z0-9_]+]] : memref<?x?xindex>)
+// CHECK-SAME:     outs(%[[ARG1:[A-Za-z0-9_]+]] : memref<?x?xindex>)
+//      CHECK:   ^bb0(%[[ARG2:[A-Za-z0-9_]+]]: index, %[[ARG3:[A-Za-z0-9_]+]]: index):
+// CHECK-NEXT:     %[[IDX0:.+]] = linalg.index 0 : index
+// CHECK-NEXT:     %[[IDX1:.+]] = linalg.index 1 : index
+// CHECK-NEXT:     %[[SUM0:.+]] = addi %[[IDX0]], %[[IDX1]] : index
+// CHECK-NEXT:     %[[SUM1:.+]] = addi %[[SUM0]], %[[ARG2]] : index
+// CHECK-NEXT:     %[[SUM2:.+]] = addi %[[SUM1]], %[[ARG3]] : index
+// CHECK-NEXT:     linalg.yield %[[SUM2]] : index

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

 }
 // CHECK-LABEL: func @fuse_indexed_generic_consumer
 // CHECK:  scf.for
 // CHECK:    scf.for
 // CHECK-NOT:  scf.for
 // CHECK:      linalg.generic
 // CHECK-NOT:    addi
 // CHECK:        addf
-// CHECK:      linalg.indexed_generic
+// CHECK:      linalg.generic
 // CHECK:        index_cast
 
 // -----
 
 #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, #id_2d, #id_2d],
^bb0( %a: f32, %b: f32):
       linalg.yield %ab : f32
     }
   }
   return
 }
 // CHECK-LABEL: func @fuse_indexed_generic_producer
 // CHECK:  scf.parallel ([[I:%.*]], [[J:%.*]]) =
 // CHECK-NOT:  scf.parallel
-// CHECK:      linalg.indexed_generic
-// CHECK:        ^bb0([[i:%.*]]: index, [[j:%.*]]: index
+// CHECK:      linalg.generic
+// CHECK:          [[i:%.*]] = linalg.index 0 : index
+// CHECK:          [[j:%.*]] = linalg.index 1 : index
 // CHECK:          [[i_new:%.*]] = addi [[i]], [[I]] : index
 // CHECK:          [[j_new:%.*]] = addi [[j]], [[J]] : index
 // CHECK:          {{.*}} = index_cast [[i_new]] : index to i32
 // CHECK:          {{.*}} = index_cast [[j_new]] : index to i32
 // CHECK:      linalg.generic
 // CHECK:          addf
 
 // -----
scf.parallel (%j) = (%c0) to (%C_Y) step (%c3) {
     scf.yield
   }
   return
 }
 // CHECK-LABEL: func @fuse_indexed_generic_producer_tile_second_dim_only
 // CHECK:  [[C0:%.*]] = constant 0 : index
 // CHECK:  scf.parallel ([[J:%.*]]) =
 // CHECK-NOT:  scf.parallel
-// CHECK:      linalg.indexed_generic
-// CHECK:        ^bb0([[i:%.*]]: index, [[j:%.*]]: index
+// CHECK:      linalg.generic
+// CHECK:          [[i:%.*]] = linalg.index 0 : index
+// CHECK:          [[j:%.*]] = linalg.index 1 : 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
 
 // -----
 

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

^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: f32, %arg5 : f32, %arg6 : f32):
       %4 = mulf %arg4, %arg5 : f32
       linalg.yield %4 : f32
     } -> tensor<5x?x?xf32>
   return %3 : tensor<5x?x?xf32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-LABEL: func @indexed_generic_op_constant_fusion
 //       CHECK:   %[[CST:.*]] = constant {{.*}} : f32
-//       CHECK:   linalg.indexed_generic
+//       CHECK:   linalg.generic
 //       CHECK:   ^{{[a-zA-Z0-9_]*}}
-//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]*]]: index
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]*]]: index
-//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]*]]: index
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9_]*]]: f32, %{{.*}}: f32)
 //       CHECK:     mulf %[[CST]], %[[ARG4]]
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2) -> ()>
 #map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 func @generic_op_zero_dim_constant_fusion(%arg0 : tensor<5x?x?xf32>)
^bb0(%arg1 : index, %arg2 : index, %arg3 : index, %arg4: f32, %arg5: f32, %arg6: f32):
       %4 = mulf %arg4, %arg5 : f32
       linalg.yield %4 : f32
     } -> tensor<5x?x?xf32>
   return %3 : tensor<5x?x?xf32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK-LABEL: func @indexed_generic_op_zero_dim_constant_fusion
 //       CHECK:   %[[CST:.*]] = constant {{.*}} : f32
-//       CHECK:   linalg.indexed_generic
+//       CHECK:   linalg.generic
 //       CHECK:   ^{{[a-zA-Z0-9_]*}}
-//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]*]]: index
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]*]]: index
-//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]*]]: index
 //  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9_]*]]: f32, %{{.*}}: f32)
 //       CHECK:     mulf %[[CST]], %[[ARG4]]
 
 // -----
 
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 func @generic_op_indexed_generic_op_fusion(%arg0: tensor<?x?xi32>,
                                            %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {
^bb0(%arg2: index, %arg3: index, %arg4: i32, %arg5: i32):       // no predecessors
       %7 = addi %arg4, %5 : i32
       %8 = subi %7, %6 : i32
       linalg.yield %8 : i32
     } -> tensor<?x?xi32>
   return %4 : tensor<?x?xi32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: func @generic_op_indexed_generic_op_fusion
-//   CHECK-NOT: linalg.generic
-//       CHECK: linalg.indexed_generic
+//   CHECK-NOT: linalg.indexed_generic
+//       CHECK: linalg.generic
 // CHECK-SAME:    indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]]
 //      CHECK: ^{{[a-zA-Z0-9_]*}}
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
 // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
 // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[ARG0:.+]] = linalg.index 0 : index
+//      CHECK:   %[[ARG1:.+]] = linalg.index 1 : index
 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ARG3]] : i32
 //      CHECK:   %[[ADD_OPERAND:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[VAL2:.+]] = addi %[[VAL1]], %[[ADD_OPERAND]] : i32
 //      CHECK:   %[[VAL3:.+]] = subi %[[VAL2]], %[[SUB_OPERAND]] : i32
 //      CHECK:   linalg.yield %[[VAL3]] : i32
 
 // -----
%4 = linalg.generic {
     ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):       // no predecessors
       %10 = addi %arg2, %arg3 : i32
       linalg.yield %10 : i32
     } -> tensor<?x?xi32>
   return %4 : tensor<?x?xi32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: func @indexed_generic_op_generic_op_fusion
-//       CHECK: linalg.indexed_generic
+//       CHECK: linalg.generic
 // CHECK-SAME:    indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]]
 //      CHECK: ^{{[a-zA-Z0-9_]*}}
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
 // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
 // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[ARG0:.+]] = linalg.index 0 : index
+//      CHECK:   %[[ARG1:.+]] = linalg.index 1 : index
 //      CHECK:   %[[ADD_OPERAND:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND]] : i32
 //      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND]] : i32
 //      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ARG3]] : i32
 //      CHECK:   linalg.yield %[[VAL3]] : i32
 //   CHECK-NOT: linalg.generic
 
^bb0(%arg2: index, %arg3: index, %arg4: i32, %arg5: i32):       // no predecessors
       %7 = addi %arg4, %5 : i32
       %8 = subi %7, %6 : i32
       linalg.yield %8 : i32
     } -> tensor<?x?xi32>
   return %4 : tensor<?x?xi32>
 }
 //   CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-LABEL: func @indexed_generic_op_fusion
-//       CHECK: linalg.indexed_generic
+//       CHECK: linalg.generic
 // CHECK-SAME:    indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
 //      CHECK: ^{{[a-zA-Z0-9_]*}}
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: index
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: index
 // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: i32
+//      CHECK:   %[[ARG0:.+]] = linalg.index 0 : index
+//      CHECK:   %[[ARG1:.+]] = linalg.index 1 : index
 //      CHECK:   %[[ADD_OPERAND1:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[SUB_OPERAND1:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND1]] : i32
 //      CHECK:   %[[VAL2:.+]] = subi %[[SUB_OPERAND1]], %[[VAL1]] : i32
 //      CHECK:   %[[ADD_OPERAND2:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND2:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ADD_OPERAND2]] : i32
 //      CHECK:   %[[VAL4:.+]] = subi %[[VAL3]], %[[SUB_OPERAND2]] : i32
 //      CHECK:   linalg.yield %[[VAL4]] : i32
-//   CHECK-NOT: linalg.indexed_generic
+//   CHECK-NOT: linalg.generic
 
 // -----
 
 // The indices of the first indexed_generic op are swapped after fusion.
 #map0 = affine_map<(d0, d1) -> (d1, d0)>
 #map1 = affine_map<(d0, d1) -> (d0, d1)>
 func @indexed_producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>)
                                                -> tensor<?x?xi32> {
%3 = linalg.generic
     } -> tensor<10xf32>
   return %3 : tensor<10xf32>
 }
 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> ()>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
 //       CHECK: func @scalar_indexed_generic_fusion
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<5x1x1xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<i32>
-//       CHECK:   %[[T0:.+]] = linalg.indexed_generic
+//       CHECK:   %[[T0:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 //  CHECK-SAME:     iterator_types = ["parallel"]
 //  CHECK-SAME:     ins(%[[ARG1]] : tensor<i32>)
 //       CHECK:     tensor.extract %[[ARG0]]
 //       CHECK:     linalg.yield
 //       CHECK   return %[[T0]]
 
 // -----

mlir/test/Dialect/Linalg/reshape_fusion.mlir

func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>,
   } -> tensor<?x?x?xi32>
   return %1 : tensor<?x?x?xi32>
 }
 
 // The generic op version of the test check for the op structure. Only
 // checking the op body here.
 //       CHECK: #[[MAP:.+]] =  affine_map<(d0, d1) -> (d0 + d1 * 4)>
 //       CHECK: func @indexed_generic_op_reshape_producer_fusion
-//       CHECK:   linalg.indexed_generic
+//       CHECK:   linalg.generic
 //       CHECK:   ^{{.*}}(
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32)
+//       CHECK:     %[[ARG2:.+]] = linalg.index 0 : index
+//       CHECK:     %[[ARG3:.+]] = linalg.index 1 : index
+//       CHECK:     %[[ARG4:.+]] = linalg.index 2 : index
+//       CHECK:     %[[ARG5:.+]] = linalg.index 3 : index
 //       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG3]], %[[ARG2]])
 //       CHECK:     %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
 //       CHECK:     %[[T5:.+]] = index_cast %[[T3]]
 //       CHECK:     %[[T6:.+]] = addi %[[T4]], %[[T5]]
 //       CHECK:     %[[T7:.+]] = index_cast %[[ARG4]]
 //       CHECK:     %[[T8:.+]] = addi %[[T6]], %[[T7]]
 //       CHECK:     %[[T9:.+]] = index_cast %[[ARG5]]
 //       CHECK:     %[[T10:.+]] = addi %[[T8]], %[[T9]]
%1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
                                  affine_map<(i, j, k, l) -> (j, k, l)>] :
     tensor<?x?xi32> into tensor<?x?x4x5xi32>
   return %1 : tensor<?x?x4x5xi32>
 }
 // The generic op version of the test check for the op structure. Only
 // checking the op body here.
 //       CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1 * 5 + d2 * 20)>
 //       CHECK: func @indexed_generic_op_reshape_consumer_fusion
-//       CHECK:   linalg.indexed_generic
+//       CHECK:   linalg.generic
 //       CHECK:   ^{{.*}}(
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32)
+//       CHECK:     %[[ARG2:.+]] = linalg.index 0 : index
+//       CHECK:     %[[ARG3:.+]] = linalg.index 1 : index
+//       CHECK:     %[[ARG4:.+]] = linalg.index 2 : index
+//       CHECK:     %[[ARG5:.+]] = linalg.index 3 : index
 //       CHECK:     %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG5]], %[[ARG4]], %[[ARG3]])
 //       CHECK:     %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]]
 //       CHECK:     %[[T5:.+]] = index_cast %[[ARG2]]
 //       CHECK:     %[[T6:.+]] = addi %[[T4]], %[[T5]]
 //       CHECK:     %[[T7:.+]] = index_cast %[[T3]]
 //       CHECK:     %[[T8:.+]] = addi %[[T6]], %[[T7]]
 //       CHECK:     linalg.yield %[[T8]]
 

 //       CHECK: func @reshape_as_consumer_permutation
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<210x6x4xi32>
 //  CHECK-SAME:   %[[ARG1:.+]]: tensor<210x4xi32>
 //   CHECK-DAG:   %[[T1:.+]] = linalg.tensor_reshape %[[ARG0]]
 //  CHECK-SAME:     [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
 //   CHECK-DAG:   %[[T2:.+]] = linalg.tensor_reshape %[[ARG1]]
 //  CHECK-SAME:     [#[[MAP3]], #[[MAP4]]]
 //   CHECK-DAG:   %[[T0:.+]] = linalg.init_tensor [2, 3, 4, 5, 6, 7]
-//       CHECK:   %[[T4:.+]] = linalg.indexed_generic
+//       CHECK:   %[[T4:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]]]
 //  CHECK-SAME:     ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>)
 //  CHECK-SAME:     outs(%[[T0]] : tensor<2x3x4x5x6x7xi32>)
 //       CHECK:   ^{{.+}}(
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG6:[a-zA-Z0-9]+]]: index, %[[ARG7:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG8:[a-zA-Z0-9]+]]: i32, %[[ARG9:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG10:[a-zA-Z0-9]+]]: i32)
+//       CHECK:     %[[ARG2:.+]] = linalg.index 0 : index
+//       CHECK:     %[[ARG3:.+]] = linalg.index 1 : index
+//       CHECK:     %[[ARG4:.+]] = linalg.index 2 : index
+//       CHECK:     %[[ARG5:.+]] = linalg.index 3 : index
+//       CHECK:     %[[ARG6:.+]] = linalg.index 4 : index
+//       CHECK:     %[[ARG7:.+]] = linalg.index 5 : index
 //   CHECK-DAG:       %[[T5:.+]] = affine.apply #[[MAP8]](%[[ARG3]], %[[ARG2]])
 //   CHECK-DAG:       %[[T6:.+]] = affine.apply #[[MAP9]](%[[ARG6]], %[[ARG5]], %[[ARG4]])
 //   CHECK-DAG:       %[[T7:.+]] = addi %[[ARG8]], %[[ARG9]]
 //       CHECK:       %[[T8:.+]] = index_cast %[[T5]]
 //       CHECK:       %[[T9:.+]] = addi %[[T7]], %[[T8]]
 //       CHECK:       %[[T10:.+]] = index_cast %[[T6]]
 //       CHECK:       %[[T11:.+]] = addi %[[T9]], %[[T10]]
 //       CHECK:       %[[T12:.+]] = index_cast %[[ARG7]]

 //   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 //   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 //   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 8)>
 //   CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
 //   CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)>
 //   CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
 //       CHECK: @reshape_as_producer_projected_permutation
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<33x8x?xi32>
-//       CHECK:   %[[RES:.+]] = linalg.indexed_generic
+//       CHECK:   %[[RES:.+]] = linalg.generic
 //  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
 //  CHECK-SAME:     ins(%[[ARG0]] : tensor<33x8x?xi32>)
 //       CHECK:   ^{{.+}}(
-//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: index,
-//  CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG5:[a-zA-Z0-9]+]]: i32,
 //  CHECK-SAME:     %[[ARG7:[a-zA-Z0-9]+]]: i32)
+//       CHECK:       %[[ARG1:.+]] = linalg.index 0 : index
+//       CHECK:       %[[ARG2:.+]] = linalg.index 1 : index
+//       CHECK:       %[[ARG3:.+]] = linalg.index 2 : index
+//       CHECK:       %[[ARG4:.+]] = linalg.index 3 : index
 //       CHECK:       %[[T0:.+]] = affine.apply #[[MAP2]](%[[ARG2]], %[[ARG1]])
 //       CHECK:       %[[T1:.+]] = index_cast %[[T0]] : index to i32
 //       CHECK:       %[[T2:.+]] = addi %[[ARG5]], %[[T1]] : i32
 //       CHECK:       %[[T3:.+]] = index_cast %[[ARG3]] : index to i32
 //       CHECK:       %[[T4:.+]] = addi %[[T2]], %[[T3]] : i32
 //       CHECK:       %[[T5:.+]] = index_cast %[[ARG4]] : index to i32
 //       CHECK:       %[[T6:.+]] = addi %[[T4]], %[[T5]] : i32
 //       CHECK:       linalg.yield %[[T6]] : i32

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

^bb0(%i: index, %operand_in: f32, %result_in: f32):
       %out = addf %operand_in, %i_float : f32
       linalg.yield %out : f32
   }
   return
 }
 // TILE-10n25-LABEL: func @indexed_generic_vector
 // TILE-10n25: %[[C10:.*]] = constant 10 : index
 // TILE-10n25: scf.for %[[J:.*]] = {{.*}} step %[[C10]]
-// TILE-10n25:   linalg.indexed_generic
-// TILE-10n25:   ^bb0(%[[I:.*]]: index, %[[IN:.*]]: f32, %[[OUT:.*]]: f32)
+// TILE-10n25:   linalg.generic
+// TILE-10n25:   ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32)
+// TILE-10n25:     %[[I:.*]] = linalg.index 0 : index
 // TILE-10n25:     %[[NEW_I:.*]] = addi %[[I]], %[[J]] : index
 // TILE-10n25:     %[[NEW_I_INT:.*]] = index_cast %[[NEW_I]] : index to i32
 // TILE-10n25:     %[[NEW_I_FLOAT:.*]] = sitofp %[[NEW_I_INT]] : i32 to f32
 // TILE-10n25:     %[[OUT:.*]] = addf %[[IN]], %[[NEW_I_FLOAT]] : f32
 
 // TILE-25n0-LABEL: func @indexed_generic_vector
 // TILE-25n0: %[[C25:.*]] = constant 25 : index
 // TILE-25n0: scf.for %[[J:.*]] = {{.*}} step %[[C25]]
-// TILE-25n0:   linalg.indexed_generic
-// TILE-25n0:   ^bb0(%[[I:.*]]: index, %[[IN:.*]]: f32, %[[OUT:.*]]: f32)
+// TILE-25n0:   linalg.generic
+// TILE-25n0:   ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32)
+// TILE-25n0:    %[[I:.*]] = linalg.index 0 : index
 // TILE-25n0:     %[[NEW_I:.*]] = addi %[[I]], %[[J]] : index
 // TILE-25n0:     %[[NEW_I_INT:.*]] = index_cast %[[NEW_I]] : index to i32
 // TILE-25n0:     %[[NEW_I_FLOAT:.*]] = sitofp %[[NEW_I_INT]] : i32 to f32
 // TILE-25n0:     %[[OUT:.*]] = addf %[[IN]], %[[NEW_I_FLOAT]] : f32
 
 // TILE-0n25-LABEL: func @indexed_generic_vector
 // TILE-0n25-NOT: scf.for %[[J:.*]] = {{.*}} step %
-// TILE-0n25: linalg.indexed_generic
+// TILE-0n25: linalg.generic
 
 #combined_indices_trait = {
   args_in = 1,
   args_out = 1,
   indexing_maps = [
     affine_map<(i, j) -> (j, i + j)>,
     affine_map<(i, j) -> (i, j)>
   ],
func @indexed_generic_matrix(%operand: memref<50x99xf32>, %result: memref<50x50xf32>) {
   }
   return
 }
 // TILE-10n25-LABEL: func @indexed_generic_matrix
 // TILE-10n25-DAG: %[[C25:.*]] = constant 25 : index
 // TILE-10n25-DAG: %[[C10:.*]] = constant 10 : index
 // TILE-10n25: scf.for %[[K:.*]] = {{.*}} step %[[C10]]
 // TILE-10n25:   scf.for %[[L:.*]] = {{.*}} step %[[C25]]
-// TILE-10n25:     linalg.indexed_generic
-// TILE-10n25:     ^bb0(%[[I:.*]]: index, %[[J:.*]]: index, %[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-10n25:     linalg.generic
+// TILE-10n25:     ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-10n25:       %[[I:.*]] = linalg.index 0 : index
+// TILE-10n25:       %[[J:.*]] = linalg.index 1 : index
 // TILE-10n25:       %[[NEW_I:.*]] = addi %[[I]], %[[K]] : index
 // TILE-10n25:       %[[NEW_J:.*]] = addi %[[J]], %[[L]] : index
 // TILE-10n25:       %[[NEW_INT_I:.*]] = index_cast %[[NEW_I]] : index to i32
 // TILE-10n25:       %[[NEW_FLOAT_I:.*]] = sitofp %[[NEW_INT_I]] : i32 to f32
 // TILE-10n25:       %[[NEW_INT_J:.*]] = index_cast %[[NEW_J]] : index to i32
 // TILE-10n25:       %[[NEW_FLOAT_J:.*]] = sitofp %[[NEW_INT_J]] : i32 to f32
 // TILE-10n25:       %[[OUT:.*]] = addf %[[NEW_FLOAT_I]], %[[NEW_FLOAT_J]] : f32
 
 // TILE-25n0-LABEL: func @indexed_generic_matrix
 // TILE-25n0: %[[C25:.*]] = constant 25 : index
 // TILE-25n0: scf.for %[[L:.*]] = {{.*}} step %[[C25]]
-// TILE-25n0:   linalg.indexed_generic
-// TILE-25n0:   ^bb0(%[[I:.*]]: index, %[[J:.*]]: index, %[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-25n0:   linalg.generic
+// TILE-25n0:   ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-25n0:     %[[I:.*]] = linalg.index 0 : index
+// TILE-25n0:     %[[J:.*]] = linalg.index 1 : index
 // TILE-25n0:     %[[NEW_I:.*]] = addi %[[I]], %[[L]] : index
 // TILE-25n0:     %[[NEW_INT_I:.*]] = index_cast %[[NEW_I]] : index to i32
 // TILE-25n0:     %[[NEW_FLOAT_I:.*]] = sitofp %[[NEW_INT_I]] : i32 to f32
 // TILE-25n0:     %[[INT_J:.*]] = index_cast %[[J]] : index to i32
 // TILE-25n0:     %[[FLOAT_J:.*]] = sitofp %[[INT_J]] : i32 to f32
 // TILE-25n0:     %[[OUT:.*]] = addf %[[NEW_FLOAT_I]], %[[FLOAT_J]] : f32
 
 // TILE-0n25-LABEL: func @indexed_generic_matrix
 // TILE-0n25: %[[C25:.*]] = constant 25 : index
 // TILE-0n25: scf.for %[[L:.*]] = {{.*}} step %[[C25]]
-// TILE-0n25:   linalg.indexed_generic
-// TILE-0n25:   ^bb0(%[[I:.*]]: index, %[[J:.*]]: index, %[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-0n25:   linalg.generic
+// TILE-0n25:   ^bb0(%[[IN:.*]]: f32, %[[OUT:.*]]: f32):
+// TILE-0n25:     %[[I:.*]] = linalg.index 0 : index
+// TILE-0n25:     %[[J:.*]] = linalg.index 1 : index
 // TILE-0n25:     %[[NEW_J:.*]] = addi %[[J]], %[[L]] : index
 // TILE-0n25:     %[[INT_I:.*]] = index_cast %[[I]] : index to i32
 // TILE-0n25:     %[[FLOAT_I:.*]] = sitofp %[[INT_I]] : i32 to f32
 // TILE-0n25:     %[[NEW_INT_J:.*]] = index_cast %[[NEW_J]] : index to i32
 // TILE-0n25:     %[[NEW_FLOAT_J:.*]] = sitofp %[[NEW_INT_J]] : i32 to f32
 // TILE-0n25:     %[[OUT:.*]] = addf %[[FLOAT_I]], %[[NEW_FLOAT_J]] : f32

mlir/test/Dialect/Linalg/tile-tensors.mlir

 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 //       CHECK:   %[[INIT:.+]] = linalg.init_tensor
 //       CHECK:   %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) {
 //       CHECK:     %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
 //       CHECK:       %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) {
 //       CHECK:       %[[STARG0:.+]] = subtensor %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
 //       CHECK:       %[[STARG1:.+]] = subtensor %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
 //       CHECK:       %[[STARG2:.+]] = subtensor %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
-//       CHECK:       %[[STRETURN:.+]] = linalg.indexed_generic
+//       CHECK:       %[[STRETURN:.+]] = linalg.generic
 //  CHECK-SAME:         ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
 //  CHECK-SAME:         outs(%[[STARG2]] : tensor<?x?x?xf32>)
 //       CHECK:       %[[TD:.+]] = subtensor_insert %[[STRETURN]] into %[[TC2]]
 //       CHECK:       scf.yield %[[TD]]
 //       CHECK:     }
 //       CHECK:     scf.yield %[[TD2]]
 //       CHECK:   }
 //       CHECK:   scf.yield %[[TD1]]