Convolution double backward registered.

Author: vbharadwaj-bk

openequivariance_extjax/src/libjax_tp_jit.cpp

@@ -529,6 +529,72 @@ ffi::Error conv_backward_impl(
     return ffi::Error::Success();
 }
 
+ffi::Error conv_double_backward_impl(
+        ffi::AnyBuffer L1_in,
+        ffi::AnyBuffer L2_in,
+        ffi::AnyBuffer W,
+        ffi::AnyBuffer L3_grad,
+        ffi::AnyBuffer L1_dgrad,
+        ffi::AnyBuffer L2_dgrad,
+        ffi::AnyBuffer W_dgrad,
+        ffi::Result<ffi::AnyBuffer> L1_grad,
+        ffi::Result<ffi::AnyBuffer> L2_grad,
+        ffi::Result<ffi::AnyBuffer> W_grad,
+        ffi::Result<ffi::AnyBuffer> L3_dgrad,
+        ffi::AnyBuffer rows,
+        ffi::AnyBuffer cols,
+        ffi::AnyBuffer workspace,
+        ffi::AnyBuffer transpose_perm,
+        cudaStream_t stream, 
+        std::string_view kernel, ffi::Dictionary forward_config, ffi::Dictionary backward_config, ffi::Dictionary double_backward_config, ffi::Dictionary kernel_prop,
+        int64_t hash) {
+    
+    auto [jit_kernel, k] = compile_conv_with_caching(
+        kernel, forward_config, backward_config, double_backward_config, kernel_prop, hash, true);
+    const int64_t nnz = rows.dimensions()[0];
+    const int64_t node_count = L1_in.dimensions()[0];
+    check_tensor(L1_in, {node_count, k.L1_dim}, k.irrep_dtype, "L1_in");
+    check_tensor(L2_in, {nnz, k.L2_dim}, k.irrep_dtype, "L2_in");
+    check_tensor(L3_grad, {node_count, k.L3_dim}, k.irrep_dtype, "L3_grad");
+    check_tensor(L1_dgrad, {node_count, k.L1_dim}, k.irrep_dtype, "L1_dgrad");
+    check_tensor(L2_dgrad, {nnz, k.L2_dim}, k.irrep_dtype, "L2_dgrad");
+    check_tensor(workspace, {k.workspace_size}, k.workspace_dtype, "workspace");
+    check_tensor(rows, {nnz}, k.idx_dtype, "rows");
+    check_tensor(cols, {nnz}, k.idx_dtype, "cols");
+
+    if (k.deterministic)
+        check_tensor(transpose_perm, {nnz}, k.idx_dtype, "transpose perm");
+    
+    if (k.shared_weights) {
+        check_tensor(W, {k.weight_numel}, k.weight_dtype, "W");
+        check_tensor(W_dgrad, {k.weight_numel}, k.weight_dtype, "W_dgrad");
+    } else {
+        check_tensor(W, {nnz, k.weight_numel}, k.weight_dtype, "W");
+        check_tensor(W_dgrad, {nnz, k.weight_numel}, k.weight_dtype, "W_dgrad");
+    }
+    if(k.shared_weights)
+        zero_buffer(*W_grad);
+    jit_kernel->double_backward(
+            data_ptr(L1_in),
+            data_ptr(L2_in),
+            data_ptr(W),
+            data_ptr(L3_grad),
+            data_ptr(L1_dgrad),
+            data_ptr(L2_dgrad),
+            data_ptr(W_dgrad),
+            data_ptr(L1_grad),
+            data_ptr(L2_grad),
+            data_ptr(W_grad),
+            data_ptr(L3_dgrad),
+            data_ptr(rows),
+            data_ptr(cols),
+            nnz, node_count,
+            data_ptr(workspace),
+            data_ptr(transpose_perm),
+            stream);
+    return ffi::Error::Success();
+}
+
 XLA_FFI_DEFINE_HANDLER_SYMBOL(
     conv_forward, conv_forward_impl,
     ffi::Ffi::Bind()
@@ -564,6 +630,30 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
         .Attr<int64_t>("hash"),
         {xla::ffi::Traits::kCmdBufferCompatible});
 
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    conv_double_backward, conv_double_backward_impl,
+    ffi::Ffi::Bind()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Ret<ffi::AnyBuffer>()
+        .Ret<ffi::AnyBuffer>()
+        .Ret<ffi::AnyBuffer>()
+        .Ret<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Ctx<ffi::PlatformStream<cudaStream_t>>()
+        .Attr<std::string_view>("kernel").Attr<ffi::Dictionary>("forward_config").Attr<ffi::Dictionary>("backward_config").Attr<ffi::Dictionary>("double_backward_config").Attr<ffi::Dictionary>("kernel_prop")
+        .Attr<int64_t>("hash"),
+        {xla::ffi::Traits::kCmdBufferCompatible});
+
+// --------------------- NB Module --------------------------
 NB_MODULE(openequivariance_extjax, m) {
     m.def("registrations", []() {
         nb::dict registrations;
@@ -572,6 +662,8 @@ NB_MODULE(openequivariance_extjax, m) {
         registrations["tp_double_backward"] = nb::capsule(reinterpret_cast<void *>(tp_double_backward));
 
         registrations["conv_forward"] = nb::capsule(reinterpret_cast<void *>(conv_forward));
+        registrations["conv_backward"] = nb::capsule(reinterpret_cast<void *>(conv_backward));
+        registrations["conv_double_backward"] = nb::capsule(reinterpret_cast<void *>(conv_double_backward));
         return registrations;
     });