This is an archive of the discontinued LLVM Phabricator instance.

Differential D102089

[mlir] Add a pattern to bufferize linalg.tensor_reshape.
ClosedPublic

Authored by pifon2a on May 7 2021, 12:22 PM.

Details

Reviewers
herhut
bkramer
nicolasvasilache
Commits
rGa3f22d020b27: [mlir] Add a pattern to bufferize linalg.tensor_reshape.

Diff Detail

Event Timeline

pifon2a created this revision.May 7 2021, 12:22 PM
Herald added subscribers: dcaballe, cota, mravishankar and 17 others. · View Herald TranscriptMay 7 2021, 12:22 PM
pifon2a requested review of this revision.May 7 2021, 12:22 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptMay 7 2021, 12:22 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer, nicolasvasilache. · View Herald Transcript
bkramer accepted this revision.May 7 2021, 12:24 PM

looks good

This revision is now accepted and ready to land.May 7 2021, 12:24 PM
This revision was landed with ongoing or failed builds.May 7 2021, 12:31 PM
Closed by commit rGa3f22d020b27: [mlir] Add a pattern to bufferize linalg.tensor_reshape. (authored by pifon2a). · Explain Why
This revision was automatically updated to reflect the committed changes.
pifon2a added a commit: rGa3f22d020b27: [mlir] Add a pattern to bufferize linalg.tensor_reshape..
Harbormaster completed remote builds in B103246: Diff 343738.May 7 2021, 1:16 PM

Revision Contents

PathSize
mlir/
lib/
Dialect/
Linalg/
Transforms/
18 lines
test/
Dialect/
Linalg/
15 lines

Diff 343740

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

Show First 20 Lines • Show All 143 Lines • ▼ Show 20 Lines matchAndRewrite(InitTensorOp op, ArrayRef<Value> operands,
linalg::InitTensorOpAdaptor adaptor(operands, op->getAttrDictionary()); linalg::InitTensorOpAdaptor adaptor(operands, op->getAttrDictionary());
rewriter.replaceOpWithNewOp<memref::AllocOp>( rewriter.replaceOpWithNewOp<memref::AllocOp>(
op, getTypeConverter()->convertType(op.getType()).cast<MemRefType>(), op, getTypeConverter()->convertType(op.getType()).cast<MemRefType>(),
adaptor.sizes()); adaptor.sizes());
return success(); return success();
} }
}; };
/// Conversion pattern that replaces `linalg.tensor_reshape` with
/// `linalg.reshape`.
class BufferizeTensorReshapeOp : public OpConversionPattern<TensorReshapeOp> {
public:
using OpConversionPattern<TensorReshapeOp>::OpConversionPattern;
LogicalResult
matchAndRewrite(TensorReshapeOp op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const final {
linalg::TensorReshapeOpAdaptor adaptor(operands, op->getAttrDictionary());
rewriter.replaceOpWithNewOp<linalg::ReshapeOp>(
op, getTypeConverter()->convertType(op.getType()).cast<MemRefType>(),
adaptor.src(), adaptor.reassociation());
return success();
}
};
/// Conversion pattern that bufferizes `linalg.fill` operation. /// Conversion pattern that bufferizes `linalg.fill` operation.
class BufferizeFillOp : public OpConversionPattern<FillOp> { class BufferizeFillOp : public OpConversionPattern<FillOp> {
public: public:
using OpConversionPattern<FillOp>::OpConversionPattern; using OpConversionPattern<FillOp>::OpConversionPattern;
LogicalResult LogicalResult
matchAndRewrite(FillOp op, ArrayRef<Value> operands, matchAndRewrite(FillOp op, ArrayRef<Value> operands,
ConversionPatternRewriter &rewriter) const final { ConversionPatternRewriter &rewriter) const final {
▲ Show 20 Lines • Show All 171 Lines • ▼ Show 20 Lines
void mlir::linalg::populateLinalgBufferizePatterns( void mlir::linalg::populateLinalgBufferizePatterns(
BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) { BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns) {
// TODO: Drop this once tensor constants work in standard. // TODO: Drop this once tensor constants work in standard.
// clang-format off // clang-format off
patterns.add< patterns.add<
BufferizeAnyLinalgOp, BufferizeAnyLinalgOp,
BufferizeFillOp, BufferizeFillOp,
BufferizeInitTensorOp, BufferizeInitTensorOp,
BufferizeTensorReshapeOp,
SubTensorOpConverter, SubTensorOpConverter,
SubTensorInsertOpConverter SubTensorInsertOpConverter
>(typeConverter, patterns.getContext()); >(typeConverter, patterns.getContext());
// clang-format on // clang-format on
} }

mlir/test/Dialect/Linalg/bufferize.mlir

Show First 20 Lines • Show All 272 Lines • ▼ Show 20 Linesfunc @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%c0 = constant 0.0 : f32 %c0 = constant 0.0 : f32
// CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32> // CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<?xf32>
// CHECK: linalg.fill(%[[MEMREF]], %cst) : memref<?xf32>, f32 // CHECK: linalg.fill(%[[MEMREF]], %cst) : memref<?xf32>, f32
// CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[MEMREF]] : memref<?xf32> // CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[MEMREF]] : memref<?xf32>
// CHECK: return %[[TENSOR]] // CHECK: return %[[TENSOR]]
%0 = linalg.fill(%arg0, %c0) : tensor<?xf32>, f32 -> tensor<?xf32> %0 = linalg.fill(%arg0, %c0) : tensor<?xf32>, f32 -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// -----
// CHECK-LABEL: func @bufferize_tensor_reshape(
// CHECK-SAME: %[[IN:.*]]: tensor<4x5xf32>
func @bufferize_tensor_reshape(%arg0: tensor<4x5xf32>) -> tensor<20xf32> {
%out = linalg.tensor_reshape %arg0 [[0, 1]] :
tensor<4x5xf32> into tensor<20xf32>
return %out : tensor<20xf32>
}
// CHECK: %[[MEMREF:.*]] = memref.buffer_cast %[[IN]] : memref<4x5xf32>
// CHECK: %[[RESHAPE:.*]] = linalg.reshape %[[MEMREF]] {{\[}}[0, 1]]
// CHECK-SAME: : memref<4x5xf32> into memref<20xf32>
// CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[RESHAPE]] : memref<20xf32>
// CHECK: return %[[TENSOR]]