Printer: prior to this PR, when printing call_tir, the printer always explicitly prints "dtype=", while doesn’t print "tir_vars=" when there is. This will make the printed script looks like
y = R.call_tir(copy, (x,), ((n * 2,)), dtype="float32", (n,))
which doesn’t conform to Python syntax.
call_tir’s handling on TIR vars: the TIR vars will be parsed as a list or a tuple of Python, while on the call_tir side the parsed tuple/list will be treated as a Expr and passed to C++ side through FFI. This will incur FFI error like
error: Check failed: (!checked_type.defined()) is false: Expected RelayExpr, but got Array
Therefore, this PR fixes both sides. On printer side, we now print tir_vars after dtype and explicitly print "tir_vars=". On call_tir side, we wrap the input tir_vars to ShapeExpr if we find it a tuple or list. One regression test that covers both issues is provided.
This PR patches both the TVMScript parser and printer on handling the
tir_varsparameter ofcall_tir.When a TIR function contains symbolic shapes, its signature can have trailing TIR vars in the param list. In such cases, when using
call_tirto call into the TIR function, we should provide the instances of those TIR vars through the paramtir_vars.tir_varsis the last parameter ofcall_tir, one behinddtype. https://github.com/tlc-pack/relax/blob/697675bf8f137488e0f38237b046ce671cdbad9f/python/tvm/relax/op/base.py#L45-L51The specific issues fixed in this PR is:
call_tir, the printer always explicitly prints"dtype=", while doesn’t print"tir_vars="when there is. This will make the printed script looks likewhich doesn’t conform to Python syntax.
call_tir’s handling on TIR vars: the TIR vars will be parsed as a list or a tuple of Python, while on thecall_tirside the parsed tuple/list will be treated as a Expr and passed to C++ side through FFI. This will incur FFI error likeTherefore, this PR fixes both sides. On printer side, we now print
tir_varsafterdtypeand explicitly print"tir_vars=". Oncall_tirside, we wrap the inputtir_varsto ShapeExpr if we find it a tuple or list. One regression test that covers both issues is provided.cc @YuchenJin