[Relax][PyTorch] Support torch.bfloat16 dtype in pytorch frontend (apache#17894)

support bfloat16 dtype in pytorch frontend

Author: mshr-h

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -58,6 +58,8 @@ def _convert_data_type(input_type: Union[str, torch.dtype], env: Optional[Dict]
             return "float32"
         elif input_type in ["float16", "torch.float16", torch.float16]:
             return "float16"
+        elif input_type in ["bfloat16", "torch.bfloat16", torch.bfloat16]:
+            return "bfloat16"
         elif input_type in ["int64", "torch.int64", torch.int64]:
             return "int64"
         elif input_type in ["int32", "torch.int32", torch.int32]:

tests/python/relax/test_frontend_from_exported_program.py

@@ -4869,5 +4869,32 @@ def main(
     verify_model(Linspace(), example_args, {}, Expected)
 
 
+def test_bfloat16():
+    # TODO(mshr-h): Add tests for all the dtypes supported in fx frontend
+    example_args = (
+        torch.randn(10, 10, dtype=torch.bfloat16),
+        torch.randn(10, 10, dtype=torch.bfloat16),
+    )
+
+    class BFloat16Model(Module):
+        def forward(self, lhs: torch.Tensor, rhs: torch.Tensor):
+            return torch.ops.aten.add(lhs, rhs)
+
+    @tvm.script.ir_module
+    class expected:
+        @R.function
+        def main(
+            lhs: R.Tensor((10, 10), dtype="bfloat16"),
+            rhs: R.Tensor((10, 10), dtype="bfloat16"),
+        ) -> R.Tuple(R.Tensor((10, 10), dtype="bfloat16")):
+            with R.dataflow():
+                lv: R.Tensor((10, 10), dtype="bfloat16") = relax.op.add(lhs, rhs)
+                gv: R.Tuple(R.Tensor((10, 10), dtype="bfloat16")) = (lv,)
+                R.output(gv)
+            return gv
+
+    verify_model(BFloat16Model(), example_args, {}, expected)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/relax/test_frontend_from_fx.py

@@ -5249,5 +5249,27 @@ def main(
         verify_model(Norm(p, dim=dim, keepdim=keepdim), input_info, {}, expected)
 
 
+def test_bfloat16():
+    # TODO(mshr-h): Add tests for all the dtypes supported in EP frontend
+    class BFloat16Model(Module):
+        def forward(self, lhs: torch.Tensor, rhs: torch.Tensor):
+            return torch.ops.aten.add(lhs, rhs)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(
+            lhs: R.Tensor((10, 10), dtype="bfloat16"),
+            rhs: R.Tensor((10, 10), dtype="bfloat16"),
+        ) -> R.Tensor((10, 10), dtype="bfloat16"):
+            with R.dataflow():
+                lv: R.Tensor((10, 10), dtype="bfloat16") = relax.op.add(lhs, rhs)
+                gv: R.Tensor((10, 10), dtype="bfloat16") = lv
+                R.output(gv)
+            return gv
+
+    verify_model(BFloat16Model(), [([10, 10], "bfloat16"), ([10, 10], "bfloat16")], {}, Expected)
+
+
 if __name__ == "__main__":
     tvm.testing.main()