Lunderberg commented on code in PR #15026:
URL: https://github.com/apache/tvm/pull/15026#discussion_r1221590440
##########
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 point regarding the mutation. Thinking on it, I'm also not sure what
the best `FuncStructInfo` would be. It could reasonably be either
`FuncStructInfo(params = [*input_tensors, *output_tensors]`, ret=None)`, which
matches the TIR function's signature, or `FuncStructInfo(params=input_tensors,
ret=relax.Tuple(output_tensors))`, which matches the exposed semantics in Relax.
The original issue I was running into was that the result of `bb.emit_te`
doesn't preserve the output struct information across mutations. If I have a
TE function that accepts dynamic shapes, but which is called using static
shapes, then the return type of the `relax::Call` should be an inferred static
shape. This works during the first usage of BlockBuilder, when a user is
calling `bb.emit_te` directly. However, when the module is mutated, any
mutation of the call node relies on the `relax::Normalizer` to regenerate the
output struct info, and it doesn't have enough information to do so.
--
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]
