[Fix][Relax] Preserve identity permute_dims in AdjustMatmulOrder

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);

tests/python/relax/test_transform_adjust_matmul_order.py

@@ -564,6 +564,68 @@ 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 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 @@ 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()