diff --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp @@ -116,30 +116,6 @@ } }; -/// Conversion pattern that replaces `linalg.tensor_reshape` with -/// `linalg.reshape`. -template -class BufferizeTensorReshapeOp : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - using ReshapeOp = typename std::conditional_t< - std::is_same::value, - memref::ExpandShapeOp, memref::CollapseShapeOp>; - - LogicalResult - matchAndRewrite(TensorReshapeOp op, Adaptor adaptor, - ConversionPatternRewriter &rewriter) const final { - rewriter.replaceOpWithNewOp(op, - this->getTypeConverter() - ->convertType(op.getType()) - .template cast(), - adaptor.src(), - adaptor.reassociation()); - return success(); - } -}; - /// Conversion pattern that bufferizes `linalg.fill` operation. class BufferizeFillOp : public OpConversionPattern { public: @@ -191,83 +167,6 @@ 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 { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(tensor::ExtractSliceOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const final { - Value sourceMemref = adaptor.source(); - assert(sourceMemref.getType().isa()); - - MemRefType subviewMemRefType = - getTypeConverter()->convertType(op.getType()).cast(); - // op.sizes() capture exactly the dynamic alloc operands matching the - // subviewMemRefType thanks to subview/slice canonicalization and - // verification. - Value alloc = rewriter.create( - op.getLoc(), subviewMemRefType, op.sizes()); - Value subView = rewriter.create( - op.getLoc(), sourceMemref, op.getMixedOffsets(), op.getMixedSizes(), - op.getMixedStrides()); - rewriter.create(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 { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(tensor::InsertSliceOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const final { - Value sourceMemRef = adaptor.source(); - assert(sourceMemRef.getType().isa()); - - // 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()); - - // Take a subview to copy the small memref. - Value subview = rewriter.create( - op.getLoc(), destMemRef, op.getMixedOffsets(), op.getMixedSizes(), - op.getMixedStrides()); - // Copy the small memref. - rewriter.create(op.getLoc(), sourceMemRef, subview); - rewriter.replaceOp(op, destMemRef); - return success(); - } -}; } // namespace namespace { @@ -283,9 +182,7 @@ target.addLegalDialect(); - target.addIllegalOp(); + target.addIllegalOp(); // Mark all Linalg operations illegal as long as they work on tensors. auto isLegalOperation = [&](Operation *op) { @@ -314,12 +211,7 @@ patterns.add< BufferizeAnyLinalgOp, BufferizeFillOp, - BufferizeInitTensorOp, - BufferizeTensorReshapeOp, - BufferizeTensorReshapeOp, - ExtractSliceOpConverter, - InsertSliceOpConverter + BufferizeInitTensorOp >(typeConverter, patterns.getContext()); // clang-format on - patterns.add(patterns.getContext()); } diff --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir --- a/mlir/test/Dialect/Linalg/bufferize.mlir +++ b/mlir/test/Dialect/Linalg/bufferize.mlir @@ -162,86 +162,6 @@ // ----- -// 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 -func @bufferize_slice(%t : tensor) -> (tensor<2x3xf32>, tensor<2x?xf32>) { - // CHECK: %[[M:.*]] = bufferization.to_memref %[[T]] : memref - - // 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 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 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 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 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 -// CHECK-SAME: %[[ST0:[0-9a-z]*]]: tensor<2x3xf32> -// CHECK-SAME: %[[ST1:[0-9a-z]*]]: tensor<2x?xf32> -func @bufferize_insert_slice(%t : tensor, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) -> - (tensor, tensor) { - // CHECK-DAG: %[[M:.*]] = bufferization.to_memref %[[T]] : memref - // 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 - // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[T]], %[[C1]] : tensor - // CHECK-NEXT: %[[M_COPY0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref - // CHECK-NEXT: memref.copy %[[M]], %[[M_COPY0]] : memref to memref - // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[M_COPY0]][0, 0] [2, 3] [1, 1] - // CHECK-SAME: memref 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 - %t0 = tensor.insert_slice %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor - - // CHECK-NEXT: %[[M_COPY1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref - // CHECK-NEXT: memref.copy %[[M]], %[[M_COPY1]] : memref to memref - // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[M_COPY1]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2] - // CHECK-SAME: memref 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 - %t1 = tensor.insert_slice %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor - - // CHECK: return %[[RT0]], %[[RT1]] - return %t0, %t1: tensor, tensor -} - -// ----- - // CHECK-LABEL: func @bufferize_fill( // CHECK-SAME: %[[IN:.*]]: tensor func @bufferize_fill(%arg0: tensor) -> tensor { @@ -256,55 +176,6 @@ // ----- -// 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) -> tensor { %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>) diff --git a/mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir b/mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir --- a/mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir +++ b/mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir @@ -1,5 +1,5 @@ -// 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 \