Fix review comments

Author: abhikran-quic

python/tvm/relax/transform/legalize_ops/manipulate.py

@@ -182,8 +182,7 @@ def te_layout_transform(data, name):
         )
 
     def set_axis_sep(axis_sep: list, sch: tir.schedule, buffer_type: str):
-        if len(axis_sep) != 0:
-            sch.set_axis_separator(primfunc_name, (buffer_type, 0), axis_separators=axis_sep)
+        sch.set_axis_separator(primfunc_name, (buffer_type, 0), axis_separators=axis_sep)
 
     index_map: tvm.tir.IndexMap = call.attrs.index_map
     pad_value = call.attrs.pad_value
@@ -199,7 +198,7 @@ def set_axis_sep(axis_sep: list, sch: tir.schedule, buffer_type: str):
     input_axis_separators: tvm.tir.IndexMap.AXIS_SEPARATOR = call.attrs.input_axis_separators
 
     # Convert to list from array
-    axis_separators = list(map(lambda x: x.value, axis_separators))
+    axis_separators = [int(sep) for sep in axis_separators]
     primfunc_name = "te_layout_transform"
     _, padding_predicate = index_map.non_surjective_inverse(call.args[0].struct_info.shape)
     if not isinstance(padding_predicate, tvm.tir.expr.IntImm):
@@ -214,7 +213,7 @@ def set_axis_sep(axis_sep: list, sch: tir.schedule, buffer_type: str):
     sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
     set_axis_sep(axis_separators, sch, "write")
     if input_axis_separators is not None:
-        input_axis_separators = list(map(lambda x: x.value, input_axis_separators))
+        input_axis_separators = [int(sep) for sep in input_axis_separators]
         set_axis_sep(input_axis_separators, sch, "read")
     gvar = bb.add_func(sch.mod["main"], primfunc_name)
     output_shape = index_map.map_shape(list(call_args[0].struct_info.shape))

python/tvm/relax/transform/transform.py

@@ -24,6 +24,7 @@
 import numpy as np  # type: ignore
 
 import tvm.ir
+from tvm.ir.container import Array
 from tvm.relax import Expr, Var, StructInfo
 from tvm.relax.dpl import DFPattern
 from tvm.runtime import NDArray, Object
@@ -1309,7 +1310,9 @@ def AlterOpImpl(
             # Extract the index_map
             if isinstance(transform, Callable):
                 transform = IndexMap.from_func_with_separators(transform)[0]
-            elif isinstance(transform, tuple) and isinstance(transform[0], IndexMap):
+            elif (isinstance(transform, tuple) or isinstance(transform, Array)) and isinstance(
+                transform[0], IndexMap
+            ):
                 transform = transform[0]
             l.append(transform)
         op_buffer_transforms[operator_name] = l

src/relax/op/tensor/manipulate.h

@@ -72,7 +72,7 @@ Expr flatten(Expr x);
  */
 Expr layout_transform(Expr x, tir::IndexMap index_map, Optional<PrimValue> pad_value,
                       Optional<Array<IntImm>> axis_separators,
-                      Optional<Array<IntImm>> input_axis_separators);
+                      Optional<Array<IntImm>> input_axis_separators = NullOpt);
 
 /*!
  * \brief Permutes the dimensions of an array.

src/relax/transform/alter_op_impl.cc

@@ -128,7 +128,8 @@ class AlterOpImplMutator : public ExprMutator {
     const auto& replacement_func = op_impl_map_[op_kind];
 
     Array<IndexMap> buffer_transforms;
-    Optional<Array<Array<IntImm>>> axis_separators, input_axis_separators;
+    Optional<Array<Array<IntImm>>> axis_separators;
+    Optional<Array<Array<IntImm>>> input_axis_separators;
     if (op_buffer_transforms__.count(op_kind)) buffer_transforms = op_buffer_transforms__[op_kind];
     if (op_buffer_axis_separators__.count(op_kind))
       axis_separators = op_buffer_axis_separators__[op_kind];
@@ -293,7 +294,8 @@ class AlterOpImplMutator : public ExprMutator {
     Array<Expr> updated_inputs;
     int index = 0;
     for (const auto& input : inputs->fields) {
-      Array<IntImm> axis_separator, input_axis_separator;
+      Array<IntImm> axis_separator;
+      Array<IntImm> input_axis_separator;
       if (axis_separators.defined()) {
         Array<Array<IntImm>> axis_separators_value = axis_separators.value();
         axis_separator = axis_separators_value[index];

tests/python/relax/test_transform_alter_op_impl.py

@@ -609,8 +609,6 @@ def relax_some_op_replacement(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer(
             for ax0, ax1 in T.grid(4, 4):
                 with T.block("T_add"):
                     v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
-                    T.reads(arg0[v_ax0, v_ax1], arg1[v_ax0, v_ax1])
-                    T.writes(output0[v_ax0, v_ax1], output1[v_ax0, v_ax1])
                     output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1]
                     output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1]
 
@@ -633,8 +631,6 @@ def some_op_2d(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float3
         for ax0, ax1 in T.grid(4, 4):
             with T.block("T_add"):
                 v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
-                T.reads(arg0[v_ax0, v_ax1], arg1[v_ax0, v_ax1])
-                T.writes(output0[v_ax0, v_ax1], output1[v_ax0, v_ax1])
                 output0[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1]
                 output1[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] - arg1[v_ax0, v_ax1]
     # fmt: on