[Mlir-commits] [mlir] 1892666 - [MLIR][SCF] Loop pipelining fails on failed predication (no assert) (#107442)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 5 11:46:21 PDT 2024
Author: SJW
Date: 2024-09-05T11:46:18-07:00
New Revision: 18926666f509104c3f478444b282291ce19fab6a
URL: https://github.com/llvm/llvm-project/commit/18926666f509104c3f478444b282291ce19fab6a
DIFF: https://github.com/llvm/llvm-project/commit/18926666f509104c3f478444b282291ce19fab6a.diff
LOG: [MLIR][SCF] Loop pipelining fails on failed predication (no assert) (#107442)
The SCFLoopPipelining allows predication on peeled or loop ops. When the
predicationFn returns a nullptr this signifies the op type is
unsupported and the pipeliner fails except in `emitPrologue` where it
asserts.
This patch fixes handling in the prologue to gracefully fail.
Added:
Modified:
mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index a34542f0161aca..7cecd4942b640f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -77,7 +77,7 @@ struct LoopPipelinerInternal {
bool initializeLoopInfo(ForOp op, const PipeliningOption &options);
/// Emits the prologue, this creates `maxStage - 1` part which will contain
/// operations from stages [0; i], where i is the part index.
- void emitPrologue(RewriterBase &rewriter);
+ LogicalResult emitPrologue(RewriterBase &rewriter);
/// Gather liverange information for Values that are used in a
diff erent stage
/// than its definition.
llvm::MapVector<Value, LiverangeInfo> analyzeCrossStageValues();
@@ -263,7 +263,7 @@ cloneAndUpdateOperands(RewriterBase &rewriter, Operation *op,
return clone;
}
-void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
+LogicalResult LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
// Initialize the iteration argument to the loop initial values.
for (auto [arg, operand] :
llvm::zip(forOp.getRegionIterArgs(), forOp.getInitsMutable())) {
@@ -311,7 +311,8 @@ void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
if (predicates[predicateIdx]) {
OpBuilder::InsertionGuard insertGuard(rewriter);
newOp = predicateFn(rewriter, newOp, predicates[predicateIdx]);
- assert(newOp && "failed to predicate op.");
+ if (newOp == nullptr)
+ return failure();
}
if (annotateFn)
annotateFn(newOp, PipeliningOption::PipelinerPart::Prologue, i);
@@ -339,6 +340,7 @@ void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
}
}
}
+ return success();
}
llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
@@ -772,7 +774,8 @@ FailureOr<ForOp> mlir::scf::pipelineForLoop(RewriterBase &rewriter, ForOp forOp,
*modifiedIR = true;
// 1. Emit prologue.
- pipeliner.emitPrologue(rewriter);
+ if (failed(pipeliner.emitPrologue(rewriter)))
+ return failure();
// 2. Track values used across stages. When a value cross stages it will
// need to be passed as loop iteration arguments.
More information about the Mlir-commits
mailing list