From d70523edb04eeacb1c5ef34f3dc9792da231f805 Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Mon, 7 Sep 2026 05:21:18 -0700 Subject: [PATCH 1/2] [Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder An explicit identity permute_dims does not transpose the inner matmul. Skip the transpose-specific reassociation when its axes preserve every dimension. The focused Relax regression covers the IR distinction from a real transpose and the CPU/LLVM integer result reported in #20277. Fixes #20277 Generated-by: OpenAI Codex (GPT-6) --- src/relax/transform/adjust_matmul_order.cc | 25 ++++++++- .../test_transform_adjust_matmul_order.py | 51 +++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/src/relax/transform/adjust_matmul_order.cc b/src/relax/transform/adjust_matmul_order.cc index 1f0cf8728156..ac3900ab43bd 100644 --- a/src/relax/transform/adjust_matmul_order.cc +++ b/src/relax/transform/adjust_matmul_order.cc @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -54,6 +55,22 @@ PrimExpr ProductDims(const ffi::Array& dims) { return product; } +bool IsIdentityPermuteDims(const Expr& expr) { + const auto* call = expr.as(); + if (call == nullptr) return false; + + const auto* attrs = call->attrs.as(); + if (attrs == nullptr || !attrs->axes.has_value()) return false; + + const auto& axes = attrs->axes.value(); + for (size_t i = 0; i < axes.size(); ++i) { + int64_t axis = axes[i]; + if (axis < 0) axis += axes.size(); + if (axis != static_cast(i)) return false; + } + return true; +} + ffi::Optional> InferBatchedMatmulBroadcastPrefix( arith::AnalyzerObj* analyzer, const ffi::Array& x1, const ffi::Array& x2) { auto infer_result = InferBinaryBroadcastShape(analyzer, x1, x2); @@ -89,8 +106,10 @@ std::tuple)>> auto pat_matmul_on_lhs = pat_matmul(pat_matmul(pat_a, pat_b), pat_c); auto pat_matmul_on_rhs = pat_matmul(pat_a, pat_matmul(pat_b, pat_c)); - auto pat_permuted_matmul_on_lhs = pat_matmul(pat_permute_dims(pat_matmul(pat_b, pat_a)), pat_c); - auto pat_permuted_matmul_on_rhs = pat_matmul(pat_a, pat_permute_dims(pat_matmul(pat_c, pat_b))); + auto pat_permuted_inner_matmul_on_lhs = pat_permute_dims(pat_matmul(pat_b, pat_a)); + auto pat_permuted_inner_matmul_on_rhs = pat_permute_dims(pat_matmul(pat_c, pat_b)); + auto pat_permuted_matmul_on_lhs = pat_matmul(pat_permuted_inner_matmul_on_lhs, pat_c); + auto pat_permuted_matmul_on_rhs = pat_matmul(pat_a, pat_permuted_inner_matmul_on_rhs); auto pat = pat_matmul_on_lhs | pat_matmul_on_rhs | pat_permuted_matmul_on_lhs | pat_permuted_matmul_on_rhs; @@ -194,12 +213,14 @@ std::tuple)>> }; if (matches.count(pat_permuted_matmul_on_lhs)) { + if (IsIdentityPermuteDims(matches[pat_permuted_inner_matmul_on_lhs])) return expr; if (shape_a.size() < 2 || shape_b.size() < 2) return expr; expr_a = permute_last_two_dims(expr_a); expr_b = permute_last_two_dims(expr_b); transpose_shape_last_two_dims(shape_a); transpose_shape_last_two_dims(shape_b); } else if (matches.count(pat_permuted_matmul_on_rhs)) { + if (IsIdentityPermuteDims(matches[pat_permuted_inner_matmul_on_rhs])) return expr; if (shape_b.size() < 2 || shape_c.size() < 2) return expr; expr_b = permute_last_two_dims(expr_b); expr_c = permute_last_two_dims(expr_c); diff --git a/tests/python/relax/test_transform_adjust_matmul_order.py b/tests/python/relax/test_transform_adjust_matmul_order.py index b5f3155248dc..d824b6a5ff07 100644 --- a/tests/python/relax/test_transform_adjust_matmul_order.py +++ b/tests/python/relax/test_transform_adjust_matmul_order.py @@ -564,6 +564,30 @@ def main( return x +class TestRHSPermuteDimsIdentity(Base): + """Do not treat an explicit identity permutation as a transpose. + + `TestRHSPermuteDims` above covers the real transpose case. Here, the + explicit axes preserve the inner matmul's order, so reassociation must not + insert transposes for its operands. + """ + + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor([2]), + A: R.Tensor([2, 1]), + B: R.Tensor([1, 2]), + ) -> R.Tensor([2]): + linear_weight: R.Tensor([2, 2]) = R.matmul(A, B) + matmul_weight: R.Tensor([2, 2]) = R.permute_dims(linear_weight, axes=[0, 1]) + out: R.Tensor([2]) = R.matmul(x, matmul_weight) + return out + + Expected = Before + + class TestRHSPermuteDimsDynamic(Base): """Prefer (x*A)*B instead of x*(A*B) @@ -852,6 +876,33 @@ def test_attention_block_numerics(self, batch, seq, dim): tvm.testing.assert_allclose(out_after, ref, rtol=1e-3, atol=1e-3) tvm.testing.assert_allclose(out_before, out_after, rtol=1e-5, atol=1e-5) + def test_identity_permute_dims_numerics(self): + bb = relax.BlockBuilder() + x = relax.Var("x", relax.TensorType((2,), "int32")) + A = relax.Var("A", relax.TensorType((2, 1), "int32")) + B = relax.Var("B", relax.TensorType((1, 2), "int32")) + with bb.function("main", [x, A, B]): + with bb.dataflow(): + linear_weight = bb.emit(relax.op.matmul(A, B)) + identity_weight = bb.emit(relax.op.permute_dims(linear_weight, axes=[0, 1])) + out = bb.emit_output(relax.op.matmul(x, identity_weight)) + bb.emit_func_output(out) + mod = bb.finalize() + mod_opt = relax.transform.AdjustMatmulOrder()(mod) + + inputs = [ + np.array([-1, 3], dtype="int32"), + np.array([[2], [-4]], dtype="int32"), + np.array([[1, -3]], dtype="int32"), + ] + expected = np.array([-14, 42], dtype="int32") + + out_before = self._run_relax_main(mod, inputs) + out_after = self._run_relax_main(mod_opt, inputs) + + np.testing.assert_array_equal(out_before, expected) + np.testing.assert_array_equal(out_after, expected) + if __name__ == "__main__": tvm.testing.main() From 27ca382079d16dccbd18672ac680cb764b5a39bc Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Mon, 7 Sep 2026 22:17:36 -0700 Subject: [PATCH 2/2] [Fix][Relax] Restrict matmul rewrite to matrix-axis swaps --- src/relax/transform/adjust_matmul_order.cc | 23 ++++++++--- .../test_transform_adjust_matmul_order.py | 38 +++++++++++++++++++ 2 files changed, 55 insertions(+), 6 deletions(-) diff --git a/src/relax/transform/adjust_matmul_order.cc b/src/relax/transform/adjust_matmul_order.cc index ac3900ab43bd..615bb6fd18a0 100644 --- a/src/relax/transform/adjust_matmul_order.cc +++ b/src/relax/transform/adjust_matmul_order.cc @@ -55,18 +55,29 @@ PrimExpr ProductDims(const ffi::Array& dims) { return product; } -bool IsIdentityPermuteDims(const Expr& expr) { +bool IsLastTwoDimsSwap(const Expr& expr) { const auto* call = expr.as(); if (call == nullptr) return false; const auto* attrs = call->attrs.as(); - if (attrs == nullptr || !attrs->axes.has_value()) return false; + const auto* input_type = GetTypeAs(call->args[0]); + if (attrs == nullptr || input_type == nullptr || input_type->ndim < 2) return false; + + size_t ndim = input_type->ndim; + if (!attrs->axes.has_value()) return ndim == 2; const auto& axes = attrs->axes.value(); + if (axes.size() != ndim) return false; for (size_t i = 0; i < axes.size(); ++i) { int64_t axis = axes[i]; - if (axis < 0) axis += axes.size(); - if (axis != static_cast(i)) return false; + if (axis < 0) axis += ndim; + size_t expected = i; + if (i == ndim - 2) { + expected = ndim - 1; + } else if (i == ndim - 1) { + expected = ndim - 2; + } + if (axis != static_cast(expected)) return false; } return true; } @@ -213,14 +224,14 @@ std::tuple)>> }; if (matches.count(pat_permuted_matmul_on_lhs)) { - if (IsIdentityPermuteDims(matches[pat_permuted_inner_matmul_on_lhs])) return expr; + if (!IsLastTwoDimsSwap(matches[pat_permuted_inner_matmul_on_lhs])) return expr; if (shape_a.size() < 2 || shape_b.size() < 2) return expr; expr_a = permute_last_two_dims(expr_a); expr_b = permute_last_two_dims(expr_b); transpose_shape_last_two_dims(shape_a); transpose_shape_last_two_dims(shape_b); } else if (matches.count(pat_permuted_matmul_on_rhs)) { - if (IsIdentityPermuteDims(matches[pat_permuted_inner_matmul_on_rhs])) return expr; + if (!IsLastTwoDimsSwap(matches[pat_permuted_inner_matmul_on_rhs])) return expr; if (shape_b.size() < 2 || shape_c.size() < 2) return expr; expr_b = permute_last_two_dims(expr_b); expr_c = permute_last_two_dims(expr_c); diff --git a/tests/python/relax/test_transform_adjust_matmul_order.py b/tests/python/relax/test_transform_adjust_matmul_order.py index d824b6a5ff07..eb886fd4d065 100644 --- a/tests/python/relax/test_transform_adjust_matmul_order.py +++ b/tests/python/relax/test_transform_adjust_matmul_order.py @@ -588,6 +588,44 @@ def main( Expected = Before +class TestRHSPermuteDimsNonMatrixAxes(Base): + """Do not rewrite permutations that move a batch axis.""" + + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor([4, 1, 4]), + A: R.Tensor([4, 4, 1]), + B: R.Tensor([4, 1, 4]), + ) -> R.Tensor([4, 1, 4]): + weight: R.Tensor([4, 4, 4]) = R.matmul(A, B) + permuted: R.Tensor([4, 4, 4]) = R.permute_dims(weight, axes=[1, 0, 2]) + out: R.Tensor([4, 1, 4]) = R.matmul(x, permuted) + return out + + Expected = Before + + +class TestLHSPermuteDimsNonMatrixAxes(Base): + """Apply the same batch-axis guard to the left-hand pattern.""" + + @I.ir_module + class Before: + @R.function + def main( + A: R.Tensor([4, 4, 1]), + B: R.Tensor([4, 1, 4]), + x: R.Tensor([4, 4, 1]), + ) -> R.Tensor([4, 4, 1]): + weight: R.Tensor([4, 4, 4]) = R.matmul(A, B) + permuted: R.Tensor([4, 4, 4]) = R.permute_dims(weight, axes=[1, 0, 2]) + out: R.Tensor([4, 4, 1]) = R.matmul(permuted, x) + return out + + Expected = Before + + class TestRHSPermuteDimsDynamic(Base): """Prefer (x*A)*B instead of x*(A*B)