[Mlir-commits] [mlir] [MLIR] Add continuous tiling to TileUsingForOp (PR #82792)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 23 04:55:41 PDT 2024
================
@@ -2305,57 +2319,105 @@ SplitOp::apply(transform::TransformRewriter &rewriter,
rewriter.getIndexAttr(getStaticSplitPoint()));
}
- // Split each target operation.
- SmallVector<Operation *> first, second;
- Operation *noSecondPart = nullptr;
- for (const auto &pair : llvm::zip(payload, splitPoints)) {
- Operation *target = std::get<0>(pair);
- auto linalgOp = dyn_cast<LinalgOp>(target);
- if (!linalgOp) {
- auto diag = emitSilenceableError() << "only applies to structured ops";
- diag.attachNote(target->getLoc()) << "target op";
- return diag;
- }
+ if (isMultiwaySplit) {
- if (getDimension() >= linalgOp.getNumLoops()) {
- auto diag = emitSilenceableError() << "dimension " << getDimension()
- << " does not exist in target op";
- diag.attachNote(target->getLoc()) << "target op";
- return diag;
+ // Split a single target operation at multiple points.
+ SmallVector<Operation *> opList;
+ Operation *head, *tail;
+ for (const auto [idx, splitPoint] : llvm::enumerate(splitPoints)) {
+
+ Operation *target;
+ if (idx == 0)
+ target = payload.front();
+ else
+ target = tail;
+
+ if (!target)
+ break;
+
+ auto linalgOp = dyn_cast<LinalgOp>(target);
+
+ if (!linalgOp) {
+ auto diag = emitSilenceableError() << "only applies to structured ops";
+ diag.attachNote(target->getLoc()) << "target op";
+ return diag;
+ }
+
+ if (getDimension() >= linalgOp.getNumLoops()) {
+ auto diag = emitSilenceableError() << "dimension " << getDimension()
+ << " does not exist in target op";
+ diag.attachNote(target->getLoc()) << "target op";
+ return diag;
+ }
----------------
muneebkhan85 wrote:
I have refactored the code so that the common checks in the if/else are not duplicated and instead now rely on two lambda functions `checkStructuredOpAndDimensions` and `checkFailureInSplitting`. The code can not be entirely moved out of the loop as in the no-multiway case there can be several different payloads and each needs to be checked separately.
https://github.com/llvm/llvm-project/pull/82792
More information about the Mlir-commits
mailing list