Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TVMScript][Relax] Allow return statement in DataflowBlock #17131

Merged

Conversation

Lunderberg
Copy link
Contributor

Prior to this commit, TVMScript required the return value of a Relax to be specified outside of any with R.dataflow() blocks. This resulted in a common pattern, where the return value of a function was first called with R.output(ret_value), to mark ret_value as a tvm::relax::Var instead of a tvm::relax::DataflowVar, followed immediately by a return ret_value statement.

This commit updates the TVMScript parser to allow a return statement inside a with R.dataflow() block. This is syntactic sugar that is equivalent to calling R.output, followed by a return.

With this change, the following two TVMScript examples are now equivalent. (Prior to this change, the return_inside_dataflow example would raise an error during parsing.)

@R.function(private=True)
def output_then_return(A: R.Tensor):
    with R.dataflow():
        B = R.add(A, A)
        C = R.multiply(B, B)
        R.output(C)

    return C

@R.function(private=True)
def return_inside_dataflow(A: R.Tensor):
    with R.dataflow():
        B = R.add(A, A)
        C = R.multiply(B, B)
        return C

@Lunderberg Lunderberg requested a review from yongwww July 1, 2024 14:17
Prior to this commit, TVMScript required the return value of a Relax
to be specified outside of any `with R.dataflow()` blocks.  This
resulted in a common pattern, where the return value of a function was
first called with `R.output(ret_value)`, to mark `ret_value` as a
`tvm::relax::Var` instead of a `tvm::relax::DataflowVar`, followed
immediately by a `return ret_value` statement.

This commit updates the TVMScript parser to allow a `return` statement
inside a `with R.dataflow()` block.  This is syntactic sugar that
is equivalent to calling `R.output`, followed by a `return`.

With this change, the following two TVMScript examples are now
equivalent.  (Prior to this change, the `return_inside_dataflow`
example would raise an error during parsing.)

```python
@R.function(private=True)
def output_then_return(A: R.Tensor):
    with R.dataflow():
        B = R.add(A, A)
        C = R.multiply(B, B)
        R.output(C)

    return C

@R.function(private=True)
def return_inside_dataflow(A: R.Tensor):
    with R.dataflow():
        B = R.add(A, A)
        C = R.multiply(B, B)
        return C
```
@Lunderberg Lunderberg force-pushed the tvmscript_relax_return_inside_dataflow_block branch from ed4c76e to dd38813 Compare September 17, 2024 14:32
@Lunderberg
Copy link
Contributor Author

Rebased onto main to resolve conflict.

@tqchen tqchen merged commit a242046 into apache:main Sep 18, 2024
17 of 18 checks passed
@Lunderberg Lunderberg deleted the tvmscript_relax_return_inside_dataflow_block branch September 18, 2024 20:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants