[GitHub] [tvm] abhikran-quic commented on a diff in pull request #15390: [Unity] Lowering of axis separator in Layout Transform

["abhikran-quic (via GitHub)" <gi...@apache.org>]

abhikran-quic commented on code in PR #15390:
URL: https://github.com/apache/tvm/pull/15390#discussion_r1274560492


##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -167,30 +170,33 @@ def _scatter_elements(bb: BlockBuilder, call: Call) -> Expr:
 
 @register_legalize("relax.layout_transform")
 def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
+    def te_layout_transform(data, name):
+        return te.compute(
+            data.shape,
+            data,
+            name=name,
+        )

Review Comment:
   IMHO, `gen_call_tir_inputs` will need a te function to be used (similar to how it's used in `bb.call_te`). This will help in generation of `call_args`, `output_sinfo`, `tir_vars` from a utility function.
   Another reason to use TE compute instead of TIR compute is the availability of dynamic shapes in TE.
   Eventually this helps to `transform_layout` and `set_axis_separator` directives to achieve the end goal.



-- 
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