Page MenuHomePhabricator

[mlir] Add a generic while/do-while loop to the SCF dialect
ClosedPublic

Authored by ftynse on Oct 27 2020, 11:20 AM.

Details

Summary

The new construct represents a generic loop with two regions: one executed
before the loop condition is verifier and another after that. This construct
can be used to express both a "while" loop and a "do-while" loop, depending on
where the main payload is located. It is intended as an intermediate
abstraction for lowering, which will be added later. This form is relatively
easy to target from higher-level abstractions and supports transformations such
as loop rotation and LICM.

Diff Detail

Event Timeline

ftynse created this revision.Oct 27 2020, 11:20 AM
ftynse requested review of this revision.Oct 27 2020, 11:20 AM
silvas added inline comments.Oct 27 2020, 11:42 AM
mlir/include/mlir/Dialect/SCF/SCFOps.td
498

Unclear what "match" means in this first sentence (match with each other? match with the outer op?). Maybe make it explicit "The two regions need not have the same argument list."

mlir/lib/Dialect/SCF/SCF.cpp
1029

A lot of this verification is subsumed by implementing RegionBranchOpInterface and calling RegionBranchOpInterface::verifyTypes (this would also reduce the testing burden, as we generally don't require testing properties verified by traits/interfaces)

Thanks for pushing on this! This is looking really good :)

Thanks for adding the while!

mlir/include/mlir/Dialect/SCF/SCFOps.td
438

.... represents a generic ...
reads better

469

Label this as "After" region

mlir/lib/Dialect/SCF/SCF.cpp
1032

You switch a bit between first/second and before/after region in your terminology. Perhaps use one name consistently throughout doc and messages?

mehdi_amini added inline comments.Oct 27 2020, 5:50 PM
mlir/include/mlir/Dialect/SCF/SCFOps.td
440–441

Random question: we tend to easily add this single-block restriction, but I'm not sure why is it the right thing to do?

rriddle added inline comments.Oct 28 2020, 12:29 AM
mlir/include/mlir/Dialect/SCF/SCFOps.td
435

Can you add the RegionControlFlowInterface to this operation?

448

nit: The two regions ?

mlir/lib/Dialect/SCF/SCF.cpp
976

parser.resolveOperands?

1005

nit: Add braces to the outer loop if the inner statement also has braces.

1043

nit: Add braces around this, it isn't really trivial.

nicolasvasilache accepted this revision.Oct 28 2020, 2:06 AM

Nice, thanks for implementing this ! :)

mlir/include/mlir/Dialect/SCF/SCFOps.td
466

You probably want to be very loud in this part of the doc to make the textual part more intuitive.

/* Compute values for "after" block*/
%foo#2 = call @compute(%arg1) : (f32) -> (i42, vector<3xf32>)

/* "Before" region. */
%condition = call @evaluate_condition(%arg1, %foo#1) : (f32, vector<3xf32>) -> i1

/* Forward the argument (as result or "after" region argument). */
scf.yield %condition, %foo#0, %foo#1, %arg1 : i1, i42, vector<3xf32>, f32
473

I'd add a comment: /* Must match enclosing scf.while operand types */ for in the same spirit of making the textual part more intuitive.

mlir/lib/Dialect/SCF/SCF.cpp
1029

+1

This revision is now accepted and ready to land.Oct 28 2020, 2:06 AM
ftynse updated this revision to Diff 301229.Oct 28 2020, 4:05 AM
ftynse marked 9 inline comments as done.

Partially address review

mlir/include/mlir/Dialect/SCF/SCFOps.td
440–441

I am considering to relax this, but this should be done with a proper RFC and apply to all SCF loops IMO. Allowing scf.yield in multiple blocks is equivalent to implementing a continue, and we may also want a break. Loop transformations currently rely on the single-block property.

466

Added another example. I want to keep the simple examples simple.

mlir/lib/Dialect/SCF/SCF.cpp
1029

Not as much actually as I would have liked. We cannot use it to verify the yield of the 'before' block because it takes the condition operand that it does not forward.

Looking at this closer, I am considering to introduce a different terminator for this region, something like scf.while_condition(%i1) %arg1, %arg2 : types to avoid breaking the "return-like" semantics of scf.yield. WDYT?

silvas added inline comments.Oct 28 2020, 12:57 PM
mlir/lib/Dialect/SCF/SCF.cpp
1029

Interesting. I thought this was supported, but perhaps not. @rriddle FYI

silvas accepted this revision.Oct 28 2020, 12:57 PM
mehdi_amini added inline comments.Oct 29 2020, 9:21 AM
mlir/lib/Dialect/SCF/SCF.cpp
1029

About the yield "return semantics": I see the terminator as always terminating the region and transferring the control to the parent op, yield fits here to me just like other loops.
The parent op then internally look at one of the result, and depending on the value forward the other results to the other region or to its control successor.

I was also expecting the region branch op interface to support this, but it may haven't been implemented yet. @jurahul worked on this at some point?

aartbik accepted this revision.Oct 29 2020, 10:08 AM

lowering next? ;-)

mlir/include/mlir/Dialect/SCF/SCFOps.td
438

The a is now at the wrong place :-)

This operation represents a generic ....

ftynse added inline comments.Oct 29 2020, 10:10 AM
mlir/lib/Dialect/SCF/SCF.cpp
1029

If I read the code correctly, there is an assumption that all operands of a ReturnLike terminator are forwarded somewhere.

I see the terminator as always terminating the region and transferring the control to the parent op,

This is what we wrote in the spec. But I interpret the RegionBranchOpInterface API as "after receiving the flow of control from region X, the op can forwarded it to region Y or to its natural control-flow successor (next op in the block)". In the former case, there is an expectation that types match exactly AFAICS.

ftynse updated this revision to Diff 301923.Fri, Oct 30, 9:13 AM

Added scf.condition and made WhileOp use RegionBranchOpInterface

ftynse marked 5 inline comments as done.Fri, Oct 30, 9:14 AM

PTAL

jurahul added inline comments.Fri, Oct 30, 9:25 AM
mlir/lib/Dialect/SCF/SCF.cpp
1029

Yes, I had spec'd an extension to RegionBranchOpInterface to support more extensive type checking, specifically around "internal" values that op's may produce and pass into attached regions. The main feature there was the interface declares, across all control flow edges, the types of values that flow across them, separately from the values. So its then becomes possible to typecheck say a region inputs when they feed not from the operation inputs but some internally generated values (which have no direct IR representation). And similarly typecheck region outputs that are consumed internally by the operation (and not directly feeding into the output).

However, I have not yet had time to actually follow that up with a discourse RFC and actual implementation.

mehdi_amini accepted this revision.Mon, Nov 2, 9:36 AM
mehdi_amini added inline comments.
mlir/lib/Dialect/SCF/SCF.cpp
1029

I don't understand how you quote from the RegionBranchOpInterface contradicts the quote you're answering from me? On the contrary it seems to say the same thing here?
In particular my comment was about the introduction of another terminator "to avoid breaking the "return-like" semantics of scf.yield", and I don't quite get the relationship with the types right now?

Can we have a TODO somewhere for the missing RegionBranchOpInterface feature?

ftynse added inline comments.Mon, Nov 2, 12:08 PM
mlir/lib/Dialect/SCF/SCF.cpp
1029

I don't think there is a contradiction on RegionBranchOpInterface.

However, there is a "ReturnLike" trait that RegionBranchOpInterface uses to understand that all operands are sent back to the parent op for dispatch. A potentially useful addition could be an interface method that specifies which operands instead of just taking all of them. That being said, "ReturnLike" seems to have some other semantics attached to it, and I am not willing to revise those semantics here or remove them from YieldOp for the purpose of this patch. Hence the new op.

mehdi_amini added inline comments.Mon, Nov 2, 12:55 PM
mlir/lib/Dialect/SCF/SCF.cpp
1029

A potentially useful addition could be an interface method that specifies which operands instead of just taking all of them. That being said, "ReturnLike" seems to have some other semantics attached to it, and I am not willing to revise those semantics here or remove them from YieldOp for the purpose of this patch. Hence the new op.

I still don't understand why yield can't be directly used here?

ftynse added inline comments.Tue, Nov 3, 1:06 AM
mlir/lib/Dialect/SCF/SCF.cpp
1029

Because it has the ReturnLike trait, which disallows to take the leading i1 operand without the next region also having the leading i1 argument. ReturnLike is also connected to other concepts, so I am reluctant to change it without a proper discussion.

This revision was landed with ongoing or failed builds.Wed, Nov 4, 12:44 AM
This revision was automatically updated to reflect the committed changes.