From 7cfac8c968f2b90fee09afc222fdbc6a350cc136 Mon Sep 17 00:00:00 2001 From: DiadkoShmek <207617272+DiadkoShmek@users.noreply.github.com> Date: Sun, 15 Nov 2026 21:37:36 +0200 Subject: [PATCH] [Fix][Relax] Preserve stack axes in ConvertLayout --- src/relax/op/tensor/manipulate.cc | 31 +++++--- .../relax/test_transform_convert_layout.py | 74 +++++++++++++++++++ 2 files changed, 96 insertions(+), 9 deletions(-) diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index fb3141b14721..a638c5199c82 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/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(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 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(attrs->axis.value()) : 0; - layout_str.insert(static_cast(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 new_attrs = ffi::make_object(*attrs); - new_attrs->axis = static_cast(FindAxis(layout->layout.value(), axis)); + new_attrs->axis = axis; return InferLayoutOutput({NLayout(input_layouts)}, output_layouts, Attrs(new_attrs)); } diff --git a/tests/python/relax/test_transform_convert_layout.py b/tests/python/relax/test_transform_convert_layout.py index 5aa5cfc40a3e..bb48297ccfad 100644 --- a/tests/python/relax/test_transform_convert_layout.py +++ b/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: