slyubomirsky commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1220408125
##########
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:
Yeah, permitting TIR calls outside of `call_tir` is something we're trying
to figure out with respect to [phase
ordering](https://discuss.tvm.apache.org/t/unity-dealing-with-phase-ordering/14709)
in Relax (see thread). I was under the impression that we did not want direct
calls to `PrimFunc`s in the front end, so we should clarify that (we could put
this on the agenda for a community meeting).
FWIW, I don't think it would be hard to give `PrimFunc`s `FuncStructInfo`,
but there is the issue that they mutate their arguments, so they should be
treated as impure (except when called via `call_tir`).
--
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]
