convolution_backward batch rule (#359)

zou3519 · web-flow · commit 7dcf2610f0fc · 2021-12-21T09:45:25.000-05:00
This is the most ridiculous batching rule we have. Featuring a guest appearance of efficient zeros tensors. We should really consider upstreaming einops.rearrange (https://einops.rocks/api/rearrange/). Many batching rules are straight up dimension manipulation and if we could specify that with strings we'd be done 10x faster. Test Plan: - run tests
diff --git a/functorch/csrc/BatchRulesConvolution.cpp b/functorch/csrc/BatchRulesConvolution.cpp
@@ -180,11 +180,327 @@ Tensor _convolution_decomp(
 //   static auto op = c10::Dispatcher::singleton()
 //     .findSchemaOrThrow("aten::cudnn_convolution_backward", "");
 //   return slow_fallback<Tensor,Tensor>(op, { self, grad_output, weight, padding, stride, dilation, groups, benchmark, deterministic, allow_tf32, output_mask });
-// }
+
+static Tensor compute_grad_bias(
+    const Tensor& grad_output_, std::array<bool, 3> output_mask) {
+  if (!output_mask[2]) {
+    return Tensor();
+  }
+  DimVector reduce_dims;
+  reduce_dims.resize(grad_output_.dim() - 1);
+  reduce_dims[0] = 0;
+  std::iota(reduce_dims.begin() + 1, reduce_dims.end(), 2);
+  return grad_output_.sum(reduce_dims);
+}
+
+// reshapes the batch_size into dim
+Tensor make_dummy(
+    const Tensor& tensor, optional<int64_t> tensor_bdim,
+    int64_t dim, int64_t batch_size) {
+  auto tensor_ = tensor_bdim ? tensor.select(*tensor_bdim, 0) : tensor;
+  auto orig_size = tensor_.size(dim);
+  tensor_ = tensor_.slice(dim, 0, 1);
+
+  DimVector expand_shape(tensor_.sizes().begin(), tensor_.sizes().end());
+  expand_shape[dim] = batch_size * orig_size;
+
+  // return tensor_.new_zeros(expand_shape);
+  return at::_efficientzerotensor(expand_shape, tensor.options());
+}
+
+std::tuple<Tensor,optional<int64_t>>
+convolution_backward_input_batch_rule(
+    const Tensor& grad_output, optional<int64_t> grad_output_bdim,
+    const Tensor& input, optional<int64_t> input_bdim,
+    const Tensor& weight, optional<int64_t> weight_bdim,
+    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, bool transposed,
+    IntArrayRef output_padding, int64_t groups) {
+  const std::array<bool, 3> mask = {true, false, false};
+  if (grad_output_bdim && weight_bdim) {
+    // regular: BNO, BOI -> N(BO), (BO)I -> N(BI)
+    // transposed: BNO, BIO -> N(BO), (BI)O -> N(BI)
+    const auto batch_size = weight.size(*weight_bdim);
+    const auto grad_output_ = reshape_dim_into(*grad_output_bdim, 1, grad_output);
+    const auto weight_ = reshape_dim_into(*weight_bdim, 0, weight);
+    auto dummy_input = make_dummy(input, input_bdim, 1, batch_size);
+    const auto result = at::convolution_backward(
+        grad_output_, dummy_input, weight_, nullopt, stride, padding,
+        dilation, transposed, output_padding, groups * batch_size, mask);
+    const auto grad_input = reshape_dim_outof(1, batch_size, std::get<0>(result));
+    return std::make_tuple(grad_input, 1);
+  } else if (grad_output_bdim && !weight_bdim) {
+    // BNO, OI -> (BN)O, OI -> (BN)I
+    // transposed is the same.
+    const auto batch_size = grad_output.size(*grad_output_bdim);
+    const auto grad_output_ = reshape_dim_into(*grad_output_bdim, 0, grad_output);
+    auto dummy_input = make_dummy(input, input_bdim, 0, batch_size);
+    const auto result = at::convolution_backward(
+        grad_output_, dummy_input, weight, nullopt, stride, padding,
+        dilation, transposed, output_padding, groups, mask);
+    const auto grad_input = reshape_dim_outof(0, batch_size, std::get<0>(result));
+    return std::make_tuple(grad_input, 0);
+  } else if (!grad_output_bdim && weight_bdim) {
+    const auto batch_size = weight.size(*weight_bdim);
+    if (groups == 1) {
+      // regular: NO, BOI -> NO, O(BI) -> N(BI)
+      // transposed: NO, BIO -> NO, (BI)O -> N(BI)
+      const auto in_ch_dim = transposed ? 0 : 1;
+      const auto weight_ = reshape_dim_into(*weight_bdim, in_ch_dim, weight);
+      auto dummy_input = make_dummy(input, input_bdim, 1, batch_size);
+      const auto result = at::convolution_backward(
+          grad_output, dummy_input, weight_, nullopt, stride, padding,
+          dilation, transposed, output_padding, groups, mask);
+      const auto grad_input = reshape_dim_outof(1, batch_size, std::get<0>(result));
+      return std::make_tuple(grad_input, 1);
+    }
+    Tensor grad_input;
+    if (!transposed) {
+      // N(GO), B(GO)I -> N(GO), (GO)(BI) -> N(GBI)
+      const auto weight_ = reshape_dim_into(*weight_bdim, 1, weight);
+      auto dummy_input = make_dummy(input, input_bdim, 1, batch_size);
+      const auto result = at::convolution_backward(
+          grad_output, dummy_input, weight_, nullopt, stride, padding,
+          dilation, transposed, output_padding, groups, mask);
+      grad_input = std::get<0>(result); // N(GBI)
+    } else {
+      // N(GO), B(GI)O -> N(GO), (GBI)O -> N(GBI)
+      auto weight_ = moveBatchDimToFront(weight, weight_bdim); // B(GI)O
+      weight_ = reshape_dim_outof(1, groups, weight_);         // BGIO
+      weight_ = weight_.transpose(0, 1);                       // GBIO
+      weight_ = weight_.flatten(0, 2);                         // (GBI)O
+      const auto dummy_input = make_dummy(input, input_bdim, 1, batch_size);
+      const auto result = at::convolution_backward(
+          grad_output, dummy_input, weight_, nullopt, stride, padding,
+          dilation, transposed, output_padding, groups, mask);
+      grad_input = std::get<0>(result); // N(GBI)
+    }
+    // N(GBI) -> NG(BI) -> NGBI -> NBGI -> NB(GI)
+    grad_input = reshape_dim_outof(1, groups, grad_input);
+    grad_input = reshape_dim_outof(2, batch_size, grad_input);
+    grad_input = grad_input.transpose(1, 2);
+    grad_input = reshape_dim_into(2, 2, grad_input);
+    return std::make_tuple(grad_input, 1);
+  } else {
+    TORCH_INTERNAL_ASSERT(input_bdim);
+    const auto dummy_input = make_dummy(input, input_bdim, 0, 1);
+    const auto result = at::convolution_backward(
+        grad_output, dummy_input, weight, nullopt, stride, padding,
+        dilation, transposed, output_padding, groups, mask);
+    return std::make_tuple(std::get<0>(result), nullopt);
+  }
+}
+std::tuple<Tensor,optional<int64_t>>
+convolution_backward_weight_batch_rule(
+    const Tensor& grad_output, optional<int64_t> grad_output_bdim,
+    const Tensor& input, optional<int64_t> input_bdim,
+    const Tensor& weight, optional<int64_t> weight_bdim,
+    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, bool transposed,
+    IntArrayRef output_padding, int64_t groups) {
+  const std::array<bool, 3> mask = {false, true, false};
+  if (grad_output_bdim && input_bdim) {
+    // BNO, BNI -> N(BO), N(BI) -> (BO)I (regular) (BI)O (transposed)
+    const auto batch_size = input.size(*input_bdim);
+    const auto grad_output_ = reshape_dim_into(*grad_output_bdim, 1, grad_output);
+    const auto input_ = reshape_dim_into(*input_bdim, 1, input);
+    const auto dummy_weight = make_dummy(weight, weight_bdim, 0, batch_size);
+    const auto result = at::convolution_backward(
+        grad_output_, input_, dummy_weight, nullopt, stride, padding,
+        dilation, transposed, output_padding, groups * batch_size, mask);
+    auto grad_weight = std::get<1>(result);
+    grad_weight = reshape_dim_outof(0, batch_size, grad_weight);
+    return std::make_tuple(grad_weight, 0);
+  } else if (grad_output_bdim && !input_bdim) {
+    const auto batch_size = grad_output.size(*grad_output_bdim);
+    if (groups == 1) {
+      // regular: BNO, NI -> N(BO), NI -> (BO)I
+      // transposed: BNO, NI -> N(BO), NI -> I(BO)
+      const auto grad_output_ = reshape_dim_into(*grad_output_bdim, 1, grad_output);
+      const auto out_ch_dim = transposed ? 1 : 0;
+      const auto dummy_weight = make_dummy(weight, weight_bdim, out_ch_dim, batch_size);
+      const auto result = at::convolution_backward(
+          grad_output_, input, dummy_weight, nullopt, stride, padding,
+          dilation, transposed, output_padding, groups, mask);
+      auto grad_weight = std::get<1>(result);
+      grad_weight = reshape_dim_outof(out_ch_dim, batch_size, grad_weight);
+      return std::make_tuple(grad_weight, out_ch_dim);
+    } else {
+      auto grad_output_ = moveBatchDimToFront(grad_output, grad_output_bdim); // BN(GO)
+      grad_output_ = reshape_dim_outof(2, groups, grad_output_);              // BNGO
+      grad_output_ = grad_output_.movedim(0, 2);                              // NGBO
+      grad_output_ = grad_output_.flatten(1, 3);                              // N(GBO)
+      if (!transposed) {
+        // BN(GO), N(GI) -> N(GBO), N(GI) -> (GBO)I
+        const auto dummy_weight = make_dummy(weight, weight_bdim, 0, batch_size);
+        const auto result = at::convolution_backward(
+            grad_output_, input, dummy_weight, nullopt, stride, padding,
+            dilation, transposed, output_padding, groups, mask);
+        auto grad_weight = std::get<1>(result);
+        grad_weight = grad_weight.unflatten(0, { groups, batch_size, -1 }); // GBOI
+        grad_weight = grad_weight.transpose(0, 1);                          // BGOI
+        grad_weight = grad_weight.flatten(1, 2);                            // B(GO)I
+        return std::make_tuple(grad_weight, 0);
+      } else {
+        // BN(GO), N(GI) -> N(GBO), N(GI) -> (GI)(BO)
+        const auto dummy_weight = make_dummy(weight, weight_bdim, 1, batch_size);
+        const auto result = at::convolution_backward(
+            grad_output_, input, dummy_weight, nullopt, stride, padding,
+            dilation, transposed, output_padding, groups, mask);
+        auto grad_weight = std::get<1>(result);
+        grad_weight = reshape_dim_outof(1, batch_size, grad_weight);
+        return std::make_tuple(grad_weight, 1);
+      }
+    }
+  } else if (!grad_output_bdim && input_bdim) {
+    const auto batch_size = input.size(*input_bdim);
+    if (groups == 1) {
+      // regular: NO, BNI -> NO, N(BI) -> O(BI)
+      // transposed: NO, BNI -> NO, N(BI) -> (BI)O
+      const auto input_ = reshape_dim_into(*input_bdim, 1, input);
+      const auto in_ch_dim = transposed ? 0 : 1;
+      const auto dummy_weight = make_dummy(weight, weight_bdim, in_ch_dim, batch_size);
+      const auto result = at::convolution_backward(
+          grad_output, input_, dummy_weight, nullopt, stride, padding,
+          dilation, transposed, output_padding, groups, mask);
+      auto grad_weight = std::get<1>(result);
+      grad_weight = reshape_dim_outof(in_ch_dim, batch_size, grad_weight);
+      return std::make_tuple(grad_weight, in_ch_dim);
+    } else {
+      auto input_ = moveBatchDimToFront(input, input_bdim); // BN(GI)
+      input_ = reshape_dim_outof(2, groups, input_);        // BNGI
+      input_ = input_.movedim(0, 2);                        // NGBI
+      input_ = input_.flatten(1, 3);                        // N(GBI)
+      if (!transposed) {
+        // regular: N(GO), BN(GI) -> N(GO), N(GBI) -> (GO)(BI)
+        const auto dummy_weight = make_dummy(weight, weight_bdim, 1, batch_size);
+        const auto result = at::convolution_backward(
+            grad_output, input_, dummy_weight, nullopt, stride, padding,
+            dilation, transposed, output_padding, groups, mask);
+        auto grad_weight = std::get<1>(result);
+        grad_weight = reshape_dim_outof(1, batch_size, grad_weight);
+        return std::make_tuple(grad_weight, 1);
+      } else {
+        // transposed: N(GO), BN(GI) -> N(GO), N(GBI) -> (GBI)O
+        const auto dummy_weight = make_dummy(weight, weight_bdim, 0, batch_size);
+        const auto result = at::convolution_backward(
+            grad_output, input_, dummy_weight, nullopt, stride, padding,
+            dilation, transposed, output_padding, groups, mask);
+        auto grad_weight = std::get<1>(result);
+        grad_weight = grad_weight.unflatten(0, { groups, batch_size, -1 }); // GBIO
+        grad_weight = grad_weight.transpose(0, 1);                          // BGIO
+        grad_weight = grad_weight.flatten(1, 2);                            // B(GI)O
+        return std::make_tuple(grad_weight, 0);
+      }
+    }
+  } else {
+    TORCH_INTERNAL_ASSERT(weight_bdim);
+    const auto dummy_weight = make_dummy(weight, weight_bdim, 0, 1);
+    const auto result = at::convolution_backward(
+        grad_output, input, dummy_weight, nullopt, stride, padding,
+        dilation, transposed, output_padding, groups, mask);
+    return std::make_tuple(std::get<1>(result), nullopt);
+
+  }
+}
+
+std::tuple<Tensor,Tensor,Tensor> convolution_backward_plumbing(
+    const Tensor& grad_output_, const Tensor& input_, const Tensor& weight_,
+    const c10::optional<IntArrayRef> bias_sizes_opt,
+    IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, bool transposed,
+    IntArrayRef output_padding, int64_t groups, std::array<bool, 3> output_mask) {
+  const auto maybe_layer = maybeCurrentDynamicLayer();
+  TORCH_INTERNAL_ASSERT(maybe_layer.has_value());
+  int64_t cur_level = maybe_layer->layerId();
+  Tensor grad_output;
+  optional<int64_t> grad_output_bdim;
+  std::tie(grad_output, grad_output_bdim) = unwrapTensorAtLevel(grad_output_, cur_level);
+  Tensor input;
+  optional<int64_t> input_bdim;
+  std::tie(input, input_bdim) = unwrapTensorAtLevel(input_, cur_level);
+  Tensor weight;
+  optional<int64_t> weight_bdim;
+  std::tie(weight, weight_bdim) = unwrapTensorAtLevel(weight_, cur_level);
+
+  const auto grad_bias = compute_grad_bias(grad_output_, output_mask);
+  output_mask[2] = false;
+
+  // TODO: A little bird says that unfold + matmul is actually faster than
+  // group convolution in many cases. We should benchmark some of
+  // the common cases and replace things with unfold + matmul as necessary.
+
+  // Notation:
+  // B - a batch dimension
+  // G - groups (sometimes omitted because it doesn't matter)
+  // NO - grad_output
+  // NI - input
+  // OI - weight
+  // "(BO)I" - we don't actually care about the values of this Tensor,
+  //           we just need to create a tensor on the same device with the
+  //           correct shape and pray that the implementation is smart enough
+  //           to not do anything with it.
+
+  // BNO, BNI, BOI
+  // AKA one of the model ensembling case
+  if (grad_output_bdim && input_bdim && weight_bdim) {
+    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
+    grad_output = reshape_dim_into(*grad_output_bdim, 1, grad_output);
+
+    // BNO, BNI, BOI -> N(BO), N(BI), (BO)I
+    const auto batch_size = weight.size(*weight_bdim);
+    input = reshape_dim_into(*input_bdim, 1, input);
+    weight = reshape_dim_into(*weight_bdim, 0, weight);
+    const auto result = at::convolution_backward(
+        grad_output, input, weight, nullopt, stride, padding, dilation,
+        transposed, output_padding, batch_size * groups, output_mask);
+    // N(BI), (BO)I -> NBI, BOI
+    const auto grad_input = output_mask[0] ?
+      reshape_dim_outof(1, batch_size, std::get<0>(result)) : Tensor();
+    const auto grad_weight = output_mask[1] ?
+      reshape_dim_outof(0, batch_size, std::get<1>(result)) : Tensor();
+    return std::make_tuple(
+        output_mask[0] ? makeBatched(grad_input, 1, cur_level) : grad_input,
+        output_mask[1] ? makeBatched(grad_weight, 0, cur_level) : grad_weight,
+        grad_bias);
+  }
+
+  Tensor grad_input;
+  if (output_mask[0]) {
+    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
+    const auto result = convolution_backward_input_batch_rule(
+        grad_output, grad_output_bdim,
+        input, input_bdim,
+        weight, weight_bdim,
+        stride, padding, dilation, transposed, output_padding, groups);
+    grad_input = makeBatched(std::get<0>(result), std::get<1>(result), cur_level);
+  }
+
+  Tensor grad_weight;
+  if (output_mask[1]) {
+    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
+    const auto result = convolution_backward_weight_batch_rule(
+        grad_output, grad_output_bdim,
+        input, input_bdim,
+        weight, weight_bdim,
+        stride, padding, dilation, transposed, output_padding, groups);
+    grad_weight = makeBatched(std::get<0>(result), std::get<1>(result), cur_level);
+  }
+  return std::make_tuple(grad_input, grad_weight, grad_bias);
+
+  // Someone's definitely going to find a problem with this batching rule so
+  // I'm leaving the following fallback if we need it back.
+  // static auto op = c10::Dispatcher::singleton()
+  //   .findSchemaOrThrow("aten::convolution_backward", "");
+  // auto result = slow_fallback<Tensor,Tensor,Tensor>(op, {
+  //   grad_output_, input_, weight_, bias_sizes_opt,
+  //   stride, padding, dilation, transposed, output_padding, groups, output_mask
+  // });
+  // return std::make_tuple(grad_input, std::get<1>(result), grad_bias);
+}
+
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("convolution", convolution_batch_rule);
   m.impl("_convolution", _convolution_decomp);
+  m.impl("convolution_backward", convolution_backward_plumbing);
 }
 
 }} // namespace at;:functorch
diff --git a/test/test_ops.py b/test/test_ops.py
@@ -790,7 +790,6 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('masked_select'),
         xfail('matrix_exp'),
         xfail('nanquantile'),
-        xfail('nn.functional.conv_transpose2d'),
         xfail('nn.functional.gelu'),
         xfail('norm', 'nuc'),
         xfail('pinverse'),
@@ -816,13 +815,11 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('cross'),
         xfail('double', 'channels_last'),
         xfail('linalg.cross'),
-        skip('nn.functional.conv1d'),
         xfail('nn.functional.gaussian_nll_loss'),
         xfail('nn.functional.hardsigmoid'),
         xfail('nn.functional.huber_loss'),
         xfail('nn.functional.instance_norm'),
         xfail('nn.functional.poisson_nll_loss'),
-        xfail('nn.functional.conv_transpose3d'),
         xfail('nn.functional.bilinear'),
         xfail('nn.functional.prelu'),
         xfail('nn.functional.glu'),
@@ -833,16 +830,12 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('nn.functional.rrelu'),
         xfail('nn.functional.embedding_bag'),
         xfail('nn.functional.softshrink'),
-        xfail('nn.functional.conv_transpose1d'),
         xfail('nn.functional.max_pool3d'),
         xfail('istft'),
         xfail('nn.functional.fractional_max_pool2d'),
         xfail('linalg.tensorsolve'),
     }))
     def test_vmapvjp_has_batch_rule(self, device, dtype, op):
-        # These are too annoying to put into the list above
-        if op.name in {'nn.functional.conv2d'}:
-            self.skipTest("Skipped! ExpectedF failures")
         if not op.supports_autograd:
             self.skipTest("Skipped! Autograd not supported.")
             return
