Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir][Bufferization] castOrReallocMemRefValue: Use BufferizationOptions #89175

Merged
merged 2 commits into from
Apr 18, 2024
Merged

[mlir][Bufferization] castOrReallocMemRefValue: Use BufferizationOptions #89175

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f9969fa
[mlir][Bufferization] castOrReallocMemRefValue: Use BufferizationOptions
mgehre-amd Apr 7, 2024
2799d51
Use default alignment of zero
mgehre-amd Apr 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,14 @@ void populateDynamicDimSizes(OpBuilder &b, Location loc, Value shapedValue,
/// This function returns `failure()` in case of unsupported casts. E.g., casts
/// with differing element types or memory spaces.
FailureOr<Value> castOrReallocMemRefValue(OpBuilder &b, Value value,
MemRefType type);
MemRefType type,
const BufferizationOptions &options);

/// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
/// to_memref op are different, a memref.cast is needed.
LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter,
ToMemrefOp toMemref);
ToMemrefOp toMemref,
const BufferizationOptions &options);

/// Add the canonicalization patterns for bufferization.dealloc to the given
/// pattern set to make them available to other passes (such as
Expand Down
29 changes: 16 additions & 13 deletions mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@ using namespace mlir::bufferization;
// Helper functions
//===----------------------------------------------------------------------===//

FailureOr<Value>
mlir::bufferization::castOrReallocMemRefValue(OpBuilder &b, Value value,
MemRefType destType) {
FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
OpBuilder &b, Value value, MemRefType destType,
const BufferizationOptions &options) {
auto srcType = llvm::cast<MemRefType>(value.getType());

// Element type, rank and memory space must match.
Expand Down Expand Up @@ -73,18 +73,21 @@ mlir::bufferization::castOrReallocMemRefValue(OpBuilder &b, Value value,
Value size = b.create<memref::DimOp>(loc, value, i);
dynamicOperands.push_back(size);
}
// TODO: Use alloc/memcpy callback from BufferizationOptions if called via
// BufferizableOpInterface impl of ToMemrefOp.
Value copy = b.create<memref::AllocOp>(loc, destType, dynamicOperands);
b.create<memref::CopyOp>(loc, value, copy);

FailureOr<Value> copy =
options.createAlloc(b, loc, destType, dynamicOperands);
if (failed(copy))
return failure();
if (failed(options.createMemCpy(b, loc, value, *copy)))
return failure();
return copy;
}

/// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the
/// to_memref op are different, a memref.cast is needed.
LogicalResult
mlir::bufferization::foldToMemrefToTensorPair(RewriterBase &rewriter,
ToMemrefOp toMemref) {
LogicalResult mlir::bufferization::foldToMemrefToTensorPair(
RewriterBase &rewriter, ToMemrefOp toMemref,
const BufferizationOptions &options) {
auto memrefToTensor = toMemref.getTensor().getDefiningOp<ToTensorOp>();
if (!memrefToTensor)
return failure();
Expand All @@ -105,7 +108,7 @@ mlir::bufferization::foldToMemrefToTensorPair(RewriterBase &rewriter,
// Ranked memref -> Ranked memref cast.
if (rankedSrcType && rankedDestType) {
FailureOr<Value> replacement = castOrReallocMemRefValue(
rewriter, memrefToTensor.getMemref(), rankedDestType);
rewriter, memrefToTensor.getMemref(), rankedDestType, options);
if (failed(replacement))
return failure();

Expand Down Expand Up @@ -795,7 +798,7 @@ struct ToMemrefToTensorFolding : public OpRewritePattern<ToMemrefOp> {

LogicalResult matchAndRewrite(ToMemrefOp toMemref,
PatternRewriter &rewriter) const final {
return foldToMemrefToTensorPair(rewriter, toMemref);
return foldToMemrefToTensorPair(rewriter, toMemref, {});
}
};

Expand Down Expand Up @@ -843,7 +846,7 @@ void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
const BufferizationOptions &options) {
// Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
(void)foldToMemrefToTensorPair(rewriter, *this);
(void)foldToMemrefToTensorPair(rewriter, *this, options);
// Note: The return value of `bufferize` indicates whether there was an error
// or not. (And not whether the pattern matched or not.)
return success();
Expand Down
6 changes: 3 additions & 3 deletions mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ BufferizeTypeConverter::BufferizeTypeConverter() {
if (!rankedDestType)
return nullptr;
FailureOr<Value> replacement =
castOrReallocMemRefValue(builder, inputs[0], rankedDestType);
castOrReallocMemRefValue(builder, inputs[0], rankedDestType, {});
mgehre-amd marked this conversation as resolved.
Show resolved Hide resolved
if (failed(replacement))
return nullptr;
return *replacement;
Expand Down Expand Up @@ -512,8 +512,8 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
// Fold all to_memref(to_tensor(x)) pairs.
for (Operation *op : toMemrefOps) {
rewriter.setInsertionPoint(op);
(void)bufferization::foldToMemrefToTensorPair(rewriter,
cast<ToMemrefOp>(op));
(void)bufferization::foldToMemrefToTensorPair(
rewriter, cast<ToMemrefOp>(op), options);
}

// Remove all dead to_tensor ops.
Expand Down
6 changes: 3 additions & 3 deletions mlir/test/Dialect/Bufferization/Transforms/finalizing-bufferize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ func.func @unable_to_convert_lone_tensor_load(%arg0: memref<f32>) {
// CHECK-SAME: %[[arg:.*]]: memref<?xf32, strided<[1], offset: ?>>)
// CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK: %[[dim:.*]] = memref.dim %[[arg]], %[[c0]]
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?xf32>
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) {{.*}} : memref<?xf32>
// CHECK: memref.copy %[[arg]], %[[alloc]]
// CHECK: return %[[alloc]]
func.func @dyn_layout_to_no_layout_cast(%m: memref<?xf32, strided<[1], offset: ?>>) -> memref<?xf32> {
Expand All @@ -48,7 +48,7 @@ func.func @dyn_layout_to_no_layout_cast(%m: memref<?xf32, strided<[1], offset: ?
// CHECK-SAME: %[[arg:.*]]: memref<?xf32, strided<[100], offset: ?>>)
// CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK: %[[dim:.*]] = memref.dim %[[arg]], %[[c0]]
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?xf32>
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) {{.*}} : memref<?xf32>
// CHECK: memref.copy %[[arg]], %[[alloc]]
// CHECK: return %[[alloc]]
func.func @fancy_layout_to_no_layout_cast(%m: memref<?xf32, strided<[100], offset: ?>>) -> memref<?xf32> {
Expand All @@ -63,7 +63,7 @@ func.func @fancy_layout_to_no_layout_cast(%m: memref<?xf32, strided<[100], offse
// CHECK-SAME: %[[arg:.*]]: memref<?xf32, strided<[1], offset: 25>>)
// CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK: %[[dim:.*]] = memref.dim %[[arg]], %[[c0]]
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?xf32>
// CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]]) {{.*}} : memref<?xf32>
// CHECK: memref.copy %[[arg]], %[[alloc]]
// CHECK: return %[[alloc]]
func.func @static_layout_to_no_layout_cast(%m: memref<?xf32, strided<[1], offset: 25>>) -> memref<?xf32> {
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-out-params.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func.func @main(%t: tensor<5xf32>) -> (f32, f32) {
// Note: This alloc is not needed, but it is inserted before the returned buffer
// is promoted to an out param to reconcile mismatching layout maps on return
// value and function signature.
// CHECK-NO-LAYOUT: %[[alloc2:.*]] = memref.alloc() : memref<2x5xf32>
// CHECK-NO-LAYOUT: %[[alloc2:.*]] = memref.alloc() {{.*}} : memref<2x5xf32>
// CHECK-NO-LAYOUT: memref.copy %[[subview]], %[[alloc2]]
// CHECK-NO-LAYOUT: memref.copy %[[alloc2]], %[[r]]

Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ func.func private @external_func_with_return_val(tensor<4xi32>) -> f32
// CHECK-NO-LAYOUT-MAP-LABEL: func @return_extract_slice(%{{.*}}) -> memref<2x?xf32>
// CHECK-NO-LAYOUT-MAP: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<20x10xf32>
// CHECK-NO-LAYOUT-MAP: %[[subview:.*]] = memref.subview {{.*}} : memref<20x10xf32> to memref<2x?xf32, strided<[10, 1], offset: ?>>
// CHECK-NO-LAYOUT-MAP: %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) : memref<2x?xf32>
// CHECK-NO-LAYOUT-MAP: %[[alloc_no_layout:.*]] = memref.alloc(%{{.*}}) {{.*}} : memref<2x?xf32>
// CHECK-NO-LAYOUT-MAP: memref.copy %[[subview]], %[[alloc_no_layout]]
// TODO: %alloc should be deallocated here, but we currently do not dealloc
// buffers that are inserted due to to_tensor/to_memref canonicalization (when
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Bufferization/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func.func @canonicalize_buffer_cast_of_tensor_load_to_copy(
// CHECK-NOT: bufferization.to_memref
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[DIM:.*]] = memref.dim %[[M]], %[[C0]] : memref<?xf32, strided<[1], offset: ?>>
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) : memref<?xf32, strided<[1], offset: 3>>
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) {{.*}} : memref<?xf32, strided<[1], offset: 3>>
// CHECK: memref.copy %[[M]], %[[ALLOC]]
// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: 3>>
// CHECK: return %[[ALLOC]]
Expand Down
Loading