[Fix][Relax] Preserve stack axes in ConvertLayout

src/relax/op/tensor/manipulate.cc

@@ -1775,21 +1775,34 @@ InferLayoutOutput InferLayoutStack(
 
   int n_tensor = nlayout.NestedArray().size();
   LayoutDecision layout = nlayout.NestedArray()[0].LeafValue();
+  int output_ndim = GetTypeAs<TensorTypeObj>(call)->ndim;
+  SLayout initial_input_layout = InitialLayout(output_ndim - 1);
+  if (layout->layout.has_value() &&
+      layout->layout.value().ndim() != layout->layout.value().ndim_primal()) {
+    // As with expand_dims, rank insertion uses an unblocked input layout.
+    layout = LayoutDecision(initial_input_layout);
+  }
+
+  int64_t axis = attrs->axis.value_or(0);
+  if (axis < 0) axis += output_ndim;
+  TVM_FFI_ICHECK(0 <= axis && axis < output_ndim);
+
+  // Relabel the input axes using the output's axis names, then apply the
+  // input permutation before inserting the new axis at its physical position.
+  std::string input_axes = InitialLayout(output_ndim).name();
+  char stack_axis = input_axes.at(axis);
+  input_axes.erase(axis, 1);
+  std::string output_axes = TransposeStrLike(input_axes, initial_input_layout, layout->layout);
+  output_axes.insert(axis, 1, stack_axis);
+
   ffi::Array<NLayout> input_layouts, output_layouts;
   for (int i = 0; i < n_tensor; ++i) {
     input_layouts.push_back(layout);
   }
-
-  // For stack, we need to adjust the output layout by inserting a new axis
-  std::string layout_str =
-      (layout->layout.has_value() ? layout->layout.value().name() : "__undef__");
-  int axis = attrs->axis.has_value() ? static_cast<int>(attrs->axis.value()) : 0;
-  layout_str.insert(static_cast<size_t>(axis), "S");  // Add stack dimension
-  SLayout output_layout = SLayout(layout_str);
-  output_layouts.push_back(LayoutDecision(output_layout));
+  output_layouts.push_back(LayoutDecision(SLayout::Create(output_axes)));
 
   ffi::ObjectPtr<StackAttrs> new_attrs = ffi::make_object<StackAttrs>(*attrs);
-  new_attrs->axis = static_cast<int64_t>(FindAxis(layout->layout.value(), axis));
+  new_attrs->axis = axis;
   return InferLayoutOutput({NLayout(input_layouts)}, output_layouts, Attrs(new_attrs));
 }
 

tests/python/relax/test_transform_convert_layout.py

@@ -16,9 +16,13 @@
 # under the License.
 # ruff: noqa: E501, F841
 
+import numpy as np
+import pytest
+
 import tvm
 import tvm_tirx
 import tvm_tirx.testing
+from tvm import relax
 from tvm.relax.transform import ConvertLayout, Normalize
 from tvm_tirx.script.parser import ir as I
 from tvm_tirx.script.parser import relax as R
@@ -863,6 +867,76 @@ def main(
     verify(Input, Expected)
 
 
+@pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="LLVM is required")
+@pytest.mark.parametrize(
+    "shape,axis", [((), -1), ((2, 3), -3), ((2, 3), -2), ((2, 3), -1), ((2, 3), 2)]
+)
+def test_stack_layout_axes(shape, axis):
+    x = relax.Var("x", relax.TensorType(shape, "float32"))
+    y = relax.Var("y", relax.TensorType(shape, "float32"))
+    bb = relax.BlockBuilder()
+    with bb.function("main", [x, y]):
+        with bb.dataflow():
+            stacked = bb.emit_output(relax.op.stack((x, y), axis=axis))
+        bb.emit_func_output(stacked)
+
+    mod = ConvertLayout({})(bb.get())
+    vm = relax.VirtualMachine(tvm.compile(mod, target="llvm"), tvm_tirx.cpu())
+    a = np.arange(max(1, int(np.prod(shape))), dtype=np.float32).reshape(shape)
+    b = np.asarray(a + 10, dtype=np.float32)
+    result = vm["main"](tvm_tirx.runtime.tensor(a), tvm_tirx.runtime.tensor(b))
+    np.testing.assert_array_equal(result.numpy(), np.stack((a, b), axis=axis), strict=True)
+
+
+def test_stack_after_conv2d_layout_conversion():
+    @I.ir_module
+    class Input:
+        @R.function
+        def main(x: R.Tensor((1, 4, 4, 5), "float32"), w: R.Tensor((4, 4, 2, 2), "float32")):
+            with R.dataflow():
+                a = R.nn.conv2d(x, w, out_dtype="float32")
+                b = R.negative(a)
+                out = R.stack((a, b), axis=-3)
+                R.output(out)
+            return out
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor((1, 4, 4, 5), "float32"), w: R.Tensor((4, 4, 2, 2), "float32")):
+            with R.dataflow():
+                image = R.permute_dims(x, axes=[0, 2, 3, 1])
+                weight = R.permute_dims(w, axes=[0, 2, 3, 1])
+                a = R.nn.conv2d(
+                    image,
+                    weight,
+                    data_layout="NHWC",
+                    kernel_layout="OHWI",
+                    out_layout="NHWC",
+                    out_dtype="float32",
+                )
+                b = R.negative(a)
+                stacked = R.stack((a, b), axis=2)
+                out = R.permute_dims(stacked, axes=[0, 4, 2, 1, 3])
+                R.output(out)
+            return out
+
+    verify(Input, Expected)
+
+    blocked = ConvertLayout({"relax.nn.conv2d": ["NCHW4c", "OIHW4o"]})(Input)
+    stack_calls = []
+
+    def collect_stack(expr):
+        if isinstance(expr, relax.Call) and expr.op == tvm_tirx.ir.Op.get("relax.stack"):
+            stack_calls.append(expr)
+
+    relax.analysis.post_order_visit(blocked["main"].body, collect_stack)
+    (stack,) = stack_calls
+    assert stack.attrs.axis == 2
+    assert [arg.ty.ndim for arg in stack.args[0].fields] == [4, 4]
+    assert [dim.value for dim in stack.ty.shape.values] == [1, 4, 2, 3, 4]
+
+
 def test_expand_dims_scalar():
     @I.ir_module
     class Input: