[GitHub] [tvm] Lunderberg commented on a diff in pull request #15026: [Unity] Provide FuncStructInfo from `bb.emit_te`

["Lunderberg (via GitHub)" <gi...@apache.org>]

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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org