This is an archive of the discontinued LLVM Phabricator instance.

Differential D75513

[mlir] support reductions in loop to std conversion
ClosedPublic

Authored by ftynse on Mar 3 2020, 4:22 AM.

Details

Reviewers
nicolasvasilache
Commits
rG8ba8ab8c95fb: [mlir] support reductions in loop to std conversion
Summary

Introduce support for converting loop.for operations with loop-carried values
to a CFG in the standard dialect. This is achieved by passing loop-carried
values as block arguments to the loop condition block. This block dominates
both the loop body and the block immediately following the loop, so the
arguments of this block are remain visible there.

Diff Detail

Event Timeline

ftynse created this revision.Mar 3 2020, 4:22 AM
Herald added a project: Restricted Project. · View Herald TranscriptMar 3 2020, 4:22 AM
Herald added subscribers: llvm-commits, Joonsoo, liufengdb and 9 others. · View Herald Transcript
ftynse added a subscriber: poechsel.Mar 3 2020, 4:23 AM
Harbormaster completed remote builds in B47895: Diff 247855.Mar 3 2020, 5:45 AM
nicolasvasilache accepted this revision.Mar 3 2020, 8:48 AM

Great, thanks Alex!

This revision is now accepted and ready to land.Mar 3 2020, 8:48 AM
Closed by commit rG8ba8ab8c95fb: [mlir] support reductions in loop to std conversion (authored by ftynse). · Explain WhyMar 3 2020, 9:41 AM
This revision was automatically updated to reflect the committed changes.

Revision Contents

PathSize
mlir/
lib/
Conversion/
LoopToStandard/
59 lines
test/
Conversion/
56 lines

Diff 247940

mlir/lib/Conversion/LoopToStandard/ConvertLoopToStandard.cpp

Show All 39 Lines
// entry and a single exit, and that the entry/exit blocks are respectively // entry and a single exit, and that the entry/exit blocks are respectively
// first/last blocks in the parent region. The original loop operation is // first/last blocks in the parent region. The original loop operation is
// replaced by the initialization operations that set up the initial value of // replaced by the initialization operations that set up the initial value of
// the loop induction variable (%iv) and computes the loop bounds that are loop- // the loop induction variable (%iv) and computes the loop bounds that are loop-
// invariant for affine loops. The operations following the original loop.for // invariant for affine loops. The operations following the original loop.for
// are split out into a separate continuation (exit) block. A condition block is // are split out into a separate continuation (exit) block. A condition block is
// created before the continuation block. It checks the exit condition of the // created before the continuation block. It checks the exit condition of the
// loop and branches either to the continuation block, or to the first block of // loop and branches either to the continuation block, or to the first block of
// the body. Induction variable modification is appended to the last block of // the body. The condition block takes as arguments the values of the induction
// the body (which is the exit block from the body subgraph thanks to the // variable followed by loop-carried values. Since it dominates both the body
// blocks and the continuation block, loop-carried values are visible in all of
// those blocks. Induction variable modification is appended to the last block
// of the body (which is the exit block from the body subgraph thanks to the
// invariant we maintain) along with a branch that loops back to the condition // invariant we maintain) along with a branch that loops back to the condition
// block. // block. Loop-carried values are the loop terminator operands, which are
// forwarded to the branch.
// //
// +---------------------------------+ // +---------------------------------+
// | <code before the ForOp> | // | <code before the ForOp> |
// | <definitions of %init...> |
// | <compute initial %iv value> | // | <compute initial %iv value> |
// | br cond(%iv) | // | br cond(%iv, %init...) |
// +---------------------------------+ // +---------------------------------+
// | // |
// -------| | // -------| |
// | v v // | v v
// | +--------------------------------+ // | +--------------------------------+
// | | cond(%iv): | // | | cond(%iv, %init...): |
// | | <compare %iv to upper bound> | // | | <compare %iv to upper bound> |
// | | cond_br %r, body, end | // | | cond_br %r, body, end |
// | +--------------------------------+ // | +--------------------------------+
// | | | // | | |
// | | -------------| // | | -------------|
// | v | // | v |
// | +--------------------------------+ | // | +--------------------------------+ |
// | | body-first: | | // | | body-first: | |
// | | <%init visible by dominance> | |
// | | <body contents> | | // | | <body contents> | |
// | +--------------------------------+ | // | +--------------------------------+ |
// | | | // | | |
// | ... | // | ... |
// | | | // | | |
// | +--------------------------------+ | // | +--------------------------------+ |
// | | body-last: | | // | | body-last: | |
// | | <body contents> | | // | | <body contents> | |
// | | <operands of yield = %yields>| |
// | | %new_iv =<add step to %iv> | | // | | %new_iv =<add step to %iv> | |
// | | br cond(%new_iv) | | // | | br cond(%new_iv, %yields) | |
// | +--------------------------------+ | // | +--------------------------------+ |
// | | | // | | |
// |----------- |-------------------- // |----------- |--------------------
// v // v
// +--------------------------------+ // +--------------------------------+
// | end: | // | end: |
// | <code after the ForOp> | // | <code after the ForOp> |
// | <%init visible by dominance> |
// +--------------------------------+ // +--------------------------------+
// //
struct ForLowering : public OpRewritePattern<ForOp> { struct ForLowering : public OpRewritePattern<ForOp> {
using OpRewritePattern<ForOp>::OpRewritePattern; using OpRewritePattern<ForOp>::OpRewritePattern;
PatternMatchResult matchAndRewrite(ForOp forOp, PatternMatchResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override; PatternRewriter &rewriter) const override;
}; };
Show All 32 Lines
// | | <else contents> | // | | <else contents> |
// | | br continue | // | | br continue |
// | +--------------------------------+ // | +--------------------------------+
// | | // | |
// ------| | // ------| |
// v v // v v
// +--------------------------------+ // +--------------------------------+
// | continue: | // | continue: |
// | <code after the IfOp> | // | <code after the IfOp> |
// +--------------------------------+ // +--------------------------------+
// //
struct IfLowering : public OpRewritePattern<IfOp> { struct IfLowering : public OpRewritePattern<IfOp> {
using OpRewritePattern<IfOp>::OpRewritePattern; using OpRewritePattern<IfOp>::OpRewritePattern;
PatternMatchResult matchAndRewrite(IfOp ifOp, PatternMatchResult matchAndRewrite(IfOp ifOp,
PatternRewriter &rewriter) const override; PatternRewriter &rewriter) const override;
}; };
Show All 12 LinesForLowering::matchAndRewrite(ForOp forOp, PatternRewriter &rewriter) const {
// Start by splitting the block containing the 'loop.for' into two parts. // Start by splitting the block containing the 'loop.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
// point. // point.
auto *initBlock = rewriter.getInsertionBlock(); auto *initBlock = rewriter.getInsertionBlock();
auto initPosition = rewriter.getInsertionPoint(); auto initPosition = rewriter.getInsertionPoint();
auto *endBlock = rewriter.splitBlock(initBlock, initPosition); auto *endBlock = rewriter.splitBlock(initBlock, initPosition);
// Use the first block of the loop body as the condition block since it is // Use the first block of the loop body as the condition block since it is the
// the block that has the induction variable as its argument. Split out // block that has the induction variable and loop-carried values as arguments.
// all operations from the first block into a new block. Move all body // Split out all operations from the first block into a new block. Move all
// blocks from the loop body region to the region containing the loop. // body blocks from the loop body region to the region containing the loop.
auto *conditionBlock = &forOp.region().front(); auto *conditionBlock = &forOp.region().front();
auto *firstBodyBlock = auto *firstBodyBlock =
rewriter.splitBlock(conditionBlock, conditionBlock->begin()); rewriter.splitBlock(conditionBlock, conditionBlock->begin());
auto *lastBodyBlock = &forOp.region().back(); auto *lastBodyBlock = &forOp.region().back();
rewriter.inlineRegionBefore(forOp.region(), endBlock); rewriter.inlineRegionBefore(forOp.region(), endBlock);
auto iv = conditionBlock->getArgument(0); auto iv = conditionBlock->getArgument(0);
// Append the induction variable stepping logic to the last body block and // Append the induction variable stepping logic to the last body block and
// branch back to the condition block. Construct an expression f : // branch back to the condition block. Loop-carried values are taken from
// (x -> x+step) and apply this expression to the induction variable. // operands of the loop terminator.
rewriter.eraseOp(lastBodyBlock->getTerminator()); Operation *terminator = lastBodyBlock->getTerminator();
rewriter.setInsertionPointToEnd(lastBodyBlock); rewriter.setInsertionPointToEnd(lastBodyBlock);
auto step = forOp.step(); auto step = forOp.step();
auto stepped = rewriter.create<AddIOp>(loc, iv, step).getResult(); auto stepped = rewriter.create<AddIOp>(loc, iv, step).getResult();
if (!stepped) if (!stepped)
return matchFailure(); return matchFailure();
rewriter.create<BranchOp>(loc, conditionBlock, stepped);
SmallVector<Value, 8> loopCarried;
loopCarried.push_back(stepped);
loopCarried.append(terminator->operand_begin(), terminator->operand_end());
rewriter.create<BranchOp>(loc, conditionBlock, loopCarried);
rewriter.eraseOp(terminator);
// Compute loop bounds before branching to the condition. // Compute loop bounds before branching to the condition.
rewriter.setInsertionPointToEnd(initBlock); rewriter.setInsertionPointToEnd(initBlock);
Value lowerBound = forOp.lowerBound(); Value lowerBound = forOp.lowerBound();
Value upperBound = forOp.upperBound(); Value upperBound = forOp.upperBound();
if (!lowerBound || !upperBound) if (!lowerBound || !upperBound)
return matchFailure(); return matchFailure();
rewriter.create<BranchOp>(loc, conditionBlock, lowerBound);
// The initial values of loop-carried values is obtained from the operands
// of the loop operation.
SmallVector<Value, 8> destOperands;
destOperands.push_back(lowerBound);
auto iterOperands = forOp.getIterOperands();
destOperands.append(iterOperands.begin(), iterOperands.end());
rewriter.create<BranchOp>(loc, conditionBlock, destOperands);
// With the body block done, we can fill in the condition block. // With the body block done, we can fill in the condition block.
rewriter.setInsertionPointToEnd(conditionBlock); rewriter.setInsertionPointToEnd(conditionBlock);
auto comparison = auto comparison =
rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, iv, upperBound); rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, iv, upperBound);
rewriter.create<CondBranchOp>(loc, comparison, firstBodyBlock, rewriter.create<CondBranchOp>(loc, comparison, firstBodyBlock,
ArrayRef<Value>(), endBlock, ArrayRef<Value>()); ArrayRef<Value>(), endBlock, ArrayRef<Value>());
// Ok, we're done! // The result of the loop operation is the values of the condition block
rewriter.eraseOp(forOp); // arguments except the induction variable on the last iteration.
rewriter.replaceOp(forOp, conditionBlock->getArguments().drop_front());
return matchSuccess(); return matchSuccess();
} }
PatternMatchResult PatternMatchResult
IfLowering::matchAndRewrite(IfOp ifOp, PatternRewriter &rewriter) const { IfLowering::matchAndRewrite(IfOp ifOp, PatternRewriter &rewriter) const {
auto loc = ifOp.getLoc(); auto loc = ifOp.getLoc();
// Start by splitting the block containing the 'loop.if' into two parts. // Start by splitting the block containing the 'loop.if' into two parts.
▲ Show 20 Lines • Show All 92 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 174 Lines • ▼ Show 20 Linesfunc @parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
%arg3 : index, %arg4 : index) { %arg3 : index, %arg4 : index) {
%step = constant 1 : index %step = constant 1 : index
loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3) loop.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
step (%arg4, %step) { step (%arg4, %step) {
%c1 = constant 1 : index %c1 = constant 1 : index
} }
return return
} }
// CHECK-LABEL: @for_yield
// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index)
// CHECK: %[[INIT0:.*]] = constant 0
// CHECK: %[[INIT1:.*]] = constant 1
// CHECK: br ^[[COND:.*]](%[[LB]], %[[INIT0]], %[[INIT1]] : index, f32, f32)
//
// CHECK: ^[[COND]](%[[ITER:.*]]: index, %[[ITER_ARG0:.*]]: f32, %[[ITER_ARG1:.*]]: f32):
// CHECK: %[[CMP:.*]] = cmpi "slt", %[[ITER]], %[[UB]] : index
// CHECK: cond_br %[[CMP]], ^[[BODY:.*]], ^[[CONTINUE:.*]]
//
// CHECK: ^[[BODY]]:
// CHECK: %[[SUM:.*]] = addf %[[ITER_ARG0]], %[[ITER_ARG1]] : f32
// CHECK: %[[STEPPED:.*]] = addi %[[ITER]], %[[STEP]] : index
// CHECK: br ^[[COND]](%[[STEPPED]], %[[SUM]], %[[SUM]] : index, f32, f32)
//
// CHECK: ^[[CONTINUE]]:
// CHECK: return %[[ITER_ARG0]], %[[ITER_ARG1]] : f32, f32
func @for_yield(%arg0 : index, %arg1 : index, %arg2 : index) -> (f32, f32) {
%s0 = constant 0.0 : f32
%s1 = constant 1.0 : f32
%result:2 = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%si = %s0, %sj = %s1) -> (f32, f32) {
%sn = addf %si, %sj : f32
loop.yield %sn, %sn : f32, f32
}
return %result#0, %result#1 : f32, f32
}
// CHECK-LABEL: @nested_for_yield
// CHECK-SAME: (%[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index)
// CHECK: %[[INIT:.*]] = constant
// CHECK: br ^[[COND_OUT:.*]](%[[LB]], %[[INIT]] : index, f32)
// CHECK: ^[[COND_OUT]](%[[ITER_OUT:.*]]: index, %[[ARG_OUT:.*]]: f32):
// CHECK: cond_br %{{.*}}, ^[[BODY_OUT:.*]], ^[[CONT_OUT:.*]]
// CHECK: ^[[BODY_OUT]]:
// CHECK: br ^[[COND_IN:.*]](%[[LB]], %[[ARG_OUT]] : index, f32)
// CHECK: ^[[COND_IN]](%[[ITER_IN:.*]]: index, %[[ARG_IN:.*]]: f32):
// CHECK: cond_br %{{.*}}, ^[[BODY_IN:.*]], ^[[CONT_IN:.*]]
// CHECK: ^[[BODY_IN]]
// CHECK: %[[RES:.*]] = addf
// CHECK: br ^[[COND_IN]](%{{.*}}, %[[RES]] : index, f32)
// CHECK: ^[[CONT_IN]]:
// CHECK: br ^[[COND_OUT]](%{{.*}}, %[[ARG_IN]] : index, f32)
// CHECK: ^[[CONT_OUT]]:
// CHECK: return %[[ARG_OUT]] : f32
func @nested_for_yield(%arg0 : index, %arg1 : index, %arg2 : index) -> f32 {
%s0 = constant 1.0 : f32
%r = loop.for %i0 = %arg0 to %arg1 step %arg2 iter_args(%iter = %s0) -> (f32) {
%result = loop.for %i1 = %arg0 to %arg1 step %arg2 iter_args(%si = %iter) -> (f32) {
%sn = addf %si, %si : f32
loop.yield %sn : f32
}
loop.yield %result : f32
}
return %r : f32
}