add bf16 support in jagged tensor ops (#1079)

Summary:
Pull Request resolved: #1079

To support bf16 training

Reviewed By: ajtulloch

Differential Revision: D35955466

fbshipit-source-id: 0f740f29074576c026005362c78f872fec80bbcc

Author: jspark1105

fbgemm_gpu/src/jagged_tensor_ops.cu

@@ -428,8 +428,9 @@ at::Tensor jagged_to_padded_dense_forward(
   Tensor padded_values_view =
       values.dim() == 1 ? padded_values.unsqueeze(-1) : padded_values;
 
-  AT_DISPATCH_ALL_TYPES_AND(
+  AT_DISPATCH_ALL_TYPES_AND2(
       at::ScalarType::Half,
+      at::ScalarType::BFloat16,
       values.scalar_type(),
       "jagged_to_padded_dense",
       [&] {
@@ -461,7 +462,9 @@ at::Tensor jagged_to_padded_dense_backward(
   auto grad_values =
       at::zeros({max_lengths[0], D}, grad_padded_values.options());
 
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
       grad_padded_values.scalar_type(),
       "jagged_2d_to_dense_backward_kernel",
       [&] {
@@ -904,7 +907,9 @@ class BatchedDenseVecJagged2DMulGPUOp
 
       AT_DISPATCH_INDEX_TYPES(
           a_offsets.scalar_type(), "dense_vec_jagged_2d_bmm_kernel_1", [&] {
-            AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+            AT_DISPATCH_FLOATING_TYPES_AND2(
+                at::ScalarType::Half,
+                at::ScalarType::BFloat16,
                 a_values.scalar_type(),
                 "dense_vec_jagged_2d_bmm_kernel_2",
                 [&] {
@@ -963,7 +968,9 @@ class BatchedDenseVecJagged2DMulGPUOp
           a_offsets.scalar_type(),
           "dense_vec_jagged_2d_bmm_baackward_kernel_1",
           [&] {
-            AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+            AT_DISPATCH_FLOATING_TYPES_AND2(
+                at::ScalarType::Half,
+                at::ScalarType::BFloat16,
                 grad_outputs[0].scalar_type(),
                 "dense_vec_jagged_2d_bmm_baackward_kernel_2",
                 [&] {

fbgemm_gpu/test/sparse_ops_test.py

@@ -1836,6 +1836,7 @@ def test_dense_to_jagged(
         outer_dense_size=st.integers(0, 5),
         inner_dense_size=st.integers(0, 5),
         padding_value=st.sampled_from([0, -1e-8]),
+        dtype=st.sampled_from([torch.float, torch.half, torch.bfloat16, torch.double]),
         use_cpu=st.booleans() if gpu_available else st.just(True),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=20, deadline=None)
@@ -1845,8 +1846,12 @@ def test_jagged_to_padded_dense(
         outer_dense_size: int,
         inner_dense_size: int,
         padding_value: float,
+        dtype: torch.dtype,
         use_cpu: bool,
     ) -> None:
+        # CPU doesn't support bfloat16
+        assume(not use_cpu or dtype != torch.bfloat16)
+
         # Testing with a basic crafted example.
         # dense representation is
         # [[[[0, 1], [ 0,  0], [0, 0]],
@@ -2006,7 +2011,7 @@ def mul_func(*args) -> torch.Tensor:
         H=st.integers(1, 3),
         max_L=st.integers(1, 32),
         D=st.integers(0, 32),
-        dtype=st.sampled_from([torch.float, torch.half, torch.double]),
+        dtype=st.sampled_from([torch.float, torch.half, torch.bfloat16, torch.double]),
         use_cpu=st.booleans() if gpu_available else st.just(True),
     )
     def test_batched_dense_vec_jagged_2d_mul(
@@ -2019,6 +2024,9 @@ def test_batched_dense_vec_jagged_2d_mul(
         use_cpu: bool,
     ) -> None:
         assume(H == 1 or B != 0)
+        # CPU doesn't support bfloat16
+        assume(not use_cpu or dtype != torch.bfloat16)
+
         device = torch.device("cpu" if use_cpu else "cuda")
         torch.backends.cuda.matmul.allow_tf32 = False