Lunderberg commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1220384886
##########
python/tvm/relax/utils.py:
##########
@@ -455,10 +462,14 @@ def _shape_with_old_tir_var(
# with old set of variables.
tir_var_inverse_map = {v: k for k, v in tir_var_map.items()}
- output_sinfo = [
- TensorStructInfo(_shape_with_old_tir_var(out.shape,
tir_var_inverse_map), out.dtype)
- for out in outs
- ]
+ def te_to_sinfo(arg):
+ return TensorStructInfo(_shape_with_old_tir_var(arg.shape,
tir_var_inverse_map), arg.dtype)
+
+ input_sinfo = [te_to_sinfo(arg) for arg in te_args]
+ output_sinfo = [te_to_sinfo(out) for out in outs]
+
+ primfunc_sinfo = FuncStructInfo([*input_sinfo, *output_sinfo],
PrimStructInfo("void"))
+ _update_struct_info(tir_func, primfunc_sinfo)
Review Comment:
Good question, and I had assumed this was intended, but would be interested
to hear on it. I had mostly assumed that a Relax function and a TIR PrimFunc
should expose the same information, so long as they have the same convention.
That is, since the callsite has no distinction between a `GlobalVar`
representing a `relax::Function` or a `tir::PrimFunc`, it seemed that the
`struct_info_` would depend only on the call sequence, and not the
implementation dialect.
Regarding `call_tir`, I think the Relax-to-TIR calls are not restricted to
the `R.call_tir` built-in, because the `LowerCallTIR` pass can output a
`relax::CallNode` with a `GlobalVar` as its operation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
