[Bug][Parsing] Mutual recursion fails to parse inside an if branch

[slyubomirsky]

I encounter a parse error on the following test case:

import tvm
from tvm.script import relax as R

@tvm.script.ir_module
class ControlFlowExample:
    @R.function
    def a(x: R.Object) -> R.Object:
        y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
        if y:
            return b(x)
        else:
            return c(x)

    @R.function
    def b(x: R.Object) -> R.Object:
        return a(x)

    @R.function
    def c(x: R.Object) -> R.Object:
        return a(x)

The resulting error complains of a violated invariant in the parser:

error:   Check failed: (!IRBuilder::Current()->FindFrame<BlockFrame>()) is false: All block frame are supposed to be popped out already
 --> /home/slyubomirsky/code/relax/tests/python/relax/test_analysis_mutual_recursion.py:293:17
     |  
 293 |                  return b(x)
     |                  ^^^^^^^^^^^

[MasterJH5574]

We now seem to have the restriction that the body of function must be LeafExpr, IIRC. Probably we can improve it with a better error msg.

[slyubomirsky]

If you mean the branches of If nodes, that would be a mistake if we did it. We require those to be SeqExprs.

Edit: If you mean that return statements should have a leaf expression, I also don't think that's something we should do, as the normalizer should be work.

[yongwww]

The restriction is that return is not allowed in true/false branch of if. Not related to recursion.
@slyubomirsky a workaround as below.

import tvm
from tvm.script import relax as R

@tvm.script.ir_module
class ControlFlowExample:
    @R.function
    def a(x: R.Object) -> R.Object:
        y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
        if y:
            r = b(x) 
        else:
            r = c(x)
        return r

    @R.function
    def b(x: R.Object) -> R.Object:
        return a(x)

    @R.function
    def c(x: R.Object) -> R.Object:
        return a(x)

[yongwww]

If we would like to support multiple return like the following example in Relax Function, one option is to introduce a tag member like Bool return_body in SeqExprNode , the tag return_body is used to tell if body of SeqExpr is a return. Any comments? cc: @tqchen @YuchenJin

    @R.function
    def foo(x: R.Tensor) -> R.Tensor:
        y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
        if y:
            return R.add(x, x)
        else:
            return R.multiply(x, x)

[tqchen]

To keep IR constrained for now, let us not update the IR. this can be sugared to assigning values to global then return. So the only question is do we want to support this syntax in parser and sugar to that IR

[yongwww]

Thanks TQ! I don't have a preference on if we need this syntax, a better err message should be good for me.

[slyubomirsky]

Okay, thanks for the update. I think the workaround is fine for my purposes. It would be nice for the parser to have syntactic sugar for it and at minimum a better error message.

[yongwww]

@tqchen @slyubomirsky I have updated the error message in apache/tvm#14123. I agree it would be nice to have this syntactic sugar (more pythonic style). I can work on adding this sugar.

[kparzysz-quic]

To keep IR constrained for now, let us not update the IR. this can be sugared to assigning values to global then return. So the only question is do we want to support this syntax in parser and sugar to that IR

What if the if part returns, but the else part doesn't?

[yongwww]

What if the if part returns, but the else part doesn't?

Then the code should look like this:

    @R.function
    def foo(x: R.Tensor) -> R.Tensor:
        y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
        if y:
            return R.add(x, x)
        else:
            r = R.multiply(x, x)
        return r

[kparzysz-quic]

Shouldn't we first define what return actually does? The draft of the Relax spec is silent on this.

[yongwww]

Shouldn't we first define what return actually does? The draft of the Relax spec is silent on this.

return is not in the IR explicitly, currently by default the body of SeqExpr of function works as return

[slyubomirsky]

Yeah, return is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.

This is actually opening a pretty tough can of worms for parsing, now that I think of it. We could indeed make our lives simpler by requiring both branches to contain a return if either does.

[yongwww]

Yeah, return is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.

This is actually opening a pretty tough can of worms for parsing, now that I think of it. We could indeed make our lives simpler by requiring both branches to contain a return if either does.

Another case is it has no false branch, it may require the member false_branch of IfNode to be optional

    @R.function
    def foo(x: R.Tensor) -> R.Tensor:
        y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
        if y:
            return R.add(x, x)
        # no false branch
        return R.multiply(x, x)

[slyubomirsky]

The Relax AST does not permit omitting the false branch, so how would we parse it? (Even without a return)

[yongwww]

The Relax AST does not permit omitting the false branch, so how would we parse it? (Even without a return)

Omitting the false branch is not allowed now, a SeqExpr will be constructed for the false branch during parsing.

[kparzysz-quic]

Yeah, return is not part of Relax's grammar; it's from Python's grammar and it is up to us to determine how to parse it.

Hmm.. So in the example below, does the return return from the if node, or from the function?

@R.function
def foo(x: R.Tensor) -> R.Tensor:
    y: R.Tensor((), dtype="bool") = R.const(True, dtype="bool")
    if y:
        return R.add(x, x)
    else:
        return R.mul(x, x)
    ...

[slyubomirsky]

I think we should parse that to If(y, R.add(x, x), R.mul(x, x)). The tricky thing is we're parsing Python statements into Relax expressions.

[kparzysz-quic]

Maybe the easiest thing to do would be to disallow return at all, except if it's the last statement in the function body, and it's "top-level", i.e. not a proper sub-statement of any other statement.

Edit: The motivation here is that the behavior of return in the script should match the behavior of it in python code.

[yongwww]

Hmm.. So in the example below, does the return return from the if node, or from the function?

from the if node.

[slyubomirsky]

Edit: The motivation here is that the behavior of return in the script should match the behavior of it in python code.

Yeah, this is a bit of the dilemma we run into with this parser. We're trying to fit a square peg (Python's statement-based grammar) into a round hole (Relax's expression-based grammar). It would be nice for common patterns from Python to work smoothly, but we do have room to say that we don't support some things. I'm fine with asking for return values to be outside the if/else block, for what it's worth.