[Bug][Relax] AdjustMatmulOrder treats an identity permute_dims (axes=[0,1]) as a transpose and produces a wrong result

[lackjava-creator]

Environment

• TVM version: release 0.26.0 and main @ 48242ec33403f2b6e4fac6e763ca7a683fb9d5df (2026-09-03, reports 0.26.dev0). Both verified 2026-09-06.

• Build / install: venv wheel for 0.26.0; source build for main, LLVM 15.0.7

• OS / Python: Ubuntu 24.04.3 LTS, x86_64, Python 3.10

• Reference implementation: onnxruntime 1.23.2 with ORT_DISABLE_ALL, cross-checked against numpy

• Target: llvm (CPU)

• Pipeline coverage (relax.get_pipeline("default_build") performs no operator fusion, so all three are reported separately). Verified 2026-09-06 on 0.26.0 and main @ 48242ec, llvm:

  • get_pipeline("default_build") (no FuseOps/FuseTIR) — correct (AdjustMatmulOrder is not in it)
  • get_pipeline("zero") (fuses) — correct (the pass is not in it either)
  • AdjustMatmulOrder applied explicitly, then get_pipeline("zero") — wrong ([[-20, 40]])
  • forced FuseOps + FuseTIR without the pass — not tested; irrelevant, since the defect is introduced by the pass itself

Minimal reproducer

a @ Transpose(b @ c, perm=[0,1]) in int32, where perm=[0,1] is the identity permutation on a rank-2 tensor. model_adjmm_identity_permute.onnx + run.py attached.

import numpy as np, onnx, onnxruntime as ort, tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
M = onnx.load("model_adjmm_identity_permute.onnx")
FEED = {"a": np.array([[-1, 3]], np.int32), "b": np.array([[2], [-4]], np.int32), "c": np.array([[1, -3]], np.int32)}
SD = {i.name: [d.dim_value for d in i.type.tensor_type.shape.dim] for i in M.graph.input}
print("numpy                 :", (FEED["a"] @ (FEED["b"] @ FEED["c"])).tolist())
def run(extra=None):
    mod = from_onnx(M, shape_dict=SD, keep_params_in_input=False)
    with tvm.transform.PassContext(opt_level=3):
        if extra is not None: mod = extra()(mod)
        ex = tvm.compile(relax.get_pipeline("zero")(mod), target="llvm")
    out = relax.VirtualMachine(ex, tvm.cpu())["main"](*[tvm.runtime.tensor(np.ascontiguousarray(FEED[n]), tvm.cpu()) for n in SD])
    return np.asarray(out.numpy()).tolist()
print("TVM, stock pipeline   :", run())
print("TVM + AdjustMatmulOrder:", run(relax.transform.AdjustMatmulOrder))

Expected vs actual

perm=[0,1] on a rank-2 tensor is the identity, so the graph computes a @ (b @ c).

	result
numpy a @ (b @ c)	[[-14, 42]]
onnxruntime, optimisations disabled	[[-14, 42]]
TVM get_pipeline("zero") (opt 0)	[[-14, 42]]
TVM get_pipeline("default_build") (opt 3)	[[-14, 42]]
TVM + relax.transform.AdjustMatmulOrder()	[[-20, 40]]

• Difference: max abs 6 on a 2-element int32 tensor, 2/2 elements wrong. Exact integers, no tolerance involved.
• [[-20, 40]] is a @ (b @ c) computed as if the permute_dims were a real transpose — i.e. the pass reassociated the chain while also honouring a transpose that does not exist.
• Identical on 0.26.0 and on main @ 48242ec.

Root cause (if known)

relax.transform.AdjustMatmulOrder folds a permute_dims into the matmul reassociation without first checking whether the permutation is the identity. Merged PR #16589 ("[Unity] Check for transpose and dynamic shape in AdjustMatmulOrder", 2024-02-16) added a transpose check to this pass; that check does not exclude the identity permutation, and the defect is still present on main two and a half years later.

Why this is a bug (not tolerance / not undefined behaviour)

Integer matmul is exact; both operands and result are small int32 values with no overflow. permute_dims(axes=[0,1]) on a rank-2 tensor is the identity by definition, so any transformation that treats it as a transpose is unsound. The pass changes the value the program computes.

Scope caveat, stated up front in the report: AdjustMatmulOrder is not part of relax.get_pipeline("zero") or "default_build", nor of tvm.relax.backend.cpu_generic.get_default_pipeline — checked on 0.26.0. So the stock tvm.compile(mod, target="llvm") path is unaffected. It is a public transform in tvm.relax.transform that a user or a custom pipeline can apply, and applying it silently changes results.

How found

Found by EquiAutomaton (equivalence-graph differential testing against onnxruntime at ORT_DISABLE_ALL), localised by applying the pass in isolation and by comparing get_pipeline("zero") / default_build against the pass-applied build.

The same blind spot — a fusion treating an identity permutation as a real transpose — is one we reported in another compiler as microsoft/onnxruntime#32228 (GemmTransposeFusion folds an identity Transpose with perm=[0,1] into Gemm). Worth citing as prior art for the class of defect.

Reproducer archive

TVM-ADJMM-reproducer.zip

Triage

• Needs triage