Relax AlterOpLayout with AXIS_SEPARATOR

Questions
#1

Hello,

I was going through relax AlterOpLayout pass by running the tests in tests/python/relax/test_transform_alter_op_impl.py .

When I try to use an index_map with AXIS_SEPARATOR, I’m not seeing any change in the IR. For example, I modified test_single_output function

diff --git a/tests/python/relax/test_transform_alter_op_impl.py b/tests/python/relax/test_transform_alter_op_impl.py
index 77e2d4e35..426c080f8 100644
--- a/tests/python/relax/test_transform_alter_op_impl.py
+++ b/tests/python/relax/test_transform_alter_op_impl.py
@@ -85,7 +85,9 @@ def test_single_output():
                 T.writes(output[v_ax0, v_ax1])
                 output[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1]
     # fmt: on
-    index_map = lambda i: (i // 4, i % 4)
+    # index_map = lambda i: (i // 4, i % 4)
+    from tvm.tir import IndexMap
+    index_map, axis_sep = IndexMap.from_func_with_separators((lambda i: (i // 4, IndexMap.AXIS_SEPARATOR, i % 4)))
     _check(

Could you please suggest if support for AXIS_SEPARATOR is present in AlterOpLayout ? If yes, could you please suggest an example on how it can be used ?

Thanks.

cc @psrivas2 @jverma @sanirudh

AXIS_SEPARATOR in Relax Tensor
#2

No, AXIS SEPARATOR is not supported yet in this pass. Although I don’t think there is any fundamental technical issue in supporting it.

1 Like
#3

Thank you @psrivas2 for your response! Sorry for coming back late on this but I was going through the source code to understand the pass better. I’m planning to add support for AXIS_SEPARATOR and have proposed below mentioned changes after going through the code. Can you please review them and share your feedback?

  1. To update inputs with index_map, TransformLayout is invoked here. I plan to add AXIS_SEPARATOR in LayoutTransformAttrs to be saved for respective buffers. We can get the AXIS_SEPARATOR from user in the same way we are getting index_map.
  2. Similarly for output buffers, TransformLayout gets invoked in TransformLayoutInverse function and AXIS_SEPARATOR support can be added in the same manner as point 1 above.
  3. The replacement_primfunc passed to AlterOpImpl pass by user will already have support for AXIS_SEPARATOR because it’s a user configurable function.

Please let me know if any point above needs any further clarification and I will be happy to answer. As always, thank you for your help! :slight_smile:

#4

The following PR is raised to enable this change: https://github.com/apache/tvm/pull/15315