This is an archive of the discontinued LLVM Phabricator instance.

Differential D90603

[mlir] Add lowering to CFG for WhileOp
ClosedPublic

Authored by ftynse on Nov 2 2020, 6:00 AM.

Details

Reviewers
aartbik
nicolasvasilache
mehdi_amini
Commits
rG4c0e255c98cc: [mlir] Add lowering to CFG for WhileOp
Summary

The lowering is a straightforward inlining of the "before" and "after" regions
connected by (conditional) branches. This plugs the WhileOp into the
progressive lowering scheme. Future commits may choose to target WhileOp
instead of CFG when lowering ForOp.

Depends On D90255

Diff Detail

Event Timeline

ftynse created this revision.Nov 2 2020, 6:00 AM
Herald added a project: Restricted Project. · View Herald TranscriptNov 2 2020, 6:00 AM
Herald added subscribers: rdzhabarov, tatianashp, msifontes and 12 others. · View Herald Transcript
ftynse requested review of this revision.Nov 2 2020, 6:00 AM
Herald added a subscriber: stephenneuendorffer. · View Herald TranscriptNov 2 2020, 6:00 AM
ftynse added a child revision: D90604: [mlir] Add a simpler lowering pattern for WhileOp representing a do-while loop.Nov 2 2020, 6:00 AM
Harbormaster completed remote builds in B77245: Diff 302258.Nov 2 2020, 6:20 AM
aartbik added inline comments.Nov 2 2020, 5:43 PM
mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
257

what do you mean by "ex-regions" here?

mlir/test/Conversion/SCFToStandard/convert-to-cfg.mlir
449

can you break this in two lines?

474

same, and below, reads a bit easier

ftynse updated this revision to Diff 302508.Nov 3 2020, 1:34 AM
ftynse marked 2 inline comments as done.

Address review

mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
257

The groups of blocks that used to form a region before inlining.

Harbormaster completed remote builds in B77376: Diff 302508.Nov 3 2020, 1:55 AM
aartbik accepted this revision.Nov 3 2020, 10:25 AM
aartbik added inline comments.
mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
257

I actually also meant to spell that out a bit more clearly in the actual comments :-)

This revision is now accepted and ready to land.Nov 3 2020, 10:25 AM
mehdi_amini accepted this revision.Nov 3 2020, 12:10 PM
mehdi_amini added inline comments.
mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
508

That's amazingly simple :)

ftynse updated this revision to Diff 302763.Nov 4 2020, 12:42 AM
ftynse marked an inline comment as done.

Address review

mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
508

The rewriter infrastructure investment starts to pay off :)

This revision was landed with ongoing or failed builds.Nov 4 2020, 12:44 AM
Closed by commit rG4c0e255c98cc: [mlir] Add lowering to CFG for WhileOp (authored by ftynse). · Explain Why
This revision was automatically updated to reflect the committed changes.
ftynse added a commit: rG4c0e255c98cc: [mlir] Add lowering to CFG for WhileOp.
Harbormaster completed remote builds in B77513: Diff 302763.Nov 4 2020, 1:06 AM

Revision Contents

PathSize
mlir/
lib/
Conversion/
SCFToStandard/
117 lines
test/
Conversion/
SCFToStandard/
113 lines

Diff 302765

mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp

Show First 20 Lines • Show All 194 Lines • ▼ Show 20 Lines
}; };
struct ParallelLowering : public OpRewritePattern<mlir::scf::ParallelOp> { struct ParallelLowering : public OpRewritePattern<mlir::scf::ParallelOp> {
using OpRewritePattern<mlir::scf::ParallelOp>::OpRewritePattern; using OpRewritePattern<mlir::scf::ParallelOp>::OpRewritePattern;
LogicalResult matchAndRewrite(mlir::scf::ParallelOp parallelOp, LogicalResult matchAndRewrite(mlir::scf::ParallelOp parallelOp,
PatternRewriter &rewriter) const override; PatternRewriter &rewriter) const override;
}; };
/// Create a CFG subgraph for this loop construct. The regions of the loop need
/// not be a single block anymore (for example, if other SCF constructs that
/// they contain have been already converted to CFG), but need to be single-exit
/// from the last block of each region. The operations following the original
/// WhileOp are split into a new continuation block. Both regions of the WhileOp
/// are inlined, and their terminators are rewritten to organize the control
/// flow implementing the loop as follows.
///
/// +---------------------------------+
/// | <code before the WhileOp> |
/// | br ^before(%operands...) |
/// +---------------------------------+
/// |
/// -------| |
/// | v v
/// | +--------------------------------+
/// | | ^before(%bargs...): |
/// | | %vals... = <some payload> |
/// | +--------------------------------+
/// | |
/// | ...
/// | |
/// | +--------------------------------+
/// | | ^before-last:
/// | | %cond = <compute condition> |
/// | | cond_br %cond, |
/// | | ^after(%vals...), ^cont |
/// | +--------------------------------+
/// | | |
/// | | -------------|
/// | v |
/// | +--------------------------------+ |
/// | | ^after(%aargs...): | |
/// | | <body contents> | |
/// | +--------------------------------+ |
/// | | |
/// | ... |
/// | | |
/// | +--------------------------------+ |
/// | | ^after-last: | |
/// | | %yields... = <some payload> | |
/// | | br ^before(%yields...) | |
/// | +--------------------------------+ |
/// | | |
/// |----------- |--------------------
/// v
/// +--------------------------------+
/// | ^cont: |
/// | <code after the WhileOp> |
/// | <%vals from 'before' region |
/// | visible by dominance> |
/// +--------------------------------+
///
/// Values are communicated between ex-regions (the groups of blocks that used
aartbikUnsubmitted
Not Done
ReplyInline Actions

what do you mean by "ex-regions" here?

aartbik: what do you mean by "ex-regions" here?
ftynseAuthorUnsubmitted
Done
ReplyInline Actions

The groups of blocks that used to form a region before inlining.

ftynse: The groups of blocks that used to form a region before inlining.
aartbikUnsubmitted
Done
ReplyInline Actions

I actually also meant to spell that out a bit more clearly in the actual comments :-)

aartbik: I actually also meant to spell that out a bit more clearly in the actual comments :-)
/// to form a region before inlining) through block arguments of their
/// entry blocks, which are visible in all other dominated blocks. Similarly,
/// the results of the WhileOp are defined in the 'before' region, which is
/// required to have a single existing block, and are therefore accessible in
/// the continuation block due to dominance.
struct WhileLowering : public OpRewritePattern<WhileOp> {
using OpRewritePattern<WhileOp>::OpRewritePattern;
LogicalResult matchAndRewrite(WhileOp whileOp,
PatternRewriter &rewriter) const override;
};
} // namespace } // namespace
LogicalResult ForLowering::matchAndRewrite(ForOp forOp, LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const { PatternRewriter &rewriter) const {
Location loc = forOp.getLoc(); Location loc = forOp.getLoc();
// Start by splitting the block containing the 'scf.for' into two parts. // Start by splitting the block containing the 'scf.for' into two parts.
// The part before will get the init code, the part after will be the end // The part before will get the init code, the part after will be the end
▲ Show 20 Lines • Show All 183 Lines • ▼ Show 20 Lines if (!yieldOperands.empty()) {
rewriter.create<scf::YieldOp>(loc, yieldOperands); rewriter.create<scf::YieldOp>(loc, yieldOperands);
} }
rewriter.replaceOp(parallelOp, loopResults); rewriter.replaceOp(parallelOp, loopResults);
return success(); return success();
} }
LogicalResult WhileLowering::matchAndRewrite(WhileOp whileOp,
PatternRewriter &rewriter) const {
OpBuilder::InsertionGuard guard(rewriter);
Location loc = whileOp.getLoc();
// Split the current block before the WhileOp to create the inlining point.
Block *currentBlock = rewriter.getInsertionBlock();
Block *continuation =
rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
// Inline both regions.
Block *after = &whileOp.after().front();
Block *afterLast = &whileOp.after().back();
Block *before = &whileOp.before().front();
Block *beforeLast = &whileOp.before().back();
rewriter.inlineRegionBefore(whileOp.after(), continuation);
rewriter.inlineRegionBefore(whileOp.before(), after);
// Branch to the "before" region.
rewriter.setInsertionPointToEnd(currentBlock);
rewriter.create<BranchOp>(loc, before, whileOp.inits());
// Replace terminators with branches. Assuming bodies are SESE, which holds
// given only the patterns from this file, we only need to look at the last
// block. This should be reconsidered if we allow break/continue in SCF.
rewriter.setInsertionPointToEnd(beforeLast);
auto condOp = cast<ConditionOp>(beforeLast->getTerminator());
rewriter.replaceOpWithNewOp<CondBranchOp>(condOp, condOp.condition(), after,
condOp.args(), continuation,
ValueRange());
rewriter.setInsertionPointToEnd(afterLast);
auto yieldOp = cast<scf::YieldOp>(afterLast->getTerminator());
rewriter.replaceOpWithNewOp<BranchOp>(yieldOp, before, yieldOp.results());
// Replace the op with values "yielded" from the "before" region, which are
// visible by dominance.
rewriter.replaceOp(whileOp, condOp.args());
return success();
}
mehdi_aminiUnsubmitted
Not Done
ReplyInline Actions

That's amazingly simple :)

mehdi_amini: That's amazingly simple :)
ftynseAuthorUnsubmitted
Done
ReplyInline Actions

The rewriter infrastructure investment starts to pay off :)

ftynse: The rewriter infrastructure investment starts to pay off :)
void mlir::populateLoopToStdConversionPatterns( void mlir::populateLoopToStdConversionPatterns(
OwningRewritePatternList &patterns, MLIRContext *ctx) { OwningRewritePatternList &patterns, MLIRContext *ctx) {
patterns.insert<ForLowering, IfLowering, ParallelLowering>(ctx); patterns.insert<ForLowering, IfLowering, ParallelLowering, WhileLowering>(
ctx);
} }
void SCFToStandardPass::runOnOperation() { void SCFToStandardPass::runOnOperation() {
OwningRewritePatternList patterns; OwningRewritePatternList patterns;
populateLoopToStdConversionPatterns(patterns, &getContext()); populateLoopToStdConversionPatterns(patterns, &getContext());
// Configure conversion to lower out scf.for, scf.if and scf.parallel. // Configure conversion to lower out scf.for, scf.if, scf.parallel and
// Anything else is fine. // scf.while. Anything else is fine.
ConversionTarget target(getContext()); ConversionTarget target(getContext());
target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp>(); target.addIllegalOp<scf::ForOp, scf::IfOp, scf::ParallelOp, scf::WhileOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
if (failed( if (failed(
applyPartialConversion(getOperation(), target, std::move(patterns)))) applyPartialConversion(getOperation(), target, std::move(patterns))))
signalPassFailure(); signalPassFailure();
} }
std::unique_ptr<Pass> mlir::createLowerToCFGPass() { std::unique_ptr<Pass> mlir::createLowerToCFGPass() {
return std::make_unique<SCFToStandardPass>(); return std::make_unique<SCFToStandardPass>();
} }

mlir/test/Conversion/SCFToStandard/convert-to-cfg.mlir

Show First 20 Lines • Show All 406 Lines • ▼ Show 20 Linesfunc @unknown_op_inside_loop(%arg0: index, %arg1: index, %arg2: index) {
// CHECK-NOT: scf.for // CHECK-NOT: scf.for
scf.for %i = %arg0 to %arg1 step %arg2 { scf.for %i = %arg0 to %arg1 step %arg2 {
// CHECK: unknown.op // CHECK: unknown.op
"unknown.op"() : () -> () "unknown.op"() : () -> ()
scf.yield scf.yield
} }
return return
} }
// CHECK-LABEL: @minimal_while
func @minimal_while() {
// CHECK: %[[COND:.*]] = "test.make_condition"() : () -> i1
// CHECK: br ^[[BEFORE:.*]]
%0 = "test.make_condition"() : () -> i1
scf.while : () -> () {
// CHECK: ^[[BEFORE]]:
// CHECK: cond_br %[[COND]], ^[[AFTER:.*]], ^[[CONT:.*]]
scf.condition(%0)
} do {
// CHECK: ^[[AFTER]]:
// CHECK: br ^[[BEFORE]]
scf.yield
}
// CHECK: ^[[CONT]]:
// CHECK: return
return
}
// CHECK-LABEL: @while_values
// CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: f32)
func @while_values(%arg0: i32, %arg1: f32) {
// CHECK: %[[COND:.*]] = "test.make_condition"() : () -> i1
%0 = "test.make_condition"() : () -> i1
%c0_i32 = constant 0 : i32
%cst = constant 0.000000e+00 : f32
// CHECK: br ^[[BEFORE:.*]](%[[ARG0]], %[[ARG1]] : i32, f32)
%1:2 = scf.while (%arg2 = %arg0, %arg3 = %arg1) : (i32, f32) -> (i64, f64) {
// CHECK: ^bb1(%[[ARG2:.*]]: i32, %[[ARG3:.]]: f32):
// CHECK: %[[VAL1:.*]] = zexti %[[ARG0]] : i32 to i64
%2 = zexti %arg0 : i32 to i64
// CHECK: %[[VAL2:.*]] = fpext %[[ARG3]] : f32 to f64
%3 = fpext %arg3 : f32 to f64
// CHECK: cond_br %[[COND]],
aartbikUnsubmitted
Done
ReplyInline Actions

can you break this in two lines?

aartbik: can you break this in two lines?
// CHECK: ^[[AFTER:.*]](%[[VAL1]], %[[VAL2]] : i64, f64),
// CHECK: ^[[CONT:.*]]
scf.condition(%0) %2, %3 : i64, f64
} do {
// CHECK: ^[[AFTER]](%[[ARG4:.*]]: i64, %[[ARG5:.*]]: f64):
^bb0(%arg2: i64, %arg3: f64): // no predecessors
// CHECK: br ^[[BEFORE]](%{{.*}}, %{{.*}} : i32, f32)
scf.yield %c0_i32, %cst : i32, f32
}
// CHECK: ^bb3:
// CHECK: return
return
}
// CHECK-LABEL: @nested_while_ops
func @nested_while_ops(%arg0: f32) -> i64 {
// CHECK: br ^[[OUTER_BEFORE:.*]](%{{.*}} : f32)
%0 = scf.while(%outer = %arg0) : (f32) -> i64 {
// CHECK: ^[[OUTER_BEFORE]](%{{.*}}: f32):
// CHECK: %[[OUTER_COND:.*]] = "test.outer_before_pre"() : () -> i1
%cond = "test.outer_before_pre"() : () -> i1
// CHECK: br ^[[INNER_BEFORE_BEFORE:.*]](%{{.*}} : f32)
%1 = scf.while(%inner = %outer) : (f32) -> i64 {
// CHECK: ^[[INNER_BEFORE_BEFORE]](%{{.*}}: f32):
// CHECK: %[[INNER1:.*]]:2 = "test.inner_before"(%{{.*}}) : (f32) -> (i1, i64)
aartbikUnsubmitted
Done
ReplyInline Actions

same, and below, reads a bit easier

aartbik: same, and below, reads a bit easier
%2:2 = "test.inner_before"(%inner) : (f32) -> (i1, i64)
// CHECK: cond_br %[[INNER1]]#0,
// CHECK: ^[[INNER_BEFORE_AFTER:.*]](%[[INNER1]]#1 : i64),
// CHECK: ^[[OUTER_BEFORE_LAST:.*]]
scf.condition(%2#0) %2#1 : i64
} do {
// CHECK: ^[[INNER_BEFORE_AFTER]](%{{.*}}: i64):
^bb0(%arg1: i64):
// CHECK: %[[INNER2:.*]] = "test.inner_after"(%{{.*}}) : (i64) -> f32
%3 = "test.inner_after"(%arg1) : (i64) -> f32
// CHECK: br ^[[INNER_BEFORE_BEFORE]](%[[INNER2]] : f32)
scf.yield %3 : f32
}
// CHECK: ^[[OUTER_BEFORE_LAST]]:
// CHECK: "test.outer_before_post"() : () -> ()
"test.outer_before_post"() : () -> ()
// CHECK: cond_br %[[OUTER_COND]],
// CHECK: ^[[OUTER_AFTER:.*]](%[[INNER1]]#1 : i64),
// CHECK: ^[[CONTINUATION:.*]]
scf.condition(%cond) %1 : i64
} do {
// CHECK: ^[[OUTER_AFTER]](%{{.*}}: i64):
^bb2(%arg2: i64):
// CHECK: "test.outer_after_pre"(%{{.*}}) : (i64) -> ()
"test.outer_after_pre"(%arg2) : (i64) -> ()
// CHECK: br ^[[INNER_AFTER_BEFORE:.*]](%{{.*}} : i64)
%4 = scf.while(%inner = %arg2) : (i64) -> f32 {
// CHECK: ^[[INNER_AFTER_BEFORE]](%{{.*}}: i64):
// CHECK: %[[INNER3:.*]]:2 = "test.inner2_before"(%{{.*}}) : (i64) -> (i1, f32)
%5:2 = "test.inner2_before"(%inner) : (i64) -> (i1, f32)
// CHECK: cond_br %[[INNER3]]#0,
// CHECK: ^[[INNER_AFTER_AFTER:.*]](%[[INNER3]]#1 : f32),
// CHECK: ^[[OUTER_AFTER_LAST:.*]]
scf.condition(%5#0) %5#1 : f32
} do {
// CHECK: ^[[INNER_AFTER_AFTER]](%{{.*}}: f32):
^bb3(%arg3: f32):
// CHECK: %{{.*}} = "test.inner2_after"(%{{.*}}) : (f32) -> i64
%6 = "test.inner2_after"(%arg3) : (f32) -> i64
// CHECK: br ^[[INNER_AFTER_BEFORE]](%{{.*}} : i64)
scf.yield %6 : i64
}
// CHECK: ^[[OUTER_AFTER_LAST]]:
// CHECK: "test.outer_after_post"() : () -> ()
"test.outer_after_post"() : () -> ()
// CHECK: br ^[[OUTER_BEFORE]](%[[INNER3]]#1 : f32)
scf.yield %4 : f32
}
// CHECK: ^[[CONTINUATION]]:
// CHECK: return %{{.*}} : i64
return %0 : i64
}