Add bfloat16 + fp16 support to fractional_max_pool for CUDA and CPU (pytorch#116950)

Adds bfloat16 to fractional_max_pool. If op supports fp32 and fp16, it really should support bf16 for the most part. Most but not all ops satisfy this, so I am adding support for the few that do not.

Pull Request resolved: pytorch#116950
Approved by: https://github.com/lezcano

Author: Skylion007

aten/src/ATen/native/FractionalMaxPool2d.cpp

@@ -321,21 +321,24 @@ TORCH_IMPL_FUNC(fractional_max_pool2d_out_cpu) (
   int64_t inputH = input.size(heightDim);
   int64_t inputW = input.size(widthDim);
 
-  AT_DISPATCH_FLOATING_TYPES(input.scalar_type(),
-  "fractional_max_pool2d_out_frame", [&] {
-    auto input_data = input.data_ptr<scalar_t>();
-    auto output_data = output.data_ptr<scalar_t>();
-    auto indices_data = indices.data_ptr<int64_t>();
-    auto randomSamples_data = randomSamples.data_ptr<scalar_t>();
-    fractional_max_pool2d_out_frame<scalar_t>(
-      input_data,
-      output_data,
-      indices_data,
-      randomSamples_data,
-      numBatch, numPlanes,
-      inputW, inputH,
-      outputW, outputH,
-      poolSizeW, poolSizeH);
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    kBFloat16,
+    kHalf,
+    input.scalar_type(),
+    "fractional_max_pool2d_out_frame", [&] {
+      auto input_data = input.data_ptr<scalar_t>();
+      auto output_data = output.data_ptr<scalar_t>();
+      auto indices_data = indices.data_ptr<int64_t>();
+      auto randomSamples_data = randomSamples.data_ptr<scalar_t>();
+      fractional_max_pool2d_out_frame<scalar_t>(
+        input_data,
+        output_data,
+        indices_data,
+        randomSamples_data,
+        numBatch, numPlanes,
+        inputW, inputH,
+        outputW, outputH,
+        poolSizeW, poolSizeH);
     }
   );
 }
@@ -375,7 +378,9 @@ TORCH_IMPL_FUNC(fractional_max_pool2d_backward_cpu) (
   auto gradOutput = gradOutput_.contiguous();
 
   /* backprop */
-  AT_DISPATCH_FLOATING_TYPES(
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    kBFloat16,
+    kHalf,
     input.scalar_type(), "fractional_max_pool2d_backward_out_frame", [&] {
       auto gradInput_data = gradInput.data_ptr<scalar_t>();
       auto gradOutput_data = gradOutput.data_ptr<scalar_t>();

aten/src/ATen/native/FractionalMaxPool3d.cpp

@@ -237,7 +237,9 @@ TORCH_IMPL_FUNC(fractional_max_pool3d_out_cpu)(
   auto input = input_.contiguous();
   auto randomSamples = randomSamples_.contiguous();
 
-  AT_DISPATCH_FLOATING_TYPES(
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    kBFloat16,
+    kHalf,
     input.scalar_type(),
     "fractional_max_pool3d_out_frame",
     [&] {
@@ -371,7 +373,9 @@ void fractional_max_pool3d_backward_out_cpu_template(
   gradInput.zero_();
 
   /* backprop */
-  AT_DISPATCH_FLOATING_TYPES(
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    kBFloat16,
+    kHalf,
     input.scalar_type(),
     "fractional_max_pool3d_backward_out_frame",
     [&]{

aten/src/ATen/native/cuda/FractionalMaxPool2d.cu

@@ -180,7 +180,10 @@ TORCH_IMPL_FUNC(fractional_max_pool2d_out_cuda) (
             input_.size(0));
   dim3 block(outputPlaneSize > 128 ? 128 : outputPlaneSize);
 
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(),
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    at::ScalarType::Half,
+    at::ScalarType::BFloat16,
+    input.scalar_type(),
     "fractional_max_pool2d_out_cuda_frame",
     [&] {
       auto devInput = input_.packed_accessor64<scalar_t, 4>();
@@ -252,7 +255,10 @@ TORCH_IMPL_FUNC(fractional_max_pool2d_backward_cuda)(
   dim3 block(outputPlaneSize > 128 ? 128 : outputPlaneSize);
 
   auto devIndices = indices_.packed_accessor64<int64_t, 4>();
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(gradOutput.scalar_type(),
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    at::ScalarType::Half,
+    at::ScalarType::BFloat16,
+    gradOutput.scalar_type(),
     "fractional_max_pool2d_backward_out_cuda_frame",
     [&] {
       auto devGradInput = gradInput_.packed_accessor64<scalar_t, 4>();

aten/src/ATen/native/cuda/FractionalMaxPool3d.cu

@@ -226,7 +226,9 @@ void fractional_max_pool3d_backward_out_cuda_template(
       gradInput_.size(0));
     dim3 block(outputPlaneSize > 128 ? 128 : outputPlaneSize);
 
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
       gradOutput.scalar_type(),
       "fractional_max_pool3d_backward_out_frame",
       [&] {
@@ -285,7 +287,9 @@ TORCH_IMPL_FUNC(fractional_max_pool3d_out_cuda) (
     input_.size(0));
   dim3 block(outputPlaneSize > 128 ? 128 : outputPlaneSize);
 
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+    at::ScalarType::Half,
+    at::ScalarType::BFloat16,
     input.scalar_type(),
     "fractional_max_pool3d_out_frame",
     [&]{

torch/testing/_internal/common_methods_invocations.py

@@ -13799,8 +13799,7 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
                wrapper_set_seed(torch.nn.functional.fractional_max_pool2d, input, *args, **kwargs),
            # vmap does not support random operations
            check_batched_forward_grad=False,
-           dtypes=floating_types(),
-           dtypesIfCUDA=floating_types_and(torch.float16),
+           dtypes=floating_types_and(torch.bfloat16, torch.float16),
            test_neg_view=False,
            sample_inputs_func=sample_inputs_fractional_max_pool2d,
            decorators=(
@@ -13820,8 +13819,7 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
                wrapper_set_seed(torch.nn.functional.fractional_max_pool3d, input, *args, **kwargs),
            # vmap does not support random operations
            check_batched_forward_grad=False,
-           dtypes=floating_types(),
-           dtypesIfCUDA=floating_types_and(torch.float16),
+           dtypes=floating_types_and(torch.bfloat16, torch.float16),
            test_neg_view=False,
            gradcheck_nondet_tol=GRADCHECK_NONDET_TOL,
            sample_inputs_func=sample_inputs_fractional_max_pool3d,