From f4a3a3af470121b6cfc894387f4dd73e254ce31a Mon Sep 17 00:00:00 2001 From: tintin1942 Date: Sun, 13 Sep 2026 22:39:44 +0800 Subject: [PATCH] [BugFix][Relax] Preserve out_dtype in matmul rewrites Forward the original MatmulAttrs::out_dtype when ExpandMatmulOfSum and ReorderTakeAfterMatmul construct replacement matmul calls. Add regression coverage for the expansion and both take-reordering paths. --- src/relax/transform/expand_matmul_of_sum.cc | 5 +- .../transform/reorder_take_after_matmul.cc | 7 +- .../test_transform_expand_matmul_of_sum.py | 27 ++++++++ ...est_transform_reorder_take_after_matmul.py | 64 +++++++++++++++++++ 4 files changed, 98 insertions(+), 5 deletions(-) diff --git a/src/relax/transform/expand_matmul_of_sum.cc b/src/relax/transform/expand_matmul_of_sum.cc index 25799747c8b6..64ca60c1d8d4 100644 --- a/src/relax/transform/expand_matmul_of_sum.cc +++ b/src/relax/transform/expand_matmul_of_sum.cc @@ -24,12 +24,12 @@ #include #include +#include #include #include #include #include -#include #include #include @@ -59,6 +59,7 @@ std::tuple)>> auto pat_matmul = IsOp("relax.matmul")(pat_lhs, pat_rhs); auto rewriter = [=](Expr expr, ffi::Map matches) -> Expr { + auto out_dtype = expr.as()->attrs.as()->out_dtype; auto lhs = matches[pat_lhs]; auto rhs_a = matches[pat_rhs_a]; auto rhs_b = matches[pat_rhs_b]; @@ -88,7 +89,7 @@ std::tuple)>> rhs_b = permute_dims(rhs_b, axes); } - return add(matmul(lhs, rhs_a, std::nullopt), matmul(lhs, rhs_b, std::nullopt)); + return add(matmul(lhs, rhs_a, out_dtype), matmul(lhs, rhs_b, out_dtype)); }; return {pat_matmul, rewriter}; diff --git a/src/relax/transform/reorder_take_after_matmul.cc b/src/relax/transform/reorder_take_after_matmul.cc index 19e30f8c6853..bbc290d5ed50 100644 --- a/src/relax/transform/reorder_take_after_matmul.cc +++ b/src/relax/transform/reorder_take_after_matmul.cc @@ -24,12 +24,12 @@ #include #include +#include #include #include #include #include -#include #include #include @@ -51,6 +51,7 @@ std::tuple)>> auto pat_matmul = IsOp("relax.matmul")(pat_lhs, pat_rhs); auto rewriter = [=](Expr expr, ffi::Map matches) -> Expr { + auto out_dtype = expr.as()->attrs.as()->out_dtype; auto lhs = matches[pat_lhs]; auto weights = matches[pat_weights]; auto indices = matches[pat_indices]; @@ -92,7 +93,7 @@ std::tuple)>> // indices.shape = [outfeatures] // out_table.shape = [*batch, table_size] - auto out_table = matmul(lhs, weights, std::nullopt); + auto out_table = matmul(lhs, weights, out_dtype); // new_output.shape = [*batch, outfeatures] auto new_output = take(out_table, indices, matmul_ty->ndim - 1, attrs->mode); @@ -116,7 +117,7 @@ std::tuple)>> auto fused_weight = reshape(reordered_weight, ShapeExpr({weight_shape[1], weight_shape[0] * weight_shape[2]})); // fused_output.shape = [batch1, batch2, table_size * outfeatures] - auto fused_output = matmul(lhs, fused_weight, std::nullopt); + auto fused_output = matmul(lhs, fused_weight, out_dtype); // indexed_output.shape = [batch1, batch2, table_size, outfeatures] auto indexed_output = reshape( fused_output, ShapeExpr({lhs_shape[0], lhs_shape[1], weight_shape[0], weight_shape[2]})); diff --git a/tests/python/relax/test_transform_expand_matmul_of_sum.py b/tests/python/relax/test_transform_expand_matmul_of_sum.py index 9e38dd68f618..07d9fbdb22b0 100644 --- a/tests/python/relax/test_transform_expand_matmul_of_sum.py +++ b/tests/python/relax/test_transform_expand_matmul_of_sum.py @@ -64,6 +64,33 @@ def main( return out +class TestPreserveOutDtype(Base): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor([16], "float16"), + A: R.Tensor([16, 32], "float16"), + B: R.Tensor([16, 32], "float16"), + ) -> R.Tensor([32], "float32"): + weight = R.add(A, B) + out = R.matmul(x, weight, out_dtype="float32") + return out + + @I.ir_module + class Expected: + @R.function + def main( + x: R.Tensor([16], "float16"), + A: R.Tensor([16, 32], "float16"), + B: R.Tensor([16, 32], "float16"), + ) -> R.Tensor([32], "float32"): + lhs = R.matmul(x, A, out_dtype="float32") + rhs = R.matmul(x, B, out_dtype="float32") + out = R.add(lhs, rhs) + return out + + class TestNoExpansionOfCompileTimeAddition(Base): """Do not expand compile-time parameters diff --git a/tests/python/relax/test_transform_reorder_take_after_matmul.py b/tests/python/relax/test_transform_reorder_take_after_matmul.py index c07a86911670..48e73410a6e6 100644 --- a/tests/python/relax/test_transform_reorder_take_after_matmul.py +++ b/tests/python/relax/test_transform_reorder_take_after_matmul.py @@ -249,5 +249,69 @@ def main( return out +class TestPreserveOutDtype(Base): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor([1, 16], "float16"), + weight_table: R.Tensor([16, 64], "float16"), + routing_table: R.Tensor([32], "int64"), + ) -> R.Tensor([1, 32], "float32"): + with R.dataflow(): + weight = R.take(weight_table, routing_table, axis=1) + out = R.matmul(x, weight, out_dtype="float32") + R.output(out) + return out + + @I.ir_module + class Expected: + @R.function + def main( + x: R.Tensor([1, 16], "float16"), + weight_table: R.Tensor([16, 64], "float16"), + routing_table: R.Tensor([32], "int64"), + ) -> R.Tensor([1, 32], "float32"): + with R.dataflow(): + out_table = R.matmul(x, weight_table, out_dtype="float32") + out = R.take(out_table, routing_table, axis=1) + R.output(out) + return out + + +class TestPreserveOutDtypeForBatchedWeights(Base): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor([128, 1, 16], "float16"), + weight_table: R.Tensor([64, 16, 32], "float16"), + routing_table: R.Tensor([128], "int64"), + ) -> R.Tensor([128, 1, 32], "float32"): + with R.dataflow(): + weight = R.take(weight_table, routing_table, axis=0) + out = R.matmul(x, weight, out_dtype="float32") + R.output(out) + return out + + @I.ir_module + class Expected: + @R.function + def main( + x: R.Tensor([128, 1, 16], "float16"), + weight_table: R.Tensor([64, 16, 32], "float16"), + routing_table: R.Tensor([128], "int64"), + ) -> R.Tensor([128, 1, 32], "float32"): + with R.dataflow(): + reordered_weight = R.permute_dims(weight_table, [1, 0, 2]) + fused_weight = R.reshape(reordered_weight, [16, 2048]) + fused_output = R.matmul(x, fused_weight, out_dtype="float32") + reordered_output = R.reshape(fused_output, [128, 1, 64, 32]) + tabular_output = R.take(reordered_output, routing_table, axis=2) + out = R.einsum([tabular_output], "ijik->ijk") + R.output(out) + return out + + if __name__ == "__main__": tvm.testing.main()