jagged_dense_elementwise_mul add Meta and Autograd backend

Summary:
Add Meta and Autograd backend for jagged_dense_elementwise_mul for dynamo and AOT autograd.

A more proper way of doing Autograd to make inductor working for jagged ops. Also added Meta tensors so we're not mercy of arbitrary zero inputs fed into by default.

Differential Revision: D41383597

fbshipit-source-id: 6b3c883980e036a9b130ee588c6186c0ca072680

Author: brad-mengchi

fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h

@@ -427,6 +427,17 @@ at::Tensor jagged_dense_dense_elementwise_add_jagged_output_forward(
     const at::Tensor& y_0,
     const at::Tensor& y_1);
 
+at::Tensor jagged_dense_elementwise_mul_forward(
+    const at::Tensor& x_values,
+    const std::vector<at::Tensor>& x_offsets,
+    const at::Tensor& y);
+
+std::tuple<at::Tensor, at::Tensor> jagged_dense_elementwise_mul_backward(
+    const at::Tensor& grad_output,
+    const std::vector<at::Tensor>& x_offsets,
+    const at::Tensor& y,
+    const at::Tensor& x_values);
+
 ///@ingroup sparse-data-cuda
 at::Tensor jagged_2d_to_dense_gpu(
     at::Tensor values,

fbgemm_gpu/src/jagged_tensor_ops.cu

@@ -1635,6 +1635,82 @@ void jagged_jagged_elementwise_dense_output_(
 #undef INVOKE_KERNEL_WITH_DIM
 }
 
+Tensor jagged_dense_elementwise_mul_forward(
+    const Tensor& x_values,
+    const std::vector<Tensor>& x_offsets,
+    const Tensor& y) {
+  at::cuda::OptionalCUDAGuard device_guard;
+  device_guard.set_index(x_values.get_device());
+
+  Tensor output = at::empty_like(x_values);
+
+  AT_DISPATCH_SWITCH(
+      x_values.scalar_type(),
+      "jagged_dense_elementwise_mul_jagged_output_forward",
+      AT_DISPATCH_CASE(
+          at::ScalarType::Half,
+          [&] {
+            jagged_dense_elementwise_jagged_output_opt_<scalar_t>(
+                x_values,
+                x_offsets,
+                y,
+                output,
+                [] __device__(scalar_t x, scalar_t y) -> scalar_t {
+                  return x * y;
+                });
+          } // lambda
+          ) // CASE
+      AT_DISPATCH_CASE_FLOATING_TYPES([&] {
+        jagged_dense_elementwise_jagged_output_<scalar_t>(
+            x_values,
+            x_offsets,
+            y,
+            output,
+            [] __device__(scalar_t x, scalar_t y) -> scalar_t {
+              return x * y;
+            });
+      } // lambda
+                                      ) // CASE_FLOATING_TYPES_AND
+  ); // SWITCH
+
+  return output;
+}
+
+std::tuple<Tensor, Tensor> jagged_dense_elementwise_mul_backward(
+    const Tensor& grad_output,
+    const std::vector<Tensor>& x_offsets,
+    const Tensor& y,
+    const Tensor& x_values) {
+  at::cuda::OptionalCUDAGuard device_guard;
+  device_guard.set_index(grad_output.get_device());
+
+  Tensor x_values_grad = at::empty_like(grad_output);
+  Tensor y_grad = at::empty_like(y);
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      x_values.scalar_type(), "jagged_scalars", [&] {
+        jagged_dense_elementwise_jagged_output_<scalar_t>(
+            grad_output,
+            x_offsets,
+            y,
+            x_values_grad,
+            [] __device__(scalar_t x, scalar_t y) -> scalar_t {
+              return x * y;
+            });
+
+        jagged_jagged_elementwise_dense_output_<scalar_t>(
+            grad_output,
+            x_offsets,
+            x_values,
+            y_grad,
+            [] __device__(scalar_t x, scalar_t y) -> scalar_t {
+              return x * y;
+            });
+      });
+
+  return {x_values_grad, y_grad};
+}
+
 class JaggedDenseMulGPUOp
     : public torch::autograd::Function<JaggedDenseMulGPUOp> {
  public:
@@ -1650,39 +1726,7 @@ class JaggedDenseMulGPUOp
     tensors_to_save.push_back(y);
     ctx->save_for_backward(tensors_to_save);
 
-    at::cuda::OptionalCUDAGuard device_guard;
-    device_guard.set_index(x_values.get_device());
-
-    Tensor output = at::empty_like(x_values);
-
-    AT_DISPATCH_SWITCH(
-        x_values.scalar_type(),
-        "jagged_dense_elementwise_mul_jagged_output_forward",
-        AT_DISPATCH_CASE(
-            at::ScalarType::Half,
-            [&] {
-              jagged_dense_elementwise_jagged_output_opt_<scalar_t>(
-                  x_values,
-                  x_offsets,
-                  y,
-                  output,
-                  [] __device__(scalar_t x, scalar_t y) -> scalar_t {
-                    return x * y;
-                  });
-            } // lambda
-            ) // CASE
-        AT_DISPATCH_CASE_FLOATING_TYPES([&] {
-          jagged_dense_elementwise_jagged_output_<scalar_t>(
-              x_values,
-              x_offsets,
-              y,
-              output,
-              [] __device__(scalar_t x, scalar_t y) -> scalar_t {
-                return x * y;
-              });
-        } // lambda
-                                        ) // CASE_FLOATING_TYPES_AND
-    ); // SWITCH
+    auto output = jagged_dense_elementwise_mul_forward(x_values, x_offsets, y);
 
     return {output};
   }
@@ -1698,34 +1742,13 @@ class JaggedDenseMulGPUOp
     Tensor y = ctx->get_saved_variables().back();
     TORCH_CHECK(grad_outputs.size() == 1);
 
-    at::cuda::OptionalCUDAGuard device_guard;
-    device_guard.set_index(grad_outputs[0].get_device());
-
-    Tensor x_values_grad = at::empty_like(grad_outputs[0]);
-    Tensor y_grad = at::empty_like(y);
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        x_values.scalar_type(), "jagged_scalars", [&] {
-          jagged_dense_elementwise_jagged_output_<scalar_t>(
-              grad_outputs[0],
-              x_offsets,
-              y,
-              x_values_grad,
-              [] __device__(scalar_t x, scalar_t y) -> scalar_t {
-                return x * y;
-              });
-
-          jagged_jagged_elementwise_dense_output_<scalar_t>(
-              grad_outputs[0],
-              x_offsets,
-              x_values,
-              y_grad,
-              [] __device__(scalar_t x, scalar_t y) -> scalar_t {
-                return x * y;
-              });
-        });
+    auto outputs = jagged_dense_elementwise_mul_backward(
+        grad_outputs[0], x_offsets, y, x_values);
 
-    return {x_values_grad, y_grad, torch::autograd::Variable()};
+    return {
+        std::get<0>(outputs),
+        std::get<1>(outputs),
+        torch::autograd::Variable()};
   }
 };
 
@@ -3006,6 +3029,12 @@ TORCH_LIBRARY_IMPL(fbgemm, CUDA, m) {
       fbgemm_gpu::jagged_dense_dense_elementwise_add_jagged_output);
   DISPATCH_TO_CUDA(
       "jagged_dense_elementwise_mul", fbgemm_gpu::jagged_dense_elementwise_mul);
+  DISPATCH_TO_CUDA(
+      "jagged_dense_elementwise_mul_forward",
+      fbgemm_gpu::jagged_dense_elementwise_mul_forward);
+  DISPATCH_TO_CUDA(
+      "jagged_dense_elementwise_mul_backward",
+      fbgemm_gpu::jagged_dense_elementwise_mul_backward);
   DISPATCH_TO_CUDA(
       "batched_dense_vec_jagged_2d_mul",
       fbgemm_gpu::batched_dense_vec_jagged_2d_mul);

fbgemm_gpu/src/jagged_tensor_ops_autograd.cpp

@@ -65,8 +65,9 @@ class JaggedToPaddedDenseAutogradOp
   }
 };
 
-class JaggedDenseDenseAddJaggedOutputOp
-    : public torch::autograd::Function<JaggedDenseDenseAddJaggedOutputOp> {
+class JaggedDenseDenseAddJaggedOutputAutogradOp
+    : public torch::autograd::Function<
+          JaggedDenseDenseAddJaggedOutputAutogradOp> {
  public:
   static torch::autograd::variable_list forward(
       torch::autograd::AutogradContext* ctx,
@@ -116,6 +117,56 @@ class JaggedDenseDenseAddJaggedOutputOp
   }
 };
 
+class JaggedDenseMulAutogradOp
+    : public torch::autograd::Function<JaggedDenseMulAutogradOp> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const Tensor& x_values,
+      const std::vector<Tensor>& x_offsets,
+      const Tensor& y) {
+    std::vector<Tensor> tensors_to_save;
+    tensors_to_save.push_back(x_values);
+    tensors_to_save.insert(
+        tensors_to_save.end(), x_offsets.begin(), x_offsets.end());
+    tensors_to_save.push_back(y);
+    ctx->save_for_backward(tensors_to_save);
+
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow(
+                "fbgemm::jagged_dense_elementwise_mul_forward", "")
+            .typed<decltype(jagged_dense_elementwise_mul_forward)>();
+    Tensor output = op.call(x_values, x_offsets, y);
+
+    return {output};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_outputs) {
+    const Tensor x_values = ctx->get_saved_variables().front();
+    std::vector<Tensor> x_offsets;
+    for (size_t i = 1; i < ctx->get_saved_variables().size() - 1; ++i) {
+      x_offsets.push_back(ctx->get_saved_variables()[i]);
+    }
+    Tensor y = ctx->get_saved_variables().back();
+    TORCH_CHECK(grad_outputs.size() == 1);
+
+    static auto op =
+        c10::Dispatcher::singleton()
+            .findSchemaOrThrow(
+                "fbgemm::jagged_dense_elementwise_mul_backward", "")
+            .typed<decltype(jagged_dense_elementwise_mul_backward)>();
+    auto outputs = op.call(grad_outputs[0], x_offsets, y, x_values);
+
+    return {
+        std::get<0>(outputs),
+        torch::autograd::Variable(),
+        std::get<1>(outputs)};
+  }
+};
+
 ///@ingroup jagged-tensor-ops-autograd
 Tensor jagged_to_padded_dense_autograd(
     const Tensor& values,
@@ -158,12 +209,22 @@ jagged_dense_dense_elementwise_add_jagged_output_autograd(
     const std::vector<Tensor>& x_offsets,
     const Tensor& y_0,
     const Tensor& y_1) {
-  auto sum_values = JaggedDenseDenseAddJaggedOutputOp::apply(
+  auto sum_values = JaggedDenseDenseAddJaggedOutputAutogradOp::apply(
       x_values, x_offsets, y_0, y_1)[0];
 
   return {sum_values, x_offsets};
 }
 
+std::tuple<Tensor, std::vector<Tensor>> jagged_dense_elementwise_mul_autograd(
+    const Tensor& x_values,
+    const std::vector<Tensor>& x_offsets,
+    const Tensor& y) {
+  // Convert to jagged
+  auto prod_values = JaggedDenseMulAutogradOp::apply(x_values, x_offsets, y)[0];
+
+  return {prod_values, x_offsets};
+}
+
 } // namespace fbgemm_gpu
 
 TORCH_LIBRARY_IMPL(fbgemm, Autograd, m) {
@@ -178,4 +239,7 @@ TORCH_LIBRARY_IMPL(fbgemm, Autograd, m) {
       "jagged_dense_dense_elementwise_add_jagged_output",
       TORCH_FN(fbgemm_gpu::
                    jagged_dense_dense_elementwise_add_jagged_output_autograd));
+  m.impl(
+      "jagged_dense_elementwise_mul",
+      TORCH_FN(fbgemm_gpu::jagged_dense_elementwise_mul_autograd));
 }