[mlir][bufferize][NFC] Remove obsolete tensor bufferization patterns from Linalg/Bufferize.cpp

[springerm]

Summary

Depends On D119935

[springerm]

update

[bkramer]

looks good

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

matchAndRewrite(InitTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     rewriter.replaceOpWithNewOp<memref::AllocOp>(
         op, getTypeConverter()->convertType(op.getType()).cast<MemRefType>(),
         adaptor.sizes());
     return success();
   }
 };
 
-/// Conversion pattern that replaces `linalg.tensor_reshape` with
-/// `linalg.reshape`.
-template <typename TensorReshapeOp,
-          typename Adaptor = typename TensorReshapeOp::Adaptor>
-class BufferizeTensorReshapeOp : public OpConversionPattern<TensorReshapeOp> {
-public:
-  using OpConversionPattern<TensorReshapeOp>::OpConversionPattern;
-  using ReshapeOp = typename std::conditional_t<
-      std::is_same<TensorReshapeOp, tensor::ExpandShapeOp>::value,
-      memref::ExpandShapeOp, memref::CollapseShapeOp>;
-
-  LogicalResult
-  matchAndRewrite(TensorReshapeOp op, Adaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const final {
-    rewriter.replaceOpWithNewOp<ReshapeOp>(op,
-                                           this->getTypeConverter()
-                                               ->convertType(op.getType())
-                                               .template cast<MemRefType>(),
-                                           adaptor.src(),
-                                           adaptor.reassociation());
-    return success();
-  }
-};
-
 /// Conversion pattern that bufferizes `linalg.fill` operation.
 class BufferizeFillOp : public OpConversionPattern<FillOp> {
 public:
   using OpConversionPattern<FillOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(FillOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
              << "Failed to allocate buffers for tensor results.";
     }
     createLinalgOpOnBuffers(rewriter, op, adaptor.inputs(), newOutputBuffers);
     // Replace the results of the old op with the new output buffers.
     rewriter.replaceOp(op, newOutputBuffers);
     return success();
   }
 };
-
-/// Convert `extract_slice %t [offsets][sizes][strides] -> %st` to an
-/// alloc + copy pattern.
-/// ```
-///   %a = alloc(sizes)
-///   %sv = subview %source [offsets][sizes][strides]
-///   memref.copy(%sv, %a)
-/// ```
-///
-/// This pattern is arguable a std pattern once memref::CopyOp becomes
-/// std::CopyOp.
-class ExtractSliceOpConverter
-    : public OpConversionPattern<tensor::ExtractSliceOp> {
-public:
-  using OpConversionPattern<tensor::ExtractSliceOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(tensor::ExtractSliceOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const final {
-    Value sourceMemref = adaptor.source();
-    assert(sourceMemref.getType().isa<MemRefType>());
-
-    MemRefType subviewMemRefType =
-        getTypeConverter()->convertType(op.getType()).cast<MemRefType>();
-    // op.sizes() capture exactly the dynamic alloc operands matching the
-    // subviewMemRefType thanks to subview/slice canonicalization and
-    // verification.
-    Value alloc = rewriter.create<memref::AllocOp>(
-        op.getLoc(), subviewMemRefType, op.sizes());
-    Value subView = rewriter.create<memref::SubViewOp>(
-        op.getLoc(), sourceMemref, op.getMixedOffsets(), op.getMixedSizes(),
-        op.getMixedStrides());
-    rewriter.create<memref::CopyOp>(op.getLoc(), subView, alloc);
-    rewriter.replaceOp(op, alloc);
-    return success();
-  }
-};
-
-/// Convert `insert_slice %source into %dest [offsets][sizes][strides] ->
-/// %t` to an buffer_cast + subview + copy + tensor_load pattern.
-/// buffer_cast and tensor_load are inserted automatically by the
-/// conversion infra:
-/// ```
-///   %sv = subview %dest [offsets][sizes][strides]
-///   memref.copy(%source, %sv)
-///   // replace with %dest
-/// ```
-///
-/// This pattern is arguable a std pattern once memref::CopyOp becomes
-/// std::CopyOp.
-class InsertSliceOpConverter
-    : public OpConversionPattern<tensor::InsertSliceOp> {
-public:
-  using OpConversionPattern<tensor::InsertSliceOp>::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(tensor::InsertSliceOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const final {
-    Value sourceMemRef = adaptor.source();
-    assert(sourceMemRef.getType().isa<MemRefType>());
-
-    // For now, be conservative and copy the converted input memref.
-    // In general, the converted input memref here could be aliased or could
-    // point into constant memory, so mutating it would lead to miscompilations.
-    Value destMemRef = cloneMemref(op.getLoc(), adaptor.dest(), rewriter);
-    assert(destMemRef.getType().isa<MemRefType>());
-
-    // Take a subview to copy the small memref.
-    Value subview = rewriter.create<memref::SubViewOp>(
-        op.getLoc(), destMemRef, op.getMixedOffsets(), op.getMixedSizes(),
-        op.getMixedStrides());
-    // Copy the small memref.
-    rewriter.create<memref::CopyOp>(op.getLoc(), sourceMemRef, subview);
-    rewriter.replaceOp(op, destMemRef);
-    return success();
-  }
-};
 } // namespace
 
 namespace {
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     bufferization::BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<arith::ArithmeticDialect, AffineDialect,
                            memref::MemRefDialect, StandardOpsDialect,
                            tensor::TensorDialect>();
-    target.addIllegalOp<InitTensorOp, tensor::PadOp, tensor::CollapseShapeOp,
-                        tensor::ExpandShapeOp, tensor::ExtractSliceOp,
-                        tensor::InsertSliceOp>();
+    target.addIllegalOp<InitTensorOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
 
     RewritePatternSet patterns(&context);

 void mlir::linalg::populateLinalgBufferizePatterns(
     bufferization::BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns) {
   // TODO: Drop this once tensor constants work in standard.
   // clang-format off
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
-      BufferizeInitTensorOp,
-      BufferizeTensorReshapeOp<tensor::ExpandShapeOp>,
-      BufferizeTensorReshapeOp<tensor::CollapseShapeOp>,
-      ExtractSliceOpConverter,
-      InsertSliceOpConverter
+      BufferizeInitTensorOp
     >(typeConverter, patterns.getContext());
   // clang-format on
-  patterns.add<GeneralizePadOpPattern>(patterns.getContext());
 }

mlir/test/Dialect/Linalg/bufferize.mlir

^bb(%v0: vector<3x4xi4>, %v1: f32) :
       linalg.yield %v1 : f32
   } -> tensor<3x2xf32>
 
   return %0 : tensor<3x2xf32>
 }
 
 // -----
 
-// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
-// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)>
-
-func private @make_index() -> index
-
-// CHECK-LABEL: func @bufferize_slice(
-//  CHECK-SAME:   %[[T:[0-9a-z]*]]: tensor<?x?xf32>
-func @bufferize_slice(%t : tensor<?x?xf32>) -> (tensor<2x3xf32>, tensor<2x?xf32>) {
-  //      CHECK: %[[M:.*]] = bufferization.to_memref %[[T]] : memref<?x?xf32>
-
-  //      CHECK: %[[IDX:.*]] = call @make_index() : () -> index
-  %i0 = call @make_index() : () -> index
-
-  // CHECK-NEXT: %[[A0:.*]] = memref.alloc() : memref<2x3xf32>
-  // CHECK-NEXT: %[[SM0:.*]] = memref.subview %[[M]][0, 0] [2, 3] [1, 1]
-  // CHECK-SAME:   memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]>
-  // CHECK-NEXT: memref.copy %[[SM0]], %[[A0]] : memref<2x3xf32, #[[$MAP0]]> to memref<2x3xf32>
-  // CHECK-NEXT: %[[RT0:.*]] = bufferization.to_tensor %[[A0]] : memref<2x3xf32>
-  %st0 = tensor.extract_slice %t[0, 0][2, 3][1, 1] : tensor<?x?xf32> to tensor<2x3xf32>
-
-  // CHECK-NEXT: %[[A1:.*]] = memref.alloc(%[[IDX]]) : memref<2x?xf32>
-  // CHECK-NEXT: %[[SM1:.*]] = memref.subview %[[M]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
-  // CHECK-SAME:   memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]>
-  // CHECK-NEXT: memref.copy %[[SM1]], %[[A1]] : memref<2x?xf32, #[[$MAP1]]> to memref<2x?xf32>
-  // CHECK-NEXT: %[[RT1:.*]] = bufferization.to_tensor %[[A1]] : memref<2x?xf32>
-  %st1 = tensor.extract_slice %t[0, %i0][2, %i0][1, 2] : tensor<?x?xf32> to tensor<2x?xf32>
-
-  // CHECK-NEXT: return %[[RT0]], %[[RT1]]
-  return %st0, %st1 : tensor<2x3xf32>, tensor<2x?xf32>
-}
-
-// -----
-
-// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
-// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)>
-
-func private @make_index() -> index
-
-// CHECK-LABEL: func @bufferize_insert_slice(
-//  CHECK-SAME:   %[[T:[0-9a-z]*]]: tensor<?x?xf32>
-//  CHECK-SAME:   %[[ST0:[0-9a-z]*]]: tensor<2x3xf32>
-//  CHECK-SAME:   %[[ST1:[0-9a-z]*]]: tensor<2x?xf32>
-func @bufferize_insert_slice(%t : tensor<?x?xf32>, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) ->
-    (tensor<?x?xf32>, tensor<?x?xf32>) {
-  // CHECK-DAG: %[[M:.*]] = bufferization.to_memref %[[T]] : memref<?x?xf32>
-  // CHECK-DAG: %[[SM0:.*]] = bufferization.to_memref %[[ST0]] : memref<2x3xf32>
-  // CHECK-DAG: %[[SM1:.*]] = bufferization.to_memref %[[ST1]] : memref<2x?xf32>
-
-  %c0 = arith.constant 0 : index
-  %c1 = arith.constant 1 : index
-  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-  %i0 = call @make_index() : () -> index
-  // CHECK: %[[IDX:.*]] = call @make_index() : () -> index
-
-
-  // CHECK-NEXT: %[[DIM0:.*]] = tensor.dim %[[T]], %[[C0]] : tensor<?x?xf32>
-  // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[T]], %[[C1]] : tensor<?x?xf32>
-  // CHECK-NEXT: %[[M_COPY0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
-  // CHECK-NEXT: memref.copy %[[M]], %[[M_COPY0]] : memref<?x?xf32> to memref<?x?xf32>
-  // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[M_COPY0]][0, 0] [2, 3] [1, 1]
-  // CHECK-SAME:   memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]>
-  // CHECK-NEXT: memref.copy %[[SM0]], %[[SUBVIEW0]] : memref<2x3xf32> to memref<2x3xf32, #[[$MAP0]]>
-  // CHECK-NEXT: %[[RT0:.*]] = bufferization.to_tensor %[[M_COPY0]] : memref<?x?xf32>
-  %t0 = tensor.insert_slice %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor<?x?xf32>
-
-  // CHECK-NEXT: %[[M_COPY1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
-  // CHECK-NEXT: memref.copy %[[M]], %[[M_COPY1]] : memref<?x?xf32> to memref<?x?xf32>
-  // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[M_COPY1]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
-  // CHECK-SAME:   memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]>
-  // CHECK-NEXT: memref.copy %[[SM1]], %[[SUBVIEW1]] : memref<2x?xf32> to memref<2x?xf32, #[[$MAP1]]>
-  // CHECK-NEXT: %[[RT1:.*]] = bufferization.to_tensor %[[M_COPY1]] : memref<?x?xf32>
-  %t1 = tensor.insert_slice %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor<?x?xf32>
-
-  //     CHECK: return %[[RT0]], %[[RT1]]
-  return %t0, %t1: tensor<?x?xf32>, tensor<?x?xf32>
-}
-
-// -----
-
 // CHECK-LABEL: func @bufferize_fill(
 // CHECK-SAME:    %[[IN:.*]]: tensor<?xf32>
 func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
   %c0 = arith.constant 0.0 : f32
   // CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<?xf32>
   // CHECK: linalg.fill(%cst, %[[MEMREF]]) : f32, memref<?xf32>
   // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<?xf32>
   // CHECK: return %[[TENSOR]]
   %0 = linalg.fill(%c0, %arg0) : f32, tensor<?xf32> -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @bufferize_tensor_collapse_shape(
-// CHECK-SAME:    %[[IN:.*]]: tensor<4x5xf32>
-func @bufferize_tensor_collapse_shape(%arg0: tensor<4x5xf32>) -> tensor<20xf32> {
-  %out = tensor.collapse_shape %arg0 [[0, 1]] :
-     tensor<4x5xf32> into tensor<20xf32>
-  return %out : tensor<20xf32>
-}
-// CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<4x5xf32>
-// CHECK: %[[RESHAPE:.*]] = memref.collapse_shape %[[MEMREF]] {{\[}}[0, 1]]
-// CHECK-SAME: : memref<4x5xf32> into memref<20xf32>
-// CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[RESHAPE]] : memref<20xf32>
-// CHECK: return %[[TENSOR]]
-
-// -----
-
-// CHECK-LABEL:   func @pad_tensor_dynamic_shape(
-// CHECK-SAME:                                   %[[IN:.*]]: tensor<4x?x2x?xf32>,
-// CHECK-SAME:                                   %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> {
-func @pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
-  %c0 = arith.constant 0 : index
-  %cst = arith.constant 0.0 : f32
-  %out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1]  {
-  ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):  
-    tensor.yield %cst : f32
-  } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
-  return %out : tensor<4x?x?x?xf32>
-}
-
-// CHECK-DAG:       %[[C3:.*]] = arith.constant 3 : index
-// CHECK-DAG:       %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:           %[[IN_MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<4x?x2x?xf32>
-// CHECK:           %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>
-// CHECK:           %[[OUT_DIM2:.*]] = arith.addi %[[OFFSET]], %[[C2]] : index
-// CHECK:           %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>
-// CHECK:           %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index
-// CHECK:           %[[FILLED:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32>
-// CHECK:           linalg.fill(%[[CST]], %[[FILLED]]) : f32, memref<4x?x?x?xf32>
-// CHECK:           %[[OUT:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32>
-// CHECK:           memref.copy %[[FILLED]], %[[OUT]] : memref<4x?x?x?xf32> to memref<4x?x?x?xf32>
-// CHECK:           %[[INTERIOR:.*]] = memref.subview %[[OUT]][0, 0, %[[OFFSET]], 0] [4, %[[DIM1]], 2, %[[DIM3]]] [1, 1, 1, 1] : memref<4x?x?x?xf32> to memref<4x?x2x?xf32, #map>
-// CHECK:           memref.copy %[[IN_MEMREF]], %[[INTERIOR]] : memref<4x?x2x?xf32> to memref<4x?x2x?xf32, #map>
-// CHECK:           %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<4x?x?x?xf32>
-// CHECK:           return %[[OUT_TENSOR]] : tensor<4x?x?x?xf32>
-// CHECK:         }
-
-// -----
-
 // CHECK-LABEL:   func @bufferize_dot
 func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> {
   %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>)
                           outs(%out : tensor<f32>) -> tensor<f32>
   return %dot : tensor<f32>
   // CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>)
   // CHECK-SAME:       outs(%[[OUT:.*]] : memref<f32>)
   // CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<f32>
   // CHECK: return %[[OUT_TENSOR]]
 }

mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir

-// RUN: mlir-opt %s -linalg-bufferize \
-// RUN: -arith-bufferize -tensor-bufferize -func-bufferize \
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns \
+// RUN: -linalg-bufferize -arith-bufferize -tensor-bufferize -func-bufferize \
 // RUN: -finalizing-bufferize -buffer-deallocation \
 // RUN: -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext,%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 
 func @main() {