[COMMON] Select the last compute op that has workgroup tilin… (#18604)
…g level
Remove the check "expected only one op with a workgroup tiling level".
diff --git a/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp b/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
index 03d2780..e722fcb 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp
@@ -53,18 +53,17 @@
static FailureOr<TilingInfo>
getTiledAndDistributionInfo(RewriterBase &rewriter,
ArrayRef<Operation *> computeOps) {
- // It is expected that at most one compute op has a workgroup tiling level.
+ // TODO: It is expected that at most one compute op has a workgroup tiling
+ // level. Currently, it selects the last compute op that has workgroup tiling
+ // level.
Operation *tilableOp = nullptr;
for (Operation *op : llvm::reverse(computeOps)) {
if (getLoweringConfig(op)) {
if (!getLoweringConfig(op).hasWorkgroupTilingLevel()) {
continue;
}
- if (tilableOp) {
- return op->emitOpError("expected only one op with a workgroup tiling"
- "level.");
- }
tilableOp = op;
+ break;
}
}
if (!tilableOp) {
