[mlir][linalg] Add bufferization pattern for `linalg.indexed_generic`.

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

for (auto shapeElement : llvm::enumerate(tensorType.getShape())) {
         return failure();
       }
     }
     resultBuffers.push_back(b.create<AllocOp>(loc, memrefType, dynOperands));
   }
   return success();
 }
 
-// Specialization for `linalg::GenericOp`.
+/// Specialization for `linalg::GenericOp` and `linalg::IndexedGenericOp`.
 /// A pattern to convert Generic Linalg operations which work on tensors to
 /// use buffers. BufferPlacement pass should be later used to move
 /// Alloc operations to the correct positions and insert the missing Dealloc
 /// operations in the correct places.
-static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
-                                     linalg::GenericOp genericOp,
-                                     ValueRange inputs, ValueRange outputs) {
+template <typename GenericOpTy>
+static void
+finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter,
+                                     GenericOpTy genericOp, ValueRange inputs,
+                                     ValueRange outputs) {
   // Generate a new linalg operation that works on buffers.
-  auto newGenericOp = rewriter.create<linalg::GenericOp>(
+  auto newGenericOp = rewriter.create<GenericOpTy>(
       genericOp.getLoc(),
       /*resultTensorTypes=*/llvm::None,
       /*inputs=*/inputs,
       /*outputBuffers=*/outputs,
       /*initTensors=*/llvm::None, genericOp.indexing_maps(),
       genericOp.iterator_types(), genericOp.docAttr(),
       genericOp.library_callAttr(), genericOp.sparseAttr());
 
for (auto &op : oldBlock->getOperations()) {
     Operation *clonedOp = rewriter.clone(op, mapping);
     mapping.map(op.getResults(), clonedOp->getResults());
   }
 
   // Replace the results of the old op with the new output buffers.
   rewriter.replaceOp(genericOp, outputs);
 }
 
-// TODO: Specialization for `linalg::IndexedGenericOp`.
-
-// Specialization for all other `linalg::LinalgOp`.
+/// Specialization for all other `linalg::LinalgOp`.
 static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
                                      linalg::LinalgOp linalgOp,
                                      ValueRange inputs, ValueRange outputs) {
   assert(!isa<linalg::GenericOp>(linalgOp.getOperation()));
   assert(!isa<linalg::IndexedGenericOp>(linalgOp.getOperation()));
   SmallVector<Value, 8> newOperands = inputs;
   newOperands.append(outputs.begin(), outputs.end());
   auto otherOperands = linalgOp.getAssumedNonShapedOperands();
if (failed(allocateBuffersForResults(loc, linalgOp, adaptor,
                                          newOutputBuffers, rewriter))) {
       linalgOp.emitOpError()
           << "Failed to allocate buffers for tensor results.";
       return failure();
     }
 
     // Delegate to the linalg generic pattern.
     if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
-      finalizeBufferAllocation(rewriter, genericOp, adaptor.inputs(),
-                               newOutputBuffers);
+      finalizeBufferAllocationForGenericOp<GenericOp>(
+          rewriter, genericOp, adaptor.inputs(), newOutputBuffers);
+      return success();
+    }
+
+    // Delegate to the linalg indexed generic pattern.
+    if (auto genericOp = dyn_cast<linalg::IndexedGenericOp>(op)) {
+      finalizeBufferAllocationForGenericOp<IndexedGenericOp>(
+          rewriter, genericOp, adaptor.inputs(), newOutputBuffers);
       return success();
     }
 
     finalizeBufferAllocation(rewriter, linalgOp, adaptor.inputs(),
                              newOutputBuffers);
     return success();
   }
 };

mlir/test/Dialect/Linalg/bufferize.mlir

 #map0 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL:   func @multiple_results
 // CHECK:           %[[RESULT0:.*]] = alloc() : memref<4xf32>
 // CHECK:           %[[RESULT1:.*]] = alloc() : memref<4xf32>
 // CHECK:           linalg.generic
 // CHECK-SAME:      ins(%{{.*}} : memref<4xf32>)
 // CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xf32>, memref<4xf32>)
+// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
 func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
     %0, %1 = linalg.generic {
       indexing_maps = [#map0, #map0, #map0],
       iterator_types = ["parallel"]
     } ins(%arg0 : tensor<4xf32>) {
       ^bb0(%gen_arg1: f32):
         %tmp1 = exp %gen_arg1 : f32
         linalg.yield %tmp1, %tmp1 : f32, f32
     } -> tensor<4xf32>, tensor<4xf32>
     return %0, %1 : tensor<4xf32>, tensor<4xf32>
 }
 
 // -----
 
+#map0 = affine_map<(d0) -> (d0)>
+
+// CHECK-LABEL:   func @multiple_results_indexed
+// CHECK:           %[[RESULT0:.*]] = alloc() : memref<4xi32>
+// CHECK:           %[[RESULT1:.*]] = alloc() : memref<4xi32>
+// CHECK:           linalg.indexed_generic
+// CHECK-SAME:      ins(%{{.*}} : memref<4xi32>)
+// CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xi32>, memref<4xi32>)
+// CHECK-NEXT: ^bb0(%{{.*}}: index, %{{.*}}: 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>) {
+      ^bb0(%i: index, %gen_arg1: i32):
+        %i_i32 = index_cast %i : index to i32
+        %tmp1 = addi %gen_arg1, %i_i32 : i32
+        linalg.yield %tmp1, %tmp1 : i32, i32
+    } -> tensor<4xi32>, tensor<4xi32>
+    return %0, %1 : tensor<4xi32>, tensor<4xi32>
+}
+
+// -----
+
 #map_2d = affine_map<(d0, d1) -> (d0, d1)>
 #map_2d_inv = affine_map<(d0, d1) -> (d1, d0)>
 
 // Check that the allocs properly consider the different shapes of the output
 // operands. The permuted indexing maps translate to different output shapes.
 
 // CHECK: #map0 = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK: #map1 = affine_map<(d0, d1) -> (d1, d0)>