Recursive visiting in `PyExprMutator`

Development unity
#1

I was trying to do something like this, basically have my own visitor that will do some work on its own and use the superclass’s visitor to do the rest:

@relax.expr_functor.mutator
class MyMutator(relax.PyExprMutator):
    ...
    def visit_dataflow_block_(self, block):
        # visit bindings...
        # Do the equivalent of C++'s PyExprMutator::visit_dataflow_block(block):
        super(relax.PyExprMutator, self).visit_dataflow_block(block)   # Did not work

    def visit_call_(self, call):
        # Not called if I have my own block visitor.

This didn’t work because the superclass didn’t have any visitor methods in it (I guess the decorator did something here). Without calling the superclass’s visitor, I don’t know how to get it to call visitors of the elements of the block.

Does PyExprMutator even support this?

1 Like
#2

I guess this is a bug of PyExprMutator (as well as visitor). We once observed the endless recursion for visit_var_binding_ and had a fix here https://github.com/apache/tvm/pull/14754. It looks to me that we can fix the issue of visit_dataflow_block in a similar way.

1 Like
#3

CC @leshenj15 the original author of PyExprMutator :slight_smile:

1 Like
#4

Thank @kparzysz for sharing your experience! PR to fix this: https://github.com/apache/tvm/pull/15189

2 Likes
#5

Thank you @leshenj15 for the quick turnaround!

1 Like
#6

got similar issue under rewrite_call

from tvm.relax.dpl import *
from tvm.script import relax as R

def test_rewrite_simple():

    @R.function
    def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
        with R.dataflow():
            y = R.multiply(x, R.const(2, "float32"))
            R.output(y)
        return y


    x = wildcard()
    y = wildcard()
    pattern = is_op("relax.multiply")(x, y)
    
    def rewriter_passed_case_1(_, matchings):
        print(_)
        t = matchings[x]
        t = R.multiply(t, R.const(2, "float32"))
        return t
    
    def rewriter_passed_case_2(_, matchings):
        print(_)
        t = matchings[x]
        t = R.layout_transform(matchings[x], index_map=lambda i, j: (i, j))
        t = R.add(t, R.const(2, "float32"))
        return t
    
    def rewriter_issue_case(_, matchings):
        print(_)
        t = matchings[x]
        t = R.layout_transform(matchings[x], index_map=lambda i, j: (i, j))
        t = R.multiply(t, R.const(2, "float32"))
        return t
    rewritten = rewrite_call(pattern, rewriter_passed_case_1, main)
    print("passed case 1: ", rewritten)
    rewritten = rewrite_call(pattern, rewriter_passed_case_2, main)
    print("passed case 2: ", rewritten)
    rewritten = rewrite_call(pattern, rewriter_issue_case, main)
    print("passed issue case: ", rewritten)
    

test_rewrite_simple()

But the same logic passed when I modified it into PyExprMutator, do rewrite_call has the same problem?