Skip to content

[mlir][TilingInterface] Early return cloned ops if tile sizes are zeros. #75410

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

Merged
merged 4 commits into from
Dec 19, 2023
Merged

[mlir][TilingInterface] Early return cloned ops if tile sizes are zeros. #75410

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c07f7e1
[mlir][TilingInterface] Early return cloned ops if tile sizes are zeros.
hanhanW Dec 13, 2023
1d1a47f
review comments
hanhanW Dec 15, 2023
29b8434
Merge branch 'main' into tiling-early-return
hanhanW Dec 15, 2023
71d7c57
add comments about returning cloned op
hanhanW Dec 18, 2023
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
12 changes: 10 additions & 2 deletions mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,14 +362,22 @@ mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
auto clonedOp = cast<TilingInterface>(
cloneOpAndUpdateDestinationArgs(rewriter, op, clonedOpDestination));

// 5b. Tile the cloned operation.
// 5b. Early return cloned op if tiling is not happening. We can not return
// the original op because it could lead to
// `rewriter.replaceOp(op, op->getResults())` and user would get crash.
if (llvm::all_of(tileSizeVector, isZeroIndex)) {
return scf::SCFTilingResult{/*tiledOps=*/{clonedOp}, /*loops=*/{},
clonedOp->getResults()};
}

// 5c. Tile the cloned operation.
FailureOr<TilingResult> tiledImplementation =
clonedOp.getTiledImplementation(rewriter, offsets, sizes);
if (failed(tiledImplementation)) {
return rewriter.notifyMatchFailure(op, "failed to tile operation");
}

// 5c. Delete the cloned operation.
// 5d. Delete the cloned operation.
rewriter.eraseOp(clonedOp);

// If loops are empty, the tiled op is used as the replacement for the untiled
Expand Down
27 changes: 27 additions & 0 deletions mlir/test/Dialect/Linalg/tile-tensors.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,33 @@ module attributes {transform.with_named_sequence} {

// -----

// CHECK-LABEL: func @matmul_tensors_with_size_zeros(
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
func.func @matmul_tensors_with_size_zeros(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> {

// CHECK: %[[RES:.*]] = linalg.matmul ins(%[[TA]], %[[TB]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[TC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: return %[[RES]]
%0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2: tensor<?x?xf32>)
-> tensor<?x?xf32>
return %0 : tensor<?x?xf32>
}

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.structured.tile_using_for %0 [0, 0, 0] : (!transform.any_op) -> (!transform.any_op)
transform.yield
}
}

// -----

func.func @generic_op_tensors(
%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
%c0 = arith.constant 0 : index
Expand Down