[MLIR] Free ReturnOp from being restricted to a FuncOp

[bondhugula]

Summary

• allow return op to be used from the top level of any op with a region (as opposed to just from the top level of a func op)
• the return operand count/types should match the op results' (whenever it's not in a FuncOp's region).

This allows std.return to be used both in a declarative (FuncOp) as well as an imperative (typical op) setting --- avoiding the need to create custom return ops wherever control is to be transferred to the enclosing operation (in line with the control flow semantics of regions) and values are to be potentially returned. This also enables the proposed execute_region and affine.graybox ops to hold a list of blocks (with return terminators and values being returned).

There is no change in the terminator restrictions on affine.for/affine.if or any other registered ops with regions. affine.for/if continue to need "affine.terminator" at the end of their blocks.

[ftynse]

We haven't reached design consensus on this AFAIK.

[bondhugula]

In D71961#1805385, @ftynse wrote:

We haven't reached design consensus on this AFAIK.

Yes, I'm awaiting responses on that thread. (I responded to Mehdi and River last - it was just before the holiday period.)
https://groups.google.com/a/tensorflow.org/d/msg/mlir/U6pxNPXiTSI/-sQjzbSJBAAJ
https://groups.google.com/a/tensorflow.org/d/msg/mlir/U6pxNPXiTSI/zFjiNUqFBAAJ

[rriddle]

lib/Dialect/StandardOps/Ops.cpp

1960 ↗

(On Diff #235476)

I'd probably just assert here to be honest.

1968 ↗

(On Diff #235476)

Is there some way that we can generalize the property of a "top level" op that ReturnOp can use? Special casing FuncOp isn't really ideal here.

[rriddle]

lib/Dialect/StandardOps/Ops.cpp

1960 ↗

(On Diff #235476)

I suppose we can keep this to prevent the verify methods from crashing. But please remove the TODO comment attached.

[bondhugula]

lib/Dialect/StandardOps/Ops.cpp

1968 ↗

(On Diff #235476)

Here, the property we are looking for is the declarative nature of the FuncOp: so we are checking the return types on the function instead of the types of the return values of the op. I guess we could use the FunctionLike trait, but the trait has to be updated to implement a method that returns the function type (or FuncOp::getType() that returns FunctionType could be moved over there.). We have this dichotomy because we turned function into an op although it's not imperative.

[bondhugula]

• fix file paths (missing mlir/ prefix)
• address review comments

[bondhugula]

@rriddle @mehdi_amini This patch also enables the std.inlined_call/execute_region op that there is consensus on. execute_region has returns at its top level just like affine.graybox.

[rriddle]

In D71961#1840721, @bondhugula wrote:

@rriddle @mehdi_amini This patch also enables the std.inlined_call/execute_region op that there is consensus on. execute_region has returns at its top level just like affine.graybox.

(Sorry for the delayed response been OOO for 1.5 weeks)

I'm +1 to generalizing ReturnOp. My main request is that we have a somewhat principled approach to what is a "top-level operation". In the current state of this revision, we won't be able to support using std.return in any other operations that are similar to FuncOp. Is there some interface that we can use
to define what we need here?

[bondhugula]

In D71961#1840829, @rriddle wrote:

In D71961#1840721, @bondhugula wrote:

@rriddle @mehdi_amini This patch also enables the std.inlined_call/execute_region op that there is consensus on. execute_region has returns at its top level just like affine.graybox.

(Sorry for the delayed response been OOO for 1.5 weeks)

I'm +1 to generalizing ReturnOp. My main request is that we have a somewhat principled approach to what is a "top-level operation". In the current state of this revision, we won't be able to support using std.return in any other operations that are similar to FuncOp. Is there some interface that we can use
to define what we need here?

Hi River, I responded to this at your inline code comment, and was waiting for your thoughts. Please see the comment thread at lib/Dialect/StandardOps/Ops.cpp near line #1968.

[rriddle]

In D71961#1840830, @bondhugula wrote:

In D71961#1840829, @rriddle wrote:

In D71961#1840721, @bondhugula wrote:

@rriddle @mehdi_amini This patch also enables the std.inlined_call/execute_region op that there is consensus on. execute_region has returns at its top level just like affine.graybox.

(Sorry for the delayed response been OOO for 1.5 weeks)

I'm +1 to generalizing ReturnOp. My main request is that we have a somewhat principled approach to what is a "top-level operation". In the current state of this revision, we won't be able to support using std.return in any other operations that are similar to FuncOp. Is there some interface that we can use
to define what we need here?

Hi River, I responded to this at your inline code comment, and was waiting for your thoughts. Please see the comment thread at lib/Dialect/StandardOps/Ops.cpp near line #1968.

Ahh yes, thank you. I am truly sorry about completely missing that!

[rriddle]

lib/Dialect/StandardOps/Ops.cpp

1968 ↗

(On Diff #235476)

I'd be +1 for using function-like if we turned it into an interface. We wouldn't be able to use 'getType' though, as it is not guaranteed to be a FunctionType(e.g. in the case of LLVMFuncOp).

[rriddle]

mlir/docs/Dialects/Standard.md

78

Can you also update the ODS description?

mlir/lib/Dialect/StandardOps/Ops.cpp

1955 ↗

(On Diff #239852)

After D73429 we should be able to use ArrayRef<Type> here.

[bondhugula]

lib/Dialect/StandardOps/Ops.cpp

1968 ↗

(On Diff #235476)

I see. We'd need an interface then with a method like getReturnTypes(). How'd you prefer this piece of code to be updated for this patch?

[bondhugula]

Fix ODS description

[bondhugula]

@rriddle Anything else here?

[rriddle]

lib/Dialect/StandardOps/Ops.cpp

1968 ↗

(On Diff #235476)

I'm fine with a TODO for now.

mlir/include/mlir/Dialect/StandardOps/Ops.td

1018 ↗

(On Diff #240475)

nit: Can you wrap this with an mlir code block.

mlir/lib/Dialect/StandardOps/Ops.cpp

1955 ↗

(On Diff #239852)

Should be possible now.

[bondhugula]

Address review comments

[flaub]

mlir/lib/Dialect/StandardOps/Ops.cpp

1945 ↗

(On Diff #241043)

typo: consistent

[bondhugula]

Fix typo.

[mehdi_amini]

Something not clear to me is what is the guideline for using std.return instead of a custom terminator? This is something I'd like clarified before we get inconsistency and this get debated over and over every time a new terminator is introduced (or every time std.return is re-used)

[rriddle]

In D71961#1847104, @mehdi_amini wrote:

Something not clear to me is what is the guideline for using std.return instead of a custom terminator? This is something I'd like clarified before we get inconsistency and this get debated over and over every time a new terminator is introduced (or every time std.return is re-used)

I thought I saw that in this revision, it may have been somewhere else. I think we should clarify that ReturnOp should only be used in "Function-Like" regions where the body of the region effectively represents a function body. Things like lambdas, execute_region, MyFuncOp etc. This way it is clear that lowering patterns don't have to "guess" what the semantics are, it is the same; e.g. the lowering of ReturnOp to LLVM is always to llvm.return.

[rriddle]

(allowing more discussion)

[bondhugula]

In D71961#1847382, @rriddle wrote:

In D71961#1847104, @mehdi_amini wrote:

Something not clear to me is what is the guideline for using std.return instead of a custom terminator? This is something I'd like clarified before we get inconsistency and this get debated over and over every time a new terminator is introduced (or every time std.return is re-used)

I thought I saw that in this revision, it may have been somewhere else. I think we should clarify that ReturnOp should only be used in "Function-Like" regions where the body of the region effectively represents a function body. Things like lambdas, execute_region, MyFuncOp etc. This way it is clear that lowering patterns don't have to "guess" what the semantics are, it is the same; e.g. the lowering of ReturnOp to LLVM is always to llvm.return.

I thought this was already clear from the documentation and ODS desc. Its semantics are now clearly defined there IMO, and so it should be used as the terminator only when exactly those semantics are desired. @mehdi_amini Where would you like additional clarification to be added if any?

[mehdi_amini]

It isn't clear to me what is a "Function-Like" region. In some way I can see every region to be a "function like": I even asked on the mailing list a while ago why wouldn't we *always* use std.return to terminate regions.

I gave the example to River today of xla.if operation which can either take a function for the conditional body or a region. In the function case it'll obviously terminate with a return, if I "inline" the function as a region without changing anything else in the IR, do I have a "function-like" region? Should we keep std.return here?

[bondhugula]

In D71961#1848687, @mehdi_amini wrote:

It isn't clear to me what is a "Function-Like" region. In some way I can see every region to be a "function like": I even asked on the mailing list a while ago why wouldn't we *always* use std.return to terminate regions.

I gave the example to River today of xla.if operation which can either take a function for the conditional body or a region. In the function case it'll obviously terminate with a return, if I "inline" the function as a region without changing anything else in the IR, do I have a "function-like" region? Should we keep std.return here?

The documentation / ODS description I have doesn't say anything about function-like. Its semantics are that it transfers control to its parent op, and if it's something declarative like FuncOp, it does <as stated in the description>. River suggested not hardcoding FuncOp here since the same concept applies to any declarative op that is FuncOp like for which we'll create an interface.

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region. For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level. (If you have an affine.for/loop.for or something else nested inside, you can't use a return inside the affine.for/loop.for itself unless of course the return's parent op allows return in its top-level like std.execute_region).

[mehdi_amini]

In D71961#1848691, @bondhugula wrote:

In D71961#1848687, @mehdi_amini wrote:

It isn't clear to me what is a "Function-Like" region. In some way I can see every region to be a "function like": I even asked on the mailing list a while ago why wouldn't we *always* use std.return to terminate regions.

I gave the example to River today of xla.if operation which can either take a function for the conditional body or a region. In the function case it'll obviously terminate with a return, if I "inline" the function as a region without changing anything else in the IR, do I have a "function-like" region? Should we keep std.return here?

The documentation / ODS description I have doesn't say anything about function-like.

No, but River's comment just above does to answer my question about the intended use and the guideline I'd like us to have.

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region.

I don't know what "top-level" means? I

For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

OK, but then I'd like to understand when to *not* use it. Can we replace affine.if to use it as well?

(If you have an affine.for/loop.for or something else nested inside, you can't use a return inside the affine.for/loop.for itself unless of course the return's parent op allows return in its top-level like std.execute_region).

I don't understand this sentence.

[bondhugula]

In D71961#1850046, @mehdi_amini wrote:

In D71961#1848691, @bondhugula wrote:

In D71961#1848687, @mehdi_amini wrote:

It isn't clear to me what is a "Function-Like" region. In some way I can see every region to be a "function like": I even asked on the mailing list a while ago why wouldn't we *always* use std.return to terminate regions.

I gave the example to River today of xla.if operation which can either take a function for the conditional body or a region. In the function case it'll obviously terminate with a return, if I "inline" the function as a region without changing anything else in the IR, do I have a "function-like" region? Should we keep std.return here?

The documentation / ODS description I have doesn't say anything about function-like.

No, but River's comment just above does to answer my question about the intended use and the guideline I'd like us to have.

Sure, I can add it.

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region.

I don't know what "top-level" means? I

"top-level" => as an operation in a block of that region, and not nested in another operation of that region.

For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

OK, but then I'd like to understand when to *not* use it. Can we replace affine.if to use it as well?

I assume xla.if can return a value? affine.if never returns a value. Strictly speaking, we can use return as affine.if's terminator but since affine.if never returns a value, the affine.terminator is just a constrained/special form of return -- avoids confusion FWIW.

(If you have an affine.for/loop.for or something else nested inside, you can't use a return inside the affine.for/loop.for itself unless of course the return's parent op allows return in its top-level like std.execute_region).

I don't understand this sentence.

It just means that return can't be misused to transfer control to a non-parent ancestor op. Since affine.for/if only accept affine.terminator as their terminators, std.return can't be the terminator of a for's or if's block. However, you could do:

affine.for ... {
  %ret = std.execute_region() -> i32 {
     %y = affine.load  ...
     return %y : i32
  }
}

[mehdi_amini]

In D71961#1850751, @bondhugula wrote:

The documentation / ODS description I have doesn't say anything about function-like.

No, but River's comment just above does to answer my question about the intended use and the guideline I'd like us to have.

Sure, I can add it.

We just first figure out what it is that needs to be added though, it isn't crystal clear to me right now.

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region.

I don't know what "top-level" means? I

"top-level" => as an operation in a block of that region, and not nested in another operation of that region.

For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

OK, but then I'd like to understand when to *not* use it. Can we replace affine.if to use it as well?

I assume xla.if can return a value? affine.if never returns a value. Strictly speaking, we can use return as affine.if's terminator but since affine.if never returns a value, the affine.terminator is just a constrained/special form of return -- avoids confusion FWIW.

So the "can return values" is the discriminator you are making?
For example a GPU kernel is like a function in many ways, but it can't return values, we shouldn't use std.return only because of that?

(If you have an affine.for/loop.for or something else nested inside, you can't use a return inside the affine.for/loop.for itself unless of course the return's parent op allows return in its top-level like std.execute_region).

I don't understand this sentence.

It just means that return can't be misused to transfer control to a non-parent ancestor op. Since affine.for/if only accept affine.terminator as their terminators, std.return can't be the terminator of a for's or if's block. However, you could do:

The "Since affine.for/if only accept affine.terminator as their terminators, std.return can't be the terminator of a for's or if's block" part seems like circular logic: we just have to update affine.for/if to only accept std.return with no operands.

[bondhugula]

In D71961#1850784, @mehdi_amini wrote:

In D71961#1850751, @bondhugula wrote:

For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

OK, but then I'd like to understand when to *not* use it. Can we replace affine.if to use it as well?

I assume xla.if can return a value? affine.if never returns a value. Strictly speaking, we can use return as affine.if's terminator but since affine.if never returns a value, the affine.terminator is just a constrained/special form of return -- avoids confusion FWIW.

So the "can return values" is the discriminator you are making?

Yes, but only as a guideline. As I said, strictly speaking, one could use the return here (while enforcing in verification that it returns 0 values). These returns would then just be erased when lowering loop.for, loop.if (just like loop.terminator is simply erased).

For example a GPU kernel is like a function in many ways, but it can't return values, we shouldn't use std.return only because of that?

Perhaps - because you won't have to verify 0 values for the return FWIW. If it returns values in the future, it would be identical in all ways to std.return.

(If you have an affine.for/loop.for or something else nested inside, you can't use a return inside the affine.for/loop.for itself unless of course the return's parent op allows return in its top-level like std.execute_region).

I don't understand this sentence.

It just means that return can't be misused to transfer control to a non-parent ancestor op. Since affine.for/if only accept affine.terminator as their terminators, std.return can't be the terminator of a for's or if's block. However, you could do:

The "Since affine.for/if only accept affine.terminator as their terminators, std.return can't be the terminator of a for's or if's block" part seems like circular logic: we just have to update affine.for/if to only accept std.return with no operands.

By "can't be", I meant "currently can't be" because the ODS for affine.if/for enforce affine.terminator as their only possible terminator.

[mehdi_amini]

In D71961#1848691, @bondhugula wrote:

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region. For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

We also have xla.while, which is structured exactly like xla.if and when the body is in a function it is using std.return. It would seem strange to me that when "inlining" the region of xla.if we would continue to use the return while we wouldn't for xla.while, don't you think?

[bondhugula]

In D71961#1852477, @mehdi_amini wrote:

In D71961#1848691, @bondhugula wrote:

The control flow semantics of a region basically means that you could always use a return from the *top-level* of your region. For your xla.if inlined region for a conditional, there isn't a need to create yet another terminator. Please use std.return - again from its top-level.

We also have xla.while, which is structured exactly like xla.if and when the body is in a function it is using std.return. It would seem strange to me that when "inlining" the region of xla.if we would continue to use the return while we wouldn't for xla.while, don't you think?

That wouldn't appear strange to me. I prefer using a 'yield' name (a name different than return) when passing something out of each xla.while iteration (for the same reasons as in the case of loop.for - the value returned doesn't map 1:1 with the result of xla.while. (I'm assuming you have similar semantics for xla.while init/yield as those proposed for loop.for). What do you think/prefer to be the guideline for return?

[bondhugula]

I'd suggest this as the guideline in the doc.

"As a guideline, std.return can be used as the terminator for either declarative function-like ops (like FuncOp) or for (imperative) ops with regions where the values being returned associate 1:1 with the ops' results."

[mehdi_amini]

I don't understand the rational right now for yield vs return: a region is behaving exactly like a closure and these terminator are all having the same semantics: giving back the control to the parent operation.

I really don't see why we should differentiate if and for in the following two cases:

func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
  %x = loop.if %arg0 {
    %x1 = addf %arg1, %arg2 : (f32, f32) -> f32
    return %x1 : f32
  } else {
    %x2 = subf %arg1, %arg2 : (f32, f32) -> f32
    return %x2 : f32
  } : (i1) -> f32
  return %x : () -> f32
}

but:

func @test(%lb ...) {
 %res:2 = loop.for %iv = %lb to %ub step %step 
      iter_args (%0 = %arg0 : tensor<32xf32>, %1 = %arg1 : tensor<16xf64>)  
      -> (tensor<32xf32>, tensor<16xf64>) {
    ...
    loop.yield %val0, %val1 : tensor<32xf32>, tensor<16xf64>
  }
  return %res#0,  %res#1
}

return instead of yield is equally confusing (or not confusing) to me in both of these cases.

Basically I find the inconsistency confusing: the lack of clear guideline and the difficulty to express it here does not leave me comfortable with the direction.

[bondhugula]

In D71961#1853008, @mehdi_amini wrote:

I don't understand the rational right now for yield vs return: a region is behaving exactly like a closure and these terminator are all having the same semantics: giving back the control to the parent operation.

I really don't see why we should differentiate if and for in the following two cases:

func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
  %x = loop.if %arg0 {
    %x1 = addf %arg1, %arg2 : (f32, f32) -> f32
    return %x1 : f32
  } else {
    %x2 = subf %arg1, %arg2 : (f32, f32) -> f32
    return %x2 : f32
  } : (i1) -> f32
  return %x : () -> f32
}

but:

func @test(%lb ...) {
 %res:2 = loop.for %iv = %lb to %ub step %step 
      iter_args (%0 = %arg0 : tensor<32xf32>, %1 = %arg1 : tensor<16xf64>)  
      -> (tensor<32xf32>, tensor<16xf64>) {
    ...
    loop.yield %val0, %val1 : tensor<32xf32>, tensor<16xf64>
  }
  return %res#0,  %res#1
}

return instead of yield is equally confusing (or not confusing) to me in both of these cases.

Basically I find the inconsistency confusing: the lack of clear guideline and the difficulty to express it here does not leave me comfortable with the direction.

Semantics-wise, they all are the same, and I see this as a naming / readability thing. If naming shouldn't become the reason to introduce a new op, then we could just recommend using std.return for all these.

[flaub]

FWIW I'm +1 with using a single op for all cases of returning control to the parent op. I personally don't have a strong preference on yield vs return.

[mehdi_amini]

If we look purely at readability, I'm wary of the return:

func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
  %x = loop.if %arg0 {
    %x1 = addf %arg1, %arg2 : (f32, f32) -> f32
    return %x1 : f32
  } else {
    %x2 = subf %arg1, %arg2 : (f32, f32) -> f32
    return %x2 : f32
  } : (i1) -> f32
  return %x : () -> f32
}

This reads "like a program" but all imperative language I know of would have the return in the if blocks returning from the function.

[bondhugula]

In D71961#1853021, @mehdi_amini wrote:

If we look purely at readability, I'm wary of the return:

func @test(%arg0: i1, %arg1: f32, %arg2: f32) {
  %x = loop.if %arg0 {
    %x1 = addf %arg1, %arg2 : (f32, f32) -> f32
    return %x1 : f32
  } else {
    %x2 = subf %arg1, %arg2 : (f32, f32) -> f32
    return %x2 : f32
  } : (i1) -> f32
  return %x : () -> f32
}

This reads "like a program" but all imperative language I know of would have the return in the if blocks returning from the function.

Yes, that's true. But I think the discussion on the last few messages was on loop.for/xla.while. Over there, I found a different name slightly better. If we just go with recommending return in all such cases consistent with a region's semantics, I don't think anything more needs to be added to this patch.

Note that unlike for loop.if/xla.if, the returned values for xla.for/loop.for don't associate with the op's result in a direct way; so a better name for readability may be desired? That's all there is.

[mehdi_amini]

Sorry, I am not sure I follow: are you saying that the example with the loop.if is just fine from a readability point of view?

[bondhugula]

In D71961#1853042, @mehdi_amini wrote:

Sorry, I am not sure I follow: are you saying that the example with the loop.if is just fine from a readability point of view?

Yes.

[bondhugula]

Pinging all reviewers again: do we need any changes to the doc/ODS desc. here at all?

[mehdi_amini]

In D71961#1856049, @bondhugula wrote:

Pinging all reviewers again: do we need any changes to the doc/ODS desc. here at all?

At this point I'm still not in full agreement with the change in the first place: the tradeoff does not seem convincing to me: the loop.if example for instance makes me worried about readability.

I also don't understand why loop.for body would be treated differently from the loop.if body: the difference you're making on having multiple iteration are not meaningful to me with respect to the region semantics.

[bondhugula]

In D71961#1856212, @mehdi_amini wrote:

In D71961#1856049, @bondhugula wrote:

Pinging all reviewers again: do we need any changes to the doc/ODS desc. here at all?

At this point I'm still not in full agreement with the change in the first place: the tradeoff does not seem convincing to me: the loop.if example for instance makes me worried about readability.

I also don't understand why loop.for body would be treated differently from the loop.if body: the difference you're making on having multiple iteration are not meaningful to me with respect to the region semantics.

But have you thought why execute_region should not be using an std.return within it? Just because you feel it's not readable for loop.if / loop.for doesn't mean it should preclude use of return where it's both readable and valid semantics wise. execute_region is the one that motivated this change in the first place (as mentioned in the commit message). Why should one create a different terminator for execute_region? I believe the discussion moved away from the central point and instead overdwelled on a less important issue.

I believe I've already responded to all your questions above -- some even multiple times! (as to why loop.for's return values are different than loop.if in some ways) -- I seriously believe you are asking the same questions again!

@rriddle -- it would be great if you or others can intervene as well -- please do take a look at the entire thread. I have no more additional arguments in favor to make here, and I strongly disagree with @mehdi_amini that this change isn't in the right direction. A guideline statement is all that the doc needs if anything.

[mehdi_amini]

In D71961#1856301, @bondhugula wrote:

I believe I've already responded to all your questions above -- some even multiple times! (as to why loop.for's return values are different than loop.if in some ways) -- I seriously believe you are asking the same questions again!

Yes, thanks for answering, but I'm still saying I disagree with your stated rational in these answer, the conclusion at this point is that I oppose this patch without a guideline.

@rriddle -- it would be great if you or others can intervene as well -- please do take a look at the entire thread. I have no more additional arguments in favor to make here, and I strongly disagree with @mehdi_amini that this change isn't in the right direction. A guideline statement is all that the doc needs if anything.

In my opinion, you haven't been able to provide a reasonable guideline here to support this change, when considering that declaring a new op is "cheap", the tradeoff exposed here does not seem worthwhile to me.

[bondhugula]

In D71961#1856319, @mehdi_amini wrote:

In D71961#1856301, @bondhugula wrote:

I believe I've already responded to all your questions above -- some even multiple times! (as to why loop.for's return values are different than loop.if in some ways) -- I seriously believe you are asking the same questions again!

Yes, thanks for answering, but I'm still saying I disagree with your stated rational in these answer, the conclusion at this point is that I oppose this patch without a guideline.

@rriddle -- it would be great if you or others can intervene as well -- please do take a look at the entire thread. I have no more additional arguments in favor to make here, and I strongly disagree with @mehdi_amini that this change isn't in the right direction. A guideline statement is all that the doc needs if anything.

In my opinion, you haven't been able to provide a reasonable guideline here to support this change, when considering that declaring a new op is "cheap", the tradeoff exposed here does not seem worthwhile to me.

@mehdi_amini You've skipped/stripped the first para of my response above which had other questions for you! - please respond to the whole thing instead of a strawman. :-) As for the guideline, I'm just copy/pasting what I posted at the discussions thread on IfOp:

"Although per the region’s semantics, irrespective of which terminator you use, the value is always returned out of the region and to the parent op to handle, readability-wise, it might be confusing to use ‘return’ in cases where the returned value(s) don’t bind 1:1 with the parent op’s result values. For loop.for, such a binding only happens for the last iteration; hence, I think it’s more readable to use something other than std.return or for that matter something other than loop.return. However, in cases, where the returned value binds 1:1 with the result value whenever something is returned (like in std.func, std.execute_region, affine.graybox, all other declarative function like ops that exist in MLIR, even loop.if), using return is both readable and obviously semantics-wise valid. Allowing reuse of std.return in cases like these is the motivation behind untying std.return from FuncOp (https://reviews.llvm.org/D71961 ), but this has been pending for over a month now! Its immediate use case is execute_region / inlined_call."

On another note, I did also go through the video recording of the meeting on 13 Feb - I find the general idea to use std.yield for single block but return for multiple block arbitrary, but let's get the usage of std.return out of the way first separately from the proposal to add a new terminator op.

[bondhugula]

@herhut @flaub @silvas @nmostafa - I'd like to get your stand here as well. Please see the comment right above for the latest.

[nmostafa]

In D71961#1896853, @bondhugula wrote:

@herhut @flaub @silvas @nmostafa - I'd like to get your stand here as well. Please see the comment right above for the latest.

I find the semantics of terminating a region for a declarative op (FuncOp) vs imperative ones (e.g. loop.IfOp) slightly different. For the former, we are not really "executing" the FuncOp on a function call, instead we are transferring control to the contained region, and an std.return implies some call-return semantics instead of returning control back to the FuncOp. Also, for imperative ops, we are actually defining values, something a FuncOp does not do.

I prefer having a terminator for function-like ops (declarative in nature, some transfer of control happens) which can be std.return, and a separate terminator for ops that actually define values and executed inline (For, If, ..etc) which can be std.yield. I also think std.yield works fine for both For/WhileOp and IfOp, since control is being passed back to the parent, it can decide to re-bind the region arguments and re-enter the region or terminate.

[mehdi_amini]

In D71961#1896969, @nmostafa wrote:

I prefer having a terminator for function-like ops (declarative in nature, some transfer of control happens) which can be std.return

This seems like a very sensible design point! We can also check it statically by ensuring that the enclosing operation for std.return defines a symbol.

[bondhugula]

In D71961#1897301, @mehdi_amini wrote:

In D71961#1896969, @nmostafa wrote:

I prefer having a terminator for function-like ops (declarative in nature, some transfer of control happens) which can be std.return

This seems like a very sensible design point! We can also check it statically by ensuring that the enclosing operation for std.return defines a symbol.

Does an op defining a symbol automatically imply declarative and function like nature? How is this future proof or even correct traits wise?

While declarative vs imperative is a key element, even among imperative ops, the distinction between an op binding the returned value to its results (like execute_region) or dealing with it differently (like *.for) is equally/more important. Using yield for imperative and return for declarative appears arbitrary to me in some ways, and using yield uniformly for all imperative ops (even where return would have been a perfect match) would also be inconsistent. IMO, the imperative vs declarative distinction doesn't warrant something different from return, but it's the treatment of the returned values by the parent op that does.

[mehdi_amini]

Does an op defining a symbol automatically imply declarative and function like nature?

Defining a symbol implies that the control-flow isn't local: you have to use an indirection to reach the region in the operation.

How is this future proof or even correct traits wise?

Maybe relying on the FunctionLike trait instead would be more accurate.

[silvas]

Since it seems the discussion is a bit stalled, maybe I can make a suggestion.

The goal of this patch can be seen from a larger perspective as reducing the need for dialect-specific terminators. From that perspective, it seems one way to approach this is to build a list of the custom terminators in the ecosystem that we think are redundant and try to classify them and see what themes emerge.

Ideally, for each distinct class of terminators we can write a verifier that guarantees some useful property (this might open up the need for new traits like “region-having control flow op with bodies of a single block”. Simply having something that you can throw on the end of the block structurally doesn’t buy much.

I can think of the following terminator classes;

• llvm-style CFG return from a function (like std.return). Ideally we can find some other return ops in the ecosystem that we can fold into this (IREE might have one, llvm dialect has one I think; probably anything with a custom FunctionLike op). Once we have some examples, tying this to FunctionLike will be an easier decision.
• structurally necessary terminator for declarative “module-like” regions
• purely imperative “this region is done” op like used for affine.if
• yield-like terminator that has operands that if the parent op’s control flow were unfolded into a CFG, would logically correspond to successor operands(arguments to phi nodes in successors). These might additionally come with the constraint that they are in a single block region.

[silvas]

FWIW, I also think that we should not use std.return for inlined_call/execute_region. I can’t think of a nontrivial verifier that we could write that would allow both that use and the more typical use in FuncOp.

[bondhugula]

In D71961#1908866, @silvas wrote:

FWIW, I also think that we should not use std.return for inlined_call/execute_region. I can’t think of a nontrivial verifier that we could write that would allow both that use and the more typical use in FuncOp.

Hi Sean, this patch already updates the verifier to work with both FuncOp and with execute_region! (see near Ops.cpp:1938) I see the change as trivial. :-) To generalize it to all func like ops, we need an op interface - but that doesn't complicate the verifier in any way.

[bondhugula]

In D71961#1896969, @nmostafa wrote:

In D71961#1896853, @bondhugula wrote:

@herhut @flaub @silvas @nmostafa - I'd like to get your stand here as well. Please see the comment right above for the latest.

I find the semantics of terminating a region for a declarative op (FuncOp) vs imperative ones (e.g. loop.IfOp) slightly different. For the former, we are not really "executing" the FuncOp on a function call, instead we are transferring control to the contained region, and an std.return implies some call-return semantics instead of returning control back to the FuncOp. Also, for imperative ops, we are actually defining values, something a FuncOp does not do.

I prefer having a terminator for function-like ops (declarative in nature, some transfer of control happens) which can be std.return, and a separate terminator for ops that actually define values and executed inline (For, If, ..etc) which can be std.yield. I also think std.yield works fine for both For/WhileOp and IfOp, since control is being passed back to the parent, it can decide to re-bind the region arguments and re-enter the region or terminate.

@nmostafa With execute_region (and also the 'if' ops and function like ops), you are relinquishing control (irreversibly) and sending the values back as the parent op's results. 'yield' in contrast often has the implication of giving way and then continuing, and in any case, doesn't imply values going back to the end of the op's execution --- 'return' OTOH does imply the latter. For the *.for ops however, one could argue that yield is better since the op doesn't finish unless the region is on the last iteration. In any case, yield vs return doesn't bring out the imperative vs declarative distinction for me.

[mehdi_amini]

In D71961#1912460, @bondhugula wrote:

@nmostafa With execute_region (and also the 'if' ops and function like ops), you are relinquishing control (irreversibly) and sending the values back as the parent op's results. 'yield' in contrast often has the implication of giving way and then continuing

Isn't "yield" in generator/scheduling/co-routine use-cases continuing from right after the yield? That seems quite different from terminating a region to me.

, and in any case, doesn't imply values going back to the end of the op's execution --- 'return' OTOH does imply the latter.

I don't have the same implication or interpretation for these ops as you do here FWIW.

[bondhugula]

In D71961#1913006, @mehdi_amini wrote:

In D71961#1912460, @bondhugula wrote:

@nmostafa With execute_region (and also the 'if' ops and function like ops), you are relinquishing control (irreversibly) and sending the values back as the parent op's results. 'yield' in contrast often has the implication of giving way and then continuing

Isn't "yield" in generator/scheduling/co-routine use-cases continuing from right after the yield? That seems quite different from terminating a region to me.

I don't immediately recognize the terms generator/scheduling/co-routine! When you hit the terminator on something like FuncOp or execute_region, it means you are done with the region *and* the op as well, of which you aren't going to execute any more.

, and in any case, doesn't imply values going back to the end of the op's execution --- 'return' OTOH does imply the latter.

I don't have the same implication or interpretation for these ops as you do here FWIW.

Okay, but why? Readability / naming wise, the choice of using 'return' (returning from a region) or 'yield' (yielding from a region) should be available when designing a new op, no matter which dialect the new op is from. You haven't really answered why it shouldn't be as far as 'return' goes. It's up to the op's semantics as to see what fits well, or if you need something new. For at least all func like ops + execute_region + affine.graybox, std.return is the right choice IMO. Creating something new would be creating an alias (redundant, lacking uniformity, and bad even for readability).

[silvas]

Hi Sean, this patch already updates the verifier to work with both FuncOp and with execute_region! (see near Ops.cpp:1938) I see the change as trivial. :-) To generalize it to all func like ops, we need an op interface - but that doesn't complicate the verifier in any way.

Sorry, I missed that (was on mobile so navigating the patch was kind of hard). Now that I look at it, it feels like kind of a hack. This verifier just pushes an "isFuncOp" check to all users of the std.return op, which I don't see as being useful. Better to have one return op for FunctionLike/declarative, and one for parent ops that have a single region and a property that the return of the region is passed identically to the pass results. I'm not sure that the latter semantics are necessarily useful enough to warrant a dedicated op for that, suggesting that we avoid putting such a terminator into std for now.

[silvas]

I'd like to suggest the following though exercise. Consider the following code in a pass:

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  ... something here ...
}

What do you put inside the ".. something here ..."?

The way I see it, with the way that this patch defines ReturnOp, it would invariably look as follows:

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  if (isFunctionLike(stdReturn.getParentOp())) {
    ... do something special for functionlike ...
  } else {
   // The parent is an "execute this region" kind of op.
   ... do something special for "execute this region" kind of op ...
  }
}

At that point, having two separate ops is better. In fact, the code might look even worse:

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  assert(isFunctionLike(stdReturn.getParentOp()) && "we shouldn't be seeing non-functionlike parents at this point");
  ... do something special for functionlike since that is all that makes sense here ...
}

Even though they are pretty sure that return ops with non-functionlike parents shouldn't be present in the IR they are processing, it's not a trivial structural guarantee so the user still feels the need to put an assertion to sanity-check themselves. Again, with two separate ops they do have a trivial structural guarantee.

[bondhugula]

In D71961#1913146, @silvas wrote:

I'd like to suggest the following though exercise. Consider the following code in a pass:

I think this is a great exercise - thanks. I'm trying to see which use cases and if there are any that would require special casing based on declarative vs. imperative as you show below. In fact, at the moment, there are none in the code base, which is why this patch had to update only the verifier! But I'm trying to see where else potentially in the future. One that comes to mind.

1. an execute_region that returns a constant can be completely erased (and its result value replaced with this constant) - but this would be implemented as a fold on execute_region (and not on the return op) and so this doesn't give you any of that casing you show below.

All other use cases I could immediately think of actually treat 'return' the same way irrespective of whether it's inside a parent func op or an imperative one -- for eg., all region local / CFG transformations. Wouldn't that actually be an upside of reusing return? (You won't have to check for two terminators FWIW) Do you have use cases for the "something here" part below that require treating imperative and declarative parent ops of std.return differently? And I think that has to be weighed against the scenarios that require the same treatment.

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  ... something here ...
}

What do you put inside the ".. something here ..."?

The way I see it, with the way that this patch defines ReturnOp, it would invariably look as follows:

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  if (isFunctionLike(stdReturn.getParentOp())) {
    ... do something special for functionlike ...
  } else {
   // The parent is an "execute this region" kind of op.
   ... do something special for "execute this region" kind of op ...
  }
}

At that point, having two separate ops is better. In fact, the code might look even worse:

if (auto stdReturn = dyn_cast<StandardOps::ReturnOp>(op)) {
  assert(isFunctionLike(stdReturn.getParentOp()) && "we shouldn't be seeing non-functionlike parents at this point");
  ... do something special for functionlike since that is all that makes sense here ...
}

Even though they are pretty sure that return ops with non-functionlike parents shouldn't be present in the IR they are processing, it's not a trivial structural guarantee so the user still feels the need to put an assertion to sanity-check themselves. Again, with two separate ops they do have a trivial structural guarantee.

[silvas]

All other use cases I could immediately think of actually treat 'return' the same way irrespective of whether it's inside a parent func op or an imperative one -- for eg., all region local / CFG transformations. Wouldn't that actually be an upside of reusing return? (You won't have to check for two terminators FWIW) Do you have use cases for the "something here" part below that require treating imperative and declarative parent ops of std.return differently? And I think that has to be weighed against the scenarios that require the same treatment.

Interesting point! Would such passes prefer to use terminatorOp->getNumSuccessors() == 0 though instead of isa<ReturnOp>(terminatorOp)? What would such passes do differently for a ReturnOp vs any terminator that has no successors?

[bondhugula]

In D71961#1913580, @silvas wrote:

All other use cases I could immediately think of actually treat 'return' the same way irrespective of whether it's inside a parent func op or an imperative one -- for eg., all region local / CFG transformations. Wouldn't that actually be an upside of reusing return? (You won't have to check for two terminators FWIW) Do you have use cases for the "something here" part below that require treating imperative and declarative parent ops of std.return differently? And I think that has to be weighed against the scenarios that require the same treatment.

Interesting point! Would such passes prefer to use terminatorOp->getNumSuccessors() == 0 though instead of isa<ReturnOp>(terminatorOp)? What would such passes do differently for a ReturnOp vs any terminator that has no successors?

The CFG passes would like to know whether it's a 'return' (as opposed to br or cond_br) and so for that purpose they would want to treat it the same way whether there are in a func op or the region held by execute_region. It's not just the number of successors, but you'd need a dyn_cast to get the return operands as well. Someone who has probably worked with CFG passes may be able to better say (not me :-)). Overall, stepping back, I could broadly classify use cases into these three:

(1) intraprocedural things/passes that are transforming the CFG - these would want to really be treating the region/cfg inside a func op the same way as the region inside execute_region. This scenario calls for a single terminator irrespective of whether the cfg region holding op is imperative or declarative since they are all doing region local things isolated from whatever's outside their parent op. Having two terminators will lead to isa/dyn_cast<ReturnOp> || isa/dyn_cast<OtherOp>(..).

(2) polyhedral passes / other passes on higher order control flow like fors/ifs: these passes really don't do anything with the terminator of either the func op or execute_region -- nor do they typically care about any block terminator; even when they do, transformation is centric on the parent op - just like an execute_region or loop.if op checking if its terminators are returning a constant (both regions for loop.if). These are canonicalization patterns on the parent op and not on the return op.

(3) interprocedural stuff: this is where the return op and its operand may get looked at less rarely -- inlining is an example. But again, even here, the same inlining code would be reused when inlining an execute_region's blocks as for a func op's (there isn't really a difference except that the former doesn't have entry block arguments to deal with). So it's only simpler if they had the same terminator FWIW.

(4) I think this just leaves "return op-centric canonicalizations", where I think two terminators help (like you showed above with the structure guarantee) - eg: if a canonicalization pattern on the return op has to propagate some information back to the call site (in a ModulePass), you'd have to check if you are dealing with a FuncLikeOp or an execute_region thing to see where you need to send stuff, but one could argue that such an op pattern should be implemented on the parent op (the execute_region op) or the call op to start with?! Do you think the ReturnOp would ever need a canonicalization pattern? It has none at the moment.

[bondhugula]

Changing reviewers based on whoever is participating.

[bondhugula]

Rebase

[mehdi_amini]

FYI I still believe that a good first step her would be to use something like the FunctionLike trait to constraint the parent of the return op.

[bondhugula]

Update op documentation.

[bondhugula]

In D71961#1948401, @mehdi_amini wrote:

FYI I still believe that a good first step her would be to use something like the FunctionLike trait to constraint the parent of the return op.

This was already discussed at the comment thread at Ops.cpp line 1968 in the first diff - (see River's/my comments). The FunctionLike trait here isn't a perfect fit here because there are ops that have the FunctionLike trait but don't model their function return type in the same way (like llvm.func) and thus you can't do a funcOp.getType().getResults() in the verifier. The right way to handle that is through an op interface (so that a FuncLikeOp::getType() provides the uniform behavior needed for verify), and so we agreed on adding a TODO comment for it there. Note that if a FunctionLike trait op needs something custom for its return type, it just shouldn't use std.return - even if it does, the verifier would fail for its needs. There are ops of course with a FunctionLike trait that could reuse std.return, and in those cases they'd better use the TODO FuncLikeOpInterface. There is still no need to restrict the parent to those implementing FuncLikeOpInterface. FYI, I've updated the op ODS description to add a clear guideline at the end. (Also, see the discussion thread with @silvas where I mention why CFG transformations would want to see the same terminator on FuncOp as well as std.execute_region - instead of having to do isa/dyn_cast<ReturnOp> || isa/dyn_cast<yield op / some other alias of return>)

[bondhugula]

This is now moot given std.execute_region suits the scf dialect and the yield op there could terminate it.

mlir/docs/Dialects/Standard.md

 ### 'return' terminator operation
 
 Syntax:
 
 ```
 operation ::= `return` (ssa-use-list `:` type-list-no-parens)?
 ```
 
-The `return` terminator operation represents the completion of a function, and
-produces the result values. The count and types of the operands must match the
-result types of the enclosing function. It is legal for multiple blocks in a
-single function to return.
+The `return` terminator operation, when enclosed immediately by a function,
+represents the completion of that function and produces its result values. The
+count and types of the operands must match the result types of the enclosing
+function. It is valid for multiple blocks in a single function to return. In the
+case that the op holding the region in which the `return` appears is not a
+function, the `return` represents a transfer of control to the enclosing op,
+whose semantics define where control flows next.
 
 ## Core Operations
 
 ### 'call' operation
 
 Syntax:
 
 ```

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 def RemFOp : FloatArithmeticOp<"remf"> {
   let summary = "floating point division remainder operation";
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
-def ReturnOp : Std_Op<"return", [NoSideEffect, HasParent<"FuncOp">,
-                                 Terminator]> {
+def ReturnOp : Std_Op<"return", [NoSideEffect, Terminator]> {
   let summary = "return operation";
   let description = [{
-    The "return" operation represents a return operation within a function.
-    The operation takes variable number of operands and produces no results.
-    The operand number and types must match the signature of the function
-    that contains the operation. For example:
+    The `return` terminator operation, when enclosed immediately by a function,
+    represents the completion of that function and produces its result values.
+    The count and types of the operands must match the result types of the
+    enclosing function. It is valid for multiple blocks in a single function to
+    return. In the case that the `return` op's parent is not a function op, the
+    `return` represents a transfer of control to its parent, whose semantics
+    determine where control flows next, and whose result values match the values
+    returned. For such an imperative parent op, as a guideline, std.return
+    should be used as an terminator only when the values being returned are
+    actually the results of that op, which is typically the case when its region
+    is being executed once.
 
-      func @foo() : (i32, f8) {
+    ```mlir
+    func @foo() : () -> (i32, f32) {
       ...
-      return %0, %1 : i32, f8
+      return %0, %1 : i32, f32
+    }
+    ```
   }];
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilder<
     "Builder *b, OperationState &result", [{ build(b, result, llvm::None); }]
   >];
 

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
   return IntegerAttr();
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReturnOp op) {
-  auto function = cast<FuncOp>(op.getParentOp());
+  auto *parentOp = op.getParentOp();
+  if (!parentOp)
+    return op.emitOpError("has no parent op");
+
+  // The operand count and types must be consistent with the parent op. When the
+  // parent is a FuncOp (which is declarative), check against its return types;
+  // for the rest, check again the number of actual SSA results.
+  // TODO: create an op interface for declarative func like ops instead of
+  // treating just the FuncOp specially.
+  ArrayRef<Type> retTypes;
+  if (auto funcOp = dyn_cast<FuncOp>(parentOp))
+    retTypes = funcOp.getType().getResults();
+  else
+    retTypes = parentOp->getResultTypes();
 
-  // The operand number and types must match the function signature.
-  const auto &results = function.getType().getResults();
-  if (op.getNumOperands() != results.size())
+  if (op.getNumOperands() != retTypes.size())
     return op.emitOpError("has ")
            << op.getNumOperands()
-           << " operands, but enclosing function returns " << results.size();
+           << " operands, but enclosing function returns " << retTypes.size();
 
-  for (unsigned i = 0, e = results.size(); i != e; ++i)
-    if (op.getOperand(i).getType() != results[i])
+  for (unsigned i = 0, e = retTypes.size(); i != e; ++i)
+    if (op.getOperand(i).getType() != retTypes[i])
       return op.emitError()
              << "type of return operand " << i << " ("
              << op.getOperand(i).getType()
-             << ") doesn't match function result type (" << results[i] << ")";
+             << ") doesn't match function result type (" << retTypes[i] << ")";
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // SelectOp
 //===----------------------------------------------------------------------===//
 

mlir/test/IR/core-ops.mlir

^bb0(%0: memref<4x4xi32>, %1: index):
   prefetch %0[%1, %1], write, locality<1>, data : memref<4x4xi32>
 
   // CHECK: prefetch %arg0[%arg1, %arg1], read, locality<3>, instr : memref<4x4xi32>
   prefetch %0[%1, %1], read, locality<3>, instr : memref<4x4xi32>
 
   return
 }
 
+// CHECK-LABEL: func @return_in_op_with_region
+func @return_in_op_with_region() {
+  "foo.region"() ({
+    %c9 = constant 9 : i32
+    return %c9 : i32
+  }): () -> (i32)
+  return
+}
+
 // Test with zero-dimensional operands using no index in load/store.
 // CHECK-LABEL: func @zero_dim_no_idx
 func @zero_dim_no_idx(%arg0 : memref<i32>, %arg1 : memref<i32>, %arg2 : memref<i32>) {
   %0 = std.load %arg0[] : memref<i32>
   std.store %0, %arg1[] : memref<i32>
   return
   // CHECK: %0 = load %{{.*}}[] : memref<i32>
   // CHECK: store %{{.*}}, %{{.*}}[] : memref<i32>

mlir/test/IR/invalid-ops.mlir

func @trunci_cast_to_same_width(%arg0 : i16) {
   %0 = trunci %arg0 : i16 to i16
   return
 }
 
 // -----
 
 func @return_not_in_function() {
   "foo.region"() ({
-    // expected-error@+1 {{'std.return' op expects parent op 'func'}}
+    // expected-error@+1 {{'std.return' op has 0 operands, but enclosing function returns 1}}
     return
-  }): () -> ()
+  }): () -> (i32)
   return
 }
 
 // -----
 
 func @invalid_splat(%v : f32) {
   splat %v : memref<8xf32>
   // expected-error@-1 {{must be vector of any type values or statically shaped tensor of any type values}}