This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 6583b2b44a [Fix][Relax] Preserve identity permute_dims in
AdjustMatmulOrder (#20287)
6583b2b44a is described below
commit 6583b2b44a29b04871b7ff2575e8bfd84966e0a0
Author: Yifan Chen <[email protected]>
AuthorDate: Tue Sep 8 08:58:05 2026 -0700
[Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder (#20287)
## Summary
`AdjustMatmulOrder` matched any `relax.permute_dims` around an inner
matmul as though it were a transpose. An explicit identity axis order,
such as `[0, 1]`, preserves the inner result and must not trigger that
transpose-specific reassociation.
This change checks explicit axes before the transpose rewrite and leaves
identity permutations unchanged. It preserves the existing behavior for
genuine transposes.
Fixes #20277.
## Validation
- `python -m pytest
tests/python/relax/test_transform_adjust_matmul_order.py -q` — 27 passed
with a local CPU/LLVM build.
- The new integer LLVM regression returns `[-14, 42]` after
`AdjustMatmulOrder`; the previous implementation returns `[-20, 40]`.
- `pre-commit run --files src/relax/transform/adjust_matmul_order.cc
tests/python/relax/test_transform_adjust_matmul_order.py`
Generated-by: OpenAI Codex (GPT-6)
---
src/relax/transform/adjust_matmul_order.cc | 36 ++++++++-
.../relax/test_transform_adjust_matmul_order.py | 89 ++++++++++++++++++++++
2 files changed, 123 insertions(+), 2 deletions(-)
diff --git a/src/relax/transform/adjust_matmul_order.cc
b/src/relax/transform/adjust_matmul_order.cc
index 1f0cf87281..615bb6fd18 100644
--- a/src/relax/transform/adjust_matmul_order.cc
+++ b/src/relax/transform/adjust_matmul_order.cc
@@ -24,6 +24,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/manipulate.h>
#include <tvm/relax/dataflow_matcher.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/expr_functor.h>
@@ -54,6 +55,33 @@ PrimExpr ProductDims(const ffi::Array<PrimExpr>& dims) {
return product;
}
+bool IsLastTwoDimsSwap(const Expr& expr) {
+ const auto* call = expr.as<CallNode>();
+ if (call == nullptr) return false;
+
+ const auto* attrs = call->attrs.as<PermuteDimsAttrs>();
+ const auto* input_type = GetTypeAs<TensorTypeNode>(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 += ndim;
+ size_t expected = i;
+ if (i == ndim - 2) {
+ expected = ndim - 1;
+ } else if (i == ndim - 1) {
+ expected = ndim - 2;
+ }
+ if (axis != static_cast<int64_t>(expected)) return false;
+ }
+ return true;
+}
+
ffi::Optional<ffi::Array<PrimExpr>> InferBatchedMatmulBroadcastPrefix(
arith::AnalyzerObj* analyzer, const ffi::Array<PrimExpr>& x1, const
ffi::Array<PrimExpr>& x2) {
auto infer_result = InferBinaryBroadcastShape(analyzer, x1, x2);
@@ -89,8 +117,10 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
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 +224,14 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
};
if (matches.count(pat_permuted_matmul_on_lhs)) {
+ 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 (!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 b5f3155248..eb886fd4d0 100644
--- a/tests/python/relax/test_transform_adjust_matmul_order.py
+++ b/tests/python/relax/test_transform_adjust_matmul_order.py
@@ -564,6 +564,68 @@ class TestRHSPermuteDims(Base):
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 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)
@@ -852,6 +914,33 @@ class TestAdjustMatmulOrderAttentionBlock:
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()
