Hi Everyone,
As part of our work on nGraph dialect, we are trying to enable tensor-level control-flow. MLIR already has a loop dialect that we favor to re-use for our purposes instead of implementing our own control-flow dialect. However, in its current state, the loop dialect is not general enough. This RFC proposes a change that aim to alleviate that.
Thanks.
Summary
Currently, Loop dialect operations do not allow any temporary SSA values defined within their regions to escape to outer regions or to live across loop back-edges. While this works well for certain types of IR lowering (e.g Affine dialect to Loop dialect), it is a limiting factor to wider use of the dialect as a general tool to model structured control-flow .
This RFC proposes a change to the Loop dialect that allows IfOp to define values. Similar change to ForOp is proposed here
The RFC proposes the following changes:
- Introduce a new YieldOp to Loop dialect to return values out of the ops regions
- Modify IfOp IR format and representation to define and yield values
- Enable lowering logic of the new representations to standard dialect
Proposed YieldOp
In general, the proposed YieldOp passes a set of values from inside the IfOp regions to the outer enclosing region. In other words, it permits a set of value to escape the inner region. The exact semantics of the operation is dictated by the enclosing Loop dialect operation. This is in agreement with region specifications. See the proposed changes below for concrete use examples.
The proposed format for the operation is
def YieldOp : Loop_Op<”yield”> {
let arguments = (ins Variadic<AnyType>:$operands);
}
Proposed Changes to IfOp
Currently the IfOp format is as follows (details elided)
def IfOp : Loop_Op<"if", [SingleBlockImplicitTerminator<"TerminatorOp">]> {
let arguments = (ins I1:$condition);
let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
}
Example with pretty-print:
loop.if %b {
...
} else {
...
}
The then-region is mandatory and must have exactly 1 block, while the else-region can be empty. The operation does not define any values, and no values defined inside its region can escape. Operations inside the regions can access outer values, but not the other way around.
The new proposed format is
def IfOp : Loop_Op<"if", [SingleBlockImplicitTerminator<"TerminatorOp">]> {
let arguments = (ins I1:$condition);
let results = (outs Variadic<AnyType>);
let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
}
The result represents all merged values from the then- and else-regions.
Example 1 - Defining and merging two scoped values
func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
%x = loop.if %arg0 {
%x1 = addf %arg1, %arg2 : (f32, f32) -> f32
loop.yield %x1 : f32
} else {
%x2 = subf %arg1, %arg2 : (f32, f32) -> f32
loop.yield %x2 : f32
} : (i1) -> f32
return %x : () -> f32
}
In the context of an IfOp, the yield in the taken execution path assigns the value directly to the destination operands.
Notes about representation:
- IfOp can define zero or more operands. The number of YieldOp operands and their types must match the IfOp destination operands.
- The YieldOp is always at the end of the basic block (before the TerminatorOp)
- If the IfOp defines no operands, the region is not allowed to have a YieldOp
- If the IfOp defines any operands, the else block is mandatory since both execution paths must yield values. In case the else-block doesn’t change the value of a variable, it should still be present to re-assign a default value from the outer regions.
Example 2 - Merging outer default value with scoped value
func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
%x_default = subf %arg1, %arg2 : (f32, f32) -> f32
%x = loop.if %arg0 {
%x1 = addf %arg1, %arg2 : (f32, f32) -> f32
loop.yield %x1 : f32
} else {
// yield default value
loop.yield %x_default : f32
} : (i1) -> f32
return %x : () -> f32
}
Lowering to Standard Dialect
For Example 1 above, the lowered standard dialect IR will be as follows
func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
cond_br %arg0, %thenBlock, %elseBlock
^thenBlock:
%x1 = addf %arg1, %arg2
br ^continueBlock(%x1)
^elseBlock:
%x2 = subf %arg1, %arg2
br ^continueBlock(%x2)
^continueBlock(%x):
return %x
}
The YieldOp translates to merging of the two values at the continue block.
