forked from OSchip/llvm-project
[mlir][linalg] Let tile and fuse fail for tile sizes zero.
Adapt `tileConsumerAndFuseProducers` to return failure if the generated tile loop nest is empty since all tile sizes are zero. Additionally, fix `LinalgTileAndFuseTensorOpsPattern` to return success if the pattern applied successfully. Reviewed By: mravishankar Differential Revision: https://reviews.llvm.org/D118878
This commit is contained in:
parent
56d46b36fc
commit
b5ea288d13
|
@ -269,10 +269,10 @@ public:
|
|||
/// Returns the loop ops generated from tiling.
|
||||
ArrayRef<scf::ForOp> getLoopOps() { return tileLoopOps; }
|
||||
|
||||
private:
|
||||
/// Returns true if the tile loop nest has no tile loops.
|
||||
bool isEmpty();
|
||||
|
||||
private:
|
||||
/// Returns true if the tile loop nest invariants are satisfied:
|
||||
/// - The `rootOp` has been tiled at least once.
|
||||
/// - The number of tile loop operations and dimensions match.
|
||||
|
|
|
@ -458,5 +458,9 @@ mlir::linalg::tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp,
|
|||
return failure();
|
||||
fuseProducersGreedily(tileLoopNest.getRootOp().getInputOperands());
|
||||
|
||||
// Exit if the tile loop nest is empty since all tile sizes are zero.
|
||||
if (tileLoopNest.isEmpty())
|
||||
return failure();
|
||||
|
||||
return tileLoopNest;
|
||||
}
|
||||
|
|
|
@ -592,10 +592,6 @@ LogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
|
|||
SmallVector<int64_t> rootTileSizes(options.tileSizes.begin(),
|
||||
options.tileSizes.begin() +
|
||||
rootOp.getNumLoops());
|
||||
if (llvm::all_of(rootTileSizes, [](int64_t ts) { return ts == 0; })) {
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "all tile sizes are zero, nothing to do");
|
||||
}
|
||||
SmallVector<int64_t> rootInterchange =
|
||||
options.tileInterchange.empty()
|
||||
? llvm::to_vector<6>(llvm::seq<int64_t>(0, rootOp.getNumLoops()))
|
||||
|
@ -603,6 +599,11 @@ LogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
|
|||
options.tileInterchange.begin() +
|
||||
rootOp.getNumLoops());
|
||||
|
||||
// Check `rootTileSizes` contains non-zero tile sizes.
|
||||
if (llvm::count(rootTileSizes, 0) == static_cast<long>(rootTileSizes.size()))
|
||||
return rewriter.notifyMatchFailure(
|
||||
op, "expect at least one non-zero tile size");
|
||||
|
||||
// Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
|
||||
// It has to be a permutation since the tiling cannot tile the same loop
|
||||
// dimension multiple times.
|
||||
|
@ -623,7 +624,7 @@ LogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
|
|||
// Apply the filter if specified.
|
||||
for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps())
|
||||
filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
/// Linalg generic interchange pattern.
|
||||
|
|
|
@ -1,4 +1,4 @@
|
|||
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.generic fuse tile-sizes=0,0 run-enable-pass=false" -cse -split-input-file | FileCheck %s
|
||||
// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.matmul fuse tile-sizes=0,0,0 run-enable-pass=false" -split-input-file | FileCheck %s
|
||||
|
||||
builtin.func @no_fuse_gemm(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
|
||||
%c0 = arith.constant 0 : index
|
||||
|
|
Loading…
Reference in New Issue