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 6b34c75a52 [Relax] Preserve out_dtype in AdjustMatmulOrder (#20296)
6b34c75a52 is described below
commit 6b34c75a52f4117e7d6a5452ae72b721e8be56c7
Author: T1Gang <[email protected]>
AuthorDate: Fri Sep 11 02:15:31 2026 +0800
[Relax] Preserve out_dtype in AdjustMatmulOrder (#20296)
Fixes #20202.
`AdjustMatmulOrder` rebuilt reassociated outer matmuls with
`out_dtype=void`, which could change the result dtype of the original
program. For example, a `float16` matmul chain whose outer matmul
explicitly
returns `float32` was rewritten to return `float16`.
This patch preserves the original outer `MatmulAttrs::out_dtype` when
constructing the replacement outer matmul. Newly introduced inner
matmuls
continue to infer their output dtype.
Tests:
- Added regression coverage for both reassociation directions.
- Covered explicit and inferred output dtypes.
- 12 new regression cases passed.
- Full `AdjustMatmulOrder` test file: 41 passed.
- `pre-commit` and `git diff --check`: passed.
---
src/relax/transform/adjust_matmul_order.cc | 10 ++--
.../relax/test_transform_adjust_matmul_order.py | 56 ++++++++++++++++++++++
2 files changed, 62 insertions(+), 4 deletions(-)
diff --git a/src/relax/transform/adjust_matmul_order.cc
b/src/relax/transform/adjust_matmul_order.cc
index 615bb6fd18..1630aeb664 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/linear_algebra.h>
#include <tvm/relax/attrs/manipulate.h>
#include <tvm/relax/dataflow_matcher.h>
#include <tvm/relax/expr.h>
@@ -167,6 +168,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
auto expr_a = matches[pat_a];
auto expr_b = matches[pat_b];
auto expr_c = matches[pat_c];
+ auto out_dtype = expr.as<CallNode>()->attrs.as<MatmulAttrs>()->out_dtype;
// If all three components are compile-time, the order doesn't
// matter as the entire expression can be lifted out and
@@ -242,9 +244,9 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
// If two of the three are compile-time, group those two values
// together, to allow them to be lifted out and pre-computed.
if (is_compile_time(expr_a) && is_compile_time(expr_b)) {
- return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c,
std::nullopt);
+ return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c, out_dtype);
} else if (is_compile_time(expr_b) && is_compile_time(expr_c)) {
- return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt),
std::nullopt);
+ return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt), out_dtype);
}
// Otherwise, select the order that reduces the total number of
@@ -319,9 +321,9 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr,
ffi::Map<DFPattern, Expr>)>>
size_N > 0 && size_R > 0 && size_M > 0 && size_B > 0);
if (analyzer->CanProve(ops_with_lhs_first < ops_with_rhs_first)) {
- return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c,
std::nullopt);
+ return matmul(matmul(expr_a, expr_b, std::nullopt), expr_c, out_dtype);
} else if (analyzer->CanProve(ops_with_rhs_first < ops_with_lhs_first)) {
- return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt),
std::nullopt);
+ return matmul(expr_a, matmul(expr_b, expr_c, std::nullopt), out_dtype);
}
// If we cannot determine which order is best, keep the existing order.
diff --git a/tests/python/relax/test_transform_adjust_matmul_order.py
b/tests/python/relax/test_transform_adjust_matmul_order.py
index eb886fd4d0..29603489f3 100644
--- a/tests/python/relax/test_transform_adjust_matmul_order.py
+++ b/tests/python/relax/test_transform_adjust_matmul_order.py
@@ -847,6 +847,62 @@ class TestBatchedSharedPrefixPreferLHSFirst(Base):
return out
[email protected]("lhs_first", [False, True])
[email protected]("compile_time,transpose", [(False, False), (False,
True), (True, False)])
[email protected]("out_dtype", [None, "float32"])
+def test_preserve_outer_out_dtype(lhs_first, compile_time, transpose,
out_dtype):
+ # Compile-time grouping should override the cheaper evaluation order.
+ shapes = [(2, 20), (20, 2), (2, 10)]
+ if lhs_first == compile_time:
+ shapes = [(10, 2), (2, 20), (20, 2)]
+ inner_indices = (1, 2) if lhs_first else (0, 1)
+ if transpose:
+ for i in inner_indices:
+ shapes[i] = shapes[i][::-1]
+
+ def build_module(expected):
+ bb = relax.BlockBuilder()
+ a, b, c = [
+ relax.Var(name, relax.TensorType(shape, "float16"))
+ for name, shape in zip("abc", shapes)
+ ]
+ params = [c, a, b] if lhs_first else [a, b, c]
+ attrs = {"num_input": 1} if compile_time else None
+ with bb.function("main", params, attrs=attrs):
+ with bb.dataflow():
+ operands = [a, b, c]
+ if expected and transpose:
+ for i in inner_indices:
+ operands[i] = relax.op.permute_dims(operands[i])
+ a, b, c = operands
+ if expected:
+ if lhs_first:
+ out = relax.op.matmul(relax.op.matmul(a, b), c,
out_dtype=out_dtype)
+ else:
+ out = relax.op.matmul(a, relax.op.matmul(b, c),
out_dtype=out_dtype)
+ else:
+ x, y = (b, c) if lhs_first else (a, b)
+ inner = bb.emit(relax.op.matmul(y, x) if transpose else
relax.op.matmul(x, y))
+ if transpose:
+ inner = bb.emit(relax.op.permute_dims(inner))
+ out = (
+ relax.op.matmul(a, inner, out_dtype=out_dtype)
+ if lhs_first
+ else relax.op.matmul(inner, c, out_dtype=out_dtype)
+ )
+ output = bb.emit_output(out)
+ bb.emit_func_output(output)
+ return bb.finalize()
+
+ before = build_module(expected=False)
+ expected = build_module(expected=True)
+ transform = relax.transform.AdjustMatmulOrder()
+ after = transform(before)
+ assert after["main"].ret_ty.dtype == (out_dtype or "float16")
+ tvm.ir.assert_structural_equal(after, expected)
+ tvm.ir.assert_structural_equal(transform(after), after)
+
+
class TestAdjustMatmulOrderAttentionBlock:
"""AdjustMatmulOrder preserves numerics on a batched attention block.
