[mlir][linalg] Remove IndexedGenericOp support from LinalgBufferize...

[gysit]

Summary

after introducing the IndexedGenericOp to GenericOp canonicalization (https://reviews.llvm.org/D101612).

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

for (auto en : llvm::enumerate(linalgOp->getResultTypes())) {
     }
 
     resultBuffers.push_back(b.create<memref::AllocOp>(
         loc, memrefType, getDynOperands(loc, resultTensor, b)));
   }
   return success();
 }
 
-/// Specialization for `linalg::GenericOp` and `linalg::IndexedGenericOp`.
+/// Specialization for `linalg::GenericOp`.
 /// 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.
-template <typename GenericOpTy>
 static void
 finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter,
-                                     GenericOpTy genericOp, ValueRange inputs,
+                                     GenericOp genericOp, ValueRange inputs,
                                      ValueRange outputs) {
   // Generate a new linalg operation that works on buffers.
-  auto newGenericOp = rewriter.create<GenericOpTy>(
+  auto newGenericOp = rewriter.create<GenericOp>(
       genericOp.getLoc(),
       /*resultTensorTypes=*/llvm::None,
       /*inputs=*/inputs,
       /*outputs=*/outputs, genericOp.indexing_maps(),
       genericOp.iterator_types(), genericOp.docAttr(),
       genericOp.library_callAttr());
 
   // Create a new block in the region of the new Generic Op.
finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter,
   rewriter.replaceOp(genericOp, outputs);
 }
 
 /// 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();
   newOperands.append(otherOperands.begin(), otherOperands.end());
   linalgOp.clone(rewriter, linalgOp.getLoc(),
                  /*resultTypes=*/ArrayRef<Type>{}, newOperands);
   // Replace the results of the old op with the new output buffers.
   rewriter.replaceOp(linalgOp, outputs);

 /// instantiating one pattern for each LinalgOp.
 class BufferizeAnyLinalgOp : public OpInterfaceConversionPattern<LinalgOp> {
 public:
   using OpInterfaceConversionPattern<LinalgOp>::OpInterfaceConversionPattern;
 
   LogicalResult
   matchAndRewrite(LinalgOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
+    // Canonicalize indexed generic operations before bufferization.
+    if (isa<IndexedGenericOp>(op))
+      return failure();
+
     // GenericOpAdaptor below expects an `operand_segment_sizes` attribute.
     if (!op->hasAttr("operand_segment_sizes"))
       return failure();
 
     // We abuse the GenericOpAdaptor here.
     // TODO: Manually create an Adaptor that captures inputs and outputs for all
     // linalg::LinalgOp interface ops.
     linalg::GenericOpAdaptor adaptor(operands, op->getAttrDictionary());
 
     Location loc = op.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
     if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
                                          newOutputBuffers, rewriter))) {
       return op.emitOpError()
              << "Failed to allocate buffers for tensor results.";
     }
 
     // Delegate to the linalg generic pattern.
     if (auto genericOp = dyn_cast<linalg::GenericOp>(*op)) {
-      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);
+      finalizeBufferAllocationForGenericOp(rewriter, genericOp,
+                                           adaptor.inputs(), newOutputBuffers);
       return success();
     }
 
     finalizeBufferAllocation(rewriter, op, adaptor.inputs(), newOutputBuffers);
     return success();
   }
 };
 

mlir/test/Dialect/Linalg/bufferize.mlir

%0, %1 = linalg.generic {
         %tmp1 = math.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:.*]] = memref.alloc() : memref<4xi32>
-// CHECK:           %[[RESULT1:.*]] = memref.alloc() : memref<4xi32>
-// CHECK:           linalg.generic
-// CHECK-SAME:      ins(%{{.*}} : memref<4xi32>)
-// CHECK-SAME:      outs(%[[RESULT0]], %[[RESULT1]] : memref<4xi32>, memref<4xi32>)
-// 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):
-        %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)>
 
 // Check that the allocs properly consider the different shapes of the output
 // operands. The permuted indexing maps translate to different output shapes.
 
 // CHECK-LABEL:   func @dynamic_results(
 // CHECK-SAME:                          %[[ARG:.*]]: tensor<?x?xf32>
 // CHECK:           %[[C0:.*]] = constant 0 : index