tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[Analysis][Bugfix] Correctly handle dataflow blocks in local functions defined inside a dataflow block #370

Closed slyubomirsky closed 1 year ago

slyubomirsky commented 1 year ago

Currently, the well-formed checker enters an invalid state if a local function that contains a dataflow block is defined inside a dataflow block, like so:

@R.function
def f(...):
    with R.dataflow():
        @R.function
        def g(...):
            with R.dataflow():
                ...

The error happens because the well-formed checker uses only a single flag is_dataflow_ for determining if the definition of a DataflowVar happens inside a dataflow block. After the inner dataflow block is finished being checked, the var is simply set to false, which results in the rest of the outer dataflow block being incorrectly treated as non-dataflow. Analogously, the set of defined dataflow vars needs to be saved and restored before recursing into an inner function just as the set of ordinary vars is saved and restored.

This PR fixes this bug by keeping a stack of dataflow block states and correctly saving and restoring the set of defined dataflow vars per scope. (Note that it is always valid to define a function inside a dataflow block, since defining a function causes no side effects.)

yongwww commented 1 year ago

thanks @slyubomirsky for the fix! Thanks @masahi @YuchenJin @MasterJH5574 for reviewing it.