Lunderberg commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1259935791
##########
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:
Summarizing our conversation from this morning:
* Shape propagation through `bb.emit_te` only works during the initial
construction of a Relax module, when the `relax::Call("relax.call_tir",...)`
node is explicitly typed. Re-derivation of the output shape is not
implemented, and so the shape information can be lost during lowering if the
arguments to `call_tir` change.
* Annotating a PrimFunc with `FuncStructInfo` to represent the output of
`call_tir` (i.e. pure function, tensor output) wouldn't be accurate, and could
cause confusion in the future.
* Annotating a PrimFunc with `FuncStructInfo` to represent the PrimFunc
itself (i.e. impure function, mutates arguments) would be accurate, but
insufficient for `call_tir` to propagate shapes, as input/output shapes are
mixed.
* Would be useful to have a purity annotations for each parameter, dividing
arguments into read-only, output, and mutate-in-place. This would allow a
PrimFunc to be accurately annotated, and would be sufficient for `call_tir` to
identify outputs for shape propagation.
--
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]
