Added grid_sample backward batch rule

Description:
- Added grid_sample backward batch rule: CPU and CUDA
- Updated tests

Notes:
I had to expand on dim 0 in most of the cases and could not use
tricks like in forward pass when batch dim is merged either with channel or H_out
due to wrong grid grads in these cases

Author: vfdev-5

codegen/codegen_outofplacebatching.py

@@ -155,7 +155,7 @@ def parse_return(return_t):
     return tuple([x.strip() for x in m.group(1).split(',')])
 
 def parse_args(args_t):
-    args = args_t.split(',')
+    args = args_t.split(', ')
     result = []
     for arg in args:
         split_idx = arg.rfind(' ')
@@ -170,8 +170,6 @@ def get_signatures(path='build/aten/src/ATen/RegistrationDeclarations.h', includ
     for line in lines:
         if 'void' in line:
             continue
-        if 'std::array' in line:
-            continue
         m = re.match(r'(.*) \w+\((.*)\); // {"schema": "aten::(\w+\.?\w*)\(.*', line)
         if m is None:
             continue

functorch/csrc/BatchRulesModules.cpp

@@ -255,6 +255,202 @@ grid_sample_batch_rule(const Tensor& input, optional<int64_t> input_bdim, const
   return result;
 }
 
+Tensor expand_reshape_dim_into(int64_t batch_size, int64_t dst, const Tensor& x) {
+  auto x_ = x.unsqueeze(0);
+  VmapDimVector new_shape(x_.sizes().begin(), x_.sizes().end());
+  new_shape[0] = batch_size;
+  x_ = x_.expand(new_shape);
+  return reshape_dim_into(0, dst, x_);
+}
+
+
+std::tuple<Tensor, Tensor, optional<int64_t>, Tensor, optional<int64_t>, int64_t>
+grid_sample_backward_helper_in(
+    const Tensor& grad_output, optional<int64_t> grad_output_bdim,
+    const Tensor& input, optional<int64_t> input_bdim,
+    const Tensor& grid, optional<int64_t> grid_bdim) {
+  auto new_grad_output = grad_output;
+  auto new_input = input;
+  auto new_grid = grid;
+
+  optional<int64_t> grad_input_out_bdim = nullopt;
+  optional<int64_t> grad_grid_out_bdim = nullopt;
+  int64_t bdim_size = 0;
+
+  if (grad_output_bdim) {
+
+    bdim_size = grad_output.sizes()[*grad_output_bdim];
+
+    if (input_bdim && grid_bdim) {
+      // case 1: (grad_output is batched, input is batched, grid is batched)
+      //        grad_output: (BN)CH_{out}W_{out}, input: (BN)CH_{in}W_{in}, grid: (BN)H_{out}W_{out}2
+      //        grad_input: (BN)CH_{in}W_{in}
+
+      new_grad_output = reshape_dim_into(*grad_output_bdim, 0, grad_output);
+      new_input = reshape_dim_into(*input_bdim, 0, input);
+      new_grid = reshape_dim_into(*grid_bdim, 0, grid);
+      grad_input_out_bdim = 0;
+      grad_grid_out_bdim = 0;
+    } else if (input_bdim && !grid_bdim) {
+      // case 2: (grad_output is batched, input is batched, grid is not batched)
+      //        IF PUT BATCH DIM TO CHANNEL -> backward produces wrong grad_grid
+      //
+      //        grad_output: (BN)CH_{out}W_{out}, input: (BN)CH_{in}W_{in}, grid: NH_{out}W_{out}2
+      //        -> grid: (BN)H_{out}W_{out}2
+      //        grad_input: (BN)CH_{in}W_{in}
+
+      new_grad_output = reshape_dim_into(*grad_output_bdim, 0, grad_output);
+      new_input = reshape_dim_into(*input_bdim, 0, input);
+      grad_input_out_bdim = 0;
+      new_grid = expand_reshape_dim_into(bdim_size, 0, grid);
+      grad_grid_out_bdim = 0;
+    } else if (!input_bdim && grid_bdim) {
+      // case 3: (grad_output is batched, input is not batched, grid is batched)
+      //        IF PUT BATCH DIM TO H_out -> backward produces wrong grad_grid
+      //
+      //        grad_output: (BN)CH_{out}W_{out}, input: NCH_{in}W_{in}, grid: (BN)H_{out}W_{out}2
+      //        -> input: (BN)CH_{in}W_{in}
+      //        grad_input: (BN)CH_{in}W_{in}
+
+      new_grad_output = reshape_dim_into(*grad_output_bdim, 0, grad_output);
+      new_grid = reshape_dim_into(*grid_bdim, 0, grid);
+      grad_grid_out_bdim = 0;
+      // expand input to (BN)H_{out}W_{out}2
+      new_input = expand_reshape_dim_into(bdim_size, 0, new_input);
+      grad_input_out_bdim = 0;
+    } else {
+      // case 4: (grad_output is batched, input is not batched, grid is not batched)
+      //        IF PUT BATCH DIM TO H_out -> backward produces wrong grad_grid
+      //
+      //        grad_output: (BN)CH_{out}W_{out}, input: NCH_{in}W_{in}, grid: NH_{out}W_{out}2
+      //        -> grid: (BN)H_{out}W_{out}2
+      //        -> input: (BN)CH_{in}W_{in}
+      //        grad_input: NCH_{in}W_{in}
+
+      new_grad_output = reshape_dim_into(*grad_output_bdim, 0, grad_output);
+      // expand grid to (BN)H_{out}W_{out}2
+      new_grid = expand_reshape_dim_into(bdim_size, 0, grid);
+      grad_grid_out_bdim = 0;
+      // expand input to (BN)CH_{in}W_{in}
+      new_input = expand_reshape_dim_into(bdim_size, 0, input);
+      grad_input_out_bdim = 0;
+    }
+  } else {
+    if (input_bdim && grid_bdim) {
+      // case 5: (grad_output is not batched, input is batched, grid is batched)
+      //        grad_output: NCH_{out}W_{out}, input: (BN)CH_{in}W_{in}, grid: (BN)H_{out}W_{out}2
+      //        -> grad_output: (BN)CH_{out}W_{out}
+      //        grad_input: (BN)CH_{in}W_{in}
+
+      bdim_size = input.sizes()[*input_bdim];
+      // expand new_grad_output to (BN)CH_{out}W_{out}
+      new_grad_output = expand_reshape_dim_into(bdim_size, 0, new_grad_output);
+      new_input = reshape_dim_into(*input_bdim, 0, input);
+      grad_input_out_bdim = 0;
+      new_grid = reshape_dim_into(*grid_bdim, 0, grid);
+      grad_grid_out_bdim = 0;
+    } else if (input_bdim && !grid_bdim) {
+      // case 6: (grad_output is not batched, input is batched, grid is not batched)
+      //        grad_output: NCH_{out}W_{out}, input: (BN)CH_{in}W_{in}, grid: NH_{out}W_{out}2
+      //        -> grad_output: (BN)CH_{out}W_{out}
+      //        -> grid: (BN)H_{out}W_{out}2
+      //        grad_input: (BN)CH_{in}W_{in}
+
+      bdim_size = input.sizes()[*input_bdim];
+      // expand new_grad_output to (BN)CH_{out}W_{out}
+      new_grad_output = expand_reshape_dim_into(bdim_size, 0, new_grad_output);
+      new_input = reshape_dim_into(*input_bdim, 0, input);
+      grad_input_out_bdim = 0;
+      // expand new_grid to (BN)H_{out}W_{out}2
+      new_grid = expand_reshape_dim_into(bdim_size, 0, grid);
+      grad_grid_out_bdim = 0;
+    } else if (!input_bdim && grid_bdim) {
+      // case 7: (grad_output is not batched, input is not batched, grid is batched)
+      //        IF PUT BATCH DIM TO H_out -> backward produces wrong grad_grid
+      //
+      //        grad_output: NCH_{out}W_{out}, input: NCH_{in}W_{in}, grid: (BN)H_{out}W_{out}2
+      //        -> grad_output: (BN)CH_{out}W_{out}
+      //        -> input: (BN)CH_{out}W_{out}
+      //        grad_input: NCH_{in}W_{in}
+
+      bdim_size = grid.sizes()[*grid_bdim];
+      // expand new_grad_output to NC(BH_{out})W_{out}
+      new_grad_output = expand_reshape_dim_into(bdim_size, 0, new_grad_output);
+      // expand new_input to (BN)CH_{in}W_{in}
+      new_input = expand_reshape_dim_into(bdim_size, 0, new_input);
+      grad_input_out_bdim = 0;
+      new_grid = reshape_dim_into(*grid_bdim, 0, grid);
+      grad_grid_out_bdim = 0;
+    } // case 8 can be ignored
+  }
+  return std::make_tuple(
+      new_grad_output, new_input, grad_input_out_bdim, new_grid, grad_grid_out_bdim, bdim_size);
+}
+
+std::tuple<Tensor, optional<int64_t>, Tensor, optional<int64_t>>
+grid_sample_backward_helper_out(
+    const std::tuple<Tensor, Tensor> & bw_out,
+    optional<int64_t> grad_input_out_bdim,
+    optional<int64_t> grad_grid_out_bdim,
+    int64_t bdim_size) {
+  auto grad_input = std::get<0>(bw_out);
+  auto grad_grid = std::get<1>(bw_out);
+  if (grad_input_out_bdim) {
+    grad_input = reshape_dim_outof(*grad_input_out_bdim, bdim_size, grad_input);
+  }
+  if (grad_grid_out_bdim) {
+    grad_grid = reshape_dim_outof(*grad_grid_out_bdim, bdim_size, grad_grid);
+  }
+  auto result = std::make_tuple(grad_input, grad_input_out_bdim, grad_grid, grad_grid_out_bdim);
+  return result;
+}
+
+
+template<typename F, F Func, typename... ExtraArgs>
+std::tuple<Tensor, optional<int64_t>, Tensor, optional<int64_t>>
+grid_sample_backward_batch_rule(
+    const Tensor& grad_output, optional<int64_t> grad_output_bdim,
+    const Tensor& input, optional<int64_t> input_bdim,
+    const Tensor& grid, optional<int64_t> grid_bdim,
+    ExtraArgs... extra_args) {
+
+  auto new_bw_input = grid_sample_backward_helper_in(
+      grad_output, grad_output_bdim, input, input_bdim, grid, grid_bdim);
+
+  auto new_grad_output = std::get<0>(new_bw_input);
+  auto new_input = std::get<1>(new_bw_input);
+  auto grad_input_out_bdim = std::get<2>(new_bw_input);
+  auto new_grid = std::get<3>(new_bw_input);
+  auto grad_grid_out_bdim = std::get<4>(new_bw_input);
+  int64_t bdim_size = std::get<5>(new_bw_input);
+
+  auto bw_out = Func(new_grad_output, new_input, new_grid, std::forward<ExtraArgs>(extra_args)...);
+
+  return grid_sample_backward_helper_out(bw_out, grad_input_out_bdim, grad_grid_out_bdim, bdim_size);
+}
+
+template<typename F, F Func>
+std::tuple<Tensor, optional<int64_t>, Tensor, optional<int64_t>>
+cudnn_grid_sample_backward_batch_rule(
+    const Tensor& input, optional<int64_t> input_bdim,
+    const Tensor& grid, optional<int64_t> grid_bdim,
+    const Tensor& grad_output, optional<int64_t> grad_output_bdim) {
+
+  auto new_bw_input = grid_sample_backward_helper_in(
+      grad_output, grad_output_bdim, input, input_bdim, grid, grid_bdim);
+
+  auto new_grad_output = std::get<0>(new_bw_input);
+  auto new_input = std::get<1>(new_bw_input);
+  auto grad_input_out_bdim = std::get<2>(new_bw_input);
+  auto new_grid = std::get<3>(new_bw_input);
+  auto grad_grid_out_bdim = std::get<4>(new_bw_input);
+  int64_t bdim_size = std::get<5>(new_bw_input);
+
+  auto bw_out = Func(new_input, new_grid, new_grad_output);
+
+  return grid_sample_backward_helper_out(bw_out, grad_input_out_bdim, grad_grid_out_bdim, bdim_size);
+}
+
 std::tuple<Tensor, optional<int64_t>> cross_batch_rule(
     const Tensor& self, optional<int64_t> self_bdim,
     const Tensor& other, optional<int64_t> other_bdim,
@@ -370,12 +566,53 @@ struct GridSampleBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
   }
 };
 
+template <typename A, A a, typename C>
+struct GridSampleBackwardBatchRuleHelper;
+
+template <typename F, F Func, typename T1, typename T2, typename T3, typename... T>
+struct GridSampleBackwardBatchRuleHelper<F, Func, typelist<T1, T2, T3, T...>> {
+  static std::tuple<Tensor, optional<int64_t>, Tensor, optional<int64_t>> apply(
+      const Tensor& grad_output, optional<int64_t> grad_output_batch_dim,
+      const Tensor& input, optional<int64_t> input_batch_dim,
+      const Tensor& grid, optional<int64_t> grid_batch_dim,
+      T... extra_args) {
+    return grid_sample_backward_batch_rule<F, Func, T...>(
+        grad_output, grad_output_batch_dim,
+        input, input_batch_dim,
+        grid, grid_batch_dim,
+        std::forward<T>(extra_args)...);
+  }
+};
+
+template <typename F, F Func>
+struct CudnnGridSampleBackwardBatchRuleHelper {
+  static std::tuple<Tensor, optional<int64_t>, Tensor, optional<int64_t>> apply(
+      const Tensor& input, optional<int64_t> input_batch_dim,
+      const Tensor& grid, optional<int64_t> grid_batch_dim,
+      const Tensor& grad_output, optional<int64_t> grad_output_batch_dim) {
+    return cudnn_grid_sample_backward_batch_rule<F, Func>(
+        input, input_batch_dim,
+        grid, grid_batch_dim,
+        grad_output, grad_output_batch_dim
+    );
+  }
+};
+
 #define GRID_SAMPLE_BATCH_RULE(fn) SINGLE_ARG(\
     GridSampleBatchRuleHelper<\
       decltype(&ATEN_FN(fn)),\
       &ATEN_FN(fn),\
       c10::guts::function_traits<decltype(ATEN_FN(fn))>::parameter_types>::apply)
 
+#define GRID_SAMPLE_BW_BATCH_RULE(fn) SINGLE_ARG(\
+    GridSampleBackwardBatchRuleHelper<\
+      decltype(&ATEN_FN(fn)),\
+      &ATEN_FN(fn),\
+      c10::guts::function_traits<decltype(ATEN_FN(fn))>::parameter_types>::apply)
+
+#define CUDNN_GRID_SAMPLE_BW_BATCH_RULE(fn)\
+    CudnnGridSampleBackwardBatchRuleHelper<decltype(&ATEN_FN(fn)), &ATEN_FN(fn)>::apply
+
 #define UPSAMPLE_BACKWARD(op, overload) VMAP_SUPPORT(#op"."#overload, SINGLE_ARG(\
     UpsampleBackwardBatchRuleHelper<\
       decltype(&ATEN_FN2(op, overload)),\
@@ -386,6 +623,12 @@ struct GridSampleBatchRuleHelper<F, Func, typelist<T1, T2, T...>> {
   EXISTING_BDIM2(op, vec); \
   EXISTING_BDIM(op);
 
+Tensor this_grid_sampler_3d_backward_cpu(const Tensor& grad_output, const Tensor& input, const Tensor& grid, int64_t interpolation_mode, int64_t padding_mode, bool align_corners) {
+  return input;
+}
+
+
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("convolution", convolution_batch_rule);
   // m.impl("conv_transpose2d", convNd_transpose_decomp);
@@ -400,7 +643,12 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   EXISTING_BDIM(im2col_backward);
 
   VMAP_SUPPORT("grid_sampler_2d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
+  VMAP_SUPPORT("grid_sampler_2d_backward", GRID_SAMPLE_BW_BATCH_RULE(grid_sampler_2d_backward));
+
   VMAP_SUPPORT("grid_sampler_3d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
+  VMAP_SUPPORT("grid_sampler_3d_backward", GRID_SAMPLE_BW_BATCH_RULE(grid_sampler_3d_backward));
+  VMAP_SUPPORT("cudnn_grid_sampler_backward", CUDNN_GRID_SAMPLE_BW_BATCH_RULE(cudnn_grid_sampler_backward));
+
   VMAP_SUPPORT("cudnn_grid_sampler", GRID_SAMPLE_BATCH_RULE(cudnn_grid_sampler));
   VMAP_SUPPORT("cross", cross_batch_rule);