add Half/BFloat16 support for grid_sample on CPU (pytorch#134812)

Fix pytorch#127224.

Pull Request resolved: pytorch#134812
Approved by: https://github.com/Skylion007, https://github.com/mingfeima

Author: jiayisunx

aten/src/ATen/cpu/vec/vec512/vec512_bfloat16.h

@@ -221,73 +221,7 @@ static_assert(
   }
   template <int64_t mask>
   static Vectorized<T> blend(const Vectorized<T>& a, const Vectorized<T>& b) {
-    __at_align__ int16_t tmp_values[size()];
-    a.store(tmp_values);
-    if (mask & 0x01)
-      tmp_values[0] = b.values[31];
-    if (mask & 0x02)
-      tmp_values[1] = b.values[30];
-    if (mask & 0x04)
-      tmp_values[2] = b.values[29];
-    if (mask & 0x08)
-      tmp_values[3] = b.values[28];
-    if (mask & 0x10)
-      tmp_values[4] = b.values[27];
-    if (mask & 0x20)
-      tmp_values[5] = b.values[26];
-    if (mask & 0x40)
-      tmp_values[6] = b.values[25];
-    if (mask & 0x80)
-      tmp_values[7] = b.values[24];
-    if (mask & 0x100)
-      tmp_values[8] = b.values[23];
-    if (mask & 0x200)
-      tmp_values[9] = b.values[22];
-    if (mask & 0x400)
-      tmp_values[10] = b.values[21];
-    if (mask & 0x800)
-      tmp_values[11] = b.values[20];
-    if (mask & 0x1000)
-      tmp_values[12] = b.values[19];
-    if (mask & 0x2000)
-      tmp_values[13] = b.values[18];
-    if (mask & 0x4000)
-      tmp_values[14] = b.values[17];
-    if (mask & 0x8000)
-      tmp_values[15] = b.values[16];
-    if (mask & 0x10000)
-      tmp_values[16] = b.values[15];
-    if (mask & 0x20000)
-      tmp_values[17] = b.values[14];
-    if (mask & 0x40000)
-      tmp_values[18] = b.values[13];
-    if (mask & 0x80000)
-      tmp_values[19] = b.values[12];
-    if (mask & 0x100000)
-      tmp_values[20] = b.values[11];
-    if (mask & 0x200000)
-      tmp_values[21] = b.values[10];
-    if (mask & 0x400000)
-      tmp_values[22] = b.values[9];
-    if (mask & 0x800000)
-      tmp_values[23] = b.values[8];
-    if (mask & 0x1000000)
-      tmp_values[24] = b.values[7];
-    if (mask & 0x2000000)
-      tmp_values[25] = b.values[6];
-    if (mask & 0x4000000)
-      tmp_values[26] = b.values[5];
-    if (mask & 0x8000000)
-      tmp_values[27] = b.values[4];
-    if (mask & 0x10000000)
-      tmp_values[28] = b.values[3];
-    if (mask & 0x20000000)
-      tmp_values[29] = b.values[2];
-    if (mask & 0x40000000)
-      tmp_values[30] = b.values[1];
-    if (mask & 0x80000000)
-      tmp_values[31] = b.values[0];
-    return loadu(tmp_values);
+    return _mm512_mask_blend_epi16(mask, a.values, b.values);
   }
   static Vectorized<T> blendv(const Vectorized<T>& a,
       const Vectorized<T>& b, const Vectorized<T>& mask) {

aten/src/ATen/native/GridSampler.cpp

@@ -930,9 +930,7 @@ Tensor grid_sampler_2d_cpu(const Tensor& input, const Tensor& grid,
   }
   // AVX gather instructions use signed 32-bit offsets to gather float values.
   // Check for possible overflow and fallback to scalar implementation
-  if (input.scalar_type() != kDouble) {
-    TORCH_CHECK(input.scalar_type() == kFloat,
-                "grid_sampler_2d_cpu not implemented for ", input.scalar_type());
+  if (input.scalar_type() == kFloat) {
     auto sizes = input.sizes();
     auto strides = input.strides();
     const auto grid_sW = grid.strides()[2];
@@ -968,7 +966,7 @@ Tensor grid_sampler_3d_cpu(const Tensor& input, const Tensor& grid,
   check_grid_sampler_common(input, grid);
   check_grid_sampler_3d(input, grid, interpolation_mode);
 
-  return AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "grid_sampler3d_cpu", [&] {
+  return AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "grid_sampler3d_cpu", [&] {
     return grid_sampler_3d_cpu_impl<scalar_t>(
       input, grid, static_cast<GridSamplerInterpolation>(interpolation_mode),
       static_cast<GridSamplerPadding>(padding_mode), align_corners);
@@ -986,9 +984,7 @@ grid_sampler_2d_backward_cpu(const Tensor& grad_output, const Tensor& input, con
 
   // AVX gather instructions use signed 32-bit offsets to gather float values.
   // Check for possible overflow and fallback to scalar implementation
-  if (input.scalar_type() != kDouble) {
-    TORCH_CHECK(input.scalar_type() == kFloat,
-                "grid_sampler_2d_backward_cpu not implemented for ", input.scalar_type());
+  if (input.scalar_type() == kFloat) {
     auto isizes = input.sizes();
     auto istrides = input.strides();
     auto gsizes = grad_output.sizes();
@@ -1033,7 +1029,7 @@ grid_sampler_3d_backward_cpu(const Tensor& grad_output, const Tensor& input, con
   check_grid_sampler_common(input, grid);
   check_grid_sampler_3d(input, grid, interpolation_mode);
 
-  return AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "grid_sampler_3d_backward_cpu", [&] {
+  return AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "grid_sampler_3d_backward_cpu", [&] {
     return grid_sampler_3d_backward_cpu_impl<scalar_t>(
       grad_output, input, grid,
       static_cast<GridSamplerInterpolation>(interpolation_mode),

aten/src/ATen/native/cpu/GridSamplerKernel.cpp

@@ -1184,7 +1184,7 @@ void grid_sampler_2d_cpu_kernel_impl(
     return;                                                                    \
   }
 
-  AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "grid_sampler_2d_cpu_kernel_impl", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "grid_sampler_2d_cpu_kernel_impl", [&] {
     auto out_acc = output.accessor<scalar_t, 4>();
     auto inp_acc = input.accessor<const scalar_t, 4>();
     auto grid_acc = grid.accessor<const scalar_t, 4>();
@@ -1272,7 +1272,7 @@ void grid_sampler_2d_backward_cpu_kernel_impl(
     return;                                                                 \
   }
 
-  AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "grid_sampler_2d_backward_cpu_kernel_impl", [&] {
+  AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "grid_sampler_2d_backward_cpu_kernel_impl", [&] {
     auto gGrid_acc = grad_grid.accessor<scalar_t, 4>();
     auto inp_acc = input.accessor<const scalar_t, 4>();
     auto grid_acc = grid.accessor<const scalar_t, 4>();

test/inductor/test_torchinductor_opinfo.py

@@ -693,7 +693,7 @@ def wrapper_noop_set_seed(op, *args, **kwargs):
     "nn.functional.cosine_similarity": {f16},
     "nn.functional.cross_entropy": {f16, f32, f64},
     "nn.functional.gaussian_nll_loss": {f16},
-    "nn.functional.grid_sample": {f32, f64},
+    "nn.functional.grid_sample": {f32, f64, f16},
     "nn.functional.interpolate.area": {f16},
     "nn.functional.nll_loss": {f16, f32, f64},
     "normal": {f16, f32, f64},

test/test_mps.py

@@ -152,7 +152,7 @@ def mps_ops_grad_modifier(ops):
 
     MACOS_12_3_XFAILLIST_GRAD = {
         # Unsupported Border padding mode, forward pass success as fallback to cpu
-        'grid_sampler_2d': [torch.float32],
+        'grid_sampler_2d': [torch.float32, torch.float16, torch.bfloat16],
         # Unimplemented
         'logaddexp2': [torch.float32],
 
@@ -165,7 +165,7 @@ def mps_ops_grad_modifier(ops):
         'masked.log_softmax': [torch.float32, torch.float16],
 
         # Unsupported Border padding mode, forward pass success as fallback to cpu
-        'grid_sampler_2d': [torch.float32],
+        'grid_sampler_2d': [torch.float32, torch.float16, torch.bfloat16],
 
         # Same issue as `argsort` and `sort` with duplicate elements (undefined behaviour).
         # Forward pass is passing since `msort` doesn't return the indices, just the values, which match the CPU.
@@ -638,7 +638,7 @@ def mps_ops_modifier(ops):
 
     MACOS_AFTER_13_1_XFAILLIST = {
         # before macOS 13.2 it falls back to cpu and pass the forward pass
-        'grid_sampler_2d': [torch.float32],  # Unsupported Border padding mode
+        'grid_sampler_2d': [torch.float32, torch.float16, torch.bfloat16],  # Unsupported Border padding mode
         # inconsistency errors between cpu and mps, max seen atol is 2
         'nn.functional.interpolatebilinear': [torch.uint8],
     }

torch/testing/_internal/common_methods_invocations.py

@@ -20811,8 +20811,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     ),
     OpInfo(
         "nn.functional.grid_sample",
-        dtypes=floating_types(),
-        dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16),
+        dtypes=floating_types_and(torch.float16, torch.bfloat16),
         supports_out=False,
         sample_inputs_func=sample_inputs_grid_sample,
         reference_inputs_func=reference_inputs_grid_sample,
@@ -20821,8 +20820,7 @@ def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs):
     # TODO: delete this OpInfo once we add meta support for grid_sampler_3d
     OpInfo(
         "grid_sampler_2d",
-        dtypes=floating_types(),
-        dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16),
+        dtypes=floating_types_and(torch.float16, torch.bfloat16),
         supports_out=False,
         sample_inputs_func=sample_inputs_grid_sampler_2d,
         supports_gradgrad=False,