[Fix][Relax] Preserve out_dtype in matmul rewrite passes

src/relax/transform/expand_matmul_of_sum.cc

@@ -24,12 +24,12 @@
 
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/linear_algebra.h>
 #include <tvm/relax/dataflow_matcher.h>
 #include <tvm/relax/expr.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/transform.h>
 
-#include <optional>
 #include <unordered_set>
 #include <vector>
 
@@ -59,6 +59,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
   auto pat_matmul = IsOp("relax.matmul")(pat_lhs, pat_rhs);
 
   auto rewriter = [=](Expr expr, ffi::Map<DFPattern, Expr> matches) -> Expr {
+    auto out_dtype = expr.as<CallNode>()->attrs.as<MatmulAttrs>()->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<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
       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};

src/relax/transform/reorder_take_after_matmul.cc

@@ -24,12 +24,12 @@
 
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/linear_algebra.h>
 #include <tvm/relax/dataflow_matcher.h>
 #include <tvm/relax/expr.h>
 #include <tvm/relax/expr_functor.h>
 #include <tvm/relax/transform.h>
 
-#include <optional>
 #include <unordered_set>
 #include <vector>
 
@@ -51,6 +51,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
   auto pat_matmul = IsOp("relax.matmul")(pat_lhs, pat_rhs);
 
   auto rewriter = [=](Expr expr, ffi::Map<DFPattern, Expr> matches) -> Expr {
+    auto out_dtype = expr.as<CallNode>()->attrs.as<MatmulAttrs>()->out_dtype;
     auto lhs = matches[pat_lhs];
     auto weights = matches[pat_weights];
     auto indices = matches[pat_indices];
@@ -92,7 +93,7 @@ std::tuple<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
       // 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<DFPattern, ffi::TypedFunction<Expr(Expr, ffi::Map<DFPattern, Expr>)>>
       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]}));

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
 

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