[Bug] left-hand side variable type not deducted automatically with call_tir.

[ganler]

from __future__ import annotations
from tvm.script import relax as R, tir as T

@R.function
def bug(x: Tensor((32, 32), "float32")) -> Tensor:
    with R.dataflow():
        lv0 = R.call_tir(my_relu, (x,), (32, 32), dtype="float32")
        R.output(lv0)
    return lv0

bug.body.blocks[0].bindings[0].var.checked_type

cc: @yongwww @YuchenJin

[yongwww]

thanks for filing this bug! the checked_type of var lv0 should be filled. Now we can manually add the annotation for the var as a workaround. Will fix it in the new parser.

[YuchenJin]

The Normalization pass only applies when parsing an IRModule: https://github.com/tlc-pack/relax/blob/relax/python/tvm/script/parser.py#L1386, so if you wrap the relax.Function in an IRModule, it will be normalized (type/shape deduced). I think when parsing a standalone relax.Function, we should also do type/shape auto deduction. cc @yongwww.

This following code works:

from __future__ import annotations
from tvm.script import relax as R, tir as T

@tvm.script.ir_module
class MyModule:
    @R.function
    def func(x: Tensor((32, 32), "float32")) -> Tensor:
        with R.dataflow():
            lv0 = R.call_tir(my_relu, (x,), (32, 32), dtype="float32")
            R.output(lv0)
        return lv0


print(MyModule["func"].body.blocks[0].bindings[0].var.checked_type)

[ganler]

@YuchenJin Thanks for the clarification! Yes, I strongly agree that we should do such deduction function-wise. It will benefit the IR rewriting a lot (where users need to create a function first and then replace it on an IRModule).

[tqchen]

Should be fixed in the latest version