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 9e5e709662 [BugFix][Relax] Skip parallel matmul fusion for mixed
output dtypes (#20208)
9e5e709662 is described below
commit 9e5e709662050f68eccf27e7bb0e3f8f2c19b0ab
Author: Midst <[email protected]>
AuthorDate: Fri Aug 28 15:44:14 2026 +0800
[BugFix][Relax] Skip parallel matmul fusion for mixed output dtypes (#20208)
`CombineParallelMatmul` currently selects the combined matmul output
dtype
from the first participating branch. When parallel branches have
different
effective output dtypes, fusion changes another branch's result dtype
and
may affect numerical behavior.
This change:
- Carries each candidate branch's inferred output dtype through split
ordering.
- Skips a fusion group when the participating branches have different
output dtypes.
- Uses the validated shared dtype when creating the combined matmul.
- Adds regression coverage for both branch orderings.
Testing:
- `git diff --check`
- C++ syntax check using TVM's generated compiler flags
- `python -m py_compile
tests/python/relax/test_transform_combine_parallel_matmul.py`
A full pytest run was not available because the Windows host does not
have
a compatible built TVM runtime library.
Fixes #20203
---
src/relax/transform/combine_parallel_matmul.cc | 15 ++++++---
.../test_transform_combine_parallel_matmul.py | 39 ++++++++++++++++++++++
2 files changed, 50 insertions(+), 4 deletions(-)
diff --git a/src/relax/transform/combine_parallel_matmul.cc
b/src/relax/transform/combine_parallel_matmul.cc
index ff85a0d360..952b7f4b62 100644
--- a/src/relax/transform/combine_parallel_matmul.cc
+++ b/src/relax/transform/combine_parallel_matmul.cc
@@ -78,6 +78,7 @@ struct SplitInfo {
ffi::Optional<Var> bias;
PrimExpr split_size;
DFPattern pattern_to_replace;
+ DLDataType out_dtype;
};
Patterns CreatePatterns(const BranchInfo& branch_info) {
@@ -163,7 +164,9 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern,
Var>, ffi::Map<Var, E
}
PrimExpr split_size = GetTensorType(rhs)->GetShape().value()[rhs_dim -
1];
DFPattern pattern_to_replace = patterns_to_replace[index];
- splits.push_back(SplitInfo{rhs, bias, split_size, pattern_to_replace});
+ DLDataType out_dtype =
+
GetTensorType(matchings[patterns.matmul[index]])->dtype.value()->dtype;
+ splits.push_back(SplitInfo{rhs, bias, split_size, pattern_to_replace,
out_dtype});
}
// At most one dynamic output shape can be part of the combined
// matmul, and it must be the last item in the split. Use
@@ -188,6 +191,12 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern,
Var>, ffi::Map<Var, E
continue;
}
+ if (std::any_of(splits.begin() + 1, splits.end(), [&](const SplitInfo&
split) {
+ return split.out_dtype != splits[0].out_dtype;
+ })) {
+ continue;
+ }
+
ffi::Array<Var> rhs;
ffi::Array<Var> bias;
for (const auto& split : splits) {
@@ -202,9 +211,7 @@ ffi::TypedFunction<ffi::Map<Var, Expr>(ffi::Map<DFPattern,
Var>, ffi::Map<Var, E
}
auto concat_rhs = concat(Tuple(rhs), rhs_dim - 1);
- DLDataType out_dtype =
-
GetTensorType(matchings[patterns.matmul[indices[0]]])->dtype.value()->dtype;
- auto matmul_combined = matmul(lhs, concat_rhs, out_dtype);
+ auto matmul_combined = matmul(lhs, concat_rhs, splits[0].out_dtype);
if (branch_info.bias_dim) {
auto bias_dim = GetTensorType(bias[0])->ndim;
diff --git a/tests/python/relax/test_transform_combine_parallel_matmul.py
b/tests/python/relax/test_transform_combine_parallel_matmul.py
index 32c0c36c0d..57be71ae88 100644
--- a/tests/python/relax/test_transform_combine_parallel_matmul.py
+++ b/tests/python/relax/test_transform_combine_parallel_matmul.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E731, F401, F841
+import pytest
+
import tvm.testing
from tvm import relax, tirx
from tvm.relax.transform import CombineParallelMatmul
@@ -694,5 +696,42 @@ def test_limit_one_dynamic_shape_in_combined_matmul():
tvm.ir.assert_structural_equal(after, expected)
[email protected]("float32_branch", [0, 1])
+def test_skip_matmuls_with_different_output_dtypes(float32_branch):
+ if float32_branch == 0:
+
+ @R.function(private=True)
+ def before(
+ x: R.Tensor((3, 4), "float16"),
+ w0: R.Tensor((4, 5), "float16"),
+ w1: R.Tensor((4, 6), "float16"),
+ ):
+ with R.dataflow():
+ y0 = R.matmul(x, w0, out_dtype="float32")
+ y1 = R.matmul(x, w1)
+ out = (y0, y1)
+ R.output(out)
+ return out
+
+ else:
+
+ @R.function(private=True)
+ def before(
+ x: R.Tensor((3, 4), "float16"),
+ w0: R.Tensor((4, 5), "float16"),
+ w1: R.Tensor((4, 6), "float16"),
+ ):
+ with R.dataflow():
+ y0 = R.matmul(x, w0)
+ y1 = R.matmul(x, w1, out_dtype="float32")
+ out = (y0, y1)
+ R.output(out)
+ return out
+
+ after = CombineParallelMatmul()(tvm.IRModule.from_expr(before))["main"]
+
+ tvm.ir.assert_structural_equal(after, before)
+
+
if __name__ == "__main__":
tvm.testing.main()
