Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/vipe/csrc/bind.cpp

@@ -0,0 +1,49 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+void pybind_droid_net_ext(py::module &m);
+void pybind_grounding_dino_ext(py::module &m);
+void pybind_utils_ext(py::module &m);
+void pybind_slam_ext(py::module &m);
+void pybind_scatter_ext(py::module &m);
+void pybind_lietorch_ext(py::module &m);
+void pybind_corr_ext(py::module &m);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    py::module m_droid_net = m.def_submodule("droid_net_ext");
+    pybind_droid_net_ext(m_droid_net);
+
+    py::module m_gdino = m.def_submodule("grounding_dino_ext");
+    pybind_grounding_dino_ext(m_gdino);
+
+    py::module m_utils = m.def_submodule("utils_ext");
+    pybind_utils_ext(m_utils);
+
+    py::module m_slam = m.def_submodule("slam_ext");
+    pybind_slam_ext(m_slam);
+
+    py::module m_scatter = m.def_submodule("scatter_ext");
+    pybind_scatter_ext(m_scatter);
+
+    py::module m_lietorch = m.def_submodule("lietorch_ext");
+    pybind_lietorch_ext(m_lietorch);
+
+    py::module m_corr = m.def_submodule("corr_ext");
+    pybind_corr_ext(m_corr);
+}

third_party/vipe/csrc/corr_ext/correlation.cpp

@@ -0,0 +1,148 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+using namespace torch;
+
+#include <vector>
+
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+template <typename scalar_t>
+static void correlate_patch(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> input2, scalar_t *dst,
+                            int kH, int kW, int dilationH, int dilationW, int u, int v, int shiftU, int shiftV) {
+    const int C = input1.size(0);
+    const int iH = input1.size(1);
+    const int iW = input1.size(2);
+    for (int c = 0; c < C; ++c) {
+        for (int i = 0; i < kH; ++i) {
+            int i1 = u + i * dilationH;
+            int i2 = i1 + shiftU;
+            if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                for (int j = 0; j < kW; ++j) {
+                    int j1 = v + j * dilationW;
+                    int j2 = j1 + shiftV;
+                    if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                        scalar_t v1 = input1[c][i1][j1];
+                        scalar_t v2 = input2[c][i2][j2];
+                        *dst += v1 * v2;
+                    }
+                }
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+static void correlate_patch_grad(TensorAccessor<scalar_t, 3> input1, TensorAccessor<scalar_t, 3> gradInput1,
+                                 TensorAccessor<scalar_t, 3> input2, TensorAccessor<scalar_t, 3> gradInput2,
+                                 scalar_t gradOutput, int kH, int kW, int dilationH, int dilationW, int u, int v,
+                                 int shiftU, int shiftV) {
+    const int C = input1.size(0);
+    const int iH = input1.size(1);
+    const int iW = input1.size(2);
+
+    for (int c = 0; c < C; ++c) {
+        for (int i = 0; i < kH; ++i) {
+            int i1 = u + i * dilationH;
+            int i2 = i1 + shiftU;
+            if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                for (int j = 0; j < kW; ++j) {
+                    int j1 = v + j * dilationW;
+                    int j2 = j1 + shiftV;
+                    if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                        scalar_t v1 = input1[c][i1][j1];
+                        scalar_t v2 = input2[c][i2][j2];
+                        gradInput2[c][i2][j2] += gradOutput * v1;
+                        gradInput1[c][i1][j1] += gradOutput * v2;
+                    }
+                }
+            }
+        }
+    }
+}
+
+torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                      int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
+                                      int dilation_patchW, int dH, int dW) {
+    const auto batch_size = input1.size(0);
+    const auto iH = input1.size(2);
+    const auto iW = input1.size(3);
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+    const int dilatedKH = (kH - 1) * dilationH + 1;
+    const int dilatedKW = (kW - 1) * dilationW + 1;
+
+    const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+    const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+    auto output = at::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+
+    int n, ph, pw, h, w;
+#pragma omp parallel for private(n, ph, pw, h, w) collapse(2)
+    for (n = 0; n < batch_size; ++n) {
+        for (ph = 0; ph < patchH; ++ph) {
+            for (pw = 0; pw < patchW; ++pw) {
+                AT_DISPATCH_FLOATING_TYPES(
+                    input1.scalar_type(), "correlation_forward_cpp", ([&] {
+                        auto input1_acc = input1.accessor<scalar_t, 4>();
+                        auto input2_acc = input2.accessor<scalar_t, 4>();
+                        auto output_acc = output.accessor<scalar_t, 5>();
+                        for (h = 0; h < oH; ++h) {
+                            for (w = 0; w < oW; ++w) {
+                                correlate_patch(input1_acc[n], input2_acc[n], &output_acc[n][ph][pw][h][w], kH, kW,
+                                                dilationH, dilationW, -padH + h * dH, -padW + w * dW,
+                                                (ph - patchRadH) * dilation_patchH, (pw - patchRadW) * dilation_patchW);
+                            }
+                        }
+                    }));
+            }
+        }
+    }
+    return output;
+}
+
+std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor input1, torch::Tensor input2,
+                                                    torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
+                                                    int padH, int padW, int dilationH, int dilationW,
+                                                    int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    const int batch_size = input1.size(0);
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+    const int oH = gradOutput.size(3);
+    const int oW = gradOutput.size(4);
+
+    auto gradInput1 = torch::zeros_like(input1);
+
+    auto gradInput2 = torch::zeros_like(input2);
+
+    int n, ph, pw, h, w;
+#pragma omp parallel for private(n, ph, pw, h, w)
+    for (n = 0; n < batch_size; ++n) {
+        AT_DISPATCH_FLOATING_TYPES(
+            input1.scalar_type(), "correlation_backward_cpp", ([&] {
+                auto input1_acc = input1.accessor<scalar_t, 4>();
+                auto gradInput1_acc = gradInput1.accessor<scalar_t, 4>();
+                auto input2_acc = input2.accessor<scalar_t, 4>();
+                auto gradInput2_acc = gradInput2.accessor<scalar_t, 4>();
+                auto gradOutput_acc = gradOutput.accessor<scalar_t, 5>();
+
+                for (ph = 0; ph < patchH; ++ph) {
+                    for (pw = 0; pw < patchW; ++pw) {
+                        for (h = 0; h < oH; ++h) {
+                            for (w = 0; w < oW; ++w) {
+                                correlate_patch_grad(input1_acc[n], gradInput1_acc[n], input2_acc[n], gradInput2_acc[n],
+                                                     gradOutput_acc[n][ph][pw][h][w], kH, kW, dilationH, dilationW,
+                                                     -padH + h * dH, -padW + w * dW, (ph - patchRadH) * dilation_patchH,
+                                                     (pw - patchRadW) * dilation_patchW);
+                            }
+                        }
+                    }
+                }
+            }));
+    }
+
+    return {gradInput1, gradInput2};
+}

third_party/vipe/csrc/corr_ext/correlation_cuda_kernel.cu

@@ -0,0 +1,286 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/types.h>
+using namespace torch;
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+#include <iostream>
+
+// Cuda tensor accessor definitions
+// restrict pointer traits piroritize speed over memory consumption
+#define TensorAcc4R PackedTensorAccessor32<scalar_t, 4, RestrictPtrTraits>
+#define TensorAcc5R PackedTensorAccessor32<scalar_t, 5, RestrictPtrTraits>
+#define WITHIN_BOUNDS(x, y, H, W) (x >= 0 && x < H && y >= 0 && y < W)
+
+#define THREADS_FORWARD 32
+#define THREADS_BACKWARD 5
+
+namespace corr {
+template <typename scalar_t>
+__global__ void correlation_cuda_forward_kernel(const TensorAcc4R rInput1, const TensorAcc4R rInput2,
+                                                TensorAcc5R output, int kH, int kW, int patchH, int patchW, int padH,
+                                                int padW, int dilationH, int dilationW, int dilation_patchH,
+                                                int dilation_patchW, int dH, int dW) {
+    const int iH = rInput1.size(1);
+    const int iW = rInput1.size(2);
+    const int C = rInput1.size(3);
+
+    const int n = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int thread = threadIdx.x;
+
+    const int start_i = -padH + h * dH;
+    const int start_j = -padW + w * dW;
+
+    const int patchRadH = dilation_patchH * (patchH - 1) / 2;
+    const int patchRadW = dilation_patchW * (patchW - 1) / 2;
+
+    __shared__ scalar_t prod_sum[THREADS_FORWARD];
+
+    for (int ph = 0; ph < patchH; ++ph) {
+        int ph_dilated = ph * dilation_patchH - patchRadH;
+        for (int pw = 0; pw < patchW; ++pw) {
+            int pw_dilated = pw * dilation_patchW - patchRadW;
+            prod_sum[thread] = 0;
+            for (int i = 0; i < kH; ++i) {
+                int i1 = start_i + i * dilationH;
+                int i2 = i1 + ph_dilated;
+                if WITHIN_BOUNDS (i1, i2, iH, iH) {
+                    for (int j = 0; j < kW; ++j) {
+                        int j1 = start_j + j * dilationW;
+                        int j2 = j1 + pw_dilated;
+                        if WITHIN_BOUNDS (j1, j2, iW, iW) {
+                            for (int c = thread; c < C; c += THREADS_FORWARD) {
+                                scalar_t v1 = rInput1[n][i1][j1][c];
+                                scalar_t v2 = rInput2[n][i2][j2][c];
+                                prod_sum[thread] += v1 * v2;
+                            }
+                        }
+                    }
+                }
+            }
+            // accumulate
+            __syncthreads();
+            if (thread == 0) {
+                scalar_t reduce_sum = 0;
+                for (int index = 0; index < THREADS_FORWARD; ++index) {
+                    reduce_sum += prod_sum[index];
+                }
+                output[n][ph][pw][h][w] = reduce_sum;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input1(const TensorAcc5R gradOutput, const TensorAcc4R input2,
+                                                        TensorAcc4R gradInput1, const int kH, const int kW,
+                                                        const int patchH, const int patchW, const int padH,
+                                                        const int padW, const int dilationH, const int dilationW,
+                                                        const int dilation_patchH, const int dilation_patchW,
+                                                        const int dH, const int dW, const int batch) {
+    const int iH = input2.size(2);
+    const int iW = input2.size(3);
+
+    const int H = gradOutput.size(3);
+    const int W = gradOutput.size(4);
+
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+
+    const int n = batch;
+    const int c = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int ph_off = threadIdx.x;
+    const int pw_off = threadIdx.y;
+
+    const int h_2 = h + padH;
+    const int w_2 = w + padW;
+    const int min_h = h_2 - kH * dilationH;
+    const int min_w = w_2 - kW * dilationW;
+
+    __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+    prod_sum[ph_off][pw_off] = 0;
+
+    for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+        int i1 = h + dilation_patchH * (ph - patchRadH);
+        for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+            int j1 = w + dilation_patchW * (pw - patchRadW);
+            if (WITHIN_BOUNDS(i1, j1, iH, iW)) {
+                scalar_t val = input2[n][c][i1][j1];
+                for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+                    int i2 = (h_3) / dH;
+                    if (i2 * dH != h_3) continue;
+                    for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+                        int j2 = (w_3) / dW;
+                        if (j2 * dW != w_3) continue;
+                        if WITHIN_BOUNDS (i2, j2, H, W) {
+                            prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    __syncthreads();
+
+    if (ph_off == 0 && pw_off == 0) {
+        scalar_t reduce_sum = 0;
+        for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
+            for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
+                reduce_sum += prod_sum[ph][pw];
+            }
+        }
+        gradInput1[n][c][h][w] = reduce_sum;
+    }
+}
+
+template <typename scalar_t>
+__global__ void correlation_cuda_backward_kernel_input2(const TensorAcc5R gradOutput, const TensorAcc4R input1,
+                                                        TensorAcc4R gradInput2, int kH, int kW, int patchH, int patchW,
+                                                        int padH, int padW, int dilationH, int dilationW,
+                                                        int dilation_patchH, int dilation_patchW, int dH, int dW,
+                                                        int batch) {
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+
+    const int patchRadH = (patchH - 1) / 2;
+    const int patchRadW = (patchW - 1) / 2;
+
+    const int H = gradOutput.size(3);
+    const int W = gradOutput.size(4);
+
+    const int dilatedKH = kH * dilationH;
+    const int dilatedKW = kW * dilationW;
+
+    const int n = batch;
+    const int c = blockIdx.x;
+    const int h = blockIdx.y;
+    const int w = blockIdx.z;
+    const int ph_off = threadIdx.x;
+    const int pw_off = threadIdx.y;
+
+    __shared__ scalar_t prod_sum[THREADS_BACKWARD][THREADS_BACKWARD];
+    prod_sum[ph_off][pw_off] = 0;
+
+    for (int ph = ph_off; ph < patchH; ph += THREADS_BACKWARD) {
+        int i1 = h - dilation_patchH * (ph - patchRadH);
+        for (int pw = pw_off; pw < patchW; pw += THREADS_BACKWARD) {
+            int j1 = w - dilation_patchW * (pw - patchRadW);
+            if WITHIN_BOUNDS (i1, j1, iH, iW) {
+                scalar_t val = input1[n][c][i1][j1];
+
+                const int h_2 = i1 + padH;
+                const int w_2 = j1 + padW;
+                const int min_h = h_2 - dilatedKH;
+                const int min_w = w_2 - dilatedKW;
+
+                for (int h_3 = h_2; h_3 > min_h; h_3 -= dilationH) {
+                    int i2 = (h_3) / dH;
+                    if (i2 * dH != h_3) continue;
+                    for (int w_3 = w_2; w_3 > min_w; w_3 -= dilationW) {
+                        int j2 = (w_3) / dW;
+                        if (j2 * dW != w_3) continue;
+                        if WITHIN_BOUNDS (i2, j2, H, W) {
+                            prod_sum[ph_off][pw_off] += gradOutput[n][ph][pw][i2][j2] * val;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    __syncthreads();
+
+    if (ph_off == 0 && pw_off == 0) {
+        scalar_t reduce_sum = 0;
+        for (int ph = 0; ph < THREADS_BACKWARD; ++ph) {
+            for (int pw = 0; pw < THREADS_BACKWARD; ++pw) {
+                reduce_sum += prod_sum[ph][pw];
+            }
+        }
+        gradInput2[n][c][h][w] = reduce_sum;
+    }
+}
+}  // namespace corr
+
+torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                       int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    const int batch_size = input1.size(0);
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+    const int dilatedKH = (kH - 1) * dilationH + 1;
+    const int dilatedKW = (kW - 1) * dilationW + 1;
+
+    const auto oH = (iH + 2 * padH - dilatedKH) / dH + 1;
+    const auto oW = (iW + 2 * padW - dilatedKW) / dW + 1;
+    auto output = torch::zeros({batch_size, patchH, patchW, oH, oW}, input1.options());
+
+    auto trInput1 = input1.permute({0, 2, 3, 1}).contiguous();
+    auto trInput2 = input2.permute({0, 2, 3, 1}).contiguous();
+
+    const int threads = THREADS_FORWARD;
+    const dim3 blocks(batch_size, oH, oW);
+
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_forward_cuda", ([&] {
+            TensorAcc4R trInput1_acc = trInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R trInput2_acc = trInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc5R output_acc = output.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
+            corr::correlation_cuda_forward_kernel<scalar_t>
+                <<<blocks, threads>>>(trInput1_acc, trInput2_acc, output_acc, kH, kW, patchH, patchW, padH, padW,
+                                      dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW);
+        }));
+
+    return output;
+}
+
+std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor input1, torch::Tensor input2,
+                                                     torch::Tensor gradOutput, int kH, int kW, int patchH, int patchW,
+                                                     int padH, int padW, int dilationH, int dilationW,
+                                                     int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    auto gradInput1 = torch::zeros_like(input1);
+    auto gradInput2 = torch::zeros_like(input2);
+
+    const int batch_size = input1.size(0);
+    const int iH = input1.size(2);
+    const int iW = input1.size(3);
+    const int C = input1.size(1);
+
+    const dim3 blocks(C, iH, iW);
+    const dim3 threads(THREADS_BACKWARD, THREADS_BACKWARD);
+
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half, at::ScalarType::BFloat16, input1.scalar_type(), "correlation_backward_cuda", ([&] {
+            TensorAcc4R input1_acc = input1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R input2_acc = input2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R gradInput1_acc = gradInput1.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc4R gradInput2_acc = gradInput2.packed_accessor32<scalar_t, 4, RestrictPtrTraits>();
+            TensorAcc5R gradOutput_acc = gradOutput.packed_accessor32<scalar_t, 5, RestrictPtrTraits>();
+
+            for (int n = 0; n < batch_size; ++n) {
+                corr::correlation_cuda_backward_kernel_input1<scalar_t>
+                    <<<blocks, threads>>>(gradOutput_acc, input2_acc, gradInput1_acc, kH, kW, patchH, patchW, padH,
+                                          padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
+            }
+
+            for (int n = 0; n < batch_size; ++n) {
+                corr::correlation_cuda_backward_kernel_input2<scalar_t>
+                    <<<blocks, threads>>>(gradOutput_acc, input1_acc, gradInput2_acc, kH, kW, patchH, patchW, padH,
+                                          padW, dilationH, dilationW, dilation_patchH, dilation_patchW, dH, dW, n);
+            }
+        }));
+
+    return {gradInput1, gradInput2};
+}

third_party/vipe/csrc/corr_ext/correlation_sampler.cpp

@@ -0,0 +1,85 @@
+/**
+ * This file includes code originally from the Pytorch correlation repository:
+ * https://github.com/ClementPinard/Pytorch-Correlation-extension、
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <vector>
+#include <iostream>
+
+// declarations
+
+torch::Tensor correlation_cpp_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                      int patchW, int padH, int padW, int dilationH, int dilationW, int dilation_patchH,
+                                      int dilation_patchW, int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cpp_backward(torch::Tensor grad_output, torch::Tensor input1,
+                                                    torch::Tensor input2, int kH, int kW, int patchH, int patchW,
+                                                    int padH, int padW, int dilationH, int dilationW,
+                                                    int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x, " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x, " must be contiguous")
+#define CHECK_INPUT(x) \
+    CHECK_CUDA(x);     \
+    CHECK_CONTIGUOUS(x)
+#define CHECK_SAME_DEVICE(x, y) TORCH_CHECK(x.device() == y.device(), #x " is not on same device as " #y)
+
+torch::Tensor correlation_cuda_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                       int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+std::vector<torch::Tensor> correlation_cuda_backward(torch::Tensor grad_output, torch::Tensor input1,
+                                                     torch::Tensor input2, int kH, int kW, int patchH, int patchW,
+                                                     int padH, int padW, int dilationH, int dilationW,
+                                                     int dilation_patchH, int dilation_patchW, int dH, int dW);
+
+// C++ interface
+
+torch::Tensor correlation_sample_forward(torch::Tensor input1, torch::Tensor input2, int kH, int kW, int patchH,
+                                         int patchW, int padH, int padW, int dilationH, int dilationW,
+                                         int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    if (input1.device().is_cuda()) {
+        CHECK_INPUT(input1);
+        CHECK_INPUT(input2);
+
+        // set device of input1 as default CUDA device
+        // https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_optional_c_u_d_a_guard.html
+        const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+        CHECK_SAME_DEVICE(input1, input2);
+
+        return correlation_cuda_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
+                                        dilation_patchH, dilation_patchW, dH, dW);
+    } else {
+        return correlation_cpp_forward(input1, input2, kH, kW, patchH, patchW, padH, padW, dilationH, dilationW,
+                                       dilation_patchH, dilation_patchW, dH, dW);
+    }
+}
+
+std::vector<torch::Tensor> correlation_sample_backward(torch::Tensor input1, torch::Tensor input2,
+                                                       torch::Tensor grad_output, int kH, int kW, int patchH,
+                                                       int patchW, int padH, int padW, int dilationH, int dilationW,
+                                                       int dilation_patchH, int dilation_patchW, int dH, int dW) {
+    if (grad_output.device().is_cuda()) {
+        CHECK_INPUT(input1);
+        CHECK_INPUT(input2);
+
+        // set device of input1 as default CUDA device
+        const at::cuda::OptionalCUDAGuard guard_input1(device_of(input1));
+        CHECK_SAME_DEVICE(input1, input2);
+        CHECK_SAME_DEVICE(input1, grad_output);
+
+        return correlation_cuda_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
+                                         dilationW, dilation_patchH, dilation_patchW, dH, dW);
+    } else {
+        return correlation_cpp_backward(input1, input2, grad_output, kH, kW, patchH, patchW, padH, padW, dilationH,
+                                        dilationW, dilation_patchH, dilation_patchW, dH, dW);
+    }
+}
+
+void pybind_corr_ext(py::module& m) {
+    m.def("forward", &correlation_sample_forward, "Spatial Correlation Sampler Forward");
+    m.def("backward", &correlation_sample_backward, "Spatial Correlation Sampler backward");
+}

third_party/vipe/csrc/droid_net_ext/altcorr_kernel.cu

@@ -0,0 +1,321 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <ATen/ATen.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+#include <ATen/native/cuda/KernelUtils.cuh>
+
+#define BLOCK_H 4
+#define BLOCK_W 8
+#define BLOCK_HW BLOCK_H *BLOCK_W
+#define CHANNEL_STRIDE 32
+
+__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
+
+template <typename scalar_t>
+__global__ void altcorr_forward_kernel(const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
+                                       const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
+                                       const torch::PackedTensorAccessor32<float, 5, torch::RestrictPtrTraits> coords,
+                                       torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr,
+                                       int r) {
+    const int b = blockIdx.x;
+    const int h0 = blockIdx.y * blockDim.x;
+    const int w0 = blockIdx.z * blockDim.y;
+    const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+    const int H1 = fmap1.size(1);
+    const int W1 = fmap1.size(2);
+    const int H2 = fmap2.size(1);
+    const int W2 = fmap2.size(2);
+    const int N = coords.size(1);
+    const int C = fmap1.size(3);
+
+    __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW];
+    __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW];
+
+    __shared__ float x2s[BLOCK_HW];
+    __shared__ float y2s[BLOCK_HW];
+
+    for (int c = 0; c < C; c += CHANNEL_STRIDE) {
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            if (within_bounds(h1, w1, H1, W1))
+                f1[c1][k1] = fmap1[b][h1][w1][c + c1];
+
+            else
+                f1[c1][k1] = 0.0;
+        }
+
+        __syncthreads();
+
+        for (int n = 0; n < N; n++) {
+            int h1 = h0 + threadIdx.x;
+            int w1 = w0 + threadIdx.y;
+            if (within_bounds(h1, w1, H1, W1)) {
+                x2s[tid] = coords[b][n][h1][w1][0];
+                y2s[tid] = coords[b][n][h1][w1][1];
+            }
+
+            float dx = x2s[tid] - floor(x2s[tid]);
+            float dy = y2s[tid] - floor(y2s[tid]);
+
+            int rd = 2 * r + 1;
+            for (int iy = 0; iy < rd + 1; iy++) {
+                for (int ix = 0; ix < rd + 1; ix++) {
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        if (within_bounds(h2, w2, H2, W2))
+                            f2[c2][k1] = fmap2[b][h2][w2][c + c2];
+
+                        else
+                            f2[c2][k1] = static_cast<scalar_t>(0.0);
+                    }
+
+                    __syncthreads();
+
+                    scalar_t s = 0.0;
+                    for (int k = 0; k < CHANNEL_STRIDE; k++) s += f1[k][tid] * f2[k][tid];
+
+                    int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
+                    int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
+                    int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
+                    int ix_se = H1 * W1 * (iy + rd * ix);
+
+                    // int ix_nw = ((iy-1) + rd*(ix-1));
+                    // int ix_ne = ((iy-1) + rd*ix);
+                    // int ix_sw = (iy + rd*(ix-1));
+                    // int ix_se = (iy + rd*ix);
+
+                    scalar_t nw = s * static_cast<scalar_t>((dy) * (dx));
+                    scalar_t ne = s * static_cast<scalar_t>((dy) * (1 - dx));
+                    scalar_t sw = s * static_cast<scalar_t>((1 - dy) * (dx));
+                    scalar_t se = s * static_cast<scalar_t>((1 - dy) * (1 - dx));
+
+                    // if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_nw][h1][w1] += nw;
+
+                    // if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_ne][h1][w1] += ne;
+
+                    // if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_sw][h1][w1] += sw;
+
+                    // if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+                    //   corr[b][n][ix_se][h1][w1] += se;
+
+                    scalar_t *corr_ptr = &corr[b][n][0][h1][w1];
+
+                    if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_nw) += nw;
+
+                    if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_ne) += ne;
+
+                    if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_sw) += sw;
+
+                    if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1)) *(corr_ptr + ix_se) += se;
+                }
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void altcorr_backward_kernel(
+    const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1,
+    const torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
+    torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap1_grad,
+    torch::PackedTensorAccessor32<scalar_t, 4, torch::RestrictPtrTraits> fmap2_grad,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> coords_grad, int r) {
+    const int b = blockIdx.x;
+    const int h0 = blockIdx.y * blockDim.x;
+    const int w0 = blockIdx.z * blockDim.y;
+    const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+    const int H1 = fmap1.size(1);
+    const int W1 = fmap1.size(2);
+    const int H2 = fmap2.size(1);
+    const int W2 = fmap2.size(2);
+    const int N = coords.size(1);
+    const int C = fmap1.size(3);
+
+    __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW + 1];
+    __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW + 1];
+
+    __shared__ scalar_t f1_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
+    __shared__ scalar_t f2_grad[CHANNEL_STRIDE][BLOCK_HW + 1];
+
+    __shared__ scalar_t x2s[BLOCK_HW];
+    __shared__ scalar_t y2s[BLOCK_HW];
+
+    for (int c = 0; c < C; c += CHANNEL_STRIDE) {
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            auto fptr = fmap1[b][h1][w1];
+            if (within_bounds(h1, w1, H1, W1))
+                f1[c1][k1] = fptr[c + c1];
+            else
+                f1[c1][k1] = 0.0;
+
+            f1_grad[c1][k1] = 0.0;
+        }
+
+        __syncthreads();
+
+        int h1 = h0 + threadIdx.x;
+        int w1 = w0 + threadIdx.y;
+
+        for (int n = 0; n < N; n++) {
+            x2s[tid] = coords[b][n][h1][w1][0];
+            y2s[tid] = coords[b][n][h1][w1][1];
+
+            scalar_t dx = x2s[tid] - floor(x2s[tid]);
+            scalar_t dy = y2s[tid] - floor(y2s[tid]);
+
+            int rd = 2 * r + 1;
+            for (int iy = 0; iy < rd + 1; iy++) {
+                for (int ix = 0; ix < rd + 1; ix++) {
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        auto fptr = fmap2[b][h2][w2];
+                        if (within_bounds(h2, w2, H2, W2))
+                            f2[c2][k1] = fptr[c + c2];
+                        else
+                            f2[c2][k1] = 0.0;
+
+                        f2_grad[c2][k1] = 0.0;
+                    }
+
+                    __syncthreads();
+
+                    const scalar_t *grad_ptr = &corr_grad[b][n][0][h1][w1];
+                    scalar_t g = 0.0;
+
+                    int ix_nw = H1 * W1 * ((iy - 1) + rd * (ix - 1));
+                    int ix_ne = H1 * W1 * ((iy - 1) + rd * ix);
+                    int ix_sw = H1 * W1 * (iy + rd * (ix - 1));
+                    int ix_se = H1 * W1 * (iy + rd * ix);
+
+                    if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_nw) * dy * dx;
+
+                    if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_ne) * dy * (1 - dx);
+
+                    if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1)) g += *(grad_ptr + ix_sw) * (1 - dy) * dx;
+
+                    if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+                        g += *(grad_ptr + ix_se) * (1 - dy) * (1 - dx);
+
+                    for (int k = 0; k < CHANNEL_STRIDE; k++) {
+                        f1_grad[k][tid] += g * f2[k][tid];
+                        f2_grad[k][tid] += g * f1[k][tid];
+                    }
+
+                    for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+                        int k1 = k + tid / CHANNEL_STRIDE;
+                        int h2 = static_cast<int>(floor(y2s[k1])) - r + iy;
+                        int w2 = static_cast<int>(floor(x2s[k1])) - r + ix;
+                        int c2 = tid % CHANNEL_STRIDE;
+
+                        scalar_t *fptr = &fmap2_grad[b][h2][w2][0];
+                        if (within_bounds(h2, w2, H2, W2)) atomicAdd(fptr + c + c2, f2_grad[c2][k1]);
+                    }
+                }
+            }
+        }
+        __syncthreads();
+
+        for (int k = 0; k < BLOCK_HW; k += BLOCK_HW / CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h1 = h0 + k1 / BLOCK_W;
+            int w1 = w0 + k1 % BLOCK_W;
+            int c1 = tid % CHANNEL_STRIDE;
+
+            scalar_t *fptr = &fmap1_grad[b][h1][w1][0];
+            if (within_bounds(h1, w1, H1, W1)) fptr[c + c1] += f1_grad[c1][k1];
+        }
+    }
+}
+
+std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                int radius) {
+    const auto B = coords.size(0);
+    const auto N = coords.size(1);
+    const auto H = coords.size(2);
+    const auto W = coords.size(3);
+
+    const auto rd = 2 * radius + 1;
+    auto opts = fmap1.options();
+    auto corr = torch::zeros({B, N, rd * rd, H, W}, opts);
+
+    const dim3 blocks(B, (H + BLOCK_H - 1) / BLOCK_H, (W + BLOCK_W - 1) / BLOCK_W);
+    const dim3 threads(BLOCK_H, BLOCK_W);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(fmap1.scalar_type(), "altcorr_forward_kernel", ([&] {
+                                            altcorr_forward_kernel<scalar_t><<<blocks, threads>>>(
+                                                fmap1.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+                                                fmap2.packed_accessor32<scalar_t, 4, torch::RestrictPtrTraits>(),
+                                                coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                                                corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {corr};
+}
+
+std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                 torch::Tensor corr_grad, int radius) {
+    const auto B = coords.size(0);
+    const auto N = coords.size(1);
+
+    const auto H1 = fmap1.size(1);
+    const auto W1 = fmap1.size(2);
+    const auto H2 = fmap2.size(1);
+    const auto W2 = fmap2.size(2);
+    const auto C = fmap1.size(3);
+
+    auto opts = fmap1.options();
+    auto fmap1_grad = torch::zeros({B, H1, W1, C}, opts);
+    auto fmap2_grad = torch::zeros({B, H2, W2, C}, opts);
+    auto coords_grad = torch::zeros({B, N, H1, W1, 2}, opts);
+
+    const dim3 blocks(B, (H1 + BLOCK_H - 1) / BLOCK_H, (W1 + BLOCK_W - 1) / BLOCK_W);
+    const dim3 threads(BLOCK_H, BLOCK_W);
+
+    altcorr_backward_kernel<float>
+        <<<blocks, threads>>>(fmap1.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              fmap2.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              coords.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                              corr_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(),
+                              fmap1_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              fmap2_grad.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                              coords_grad.packed_accessor32<float, 5, torch::RestrictPtrTraits>(), radius);
+
+    return {fmap1_grad, fmap2_grad, coords_grad};
+}

third_party/vipe/csrc/droid_net_ext/correlation_kernels.cu

@@ -0,0 +1,159 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <ATen/ATen.h>
+#include <ATen/NativeFunctions.h>
+#include <ATen/Parallel.h>
+
+#define BLOCK 16
+
+__forceinline__ __device__ bool within_bounds(int h, int w, int H, int W) { return h >= 0 && h < H && w >= 0 && w < W; }
+
+template <typename scalar_t>
+__global__ void corr_index_forward_kernel(
+    const torch::PackedTensorAccessor64<scalar_t, 5, torch::RestrictPtrTraits> volume,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr, int r) {
+    // batch index
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    const int n = blockIdx.z;
+
+    const int h1 = volume.size(1);
+    const int w1 = volume.size(2);
+    const int h2 = volume.size(3);
+    const int w2 = volume.size(4);
+
+    if (!within_bounds(y, x, h1, w1)) {
+        return;
+    }
+
+    float x0 = coords[n][0][y][x];
+    float y0 = coords[n][1][y][x];
+
+    float dx = x0 - floor(x0);
+    float dy = y0 - floor(y0);
+
+    int rd = 2 * r + 1;
+    for (int i = 0; i < rd + 1; i++) {
+        for (int j = 0; j < rd + 1; j++) {
+            int x1 = static_cast<int>(floor(x0)) - r + i;
+            int y1 = static_cast<int>(floor(y0)) - r + j;
+
+            if (within_bounds(y1, x1, h2, w2)) {
+                scalar_t s = volume[n][y][x][y1][x1];
+
+                if (i > 0 && j > 0) corr[n][i - 1][j - 1][y][x] += s * scalar_t(dx * dy);
+
+                if (i > 0 && j < rd) corr[n][i - 1][j][y][x] += s * scalar_t(dx * (1.0f - dy));
+
+                if (i < rd && j > 0) corr[n][i][j - 1][y][x] += s * scalar_t((1.0f - dx) * dy);
+
+                if (i < rd && j < rd) corr[n][i][j][y][x] += s * scalar_t((1.0f - dx) * (1.0f - dy));
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void corr_index_backward_kernel(
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> coords,
+    const torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> corr_grad,
+    torch::PackedTensorAccessor32<scalar_t, 5, torch::RestrictPtrTraits> volume_grad, int r) {
+    // batch index
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    const int n = blockIdx.z;
+
+    const int h1 = volume_grad.size(1);
+    const int w1 = volume_grad.size(2);
+    const int h2 = volume_grad.size(3);
+    const int w2 = volume_grad.size(4);
+
+    if (!within_bounds(y, x, h1, w1)) {
+        return;
+    }
+
+    float x0 = coords[n][0][y][x];
+    float y0 = coords[n][1][y][x];
+
+    float dx = x0 - floor(x0);
+    float dy = y0 - floor(y0);
+
+    int rd = 2 * r + 1;
+    for (int i = 0; i < rd + 1; i++) {
+        for (int j = 0; j < rd + 1; j++) {
+            int x1 = static_cast<int>(floor(x0)) - r + i;
+            int y1 = static_cast<int>(floor(y0)) - r + j;
+
+            if (within_bounds(y1, x1, h2, w2)) {
+                scalar_t g = 0.0;
+                if (i > 0 && j > 0) g += corr_grad[n][i - 1][j - 1][y][x] * scalar_t(dx * dy);
+
+                if (i > 0 && j < rd) g += corr_grad[n][i - 1][j][y][x] * scalar_t(dx * (1.0f - dy));
+
+                if (i < rd && j > 0) g += corr_grad[n][i][j - 1][y][x] * scalar_t((1.0f - dx) * dy);
+
+                if (i < rd && j < rd) g += corr_grad[n][i][j][y][x] * scalar_t((1.0f - dx) * (1.0f - dy));
+
+                volume_grad[n][y][x][y1][x1] += g;
+            }
+        }
+    }
+}
+
+std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
+    const auto batch_size = volume.size(0);
+    const auto ht = volume.size(1);
+    const auto wd = volume.size(2);
+
+    const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
+
+    const dim3 threads(BLOCK, BLOCK);
+
+    auto opts = volume.options();
+    torch::Tensor corr = torch::zeros({batch_size, 2 * radius + 1, 2 * radius + 1, ht, wd}, opts);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_forward_kernel", ([&] {
+                                            corr_index_forward_kernel<scalar_t><<<blocks, threads>>>(
+                                                volume.packed_accessor64<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                                                corr.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {corr};
+}
+
+std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                                    int radius) {
+    const auto batch_size = volume.size(0);
+    const auto ht = volume.size(1);
+    const auto wd = volume.size(2);
+
+    auto volume_grad = torch::zeros_like(volume);
+
+    const dim3 blocks((wd + BLOCK - 1) / BLOCK, (ht + BLOCK - 1) / BLOCK, batch_size);
+
+    const dim3 threads(BLOCK, BLOCK);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(volume.scalar_type(), "sampler_backward_kernel", ([&] {
+                                            corr_index_backward_kernel<scalar_t><<<blocks, threads>>>(
+                                                coords.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+                                                corr_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                volume_grad.packed_accessor32<scalar_t, 5, torch::RestrictPtrTraits>(),
+                                                radius);
+                                        }));
+
+    return {volume_grad};
+}

third_party/vipe/csrc/droid_net_ext/droid.cpp

@@ -0,0 +1,63 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> corr_index_cuda_forward(torch::Tensor volume, torch::Tensor coords, int radius);
+std::vector<torch::Tensor> corr_index_cuda_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                                    int radius);
+std::vector<torch::Tensor> altcorr_cuda_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                int radius);
+std::vector<torch::Tensor> altcorr_cuda_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                                 torch::Tensor corr_grad, int radius);
+
+// c++ python binding
+std::vector<torch::Tensor> corr_index_forward(torch::Tensor volume, torch::Tensor coords, int radius) {
+    CHECK_INPUT(volume);
+    CHECK_INPUT(coords);
+
+    return corr_index_cuda_forward(volume, coords, radius);
+}
+
+std::vector<torch::Tensor> corr_index_backward(torch::Tensor volume, torch::Tensor coords, torch::Tensor corr_grad,
+                                               int radius) {
+    CHECK_INPUT(volume);
+    CHECK_INPUT(coords);
+    CHECK_INPUT(corr_grad);
+
+    auto volume_grad = corr_index_cuda_backward(volume, coords, corr_grad, radius);
+    return {volume_grad};
+}
+
+std::vector<torch::Tensor> altcorr_forward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords, int radius) {
+    CHECK_INPUT(fmap1);
+    CHECK_INPUT(fmap2);
+    CHECK_INPUT(coords);
+
+    return altcorr_cuda_forward(fmap1, fmap2, coords, radius);
+}
+
+std::vector<torch::Tensor> altcorr_backward(torch::Tensor fmap1, torch::Tensor fmap2, torch::Tensor coords,
+                                            torch::Tensor corr_grad, int radius) {
+    CHECK_INPUT(fmap1);
+    CHECK_INPUT(fmap2);
+    CHECK_INPUT(coords);
+    CHECK_INPUT(corr_grad);
+
+    return altcorr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
+}
+
+void pybind_droid_net_ext(py::module& m) {
+    // correlation volume kernels
+    m.def("altcorr_forward", &altcorr_forward, "ALTCORR forward");
+    m.def("altcorr_backward", &altcorr_backward, "ALTCORR backward");
+    m.def("corr_index_forward", &corr_index_forward, "INDEX forward");
+    m.def("corr_index_backward", &corr_index_backward, "INDEX backward");
+}

third_party/vipe/csrc/grounding_dino_ext/ms_deform_attn_cuda.cu

@@ -0,0 +1,136 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#include <vector>
+#include "ms_deform_im2col_cuda.cuh"
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                       const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                       const at::Tensor &attn_weight, const int im2col_step) {
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+    auto output = at::zeros({batch, num_query, num_heads, channels}, value.options());
+
+    const int batch_n = im2col_step_;
+    auto output_n = output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto columns = output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(
+            value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] {
+                ms_deformable_im2col_cuda(
+                    at::cuda::getCurrentCUDAStream(), value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                    spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+                    sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                    attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size, batch_n, spatial_size,
+                    num_heads, channels, num_levels, num_query, num_point, columns.data<scalar_t>());
+            }));
+    }
+
+    output = output.view({batch, num_query, num_heads * channels});
+
+    return output;
+}
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                     const at::Tensor &level_start_index,
+                                                     const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
+                                                     const at::Tensor &grad_output, const int im2col_step) {
+    AT_ASSERTM(value.is_contiguous(), "value tensor has to be contiguous");
+    AT_ASSERTM(spatial_shapes.is_contiguous(), "spatial_shapes tensor has to be contiguous");
+    AT_ASSERTM(level_start_index.is_contiguous(), "level_start_index tensor has to be contiguous");
+    AT_ASSERTM(sampling_loc.is_contiguous(), "sampling_loc tensor has to be contiguous");
+    AT_ASSERTM(attn_weight.is_contiguous(), "attn_weight tensor has to be contiguous");
+    AT_ASSERTM(grad_output.is_contiguous(), "grad_output tensor has to be contiguous");
+
+    AT_ASSERTM(value.is_cuda(), "value must be a CUDA tensor");
+    AT_ASSERTM(spatial_shapes.is_cuda(), "spatial_shapes must be a CUDA tensor");
+    AT_ASSERTM(level_start_index.is_cuda(), "level_start_index must be a CUDA tensor");
+    AT_ASSERTM(sampling_loc.is_cuda(), "sampling_loc must be a CUDA tensor");
+    AT_ASSERTM(attn_weight.is_cuda(), "attn_weight must be a CUDA tensor");
+    AT_ASSERTM(grad_output.is_cuda(), "grad_output must be a CUDA tensor");
+
+    const int batch = value.size(0);
+    const int spatial_size = value.size(1);
+    const int num_heads = value.size(2);
+    const int channels = value.size(3);
+
+    const int num_levels = spatial_shapes.size(0);
+
+    const int num_query = sampling_loc.size(1);
+    const int num_point = sampling_loc.size(4);
+
+    const int im2col_step_ = std::min(batch, im2col_step);
+
+    AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_);
+
+    auto grad_value = at::zeros_like(value);
+    auto grad_sampling_loc = at::zeros_like(sampling_loc);
+    auto grad_attn_weight = at::zeros_like(attn_weight);
+
+    const int batch_n = im2col_step_;
+    auto per_value_size = spatial_size * num_heads * channels;
+    auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2;
+    auto per_attn_weight_size = num_query * num_heads * num_levels * num_point;
+    auto grad_output_n = grad_output.view({batch / im2col_step_, batch_n, num_query, num_heads, channels});
+
+    for (int n = 0; n < batch / im2col_step_; ++n) {
+        auto grad_output_g = grad_output_n.select(0, n);
+        AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] {
+                                       ms_deformable_col2im_cuda(
+                                           at::cuda::getCurrentCUDAStream(), grad_output_g.data<scalar_t>(),
+                                           value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                                           spatial_shapes.data<int64_t>(), level_start_index.data<int64_t>(),
+                                           sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                           attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size,
+                                           batch_n, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                                           grad_value.data<scalar_t>() + n * im2col_step_ * per_value_size,
+                                           grad_sampling_loc.data<scalar_t>() + n * im2col_step_ * per_sample_loc_size,
+                                           grad_attn_weight.data<scalar_t>() + n * im2col_step_ * per_attn_weight_size);
+                                   }));
+    }
+
+    return {grad_value, grad_sampling_loc, grad_attn_weight};
+}
+
+}  // namespace groundingdino

third_party/vipe/csrc/grounding_dino_ext/ms_deform_attn_cuda.h

@@ -0,0 +1,25 @@
+/*!
+**************************************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************************************
+* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
+**************************************************************************************************
+*/
+
+#pragma once
+#include <torch/extension.h>
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_cuda_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                       const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                       const at::Tensor &attn_weight, const int im2col_step);
+
+std::vector<at::Tensor> ms_deform_attn_cuda_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                     const at::Tensor &level_start_index,
+                                                     const at::Tensor &sampling_loc, const at::Tensor &attn_weight,
+                                                     const at::Tensor &grad_output, const int im2col_step);
+
+}  // namespace groundingdino

third_party/vipe/csrc/grounding_dino_ext/ms_deform_im2col_cuda.cuh

@@ -0,0 +1,914 @@
+/*!
+**************************************************************************
+* Deformable DETR
+* Copyright (c) 2020 SenseTime. All Rights Reserved.
+* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+**************************************************************************
+* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
+* Copyright (c) 2018 Microsoft
+**************************************************************************
+*/
+
+#include <cstdio>
+#include <algorithm>
+#include <cstring>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <THC/THCAtomics.cuh>
+
+#define CUDA_KERNEL_LOOP(i, n) for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N, const int num_threads) { return (N + num_threads - 1) / num_threads; }
+
+template <typename scalar_t>
+__device__ scalar_t ms_deform_attn_im2col_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
+                                                   const int &nheads, const int &channels, const scalar_t &h,
+                                                   const scalar_t &w, const int &m, const int &c) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+    }
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    return val;
+}
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear(const scalar_t *&bottom_data, const int &height, const int &width,
+                                               const int &nheads, const int &channels, const scalar_t &h,
+                                               const scalar_t &w, const int &m, const int &c, const scalar_t &top_grad,
+                                               const scalar_t &attn_weight, scalar_t *&grad_value,
+                                               scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+    const scalar_t top_grad_value = top_grad * attn_weight;
+    scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+        grad_h_weight -= hw * v1;
+        grad_w_weight -= hh * v1;
+        atomicAdd(grad_value + ptr1, w1 * top_grad_value);
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+        grad_h_weight -= lw * v2;
+        grad_w_weight += hh * v2;
+        atomicAdd(grad_value + ptr2, w2 * top_grad_value);
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+        grad_h_weight += hw * v3;
+        grad_w_weight -= lh * v3;
+        atomicAdd(grad_value + ptr3, w3 * top_grad_value);
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+        grad_h_weight += lw * v4;
+        grad_w_weight += lh * v4;
+        atomicAdd(grad_value + ptr4, w4 * top_grad_value);
+    }
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    *grad_attn_weight = top_grad * val;
+    *grad_sampling_loc = width * grad_w_weight * top_grad_value;
+    *(grad_sampling_loc + 1) = height * grad_h_weight * top_grad_value;
+}
+
+template <typename scalar_t>
+__device__ void ms_deform_attn_col2im_bilinear_gm(const scalar_t *&bottom_data, const int &height, const int &width,
+                                                  const int &nheads, const int &channels, const scalar_t &h,
+                                                  const scalar_t &w, const int &m, const int &c,
+                                                  const scalar_t &top_grad, const scalar_t &attn_weight,
+                                                  scalar_t *&grad_value, scalar_t *grad_sampling_loc,
+                                                  scalar_t *grad_attn_weight) {
+    const int h_low = floor(h);
+    const int w_low = floor(w);
+    const int h_high = h_low + 1;
+    const int w_high = w_low + 1;
+
+    const scalar_t lh = h - h_low;
+    const scalar_t lw = w - w_low;
+    const scalar_t hh = 1 - lh, hw = 1 - lw;
+
+    const int w_stride = nheads * channels;
+    const int h_stride = width * w_stride;
+    const int h_low_ptr_offset = h_low * h_stride;
+    const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
+    const int w_low_ptr_offset = w_low * w_stride;
+    const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
+    const int base_ptr = m * channels + c;
+
+    const scalar_t w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;
+    const scalar_t top_grad_value = top_grad * attn_weight;
+    scalar_t grad_h_weight = 0, grad_w_weight = 0;
+
+    scalar_t v1 = 0;
+    if (h_low >= 0 && w_low >= 0) {
+        const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
+        v1 = bottom_data[ptr1];
+        grad_h_weight -= hw * v1;
+        grad_w_weight -= hh * v1;
+        atomicAdd(grad_value + ptr1, w1 * top_grad_value);
+    }
+    scalar_t v2 = 0;
+    if (h_low >= 0 && w_high <= width - 1) {
+        const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
+        v2 = bottom_data[ptr2];
+        grad_h_weight -= lw * v2;
+        grad_w_weight += hh * v2;
+        atomicAdd(grad_value + ptr2, w2 * top_grad_value);
+    }
+    scalar_t v3 = 0;
+    if (h_high <= height - 1 && w_low >= 0) {
+        const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
+        v3 = bottom_data[ptr3];
+        grad_h_weight += hw * v3;
+        grad_w_weight -= lh * v3;
+        atomicAdd(grad_value + ptr3, w3 * top_grad_value);
+    }
+    scalar_t v4 = 0;
+    if (h_high <= height - 1 && w_high <= width - 1) {
+        const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
+        v4 = bottom_data[ptr4];
+        grad_h_weight += lw * v4;
+        grad_w_weight += lh * v4;
+        atomicAdd(grad_value + ptr4, w4 * top_grad_value);
+    }
+
+    const scalar_t val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+    atomicAdd(grad_attn_weight, top_grad * val);
+    atomicAdd(grad_sampling_loc, width * grad_w_weight * top_grad_value);
+    atomicAdd(grad_sampling_loc + 1, height * grad_h_weight * top_grad_value);
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_im2col_gpu_kernel(const int n, const scalar_t *data_value,
+                                                const int64_t *data_spatial_shapes,
+                                                const int64_t *data_level_start_index,
+                                                const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                                                const int batch_size, const int spatial_size, const int num_heads,
+                                                const int channels, const int num_levels, const int num_query,
+                                                const int num_point, scalar_t *data_col) {
+    CUDA_KERNEL_LOOP(index, n) {
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        scalar_t *data_col_ptr = data_col + index;
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+        scalar_t col = 0;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels,
+                                                          h_im, w_im, m_col, c_col) *
+                           weight;
+                }
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+            }
+        }
+        *data_col_ptr = col;
+    }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+        __shared__ scalar_t cache_grad_attn_weight[blockSize];
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+                if (tid == 0) {
+                    scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
+                             _grad_a = cache_grad_attn_weight[0];
+                    int sid = 2;
+                    for (unsigned int tid = 1; tid < blockSize; ++tid) {
+                        _grad_w += cache_grad_sampling_loc[sid];
+                        _grad_h += cache_grad_sampling_loc[sid + 1];
+                        _grad_a += cache_grad_attn_weight[tid];
+                        sid += 2;
+                    }
+
+                    *grad_sampling_loc = _grad_w;
+                    *(grad_sampling_loc + 1) = _grad_h;
+                    *grad_attn_weight = _grad_a;
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t, unsigned int blockSize>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        __shared__ scalar_t cache_grad_sampling_loc[blockSize * 2];
+        __shared__ scalar_t cache_grad_attn_weight[blockSize];
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockSize / 2; s > 0; s >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    *grad_sampling_loc = cache_grad_sampling_loc[0];
+                    *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+                    *grad_attn_weight = cache_grad_attn_weight[0];
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v1(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+                if (tid == 0) {
+                    scalar_t _grad_w = cache_grad_sampling_loc[0], _grad_h = cache_grad_sampling_loc[1],
+                             _grad_a = cache_grad_attn_weight[0];
+                    int sid = 2;
+                    for (unsigned int tid = 1; tid < blockDim.x; ++tid) {
+                        _grad_w += cache_grad_sampling_loc[sid];
+                        _grad_h += cache_grad_sampling_loc[sid + 1];
+                        _grad_a += cache_grad_attn_weight[tid];
+                        sid += 2;
+                    }
+
+                    *grad_sampling_loc = _grad_w;
+                    *(grad_sampling_loc + 1) = _grad_h;
+                    *grad_attn_weight = _grad_a;
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                        if (tid + (s << 1) < spre) {
+                            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+                            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+                            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+                        }
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    *grad_sampling_loc = cache_grad_sampling_loc[0];
+                    *(grad_sampling_loc + 1) = cache_grad_sampling_loc[1];
+                    *grad_attn_weight = cache_grad_attn_weight[0];
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks(
+    const int n, const scalar_t *grad_col, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+    const int64_t *data_level_start_index, const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+    const int batch_size, const int spatial_size, const int num_heads, const int channels, const int num_levels,
+    const int num_query, const int num_point, scalar_t *grad_value, scalar_t *grad_sampling_loc,
+    scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        extern __shared__ int _s[];
+        scalar_t *cache_grad_sampling_loc = (scalar_t *)_s;
+        scalar_t *cache_grad_attn_weight = cache_grad_sampling_loc + 2 * blockDim.x;
+        unsigned int tid = threadIdx.x;
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                *(cache_grad_sampling_loc + (threadIdx.x << 1)) = 0;
+                *(cache_grad_sampling_loc + ((threadIdx.x << 1) + 1)) = 0;
+                *(cache_grad_attn_weight + threadIdx.x) = 0;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                   w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                   cache_grad_sampling_loc + (threadIdx.x << 1),
+                                                   cache_grad_attn_weight + threadIdx.x);
+                }
+
+                __syncthreads();
+
+                for (unsigned int s = blockDim.x / 2, spre = blockDim.x; s > 0; s >>= 1, spre >>= 1) {
+                    if (tid < s) {
+                        const unsigned int xid1 = tid << 1;
+                        const unsigned int xid2 = (tid + s) << 1;
+                        cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + s];
+                        cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2];
+                        cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1];
+                        if (tid + (s << 1) < spre) {
+                            cache_grad_attn_weight[tid] += cache_grad_attn_weight[tid + (s << 1)];
+                            cache_grad_sampling_loc[xid1] += cache_grad_sampling_loc[xid2 + (s << 1)];
+                            cache_grad_sampling_loc[xid1 + 1] += cache_grad_sampling_loc[xid2 + 1 + (s << 1)];
+                        }
+                    }
+                    __syncthreads();
+                }
+
+                if (tid == 0) {
+                    atomicAdd(grad_sampling_loc, cache_grad_sampling_loc[0]);
+                    atomicAdd(grad_sampling_loc + 1, cache_grad_sampling_loc[1]);
+                    atomicAdd(grad_attn_weight, cache_grad_attn_weight[0]);
+                }
+                __syncthreads();
+
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+__global__ void ms_deformable_col2im_gpu_kernel_gm(const int n, const scalar_t *grad_col, const scalar_t *data_value,
+                                                   const int64_t *data_spatial_shapes,
+                                                   const int64_t *data_level_start_index,
+                                                   const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                                                   const int batch_size, const int spatial_size, const int num_heads,
+                                                   const int channels, const int num_levels, const int num_query,
+                                                   const int num_point, scalar_t *grad_value,
+                                                   scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    CUDA_KERNEL_LOOP(index, n) {
+        int _temp = index;
+        const int c_col = _temp % channels;
+        _temp /= channels;
+        const int sampling_index = _temp;
+        const int m_col = _temp % num_heads;
+        _temp /= num_heads;
+        const int q_col = _temp % num_query;
+        _temp /= num_query;
+        const int b_col = _temp;
+
+        const scalar_t top_grad = grad_col[index];
+
+        int data_weight_ptr = sampling_index * num_levels * num_point;
+        int data_loc_w_ptr = data_weight_ptr << 1;
+        const int grad_sampling_ptr = data_weight_ptr;
+        grad_sampling_loc += grad_sampling_ptr << 1;
+        grad_attn_weight += grad_sampling_ptr;
+        const int grad_weight_stride = 1;
+        const int grad_loc_stride = 2;
+        const int qid_stride = num_heads * channels;
+        const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
+
+        for (int l_col = 0; l_col < num_levels; ++l_col) {
+            const int level_start_id = data_level_start_index[l_col];
+            const int spatial_h_ptr = l_col << 1;
+            const int spatial_h = data_spatial_shapes[spatial_h_ptr];
+            const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
+            const int value_ptr_offset = data_value_ptr_init_offset + level_start_id * qid_stride;
+            const scalar_t *data_value_ptr = data_value + value_ptr_offset;
+            scalar_t *grad_value_ptr = grad_value + value_ptr_offset;
+
+            for (int p_col = 0; p_col < num_point; ++p_col) {
+                const scalar_t loc_w = data_sampling_loc[data_loc_w_ptr];
+                const scalar_t loc_h = data_sampling_loc[data_loc_w_ptr + 1];
+                const scalar_t weight = data_attn_weight[data_weight_ptr];
+
+                const scalar_t h_im = loc_h * spatial_h - 0.5;
+                const scalar_t w_im = loc_w * spatial_w - 0.5;
+                if (h_im > -1 && w_im > -1 && h_im < spatial_h && w_im < spatial_w) {
+                    ms_deform_attn_col2im_bilinear_gm(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im,
+                                                      w_im, m_col, c_col, top_grad, weight, grad_value_ptr,
+                                                      grad_sampling_loc, grad_attn_weight);
+                }
+                data_weight_ptr += 1;
+                data_loc_w_ptr += 2;
+                grad_attn_weight += grad_weight_stride;
+                grad_sampling_loc += grad_loc_stride;
+            }
+        }
+    }
+}
+
+template <typename scalar_t>
+void ms_deformable_im2col_cuda(cudaStream_t stream, const scalar_t *data_value, const int64_t *data_spatial_shapes,
+                               const int64_t *data_level_start_index, const scalar_t *data_sampling_loc,
+                               const scalar_t *data_attn_weight, const int batch_size, const int spatial_size,
+                               const int num_heads, const int channels, const int num_levels, const int num_query,
+                               const int num_point, scalar_t *data_col) {
+    const int num_kernels = batch_size * num_query * num_heads * channels;
+    const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+    const int num_threads = CUDA_NUM_THREADS;
+    ms_deformable_im2col_gpu_kernel<scalar_t><<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+        num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight,
+        batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("error in ms_deformable_im2col_cuda: %s\n", cudaGetErrorString(err));
+    }
+}
+
+template <typename scalar_t>
+void ms_deformable_col2im_cuda(cudaStream_t stream, const scalar_t *grad_col, const scalar_t *data_value,
+                               const int64_t *data_spatial_shapes, const int64_t *data_level_start_index,
+                               const scalar_t *data_sampling_loc, const scalar_t *data_attn_weight,
+                               const int batch_size, const int spatial_size, const int num_heads, const int channels,
+                               const int num_levels, const int num_query, const int num_point, scalar_t *grad_value,
+                               scalar_t *grad_sampling_loc, scalar_t *grad_attn_weight) {
+    const int num_threads = (channels > CUDA_NUM_THREADS) ? CUDA_NUM_THREADS : channels;
+    const int num_kernels = batch_size * num_query * num_heads * channels;
+    const int num_actual_kernels = batch_size * num_query * num_heads * channels;
+    if (channels > 1024) {
+        if ((channels & 1023) == 0) {
+            ms_deformable_col2im_gpu_kernel_shm_reduce_v2_multi_blocks<scalar_t>
+                <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                   stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                             data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                             num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+        } else {
+            ms_deformable_col2im_gpu_kernel_gm<scalar_t>
+                <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                    num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc,
+                    data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point,
+                    grad_value, grad_sampling_loc, grad_attn_weight);
+        }
+    } else {
+        switch (channels) {
+            case 1:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 1>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 2:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 2>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 4:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 4>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 8:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 8>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 16:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 16>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 32:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v1<scalar_t, 32>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 64:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 64>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 128:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 128>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 256:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 256>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 512:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 512>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            case 1024:
+                ms_deformable_col2im_gpu_kernel_shm_blocksize_aware_reduce_v2<scalar_t, 1024>
+                    <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, 0, stream>>>(
+                        num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                        data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels, num_levels,
+                        num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                break;
+            default:
+                if (channels < 64) {
+                    ms_deformable_col2im_gpu_kernel_shm_reduce_v1<scalar_t>
+                        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                           stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                                     data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                                     num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                } else {
+                    ms_deformable_col2im_gpu_kernel_shm_reduce_v2<scalar_t>
+                        <<<GET_BLOCKS(num_actual_kernels, num_threads), num_threads, num_threads * 3 * sizeof(scalar_t),
+                           stream>>>(num_kernels, grad_col, data_value, data_spatial_shapes, data_level_start_index,
+                                     data_sampling_loc, data_attn_weight, batch_size, spatial_size, num_heads, channels,
+                                     num_levels, num_query, num_point, grad_value, grad_sampling_loc, grad_attn_weight);
+                }
+        }
+    }
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("error in ms_deformable_col2im_cuda: %s\n", cudaGetErrorString(err));
+    }
+}

third_party/vipe/csrc/grounding_dino_ext/vision.cpp

@@ -0,0 +1,33 @@
+// Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+#include "ms_deform_attn_cuda.h"
+
+namespace groundingdino {
+
+at::Tensor ms_deform_attn_forward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                  const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                  const at::Tensor &attn_weight, const int im2col_step) {
+    if (value.is_cuda()) {
+        return ms_deform_attn_cuda_forward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
+                                           im2col_step);
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+std::vector<at::Tensor> ms_deform_attn_backward(const at::Tensor &value, const at::Tensor &spatial_shapes,
+                                                const at::Tensor &level_start_index, const at::Tensor &sampling_loc,
+                                                const at::Tensor &attn_weight, const at::Tensor &grad_output,
+                                                const int im2col_step) {
+    if (value.is_cuda()) {
+        return ms_deform_attn_cuda_backward(value, spatial_shapes, level_start_index, sampling_loc, attn_weight,
+                                            grad_output, im2col_step);
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+
+}  // namespace groundingdino
+
+void pybind_grounding_dino_ext(py::module &m) {
+    m.def("ms_deform_attn_forward", &groundingdino::ms_deform_attn_forward, "ms_deform_attn_forward");
+    m.def("ms_deform_attn_backward", &groundingdino::ms_deform_attn_backward, "ms_deform_attn_backward");
+}

third_party/vipe/csrc/lietorch_ext/common.h

@@ -0,0 +1,16 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef COMMON_H
+#define COMMON_H
+
+#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
+#define EIGEN_RUNTIME_NO_MALLOC
+
+#define EPS 1e-6
+#define PI 3.14159265358979323846
+
+#endif

third_party/vipe/csrc/lietorch_ext/dispatch.h

@@ -0,0 +1,52 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef DISPATCH_H
+#define DISPATCH_H
+
+#include <torch/extension.h>
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+#define PRIVATE_CASE_TYPE(group_index, enum_type, type, ...) \
+    case enum_type: {                                        \
+        using scalar_t = type;                               \
+        switch (group_index) {                               \
+            case 1: {                                        \
+                using group_t = SO3<type>;                   \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 2: {                                        \
+                using group_t = RxSO3<type>;                 \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 3: {                                        \
+                using group_t = SE3<type>;                   \
+                return __VA_ARGS__();                        \
+            }                                                \
+            case 4: {                                        \
+                using group_t = Sim3<type>;                  \
+                return __VA_ARGS__();                        \
+            }                                                \
+        }                                                    \
+    }
+
+#define DISPATCH_GROUP_AND_FLOATING_TYPES(GROUP_INDEX, TYPE, NAME, ...)                 \
+    [&] {                                                                               \
+        const auto& the_type = TYPE;                                                    \
+        /* don't use TYPE again in case it is an expensive or side-effect op */         \
+        switch (the_type) {                                                             \
+            PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Double, double, __VA_ARGS__) \
+            PRIVATE_CASE_TYPE(GROUP_INDEX, at::ScalarType::Float, float, __VA_ARGS__)   \
+            default:                                                                    \
+                break;                                                                  \
+        }                                                                               \
+    }()
+
+#endif

third_party/vipe/csrc/lietorch_ext/lietorch.cpp

@@ -0,0 +1,336 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+#include "lietorch_gpu.h"
+#include "lietorch_cpu.h"
+#include <c10/cuda/CUDAGuard.h>
+
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+/* Interface for cuda and c++ group operations
+
+    enum group_t { SO3=1, SE3=2, Sim3=3 };
+    X, Y, Z: (uppercase) Lie Group Elements
+    a, b, c: (lowercase) Lie Algebra Elements
+*/
+
+// Unary operations
+torch::Tensor expm(int group_index, torch::Tensor a) {
+    CHECK_CONTIGUOUS(a);
+    if (a.device().type() == torch::DeviceType::CPU) {
+        return exp_forward_cpu(group_index, a);
+
+    } else if (a.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+        return exp_forward_gpu(group_index, a);
+    }
+
+    return a;
+}
+
+std::vector<torch::Tensor> expm_backward(int group_index, torch::Tensor grad, torch::Tensor a) {
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+    if (a.device().type() == torch::DeviceType::CPU) {
+        return exp_backward_cpu(group_index, grad, a);
+
+    } else if (a.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+        return exp_backward_gpu(group_index, grad, a);
+    }
+
+    return {};
+}
+
+torch::Tensor logm(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return log_forward_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return log_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> logm_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return log_backward_cpu(group_index, grad, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return log_backward_gpu(group_index, grad, X);
+    }
+
+    return {};
+}
+
+torch::Tensor inv(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return inv_forward_cpu(group_index, X);
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return inv_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> inv_backward(int group_index, torch::Tensor grad, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return inv_backward_cpu(group_index, grad, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return inv_backward_gpu(group_index, grad, X);
+    }
+
+    return {};
+}
+
+// Binary operations
+
+torch::Tensor mul(int group_index, torch::Tensor X, torch::Tensor Y) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(Y);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return mul_forward_cpu(group_index, X, Y);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return mul_forward_gpu(group_index, X, Y);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> mul_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(Y);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return mul_backward_cpu(group_index, grad, X, Y);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return mul_backward_gpu(group_index, grad, X, Y);
+    }
+
+    return {};
+}
+
+torch::Tensor adj(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adj_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adj_forward_gpu(group_index, X, a);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> adj_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adj_backward_cpu(group_index, grad, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adj_backward_gpu(group_index, grad, X, a);
+    }
+
+    return {};
+}
+
+torch::Tensor adjT(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adjT_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adjT_forward_gpu(group_index, X, a);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> adjT_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return adjT_backward_cpu(group_index, grad, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return adjT_backward_gpu(group_index, grad, X, a);
+    }
+
+    return {};
+}
+
+torch::Tensor act(int group_index, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act_forward_cpu(group_index, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act_forward_gpu(group_index, X, p);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> act_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act_backward_cpu(group_index, grad, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act_backward_gpu(group_index, grad, X, p);
+    }
+
+    return {};
+}
+
+torch::Tensor act4(int group_index, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act4_forward_cpu(group_index, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act4_forward_gpu(group_index, X, p);
+    }
+
+    return X;
+}
+
+std::vector<torch::Tensor> act4_backward(int group_index, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(p);
+    CHECK_CONTIGUOUS(grad);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return act4_backward_cpu(group_index, grad, X, p);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return act4_backward_gpu(group_index, grad, X, p);
+    }
+
+    return {};
+}
+
+torch::Tensor projector(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return orthogonal_projector_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return orthogonal_projector_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+torch::Tensor as_matrix(int group_index, torch::Tensor X) {
+    CHECK_CONTIGUOUS(X);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return as_matrix_forward_cpu(group_index, X);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return as_matrix_forward_gpu(group_index, X);
+    }
+
+    return X;
+}
+
+torch::Tensor Jinv(int group_index, torch::Tensor X, torch::Tensor a) {
+    CHECK_CONTIGUOUS(X);
+    CHECK_CONTIGUOUS(a);
+
+    if (X.device().type() == torch::DeviceType::CPU) {
+        return jleft_forward_cpu(group_index, X, a);
+
+    } else if (X.device().type() == torch::DeviceType::CUDA) {
+        const at::cuda::OptionalCUDAGuard device_guard(device_of(X));
+        return jleft_forward_gpu(group_index, X, a);
+    }
+
+    return a;
+}
+
+void pybind_lietorch_ext(py::module &m) {
+    // {exp, log, inv, mul, adj, adjT, act, act4} forward/backward bindings
+    m.def("expm", &expm, "exp map forward");
+    m.def("expm_backward", &expm_backward, "exp map backward");
+
+    m.def("logm", &logm, "log map forward");
+    m.def("logm_backward", &logm_backward, "log map backward");
+
+    m.def("inv", &inv, "inverse operator");
+    m.def("inv_backward", &inv_backward, "inverse operator backward");
+
+    m.def("mul", &mul, "group operator");
+    m.def("mul_backward", &mul_backward, "group operator backward");
+
+    m.def("adj", &adj, "adjoint operator");
+    m.def("adj_backward", &adj_backward, "adjoint operator backward");
+
+    m.def("adjT", &adjT, "transposed adjoint operator");
+    m.def("adjT_backward", &adjT_backward, "transposed adjoint operator backward");
+
+    m.def("act", &act, "action on point");
+    m.def("act_backward", &act_backward, "action on point backward");
+
+    m.def("act4", &act4, "action on homogeneous point");
+    m.def("act4_backward", &act4_backward, "action on homogeneous point backward");
+
+    // functions with no gradient
+    m.def("as_matrix", &as_matrix, "convert to matrix");
+    m.def("projector", &projector, "orthogonal projection matrix");
+    m.def("Jinv", &Jinv, "left inverse jacobian operator");
+}

third_party/vipe/csrc/lietorch_ext/lietorch_cpu.cpp

@@ -0,0 +1,610 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "lietorch_cpu.h"
+#include <eigen3/Eigen/Dense>
+
+#include <iostream>
+#include "common.h"
+#include "dispatch.h"
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+template <typename Group, typename scalar_t>
+void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int batch_size) {
+    // exponential map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int batch_size) {
+    // exponential map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a(a_ptr + i * Group::K);
+            Grad dX(grad + i * Group::N);
+            Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int batch_size) {
+    // logarithm map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a = Group(X_ptr + i * Group::N).Log();
+            Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+    // logarithm map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Tangent a = Group(X_ptr + i * Group::N).Log();
+            Grad da(grad + i * Group::K);
+            Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+    // group inverse backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group Y = Group(X_ptr + i * Group::N).inv();
+            Grad dY(grad + i * Group::N);
+            Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int batch_size) {
+    // group multiplication forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
+            Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
+        }
+    });
+}
+
+template <class Group, typename scalar_t>
+void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX, scalar_t* dY,
+                         int batch_size) {
+    // group multiplication backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Grad dZ(grad + i * Group::N);
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Grad>(dX + i * Group::N) = dZ;
+            Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da,
+                         int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Grad db(grad + i * Group::K);
+
+            Tangent a(a_ptr + i * Group::K);
+            Tangent b = X.Adj() * a;
+
+            Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
+            Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                          scalar_t* da, int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent db(grad + i * Group::K);
+            Grad a(a_ptr + i * Group::K);
+
+            Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
+            Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 3);
+            Eigen::Map<Point>(q_ptr + i * 3) = X * p;
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp,
+                         int batch_size) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 3);
+            PointGrad dq(grad + i * 3);
+
+            Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix().template block<3, 3>(0, 0);
+            Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
+        }
+    });
+}
+
+// template <typename Group, typename scalar_t>
+// void tovec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+//     // group inverse forward kernel
+//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
+//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
+
+//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+//         for (int64_t i=start; i<end; i++) {
+//             Group X(X_ptr + i*Group::N);
+//             Grad g(grad + i*Group::N);
+//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
+//         }
+//     });
+// }
+
+// template <typename Group, typename scalar_t>
+// void fromvec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
+//     // group inverse forward kernel
+//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
+//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;
+
+//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+//         for (int64_t i=start; i<end; i++) {
+//             Group X(X_ptr + i*Group::N);
+//             Grad g(grad + i*Group::N);
+//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
+//         }
+//     });
+// }
+
+template <typename Group, typename scalar_t>
+void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
+    // action on homogeneous point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 4);
+            Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                          scalar_t* dp, int batch_size) {
+    // action on homogeneous point backward kernel
+
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Point p(p_ptr + i * 4);
+            PointGrad dq(grad + i * 4);
+
+            Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
+            const Point q = X.act4(p);
+            Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int batch_size) {
+    // group inverse forward kernel
+    using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
+        }
+    });
+}
+
+template <typename Group, typename scalar_t>
+void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
+    // left-jacobian inverse action
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
+        for (int64_t i = start; i < end; i++) {
+            Group X(X_ptr + i * Group::N);
+            Tangent a(a_ptr + i * Group::K);
+            Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
+            Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
+        }
+    });
+}
+
+// unary operations
+
+torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor X;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
+                                          X = torch::zeros({batch_size, group_t::N}, a.options());
+                                          exp_forward_kernel<group_t, scalar_t>(a.data_ptr<scalar_t>(),
+                                                                                X.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return X;
+}
+
+std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
+                                          exp_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 a.data_ptr<scalar_t>(),
+                                                                                 da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {da};
+}
+
+torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor a;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
+                                          a = torch::zeros({batch_size, group_t::K}, X.options());
+                                          log_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                a.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return a;
+}
+
+std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_backward_kernel", ([&] {
+                                          log_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 X.data_ptr<scalar_t>(),
+                                                                                 dX.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX};
+}
+
+torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor Y = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "inv_forward_kernel",
+        ([&] { inv_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size); }));
+
+    return Y;
+}
+
+std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
+                                          inv_backward_kernel<group_t, scalar_t>(grad.data_ptr<scalar_t>(),
+                                                                                 X.data_ptr<scalar_t>(),
+                                                                                 dX.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX};
+}
+
+// binary operations
+torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor Z = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
+                                          mul_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                Y.data_ptr<scalar_t>(),
+                                                                                Z.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return Z;
+}
+
+std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
+                                          mul_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dY.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dY};
+}
+
+torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
+                                          adj_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                a.data_ptr<scalar_t>(),
+                                                                                b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
+                                          adj_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, da};
+}
+
+torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
+                                          adjT_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                 a.data_ptr<scalar_t>(),
+                                                                                 b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
+                                          adjT_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, da};
+}
+
+torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_forward_kernel", ([&] {
+                                          act_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                p.data_ptr<scalar_t>(),
+                                                                                q.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act_backward_kernel", ([&] {
+                                          act_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dp};
+}
+
+torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
+                                          act4_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                 p.data_ptr<scalar_t>(),
+                                                                                 q.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
+                                          act4_backward_kernel<group_t, scalar_t>(
+                                              grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(),
+                                              dX.data_ptr<scalar_t>(), dp.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return {dX, dp};
+}
+
+torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
+                                          as_matrix_forward_kernel<group_t, scalar_t>(
+                                              X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return T4x4;
+}
+
+torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor P;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
+                                          P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
+                                          orthogonal_projector_kernel<group_t, scalar_t>(
+                                              X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return P;
+}
+
+torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
+                                          jleft_forward_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(),
+                                                                                  a.data_ptr<scalar_t>(),
+                                                                                  b.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return b;
+}

third_party/vipe/csrc/lietorch_ext/lietorch_cpu.h

@@ -0,0 +1,50 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef LIETORCH_CPU_H_
+#define LIETORCH_CPU_H_
+
+#include <vector>
+#include <torch/extension.h>
+
+// unary operations
+torch::Tensor exp_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> exp_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor log_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> log_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor inv_forward_cpu(int, torch::Tensor);
+std::vector<torch::Tensor> inv_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+// binary operations
+torch::Tensor mul_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> mul_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adj_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adj_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adjT_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adjT_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act4_forward_cpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act4_backward_cpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+// conversion operations
+// std::vector<torch::Tensor> to_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
+// std::vector<torch::Tensor> from_vec_backward_cpu(int, torch::Tensor, torch::Tensor);
+
+// utility operations
+torch::Tensor orthogonal_projector_cpu(int, torch::Tensor);
+
+torch::Tensor as_matrix_forward_cpu(int, torch::Tensor);
+
+torch::Tensor jleft_forward_cpu(int, torch::Tensor, torch::Tensor);
+
+#endif

third_party/vipe/csrc/lietorch_ext/lietorch_gpu.cu

@@ -0,0 +1,562 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "lietorch_gpu.h"
+#include <eigen3/Eigen/Dense>
+
+#include "common.h"
+#include "dispatch.h"
+
+#include "so3.h"
+#include "rxso3.h"
+#include "se3.h"
+#include "sim3.h"
+
+#define GPU_1D_KERNEL_LOOP(i, n) \
+    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
+
+#define NUM_THREADS 256
+#define NUM_BLOCKS(batch_size) ((batch_size + NUM_THREADS - 1) / NUM_THREADS)
+
+template <typename Group, typename scalar_t>
+__global__ void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int num_threads) {
+    // exponential map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Data>(X_ptr + i * Group::N) = Group::Exp(a).data();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int num_threads) {
+    // exponential map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a(a_ptr + i * Group::K);
+        Grad dX(grad + i * Group::N);
+        Eigen::Map<Grad>(da + i * Group::K) = dX * Group::left_jacobian(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int num_threads) {
+    // logarithm map forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a = Group(X_ptr + i * Group::N).Log();
+        Eigen::Map<Tangent>(a_ptr + i * Group::K) = a;
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
+    // logarithm map backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Tangent a = Group(X_ptr + i * Group::N).Log();
+        Grad da(grad + i * Group::K);
+        Eigen::Map<Grad>(dX + i * Group::N) = da * Group::left_jacobian_inverse(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int num_threads) {
+    // group inverse forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Data>(Y_ptr + i * Group::N) = X.inv().data();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int num_threads) {
+    // group inverse backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group Y = Group(X_ptr + i * Group::N).inv();
+        Grad dY(grad + i * Group::N);
+        Eigen::Map<Grad>(dX + i * Group::N) = -dY * Y.Adj();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int num_threads) {
+    // group multiplication forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group Z = Group(X_ptr + i * Group::N) * Group(Y_ptr + i * Group::N);
+        Eigen::Map<Data>(Z_ptr + i * Group::N) = Z.data();
+    }
+}
+
+template <class Group, typename scalar_t>
+__global__ void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX,
+                                    scalar_t* dY, int num_threads) {
+    // group multiplication backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Grad dZ(grad + i * Group::N);
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Grad>(dX + i * Group::N) = dZ;
+        Eigen::Map<Grad>(dY + i * Group::N) = dZ * X.Adj();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.Adj(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                                    scalar_t* da, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Grad db(grad + i * Group::K);
+
+        Tangent a(a_ptr + i * Group::K);
+        Tangent b = X.Adj() * a;
+
+        Eigen::Map<Grad>(da + i * Group::K) = db * X.Adj();
+        Eigen::Map<Grad>(dX + i * Group::N) = -db * Group::adj(b);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // adjoint forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = X.AdjT(a);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX,
+                                     scalar_t* da, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent db(grad + i * Group::K);
+        Grad a(a_ptr + i * Group::K);
+
+        Eigen::Map<Tangent>(da + i * Group::K) = X.Adj(db);
+        Eigen::Map<Grad>(dX + i * Group::N) = -a * Group::adj(X.Adj(db));
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 3);
+        Eigen::Map<Point>(q_ptr + i * 3) = X * p;
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                                    scalar_t* dp, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 3, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 3>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 3);
+        PointGrad dq(grad + i * 3);
+
+        Eigen::Map<PointGrad>(dp + i * 3) = dq * X.Matrix4x4().block<3, 3>(0, 0);
+        Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act_jacobian(X * p);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int num_threads) {
+    // action on point forward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 4);
+        Eigen::Map<Point>(q_ptr + i * 4) = X.act4(p);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX,
+                                     scalar_t* dp, int num_threads) {
+    // adjoint backward kernel
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Grad = Eigen::Matrix<scalar_t, 1, Group::K>;
+    using Point = Eigen::Matrix<scalar_t, 4, 1>;
+    using PointGrad = Eigen::Matrix<scalar_t, 1, 4>;
+    using Transformation = Eigen::Matrix<scalar_t, 4, 4>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Point p(p_ptr + i * 4);
+        PointGrad dq(grad + i * 4);
+
+        Eigen::Map<PointGrad>(dp + i * 4) = dq * X.Matrix4x4();
+        const Point q = X.act4(p);
+        Eigen::Map<Grad>(dX + i * Group::N) = dq * Group::act4_jacobian(q);
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int num_threads) {
+    // convert to 4x4 matrix representation
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+    using Matrix4 = Eigen::Matrix<scalar_t, 4, 4, Eigen::RowMajor>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Matrix4>(T_ptr + i * 16) = X.Matrix4x4();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int num_threads) {
+    // orthogonal projection matrix
+    using Proj = Eigen::Matrix<scalar_t, Group::N, Group::N, Eigen::RowMajor>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Eigen::Map<Proj>(P_ptr + i * Group::N * Group::N) = X.orthogonal_projector();
+    }
+}
+
+template <typename Group, typename scalar_t>
+__global__ void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int num_threads) {
+    // left jacobian inverse action
+    using Tangent = Eigen::Matrix<scalar_t, Group::K, 1>;
+    using Data = Eigen::Matrix<scalar_t, Group::N, 1>;
+
+    GPU_1D_KERNEL_LOOP(i, num_threads) {
+        Group X(X_ptr + i * Group::N);
+        Tangent a(a_ptr + i * Group::K);
+        Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
+        Eigen::Map<Tangent>(b_ptr + i * Group::K) = b;
+    }
+}
+
+// unary operations
+
+torch::Tensor exp_forward_gpu(int group_id, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor X;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.scalar_type(), "exp_forward_kernel", ([&] {
+                                          X = torch::zeros({batch_size, group_t::N}, a.options());
+                                          exp_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  a.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return X;
+}
+
+std::vector<torch::Tensor> exp_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor a) {
+    int batch_size = a.size(0);
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, a.scalar_type(), "exp_backward_kernel", ([&] {
+            exp_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {da};
+}
+
+torch::Tensor log_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor a;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "log_forward_kernel", ([&] {
+                                          a = torch::zeros({batch_size, group_t::K}, X.options());
+                                          log_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return a;
+}
+
+std::vector<torch::Tensor> log_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "log_backward_kernel", ([&] {
+            log_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX};
+}
+
+torch::Tensor inv_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor Y = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "inv_forward_kernel", ([&] {
+                                          inv_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return Y;
+}
+
+std::vector<torch::Tensor> inv_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "inv_backward_kernel", ([&] {
+            inv_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX};
+}
+
+// binary operations
+torch::Tensor mul_forward_gpu(int group_id, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor Z = torch::zeros_like(X);
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "mul_forward_kernel", ([&] {
+            mul_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), Z.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return Z;
+}
+
+std::vector<torch::Tensor> mul_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "mul_backward_kernel", ([&] {
+            mul_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), Y.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dY.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dY};
+}
+
+torch::Tensor adj_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adj_forward_kernel", ([&] {
+            adj_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adj_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adj_backward_kernel", ([&] {
+            adj_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, da};
+}
+
+torch::Tensor adjT_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adjT_forward_kernel", ([&] {
+            adjT_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}
+
+std::vector<torch::Tensor> adjT_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor da = torch::zeros(a.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "adjT_backward_kernel", ([&] {
+            adjT_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                da.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, da};
+}
+
+torch::Tensor act_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act_forward_kernel", ([&] {
+            act_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act_backward_kernel", ([&] {
+            act_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dp.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dp};
+}
+
+torch::Tensor act4_forward_gpu(int group_id, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor q = torch::zeros(p.sizes(), p.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act4_forward_kernel", ([&] {
+            act4_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), q.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return q;
+}
+
+std::vector<torch::Tensor> act4_backward_gpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
+    int batch_size = X.size(0);
+    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
+    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "act4_backward_kernel", ([&] {
+            act4_backward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                grad.data_ptr<scalar_t>(), X.data_ptr<scalar_t>(), p.data_ptr<scalar_t>(), dX.data_ptr<scalar_t>(),
+                dp.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return {dX, dp};
+}
+
+torch::Tensor as_matrix_forward_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "as_matrix_forward_kernel", ([&] {
+                                          as_matrix_forward_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), T4x4.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return T4x4;
+}
+
+torch::Tensor orthogonal_projector_gpu(int group_id, torch::Tensor X) {
+    int batch_size = X.size(0);
+    torch::Tensor P;
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.scalar_type(), "orthogonal_projector_kernel", ([&] {
+                                          P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
+                                          orthogonal_projector_kernel<group_t, scalar_t>
+                                              <<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                                                  X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
+                                      }));
+
+    return P;
+}
+
+torch::Tensor jleft_forward_gpu(int group_id, torch::Tensor X, torch::Tensor a) {
+    int batch_size = X.size(0);
+    torch::Tensor b = torch::zeros(a.sizes(), a.options());
+
+    DISPATCH_GROUP_AND_FLOATING_TYPES(
+        group_id, X.scalar_type(), "jleft_forward_kernel", ([&] {
+            jleft_forward_kernel<group_t, scalar_t><<<NUM_BLOCKS(batch_size), NUM_THREADS>>>(
+                X.data_ptr<scalar_t>(), a.data_ptr<scalar_t>(), b.data_ptr<scalar_t>(), batch_size);
+        }));
+
+    return b;
+}

third_party/vipe/csrc/lietorch_ext/lietorch_gpu.h

@@ -0,0 +1,52 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef LIETORCH_GPU_H_
+#define LIETORCH_GPU_H_
+
+#include <vector>
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+// unary operations
+torch::Tensor exp_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> exp_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor log_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> log_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+torch::Tensor inv_forward_gpu(int, torch::Tensor);
+std::vector<torch::Tensor> inv_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+// binary operations
+torch::Tensor mul_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> mul_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adj_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adj_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor adjT_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> adjT_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+torch::Tensor act4_forward_gpu(int, torch::Tensor, torch::Tensor);
+std::vector<torch::Tensor> act4_backward_gpu(int, torch::Tensor, torch::Tensor, torch::Tensor);
+
+// conversion operations
+// std::vector<torch::Tensor> to_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
+// std::vector<torch::Tensor> from_vec_backward_gpu(int, torch::Tensor, torch::Tensor);
+
+// utility operators
+torch::Tensor orthogonal_projector_gpu(int, torch::Tensor);
+
+torch::Tensor as_matrix_forward_gpu(int, torch::Tensor);
+
+torch::Tensor jleft_forward_gpu(int, torch::Tensor, torch::Tensor);
+
+#endif

third_party/vipe/csrc/lietorch_ext/rxso3.h

@@ -0,0 +1,313 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef RxSO3_HEADER
+#define RxSO3_HEADER
+
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+
+template <typename Scalar>
+class RxSO3 {
+   public:
+    const static int constexpr K = 4;  // manifold dimension
+    const static int constexpr N = 5;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+
+    using Quaternion = Eigen::Quaternion<Scalar>;
+    using Transformation = Eigen::Matrix<Scalar, 3, 3>;
+    using Adjoint = Eigen::Matrix<Scalar, 4, 4>;
+
+    EIGEN_DEVICE_FUNC RxSO3(Quaternion const& q, Scalar const s) : unit_quaternion(q), scale(s) {
+        unit_quaternion.normalize();
+    };
+
+    EIGEN_DEVICE_FUNC RxSO3(const Scalar* data) : unit_quaternion(data), scale(data[4]) {
+        unit_quaternion.normalize();
+    };
+
+    EIGEN_DEVICE_FUNC RxSO3() {
+        unit_quaternion = Quaternion::Identity();
+        scale = Scalar(1.0);
+    }
+
+    EIGEN_DEVICE_FUNC RxSO3<Scalar> inv() { return RxSO3<Scalar>(unit_quaternion.conjugate(), 1.0 / scale); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << unit_quaternion.coeffs(), scale;
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC RxSO3<Scalar> operator*(RxSO3<Scalar> const& other) {
+        return RxSO3<Scalar>(unit_quaternion * other.unit_quaternion, scale * other.scale);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
+        const Quaternion& q = unit_quaternion;
+        Point uv = q.vec().cross(p);
+        uv += uv;
+        return scale * (p + q.w() * uv + q.vec().cross(uv));
+    }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << this->operator*(p.template segment<3>(0)), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Adjoint Ad = Adjoint::Identity();
+        Ad.template block<3, 3>(0, 0) = unit_quaternion.toRotationMatrix();
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const { return scale * unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
+        Eigen::Matrix<Scalar, 4, 4> T;
+        T = Eigen::Matrix<Scalar, 4, 4>::Identity();
+        T.template block<3, 3>(0, 0) = Matrix();
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 5, 5> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 5, 5> J = Eigen::Matrix<Scalar, 5, 5>::Zero();
+
+        J.template block<3, 3>(0, 0) =
+            0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
+
+        J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
+
+        // scale
+        J(4, 3) = scale;
+
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Rotation() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi_sigma) {
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        return SO3<Scalar>::hat(phi) + phi(3) * Transformation::Identity();
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi_sigma) {
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+
+        Adjoint ad = Adjoint::Zero();
+        ad.template block<3, 3>(0, 0) = Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        using std::abs;
+        using std::atan;
+        using std::sqrt;
+
+        Scalar squared_n = unit_quaternion.vec().squaredNorm();
+        Scalar w = unit_quaternion.w();
+        Scalar two_atan_nbyw_by_n;
+
+        /// Atan-based log thanks to
+        ///
+        /// C. Hertzberg et al.:
+        /// "Integrating Generic Sensor Fusion Algorithms with Sound State
+        /// Representation through Encapsulation of Manifolds"
+        /// Information Fusion, 2011
+
+        if (squared_n < EPS * EPS) {
+            two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * w * w);
+        } else {
+            Scalar n = sqrt(squared_n);
+            if (abs(w) < EPS) {
+                if (w > Scalar(0)) {
+                    two_atan_nbyw_by_n = PI / n;
+                } else {
+                    two_atan_nbyw_by_n = -PI / n;
+                }
+            } else {
+                two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
+            }
+        }
+
+        Tangent phi_sigma;
+        phi_sigma << two_atan_nbyw_by_n * unit_quaternion.vec(), log(scale);
+
+        return phi_sigma;
+    }
+
+    EIGEN_DEVICE_FUNC static RxSO3<Scalar> Exp(Tangent const& phi_sigma) {
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        Scalar scale = exp(phi_sigma(3));
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar imag_factor;
+        Scalar real_factor;
+
+        if (theta < EPS) {
+            Scalar theta4 = theta2 * theta2;
+            imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
+            real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
+        } else {
+            imag_factor = sin(.5 * theta) / theta;
+            real_factor = cos(.5 * theta);
+        }
+
+        Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
+        return RxSO3<Scalar>(q, scale);
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcW(Tangent const& phi_sigma) {
+        // left jacobian
+        using std::abs;
+        Matrix3 const I = Matrix3::Identity();
+        Scalar const one(1);
+        Scalar const half(0.5);
+
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Scalar const sigma = phi_sigma(3);
+        Scalar const theta = phi.norm();
+
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+        Matrix3 const Phi2 = Phi * Phi;
+        Scalar const scale = exp(sigma);
+
+        Scalar A, B, C;
+        if (abs(sigma) < EPS) {
+            C = one;
+            if (abs(theta) < EPS) {
+                A = half;
+                B = Scalar(1. / 6.);
+            } else {
+                Scalar theta_sq = theta * theta;
+                A = (one - cos(theta)) / theta_sq;
+                B = (theta - sin(theta)) / (theta_sq * theta);
+            }
+        } else {
+            C = (scale - one) / sigma;
+            if (abs(theta) < EPS) {
+                Scalar sigma_sq = sigma * sigma;
+                A = ((sigma - one) * scale + one) / sigma_sq;
+                B = (scale * half * sigma_sq + scale - one - sigma * scale) / (sigma_sq * sigma);
+            } else {
+                Scalar theta_sq = theta * theta;
+                Scalar a = scale * sin(theta);
+                Scalar b = scale * cos(theta);
+                Scalar c = theta_sq + sigma * sigma;
+                A = (a * sigma + (one - b) * theta) / (theta * c);
+                B = (C - ((b - one) * sigma + a * theta) / (c)) * one / (theta_sq);
+            }
+        }
+        return A * Phi + B * Phi2 + C * I;
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcWInv(Tangent const& phi_sigma) {
+        // left jacobian inverse
+        Matrix3 const I = Matrix3::Identity();
+        Scalar const half(0.5);
+        Scalar const one(1);
+        Scalar const two(2);
+
+        Vector3 const phi = phi_sigma.template segment<3>(0);
+        Scalar const sigma = phi_sigma(3);
+        Scalar const theta = phi.norm();
+        Scalar const scale = exp(sigma);
+
+        Matrix3 const Phi = SO3<Scalar>::hat(phi);
+        Matrix3 const Phi2 = Phi * Phi;
+        Scalar const scale_sq = scale * scale;
+        Scalar const theta_sq = theta * theta;
+        Scalar const sin_theta = sin(theta);
+        Scalar const cos_theta = cos(theta);
+
+        Scalar a, b, c;
+        if (abs(sigma * sigma) < EPS) {
+            c = one - half * sigma;
+            a = -half;
+            if (abs(theta_sq) < EPS) {
+                b = Scalar(1. / 12.);
+            } else {
+                b = (theta * sin_theta + two * cos_theta - two) / (two * theta_sq * (cos_theta - one));
+            }
+        } else {
+            Scalar const scale_cu = scale_sq * scale;
+            c = sigma / (scale - one);
+            if (abs(theta_sq) < EPS) {
+                a = (-sigma * scale + scale - one) / ((scale - one) * (scale - one));
+                b = (scale_sq * sigma - two * scale_sq + scale * sigma + two * scale) /
+                    (two * scale_cu - Scalar(6) * scale_sq + Scalar(6) * scale - two);
+            } else {
+                Scalar const s_sin_theta = scale * sin_theta;
+                Scalar const s_cos_theta = scale * cos_theta;
+                a = (theta * s_cos_theta - theta - sigma * s_sin_theta) /
+                    (theta * (scale_sq - two * s_cos_theta + one));
+                b = -scale *
+                    (theta * s_sin_theta - theta * sin_theta + sigma * s_cos_theta - scale * sigma + sigma * cos_theta -
+                     sigma) /
+                    (theta_sq * (scale_cu - two * scale * s_cos_theta - scale_sq + two * s_cos_theta + scale - one));
+            }
+        }
+        return a * Phi + b * Phi2 + c * I;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi_sigma) {
+        // left jacobian
+        Adjoint J = Adjoint::Identity();
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian(phi);
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi_sigma) {
+        // left jacobian inverse
+        Adjoint Jinv = Adjoint::Identity();
+        Vector3 phi = phi_sigma.template segment<3>(0);
+        Jinv.template block<3, 3>(0, 0) = SO3<Scalar>::left_jacobian_inverse(phi);
+        return Jinv;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 4> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 4> Ja;
+        Ja << SO3<Scalar>::hat(-p), p;
+        return Ja;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 4> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        J.template block<3, 1>(0, 3) = p.template segment<3>(0);
+        return J;
+    }
+
+   private:
+    Quaternion unit_quaternion;
+    Scalar scale;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/se3.h

@@ -0,0 +1,214 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef SE3_HEADER
+#define SE3_HEADER
+
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+#include "so3.h"
+
+template <typename Scalar>
+class SE3 {
+   public:
+    const static int constexpr K = 6;  // manifold dimension
+    const static int constexpr N = 7;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 4, 4>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+
+    EIGEN_DEVICE_FUNC SE3() { translation = Vector3::Zero(); }
+
+    EIGEN_DEVICE_FUNC SE3(SO3<Scalar> const& so3, Vector3 const& t) : so3(so3), translation(t) {};
+
+    EIGEN_DEVICE_FUNC SE3(const Scalar* data) : translation(data), so3(data + 3) {};
+
+    EIGEN_DEVICE_FUNC SE3<Scalar> inv() { return SE3(so3.inv(), -(so3.inv() * translation)); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << translation, so3.data();
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC SE3<Scalar> operator*(SE3<Scalar> const& other) {
+        return SE3(so3 * other.so3, translation + so3 * other.translation);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return so3 * p + translation; }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << so3 * p.template segment<3>(0) + translation * p(3), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Matrix3 R = so3.Matrix();
+        Matrix3 tx = SO3<Scalar>::hat(translation);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint Ad;
+        Ad << R, tx * R, Zer, R;
+
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = so3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix4x4() const { return Matrix(); }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        Transformation TauPhi = Transformation::Zero();
+        TauPhi.template block<3, 3>(0, 0) = SO3<Scalar>::hat(phi);
+        TauPhi.template block<3, 1>(0, 3) = tau;
+
+        return TauPhi;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint ad;
+        ad << Phi, Tau, Zer, Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 7, 7> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 7, 7> J = Eigen::Matrix<Scalar, 7, 7>::Zero();
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
+        J.template block<4, 4>(3, 3) = so3.orthogonal_projector();
+
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        Vector3 phi = so3.Log();
+        Matrix3 Vinv = SO3<Scalar>::left_jacobian_inverse(phi);
+
+        Tangent tau_phi;
+        tau_phi << Vinv * translation, phi;
+
+        return tau_phi;
+    }
+
+    EIGEN_DEVICE_FUNC static SE3<Scalar> Exp(Tangent const& tau_phi) {
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+
+        SO3<Scalar> so3 = SO3<Scalar>::Exp(phi);
+        Vector3 t = SO3<Scalar>::left_jacobian(phi) * tau;
+
+        return SE3<Scalar>(so3, t);
+    }
+
+    EIGEN_DEVICE_FUNC static Matrix3 calcQ(Tangent const& tau_phi) {
+        // Q matrix
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+
+        Scalar theta = phi.norm();
+        Scalar theta_pow2 = theta * theta;
+        Scalar theta_pow4 = theta_pow2 * theta_pow2;
+
+        Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta_pow2
+                                     : (theta - sin(theta)) / (theta_pow2 * theta);
+
+        Scalar coef2 = (theta < EPS) ? Scalar(1.0 / 24.0) - Scalar(1.0 / 720.0) * theta_pow2
+                                     : (theta_pow2 + 2 * cos(theta) - 2) / (2 * theta_pow4);
+
+        Scalar coef3 = (theta < EPS) ? Scalar(1.0 / 120.0) - Scalar(1.0 / 2520.0) * theta_pow2
+                                     : (2 * theta - 3 * sin(theta) + theta * cos(theta)) / (2 * theta_pow4 * theta);
+
+        Matrix3 Q = Scalar(0.5) * Tau + coef1 * (Phi * Tau + Tau * Phi + Phi * Tau * Phi) +
+                    coef2 * (Phi * Phi * Tau + Tau * Phi * Phi - 3 * Phi * Tau * Phi) +
+                    coef3 * (Phi * Tau * Phi * Phi + Phi * Phi * Tau * Phi);
+
+        return Q;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi) {
+        // left jacobian
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 J = SO3<Scalar>::left_jacobian(phi);
+        Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint J6x6;
+        J6x6 << J, Q, Zer, J;
+
+        return J6x6;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi) {
+        // left jacobian inverse
+        Vector3 tau = tau_phi.template segment<3>(0);
+        Vector3 phi = tau_phi.template segment<3>(3);
+        Matrix3 Jinv = SO3<Scalar>::left_jacobian_inverse(phi);
+        Matrix3 Q = SE3<Scalar>::calcQ(tau_phi);
+        Matrix3 Zer = Matrix3::Zero();
+
+        Adjoint J6x6;
+        J6x6 << Jinv, -Jinv * Q * Jinv, Zer, Jinv;
+
+        return J6x6;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 6> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 6> J;
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 6> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 6> J = Eigen::Matrix<Scalar, 4, 6>::Zero();
+        J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        return J;
+    }
+
+   private:
+    SO3<Scalar> so3;
+    Vector3 translation;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/sim3.h

@@ -0,0 +1,203 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef Sim3_HEADER
+#define Sim3_HEADER
+
+#include <stdio.h>
+#include <iostream>
+
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+#include "so3.h"
+#include "rxso3.h"
+
+template <typename Scalar>
+class Sim3 {
+   public:
+    const static int constexpr K = 7;  // manifold dimension
+    const static int constexpr N = 8;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 4, 4>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+
+    EIGEN_DEVICE_FUNC Sim3() { translation = Vector3::Zero(); }
+
+    EIGEN_DEVICE_FUNC Sim3(RxSO3<Scalar> const& rxso3, Vector3 const& t) : rxso3(rxso3), translation(t) {};
+
+    EIGEN_DEVICE_FUNC Sim3(const Scalar* data) : translation(data), rxso3(data + 3) {};
+
+    EIGEN_DEVICE_FUNC Sim3<Scalar> inv() { return Sim3<Scalar>(rxso3.inv(), -(rxso3.inv() * translation)); }
+
+    EIGEN_DEVICE_FUNC Data data() const {
+        Data data_vec;
+        data_vec << translation, rxso3.data();
+        return data_vec;
+    }
+
+    EIGEN_DEVICE_FUNC Sim3<Scalar> operator*(Sim3<Scalar> const& other) {
+        return Sim3(rxso3 * other.rxso3, translation + rxso3 * other.translation);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const { return (rxso3 * p) + translation; }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << rxso3 * p.template segment<3>(0) + p(3) * translation, p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = rxso3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix4x4() const {
+        Transformation T = Transformation::Identity();
+        T.template block<3, 3>(0, 0) = rxso3.Matrix();
+        T.template block<3, 1>(0, 3) = translation;
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 8, 8> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 8, 8> J = Eigen::Matrix<Scalar, 8, 8>::Zero();
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-translation);
+        J.template block<3, 1>(0, 6) = translation;
+        J.template block<5, 5>(3, 3) = rxso3.orthogonal_projector();
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const {
+        Adjoint Ad = Adjoint::Identity();
+        Matrix3 sR = rxso3.Matrix();
+        Matrix3 tx = SO3<Scalar>::hat(translation);
+        Matrix3 R = rxso3.Rotation();
+
+        Ad.template block<3, 3>(0, 0) = sR;
+        Ad.template block<3, 3>(0, 3) = tx * R;
+        Ad.template block<3, 1>(0, 6) = -translation;
+        Ad.template block<3, 3>(3, 3) = R;
+
+        return Ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& tau_phi_sigma) {
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector3 phi = tau_phi_sigma.template segment<3>(3);
+        Scalar sigma = tau_phi_sigma(6);
+
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 I = Matrix3::Identity();
+
+        Transformation Omega = Transformation::Zero();
+        Omega.template block<3, 3>(0, 0) = Phi + sigma * I;
+        Omega.template block<3, 1>(0, 3) = tau;
+
+        return Omega;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& tau_phi_sigma) {
+        Adjoint ad = Adjoint::Zero();
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector3 phi = tau_phi_sigma.template segment<3>(3);
+        Scalar sigma = tau_phi_sigma(6);
+
+        Matrix3 Tau = SO3<Scalar>::hat(tau);
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 I = Matrix3::Identity();
+
+        ad.template block<3, 3>(0, 0) = Phi + sigma * I;
+        ad.template block<3, 3>(0, 3) = Tau;
+        ad.template block<3, 1>(0, 6) = -tau;
+        ad.template block<3, 3>(3, 3) = Phi;
+
+        return ad;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        // logarithm map
+        Vector4 phi_sigma = rxso3.Log();
+        Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
+
+        Tangent tau_phi_sigma;
+        tau_phi_sigma << W.inverse() * translation, phi_sigma;
+
+        return tau_phi_sigma;
+    }
+
+    EIGEN_DEVICE_FUNC static Sim3<Scalar> Exp(Tangent const& tau_phi_sigma) {
+        // exponential map
+        Vector3 tau = tau_phi_sigma.template segment<3>(0);
+        Vector4 phi_sigma = tau_phi_sigma.template segment<4>(3);
+
+        RxSO3<Scalar> rxso3 = RxSO3<Scalar>::Exp(phi_sigma);
+        Matrix3 W = RxSO3<Scalar>::calcW(phi_sigma);
+
+        return Sim3<Scalar>(rxso3, W * tau);
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& tau_phi_sigma) {
+        // left jacobian
+        Adjoint const Xi = adj(tau_phi_sigma);
+        Adjoint const Xi2 = Xi * Xi;
+        Adjoint const Xi4 = Xi2 * Xi2;
+
+        return Adjoint::Identity() + Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 6.0) * Xi2 + Scalar(1.0 / 24.0) * Xi * Xi2 +
+               Scalar(1.0 / 120.0) * Xi4;
+        +Scalar(1.0 / 720.0) * Xi* Xi4;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& tau_phi_sigma) {
+        // left jacobian inverse
+        Adjoint const Xi = adj(tau_phi_sigma);
+        Adjoint const Xi2 = Xi * Xi;
+        Adjoint const Xi4 = Xi2 * Xi2;
+
+        return Adjoint::Identity() - Scalar(1.0 / 2.0) * Xi + Scalar(1.0 / 12.0) * Xi2 - Scalar(1.0 / 720.0) * Xi4;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 7> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 3, 7> J;
+        J.template block<3, 3>(0, 0) = Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p);
+        J.template block<3, 1>(0, 6) = p;
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 7> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 7> J = Eigen::Matrix<Scalar, 4, 7>::Zero();
+        J.template block<3, 3>(0, 0) = p(3) * Matrix3::Identity();
+        J.template block<3, 3>(0, 3) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        J.template block<3, 1>(0, 6) = p.template segment<3>(0);
+        return J;
+    }
+
+   private:
+    Vector3 translation;
+    RxSO3<Scalar> rxso3;
+};
+
+#endif

third_party/vipe/csrc/lietorch_ext/so3.h

@@ -0,0 +1,202 @@
+/**
+ * This file includes code originally from the lietorch repository:
+ * https://github.com/princeton-vl/lietorch
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#ifndef SO3_HEADER
+#define SO3_HEADER
+
+#include <cuda.h>
+#include <stdio.h>
+#include <eigen3/Eigen/Dense>
+#include <eigen3/Eigen/Geometry>
+
+#include "common.h"
+
+template <typename Scalar>
+class SO3 {
+   public:
+    const static int constexpr K = 3;  // manifold dimension
+    const static int constexpr N = 4;  // embedding dimension
+
+    using Vector3 = Eigen::Matrix<Scalar, 3, 1>;
+    using Vector4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Matrix3 = Eigen::Matrix<Scalar, 3, 3>;
+
+    using Tangent = Eigen::Matrix<Scalar, K, 1>;
+    using Data = Eigen::Matrix<Scalar, N, 1>;
+
+    using Point = Eigen::Matrix<Scalar, 3, 1>;
+    using Point4 = Eigen::Matrix<Scalar, 4, 1>;
+    using Transformation = Eigen::Matrix<Scalar, 3, 3>;
+    using Adjoint = Eigen::Matrix<Scalar, K, K>;
+    using Quaternion = Eigen::Quaternion<Scalar>;
+
+    EIGEN_DEVICE_FUNC SO3(Quaternion const& q) : unit_quaternion(q) { unit_quaternion.normalize(); };
+
+    EIGEN_DEVICE_FUNC SO3(const Scalar* data) : unit_quaternion(data) { unit_quaternion.normalize(); };
+
+    EIGEN_DEVICE_FUNC SO3() { unit_quaternion = Quaternion::Identity(); }
+
+    EIGEN_DEVICE_FUNC SO3<Scalar> inv() { return SO3<Scalar>(unit_quaternion.conjugate()); }
+
+    EIGEN_DEVICE_FUNC Data data() const { return unit_quaternion.coeffs(); }
+
+    EIGEN_DEVICE_FUNC SO3<Scalar> operator*(SO3<Scalar> const& other) {
+        return SO3(unit_quaternion * other.unit_quaternion);
+    }
+
+    EIGEN_DEVICE_FUNC Point operator*(Point const& p) const {
+        const Quaternion& q = unit_quaternion;
+        Point uv = q.vec().cross(p);
+        uv += uv;
+        return p + q.w() * uv + q.vec().cross(uv);
+    }
+
+    EIGEN_DEVICE_FUNC Point4 act4(Point4 const& p) const {
+        Point4 p1;
+        p1 << this->operator*(p.template segment<3>(0)), p(3);
+        return p1;
+    }
+
+    EIGEN_DEVICE_FUNC Adjoint Adj() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Transformation Matrix() const { return unit_quaternion.toRotationMatrix(); }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> Matrix4x4() const {
+        Eigen::Matrix<Scalar, 4, 4> T = Eigen::Matrix<Scalar, 4, 4>::Identity();
+        T.template block<3, 3>(0, 0) = Matrix();
+        return T;
+    }
+
+    EIGEN_DEVICE_FUNC Eigen::Matrix<Scalar, 4, 4> orthogonal_projector() const {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 4> J = Eigen::Matrix<Scalar, 4, 4>::Zero();
+        J.template block<3, 3>(0, 0) =
+            0.5 * (unit_quaternion.w() * Matrix3::Identity() + SO3<Scalar>::hat(-unit_quaternion.vec()));
+
+        J.template block<1, 3>(3, 0) = 0.5 * (-unit_quaternion.vec());
+        return J;
+    }
+
+    EIGEN_DEVICE_FUNC Tangent Adj(Tangent const& a) const { return Adj() * a; }
+
+    EIGEN_DEVICE_FUNC Tangent AdjT(Tangent const& a) const { return Adj().transpose() * a; }
+
+    EIGEN_DEVICE_FUNC static Transformation hat(Tangent const& phi) {
+        Transformation Phi;
+        Phi << 0.0, -phi(2), phi(1), phi(2), 0.0, -phi(0), -phi(1), phi(0), 0.0;
+
+        return Phi;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint adj(Tangent const& phi) { return SO3<Scalar>::hat(phi); }
+
+    EIGEN_DEVICE_FUNC Tangent Log() const {
+        using std::abs;
+        using std::atan;
+        using std::sqrt;
+        Scalar squared_n = unit_quaternion.vec().squaredNorm();
+        Scalar w = unit_quaternion.w();
+
+        Scalar two_atan_nbyw_by_n;
+
+        /// Atan-based log thanks to
+        ///
+        /// C. Hertzberg et al.:
+        /// "Integrating Generic Sensor Fusion Algorithms with Sound State
+        /// Representation through Encapsulation of Manifolds"
+        /// Information Fusion, 2011
+
+        if (squared_n < EPS * EPS) {
+            // If quaternion is normalized and n=0, then w should be 1;
+            // w=0 should never happen here!
+            Scalar squared_w = w * w;
+            two_atan_nbyw_by_n = Scalar(2) / w - Scalar(2.0 / 3.0) * (squared_n) / (w * squared_w);
+        } else {
+            Scalar n = sqrt(squared_n);
+            if (abs(w) < EPS) {
+                if (w > Scalar(0)) {
+                    two_atan_nbyw_by_n = Scalar(PI) / n;
+                } else {
+                    two_atan_nbyw_by_n = -Scalar(PI) / n;
+                }
+            } else {
+                two_atan_nbyw_by_n = Scalar(2) * atan(n / w) / n;
+            }
+        }
+
+        return two_atan_nbyw_by_n * unit_quaternion.vec();
+    }
+
+    EIGEN_DEVICE_FUNC static SO3<Scalar> Exp(Tangent const& phi) {
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar imag_factor;
+        Scalar real_factor;
+
+        if (theta < EPS) {
+            Scalar theta4 = theta2 * theta2;
+            imag_factor = Scalar(0.5) - Scalar(1.0 / 48.0) * theta2 + Scalar(1.0 / 3840.0) * theta4;
+            real_factor = Scalar(1) - Scalar(1.0 / 8.0) * theta2 + Scalar(1.0 / 384.0) * theta4;
+        } else {
+            imag_factor = sin(.5 * theta) / theta;
+            real_factor = cos(.5 * theta);
+        }
+
+        Quaternion q(real_factor, imag_factor * phi.x(), imag_factor * phi.y(), imag_factor * phi.z());
+        return SO3<Scalar>(q);
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian(Tangent const& phi) {
+        // left jacobian
+        Matrix3 I = Matrix3::Identity();
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Phi2 = Phi * Phi;
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+
+        Scalar coef1 = (theta < EPS) ? Scalar(1.0 / 2.0) - Scalar(1.0 / 24.0) * theta2 : (1.0 - cos(theta)) / theta2;
+
+        Scalar coef2 =
+            (theta < EPS) ? Scalar(1.0 / 6.0) - Scalar(1.0 / 120.0) * theta2 : (theta - sin(theta)) / (theta2 * theta);
+
+        return I + coef1 * Phi + coef2 * Phi2;
+    }
+
+    EIGEN_DEVICE_FUNC static Adjoint left_jacobian_inverse(Tangent const& phi) {
+        // left jacobian inverse
+        Matrix3 I = Matrix3::Identity();
+        Matrix3 Phi = SO3<Scalar>::hat(phi);
+        Matrix3 Phi2 = Phi * Phi;
+
+        Scalar theta2 = phi.squaredNorm();
+        Scalar theta = sqrt(theta2);
+        Scalar half_theta = Scalar(.5) * theta;
+
+        Scalar coef2 = (theta < EPS)
+                           ? Scalar(1.0 / 12.0)
+                           : (Scalar(1) - theta * cos(half_theta) / (Scalar(2) * sin(half_theta))) / (theta * theta);
+
+        return I + Scalar(-0.5) * Phi + coef2 * Phi2;
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 3, 3> act_jacobian(Point const& p) {
+        // jacobian action on a point
+        return SO3<Scalar>::hat(-p);
+    }
+
+    EIGEN_DEVICE_FUNC static Eigen::Matrix<Scalar, 4, 3> act4_jacobian(Point4 const& p) {
+        // jacobian action on a point
+        Eigen::Matrix<Scalar, 4, 3> J = Eigen::Matrix<Scalar, 4, 3>::Zero();
+        J.template block<3, 3>(0, 0) = SO3<Scalar>::hat(-p.template segment<3>(0));
+        return J;
+    }
+
+   private:
+    Quaternion unit_quaternion;
+};
+
+#endif

third_party/vipe/csrc/scatter_ext/cuda/atomics.cuh

@@ -0,0 +1,300 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#define ATOMIC(NAME)                                                                                                  \
+    template <typename scalar, size_t size>                                                                           \
+    struct Atomic##NAME##IntegerImpl;                                                                                 \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 1> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)(address - ((size_t)address & 3));                                  \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t shift = ((size_t)address & 3) * 8;                                                               \
+            uint32_t sum;                                                                                             \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                sum = OP(val, scalar((old >> shift) & 0xff));                                                         \
+                old = (old & ~(0x000000ff << shift)) | (sum << shift);                                                \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 2> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)((char *)address - ((size_t)address & 2));                          \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t sum;                                                                                             \
+            uint32_t newval;                                                                                          \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                sum = OP(val, (size_t)address & 2 ? scalar(old >> 16) : scalar(old & 0xffff));                        \
+                newval = (size_t)address & 2 ? (old & 0xffff) | (sum << 16) : (old & 0xffff0000) | sum;               \
+                old = atomicCAS(address_as_ui, assumed, newval);                                                      \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 4> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            uint32_t *address_as_ui = (uint32_t *)address;                                                            \
+            uint32_t old = *address_as_ui;                                                                            \
+            uint32_t assumed;                                                                                         \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_ui, assumed, OP(val, (scalar)old));                                        \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##IntegerImpl<scalar, 8> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            unsigned long long *address_as_ull = (unsigned long long *)address;                                       \
+            unsigned long long old = *address_as_ull;                                                                 \
+            unsigned long long assumed;                                                                               \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_ull, assumed, OP(val, (scalar)old));                                       \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar, size_t size>                                                                           \
+    struct Atomic##NAME##DecimalImpl;                                                                                 \
+                                                                                                                      \
+    template <>                                                                                                       \
+    struct Atomic##NAME##DecimalImpl<at::Half, 2> {                                                                   \
+        inline __device__ void operator()(at::Half *address, at::Half val) {                                          \
+            unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2));                  \
+            unsigned int old = *address_as_ui;                                                                        \
+            unsigned int assumed;                                                                                     \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                at::Half hsum;                                                                                        \
+                hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);                                          \
+                hsum = OP(hsum, val);                                                                                 \
+                old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x;            \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <>                                                                                                       \
+    struct Atomic##NAME##DecimalImpl<at::BFloat16, 2> {                                                               \
+        inline __device__ void operator()(at::BFloat16 *address, at::BFloat16 val) {                                  \
+            unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2));                  \
+            unsigned int old = *address_as_ui;                                                                        \
+            unsigned int assumed;                                                                                     \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                at::BFloat16 hsum;                                                                                    \
+                hsum.x = (size_t)address & 2 ? (old >> 16) : (old & 0xffff);                                          \
+                hsum = OP(hsum, val);                                                                                 \
+                old = (size_t)address & 2 ? (old & 0xffff) | (hsum.x << 16) : (old & 0xffff0000) | hsum.x;            \
+                old = atomicCAS(address_as_ui, assumed, old);                                                         \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##DecimalImpl<scalar, 4> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            int *address_as_i = (int *)address;                                                                       \
+            int old = *address_as_i;                                                                                  \
+            int assumed;                                                                                              \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old = atomicCAS(address_as_i, assumed, __float_as_int(OP(val, __int_as_float(assumed))));             \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };                                                                                                                \
+                                                                                                                      \
+    template <typename scalar>                                                                                        \
+    struct Atomic##NAME##DecimalImpl<scalar, 8> {                                                                     \
+        inline __device__ void operator()(scalar *address, scalar val) {                                              \
+            unsigned long long int *address_as_ull = (unsigned long long int *)address;                               \
+            unsigned long long int old = *address_as_ull;                                                             \
+            unsigned long long int assumed;                                                                           \
+                                                                                                                      \
+            do {                                                                                                      \
+                assumed = old;                                                                                        \
+                old =                                                                                                 \
+                    atomicCAS(address_as_ull, assumed, __double_as_longlong(OP(val, __longlong_as_double(assumed)))); \
+            } while (assumed != old);                                                                                 \
+        }                                                                                                             \
+    };
+
+#define OP(X, Y) Y + X
+ATOMIC(Add)
+#undef OP
+static inline __device__ void atomAdd(uint8_t *address, uint8_t val) {
+    AtomicAddIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int8_t *address, int8_t val) {
+    AtomicAddIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int16_t *address, int16_t val) {
+    AtomicAddIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomAdd(int32_t *address, int32_t val) { atomicAdd(address, val); }
+static inline __device__ void atomAdd(int64_t *address, int64_t val) {
+    AtomicAddIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+#if defined(USE_ROCM) || (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700 || CUDA_VERSION < 10000))
+static inline __device__ void atomAdd(at::Half *address, at::Half val) {
+    AtomicAddDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+#else
+static inline __device__ void atomAdd(at::Half *address, at::Half val) {
+    atomicAdd(reinterpret_cast<__half *>(address), val);
+}
+#endif
+static inline __device__ void atomAdd(float *address, float val) { atomicAdd(address, val); }
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600 || CUDA_VERSION < 8000)
+static inline __device__ void atomAdd(double *address, double val) {
+    AtomicAddDecimalImpl<double, sizeof(double)>()(address, val);
+}
+#else
+static inline __device__ void atomAdd(double *address, double val) { atomicAdd(address, val); }
+#endif
+static inline __device__ void atomAdd(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicAddDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) Y *X
+ATOMIC(Mul)
+#undef OP
+static inline __device__ void atomMul(uint8_t *address, uint8_t val) {
+    AtomicMulIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMul(int8_t *address, int8_t val) {
+    AtomicMulIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMul(int16_t *address, int16_t val) {
+    AtomicMulIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMul(int32_t *address, int32_t val) {
+    AtomicMulIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
+}
+static inline __device__ void atomMul(int64_t *address, int64_t val) {
+    AtomicMulIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMul(float *address, float val) {
+    AtomicMulDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMul(at::Half *address, at::Half val) {
+    AtomicMulDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMul(double *address, double val) {
+    AtomicMulDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMul(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMulDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) Y / X
+ATOMIC(Div)
+#undef OP
+static inline __device__ void atomDiv(uint8_t *address, uint8_t val) {
+    AtomicDivIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int8_t *address, int8_t val) {
+    AtomicDivIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int16_t *address, int16_t val) {
+    AtomicDivIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int32_t *address, int32_t val) {
+    AtomicDivIntegerImpl<int32_t, sizeof(int32_t)>()(address, val);
+}
+static inline __device__ void atomDiv(int64_t *address, int64_t val) {
+    AtomicDivIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomDiv(at::Half *address, at::Half val) {
+    AtomicDivDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomDiv(float *address, float val) {
+    AtomicDivDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomDiv(double *address, double val) {
+    AtomicDivDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomDiv(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicDivDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) max(Y, X)
+ATOMIC(Max)
+#undef OP
+static inline __device__ void atomMax(uint8_t *address, uint8_t val) {
+    AtomicMaxIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMax(int8_t *address, int8_t val) {
+    AtomicMaxIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMax(int16_t *address, int16_t val) {
+    AtomicMaxIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMax(int32_t *address, int32_t val) { atomicMax(address, val); }
+static inline __device__ void atomMax(int64_t *address, int64_t val) {
+    AtomicMaxIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMax(at::Half *address, at::Half val) {
+    AtomicMaxDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMax(float *address, float val) {
+    AtomicMaxDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMax(double *address, double val) {
+    AtomicMaxDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMax(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMaxDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}
+
+#define OP(X, Y) min(Y, X)
+ATOMIC(Min)
+#undef OP
+static inline __device__ void atomMin(uint8_t *address, uint8_t val) {
+    AtomicMinIntegerImpl<uint8_t, sizeof(uint8_t)>()(address, val);
+}
+static inline __device__ void atomMin(int8_t *address, int8_t val) {
+    AtomicMinIntegerImpl<int8_t, sizeof(int8_t)>()(address, val);
+}
+static inline __device__ void atomMin(int16_t *address, int16_t val) {
+    AtomicMinIntegerImpl<int16_t, sizeof(int16_t)>()(address, val);
+}
+static inline __device__ void atomMin(int32_t *address, int32_t val) { atomicMin(address, val); }
+static inline __device__ void atomMin(int64_t *address, int64_t val) {
+    AtomicMinIntegerImpl<int64_t, sizeof(int64_t)>()(address, val);
+}
+static inline __device__ void atomMin(at::Half *address, at::Half val) {
+    AtomicMinDecimalImpl<at::Half, sizeof(at::Half)>()(address, val);
+}
+static inline __device__ void atomMin(float *address, float val) {
+    AtomicMinDecimalImpl<float, sizeof(float)>()(address, val);
+}
+static inline __device__ void atomMin(double *address, double val) {
+    AtomicMinDecimalImpl<double, sizeof(double)>()(address, val);
+}
+static inline __device__ void atomMin(at::BFloat16 *address, at::BFloat16 val) {
+    AtomicMinDecimalImpl<at::BFloat16, sizeof(at::BFloat16)>()(address, val);
+}

third_party/vipe/csrc/scatter_ext/cuda/reducer.cuh

@@ -0,0 +1,115 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#include <limits>
+#include <map>
+
+#include "atomics.cuh"
+
+enum ReductionType { SUM, MEAN, MUL, DIV, MIN, MAX };
+
+const std::map<std::string, ReductionType> reduce2REDUCE = {
+    {"sum", SUM}, {"mean", MEAN}, {"mul", MUL}, {"div", DIV}, {"min", MIN}, {"max", MAX},
+};
+
+#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...)              \
+    [&] {                                                     \
+        switch (reduce2REDUCE.at(reduce)) {                   \
+            case SUM: {                                       \
+                static constexpr ReductionType REDUCE = SUM;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MEAN: {                                      \
+                static constexpr ReductionType REDUCE = MEAN; \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MUL: {                                       \
+                static constexpr ReductionType REDUCE = MUL;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case DIV: {                                       \
+                static constexpr ReductionType REDUCE = DIV;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MIN: {                                       \
+                static constexpr ReductionType REDUCE = MIN;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+            case MAX: {                                       \
+                static constexpr ReductionType REDUCE = MAX;  \
+                return __VA_ARGS__();                         \
+            }                                                 \
+        }                                                     \
+    }()
+
+template <typename scalar_t, ReductionType REDUCE>
+struct Reducer {
+    static inline __host__ __device__ scalar_t init() {
+        if (REDUCE == MUL || REDUCE == DIV)
+            return (scalar_t)1;
+        else if (REDUCE == MIN)
+            return std::numeric_limits<scalar_t>::max();
+        else if (REDUCE == MAX)
+            return std::numeric_limits<scalar_t>::lowest();
+        else
+            return (scalar_t)0;
+    }
+
+    static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            *val = *val + new_val;
+        else if (REDUCE == MUL)
+            *val = *val * new_val;
+        else if (REDUCE == DIV)
+            *val = *val / new_val;
+        else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
+            *val = new_val;
+        }
+    }
+
+    static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val, int64_t *arg, int64_t new_arg) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            *val = *val + new_val;
+        else if (REDUCE == MUL)
+            *val = *val * new_val;
+        else if (REDUCE == DIV)
+            *val = *val / new_val;
+        else if ((REDUCE == MIN && new_val < *val) || (REDUCE == MAX && new_val > *val)) {
+            *val = new_val;
+            *arg = new_arg;
+        }
+    }
+
+    static inline __host__ __device__ void write(scalar_t *address, scalar_t val, int64_t *arg_address, int64_t arg,
+                                                 int count) {
+        if (REDUCE == SUM || REDUCE == MUL || REDUCE == DIV)
+            *address = val;
+        else if (REDUCE == MEAN)
+            *address = val / (scalar_t)(count > 0 ? count : 1);
+        else if (REDUCE == MIN || REDUCE == MAX) {
+            if (count > 0) {
+                *address = val;
+                *arg_address = arg;
+            } else
+                *address = (scalar_t)0;
+        }
+    }
+
+    static inline __device__ void atomic_write(scalar_t *address, scalar_t val) {
+        if (REDUCE == SUM || REDUCE == MEAN)
+            atomAdd(address, val);
+        else if (REDUCE == MUL)
+            atomMul(address, val);
+        else if (REDUCE == DIV)
+            atomDiv(address, val);
+        else if (REDUCE == MIN)
+            atomMin(address, val);
+        else if (REDUCE == MAX)
+            atomMax(address, val);
+    }
+};

third_party/vipe/csrc/scatter_ext/cuda/scatter_cuda.cu

@@ -0,0 +1,128 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include "scatter_cuda.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/Dispatch.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/detail/TensorInfo.cuh>
+
+#include "reducer.cuh"
+#include "utils.cuh"
+
+#define THREADS 256
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+template <typename scalar_t, ReductionType REDUCE>
+__global__ void scatter_kernel(const scalar_t *src_data, const at::cuda::detail::TensorInfo<int64_t, int> index_info,
+                               scalar_t *out_data, int E, int K, int N, int numel) {
+    int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    int b = thread_idx / (E * K);
+    int k = thread_idx % K;
+
+    if (thread_idx < numel) {
+        int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
+        int64_t idx = index_info.data[offset];
+
+        Reducer<scalar_t, REDUCE>::atomic_write(out_data + b * N * K + idx * K + k, src_data[thread_idx]);
+    }
+}
+
+template <typename scalar_t>
+__global__ void scatter_arg_kernel(const scalar_t *src_data,
+                                   const at::cuda::detail::TensorInfo<int64_t, int> index_info,
+                                   const scalar_t *out_data, int64_t *arg_out_data, int E, int K, int N, int numel) {
+    int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    int b = thread_idx / (E * K);
+    int e = (thread_idx / K) % E;
+    int k = thread_idx % K;
+
+    if (thread_idx < numel) {
+        int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(thread_idx, index_info);
+        int64_t idx = index_info.data[offset];
+
+        if (src_data[thread_idx] == out_data[b * N * K + idx * K + k]) {
+            arg_out_data[b * N * K + idx * K + k] = e;
+        }
+    }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
+                                                                     int64_t dim,
+                                                                     std::optional<torch::Tensor> optional_out,
+                                                                     std::optional<int64_t> dim_size,
+                                                                     std::string reduce) {
+    CHECK_CUDA(src);
+    CHECK_CUDA(index);
+    if (optional_out.has_value()) CHECK_CUDA(optional_out.value());
+    c10::cuda::MaybeSetDevice(src.get_device());
+
+    CHECK_INPUT(src.dim() == index.dim());
+    for (auto i = 0; i < index.dim() - 1; i++) CHECK_INPUT(src.size(i) >= index.size(i));
+
+    src = src.contiguous();
+
+    torch::Tensor out;
+    if (optional_out.has_value()) {
+        out = optional_out.value().contiguous();
+        for (auto i = 0; i < out.dim(); i++)
+            if (i != dim) CHECK_INPUT(src.size(i) == out.size(i));
+    } else {
+        auto sizes = src.sizes().vec();
+        if (dim_size.has_value())
+            sizes[dim] = dim_size.value();
+        else if (index.numel() == 0)
+            sizes[dim] = 0;
+        else {
+            sizes[dim] = 1 + index.max().cpu().data_ptr<int64_t>()[0];
+        }
+        out = torch::empty(sizes, src.options());
+    }
+
+    std::optional<torch::Tensor> arg_out = std::nullopt;
+    int64_t *arg_out_data = nullptr;
+    if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
+        arg_out = torch::full_like(out, src.size(dim), index.options());
+        arg_out_data = arg_out.value().data_ptr<int64_t>();
+    }
+
+    if (src.numel() == 0) {
+        if (!optional_out.has_value()) out.fill_(0);
+        return std::make_tuple(out, arg_out);
+    }
+
+    auto B = 1;
+    for (auto i = 0; i < dim; i++) B *= src.size(i);
+    auto E = src.size(dim);
+    auto K = src.numel() / (B * E);
+    auto N = out.size(dim);
+
+    auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, src.scalar_type(), "_", [&] {
+        auto src_data = src.data_ptr<scalar_t>();
+        auto out_data = out.data_ptr<scalar_t>();
+
+        AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
+            if (!optional_out.has_value()) out.fill_(Reducer<scalar_t, REDUCE>::init());
+
+            scatter_kernel<scalar_t, REDUCE>
+                <<<BLOCKS(src.numel()), THREADS, 0, stream>>>(src_data, index_info, out_data, E, K, N, src.numel());
+
+            if (!optional_out.has_value() && (REDUCE == MIN || REDUCE == MAX))
+                out.masked_fill_(out == Reducer<scalar_t, REDUCE>::init(), (scalar_t)0);
+
+            if (REDUCE == MIN || REDUCE == MAX)
+                scatter_arg_kernel<scalar_t><<<BLOCKS(src.numel()), THREADS, 0, stream>>>(
+                    src_data, index_info, out_data, arg_out_data, E, K, N, src.numel());
+        });
+    });
+
+    return std::make_tuple(out, arg_out);
+}

third_party/vipe/csrc/scatter_ext/cuda/scatter_cuda.h

@@ -0,0 +1,15 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#include <torch/extension.h>
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_cuda(torch::Tensor src, torch::Tensor index,
+                                                                     int64_t dim,
+                                                                     std::optional<torch::Tensor> optional_out,
+                                                                     std::optional<int64_t> dim_size,
+                                                                     std::string reduce);

third_party/vipe/csrc/scatter_ext/cuda/utils.cuh

@@ -0,0 +1,29 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#pragma once
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
+#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
+
+__device__ __inline__ at::Half __shfl_up_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
+    return __shfl_up_sync(mask, var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_down_sync(const unsigned mask, const at::Half var, const unsigned int delta) {
+    return __shfl_down_sync(mask, var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_up(const at::Half var, const unsigned int delta) {
+    return __shfl_up(var.operator __half(), delta);
+}
+
+__device__ __inline__ at::Half __shfl_down(const at::Half var, const unsigned int delta) {
+    return __shfl_down(var.operator __half(), delta);
+}
+
+#define SHFL_UP_SYNC __shfl_up_sync
+#define SHFL_DOWN_SYNC __shfl_down_sync

third_party/vipe/csrc/scatter_ext/scatter.cpp

@@ -0,0 +1,238 @@
+/**
+ * This file includes code originally from the Pytorch Scatter repository:
+ * https://github.com/rusty1s/pytorch_scatter
+ * Licensed under the MIT License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+#include "cuda/scatter_cuda.h"
+
+inline std::vector<int64_t> list2vec(const c10::List<int64_t> list) {
+    std::vector<int64_t> result;
+    result.reserve(list.size());
+    for (size_t i = 0; i < list.size(); i++) result.push_back(list[i]);
+    return result;
+}
+
+torch::Tensor broadcast(torch::Tensor src, torch::Tensor other, int64_t dim) {
+    if (src.dim() == 1)
+        for (auto i = 0; i < dim; i++) src = src.unsqueeze(0);
+    for (auto i = src.dim(); i < other.dim(); i++) src = src.unsqueeze(-1);
+    src = src.expand(other.sizes().vec());
+    return src;
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> scatter_fw(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                                   std::optional<torch::Tensor> optional_out,
+                                                                   std::optional<int64_t> dim_size,
+                                                                   std::string reduce) {
+    return scatter_cuda(src, index, dim, optional_out, dim_size, reduce);
+}
+
+using torch::autograd::AutogradContext;
+using torch::autograd::Variable;
+using torch::autograd::variable_list;
+
+class ScatterSum : public torch::autograd::Function<ScatterSum> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
+        auto out = std::get<0>(result);
+        ctx->save_for_backward({index});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        auto grad_in = torch::gather(grad_out, dim, index, false);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMul : public torch::autograd::Function<ScatterMul> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "mul");
+        auto out = std::get<0>(result);
+        ctx->save_for_backward({src, index, out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto src = saved[0];
+        auto index = saved[1];
+        auto out = saved[2];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        auto grad_in = torch::gather(grad_out * out, dim, index, false).div_(src);
+        grad_in.masked_fill_(grad_in.isnan(), 0);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMean : public torch::autograd::Function<ScatterMean> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        auto old_index = index;
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "sum");
+        auto out = std::get<0>(result);
+
+        auto ones = torch::ones(old_index.sizes(), src.options());
+        result = scatter_fw(ones, old_index, old_index.dim() <= dim ? old_index.dim() - 1 : dim, std::nullopt,
+                            out.size(dim), "sum");
+        auto count = std::get<0>(result);
+        count.masked_fill_(count < 1, 1);
+        count = broadcast(count, out, dim);
+        if (out.is_floating_point())
+            out.true_divide_(count);
+        else
+            out.div_(count, "floor");
+
+        ctx->save_for_backward({index, count});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto count = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        count = torch::gather(count, dim, index, false);
+        auto grad_in = torch::gather(grad_out, dim, index, false);
+        grad_in.true_divide_(count);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMin : public torch::autograd::Function<ScatterMin> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "min");
+        auto out = std::get<0>(result);
+        auto arg_out = std::get<1>(result).value();
+        ctx->save_for_backward({index, arg_out});
+        ctx->mark_non_differentiable({arg_out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out, arg_out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto arg_out = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        src_shape[dim] += 1;
+        auto grad_in = torch::zeros(src_shape, grad_out.options());
+        grad_in.scatter_(dim, arg_out, grad_out);
+        grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+class ScatterMax : public torch::autograd::Function<ScatterMax> {
+   public:
+    static variable_list forward(AutogradContext *ctx, Variable src, Variable index, int64_t dim,
+                                 std::optional<Variable> optional_out, std::optional<int64_t> dim_size) {
+        dim = dim < 0 ? src.dim() + dim : dim;
+        ctx->saved_data["dim"] = dim;
+        ctx->saved_data["src_shape"] = src.sizes();
+
+        index = broadcast(index, src, dim);
+        auto result = scatter_fw(src, index, dim, optional_out, dim_size, "max");
+        auto out = std::get<0>(result);
+        auto arg_out = std::get<1>(result).value();
+        ctx->save_for_backward({index, arg_out});
+        ctx->mark_non_differentiable({arg_out});
+        if (optional_out.has_value()) ctx->mark_dirty({optional_out.value()});
+        return {out, arg_out};
+    }
+
+    static variable_list backward(AutogradContext *ctx, variable_list grad_outs) {
+        auto grad_out = grad_outs[0];
+        auto saved = ctx->get_saved_variables();
+        auto index = saved[0];
+        auto arg_out = saved[1];
+        auto dim = ctx->saved_data["dim"].toInt();
+        auto src_shape = list2vec(ctx->saved_data["src_shape"].toIntList());
+        src_shape[dim] += 1;
+        auto grad_in = torch::zeros(src_shape, grad_out.options());
+        grad_in.scatter_(dim, arg_out, grad_out);
+        grad_in = grad_in.narrow(dim, 0, src_shape[dim] - 1);
+        return {grad_in, Variable(), Variable(), Variable(), Variable()};
+    }
+};
+
+torch::Tensor scatter_sum(torch::Tensor src, torch::Tensor index, int64_t dim,
+                          std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterSum::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+torch::Tensor scatter_mul(torch::Tensor src, torch::Tensor index, int64_t dim,
+                          std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterMul::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+torch::Tensor scatter_mean(torch::Tensor src, torch::Tensor index, int64_t dim,
+                           std::optional<torch::Tensor> optional_out, std::optional<int64_t> dim_size) {
+    return ScatterMean::apply(src, index, dim, optional_out, dim_size)[0];
+}
+
+std::tuple<torch::Tensor, torch::Tensor> scatter_min(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                     std::optional<torch::Tensor> optional_out,
+                                                     std::optional<int64_t> dim_size) {
+    auto result = ScatterMin::apply(src, index, dim, optional_out, dim_size);
+    return std::make_tuple(result[0], result[1]);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> scatter_max(torch::Tensor src, torch::Tensor index, int64_t dim,
+                                                     std::optional<torch::Tensor> optional_out,
+                                                     std::optional<int64_t> dim_size) {
+    auto result = ScatterMax::apply(src, index, dim, optional_out, dim_size);
+    return std::make_tuple(result[0], result[1]);
+}
+
+void pybind_scatter_ext(py::module &m) {
+    m.def("scatter_sum", &scatter_sum, "Scatter sum operation");
+    m.def("scatter_mul", &scatter_mul, "Scatter multiplication operation");
+    m.def("scatter_mean", &scatter_mean, "Scatter mean operation");
+    m.def("scatter_min", &scatter_min, "Scatter min operation");
+    m.def("scatter_max", &scatter_max, "Scatter max operation");
+}

third_party/vipe/csrc/slam_ext/slam.cpp

@@ -0,0 +1,37 @@
+/**
+ * This file includes code originally from the DROID-SLAM repository:
+ * https://github.com/princeton-vl/DROID-SLAM
+ * Licensed under the BSD-3 License. See THIRD_PARTY_LICENSES.md for details.
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+namespace slam_ext {
+
+torch::Tensor depth_filter_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor ix,
+                                torch::Tensor thresh);
+
+torch::Tensor frame_distance_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics, torch::Tensor pi,
+                                  torch::Tensor pj, torch::Tensor qi, torch::Tensor qj, torch::Tensor di,
+                                  const float beta);
+
+std::vector<torch::Tensor> projmap_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
+                                        torch::Tensor ii, torch::Tensor jj);
+
+torch::Tensor iproj_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics);
+
+std::vector<torch::Tensor> ba_cuda(torch::Tensor poses, torch::Tensor disps, torch::Tensor intrinsics,
+                                   torch::Tensor disps_sens, torch::Tensor targets, torch::Tensor weights,
+                                   torch::Tensor eta, torch::Tensor ii, torch::Tensor jj, const int t0, const int t1,
+                                   const int iterations, const float lm, const float ep, const bool motion_only);
+
+}  // namespace slam_ext
+
+void pybind_slam_ext(py::module &m) {
+    m.def("ba", &slam_ext::ba_cuda, "bundle adjustment");
+    m.def("frame_distance", &slam_ext::frame_distance_cuda, "frame_distance");
+    m.def("projmap", &slam_ext::projmap_cuda, "projmap");
+    m.def("depth_filter", &slam_ext::depth_filter_cuda, "depth_filter");
+    m.def("iproj", &slam_ext::iproj_cuda, "back projection");
+}

third_party/vipe/csrc/utils_ext/cuda_kdtree.cuh

@@ -0,0 +1,137 @@
+/***********************************************************************
+ * The code in this file is modified from the original code in the
+ * https://github.com/jlblancoc/nanoflann repository.
+ *
+ * Software License Agreement (BSD License)
+ *
+ * Copyright 2008-2009  Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
+ * Copyright 2008-2009  David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
+ * Copyright 2011-2025  Jose Luis Blanco (joseluisblancoc@gmail.com).
+ *   All rights reserved.
+ *
+ * THE BSD LICENSE
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions
+ * are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
+ * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+ * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
+ * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
+ * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+ * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *************************************************************************/
+
+#ifndef FLANN_CUDA_KD_TREE_BUILDER_H_
+#define FLANN_CUDA_KD_TREE_BUILDER_H_
+
+#include <thrust/host_vector.h>
+#include <thrust/device_vector.h>
+#include <thrust/sort.h>
+#include <thrust/partition.h>
+#include <thrust/unique.h>
+#include <thrust/scan.h>
+#include "math_util.h"
+#include <cstdlib>
+#include <map>
+
+namespace tinyflann {
+
+// Distance types.
+struct CudaL1;
+struct CudaL2;
+
+// Parameters
+struct KDTreeCuda3dIndexParams {
+    int leaf_max_size = 64;
+};
+
+struct SearchParams {
+    SearchParams(int checks_ = 32, float eps_ = 0.0, bool sorted_ = true)
+        : checks(checks_), eps(eps_), sorted(sorted_) {
+        max_neighbors = -1;
+        use_heap = true;
+    }
+
+    // how many leafs to visit when searching for neighbours (-1 for unlimited)
+    int checks;
+    // search for eps-approximate neighbours (default: 0)
+    float eps;
+    // only for radius search, require neighbours sorted by distance (default: true)
+    bool sorted;
+    // maximum number of neighbors radius search should return (-1 for unlimited)
+    int max_neighbors;
+    // use a heap to manage the result set (default: FLANN_Undefined)
+    bool use_heap;
+};
+
+template <typename Distance>
+class KDTreeCuda3dIndex {
+   public:
+    int visited_leafs;
+    KDTreeCuda3dIndex(const float* input_data, size_t input_count,
+                      const KDTreeCuda3dIndexParams& params = KDTreeCuda3dIndexParams())
+        : dataset_(input_data), leaf_count_(0), visited_leafs(0), node_count_(0), current_node_count_(0) {
+        size_ = input_count;
+        leaf_max_size_ = params.leaf_max_size;
+        gpu_helper_ = 0;
+    }
+
+    /**
+     * Standard destructor
+     */
+    ~KDTreeCuda3dIndex() { clearGpuBuffers(); }
+
+    /**
+     * Builds the index
+     */
+    void buildIndex() {
+        leaf_count_ = 0;
+        node_count_ = 0;
+        uploadTreeToGpu();
+    }
+
+    /**
+     * queries: (N, p) flattened float cuda array where only first 3 elements are used.
+     * n_query: N
+     * n_query_stride: p
+     * indices: (N, knn) int cuda array
+     * dists: (N, knn) float cuda array
+     */
+    void knnSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, size_t knn,
+                   const SearchParams& params = SearchParams()) const;
+    int radiusSearch(const float* queries, size_t n_query, int n_query_stride, int* indices, float* dists, float radius,
+                     const SearchParams& params = SearchParams()) const;
+
+   private:
+    void uploadTreeToGpu();
+    void clearGpuBuffers();
+
+   private:
+    struct GpuHelper;
+    GpuHelper* gpu_helper_;
+
+    const float* dataset_;
+    int leaf_max_size_;
+    int leaf_count_;
+    int node_count_;
+    //! used by convertTreeToGpuFormat
+    int current_node_count_;
+    size_t size_;
+
+};  // class KDTreeCuda3dIndex
+
+}  // namespace tinyflann
+
+#endif

third_party/vipe/csrc/utils_ext/knn.cu

@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAException.h>
+
+#include "cuda_kdtree.cuh"
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_IS_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor")
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn) {
+    CHECK_CUDA(query);
+    CHECK_IS_FLOAT(query);
+    CHECK_CUDA(tree);
+    CHECK_IS_FLOAT(tree);
+    TORCH_CHECK(tree.size(0) >= knn, "knn is too small compared to the size of point cloud!");
+
+    torch::Tensor strided_tree = tree;
+    torch::Tensor strided_query = query;
+    torch::Device device = tree.device();
+    long n_queries = query.size(0);
+    long n_tree = tree.size(0);
+
+    if (strided_tree.stride(0) != 4) {
+        strided_tree = torch::zeros({n_tree, 4}, torch::dtype(torch::kFloat32).device(device));
+        strided_tree.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, tree.size(1))}, tree);
+    }
+
+    if (strided_query.stride(0) != 3) {
+        strided_query = torch::zeros({n_queries, 3}, torch::dtype(torch::kFloat32).device(device));
+        strided_query.index_put_(
+            {torch::indexing::Ellipsis, torch::indexing::Slice(torch::indexing::None, query.size(1))}, query);
+    }
+
+    auto *knn_index = new tinyflann::KDTreeCuda3dIndex<tinyflann::CudaL2>(strided_tree.data_ptr<float>(), n_tree);
+    knn_index->buildIndex();
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    torch::Tensor dist = torch::empty(n_queries * knn, torch::dtype(torch::kFloat32).device(device));
+    torch::Tensor indices = torch::empty(n_queries * knn, torch::dtype(torch::kInt32).device(device));
+
+    knn_index->knnSearch(strided_query.data_ptr<float>(), n_queries, 3, indices.data_ptr<int>(), dist.data_ptr<float>(),
+                         knn);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    delete knn_index;
+
+    return {dist.reshape({n_queries, knn}), indices.reshape({n_queries, knn})};
+}

third_party/vipe/csrc/utils_ext/platform.h

@@ -0,0 +1,50 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CODE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_
+#endif
+
+#if defined(__CUDACC__)
+#define _CPU_AND_GPU_CODE_TEMPLATE_ __host__ __device__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CODE_TEMPLATE_
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#define _CPU_AND_GPU_CONSTANT_ __constant__  // for CUDA device code
+#else
+#define _CPU_AND_GPU_CONSTANT_ const
+#endif
+
+#if defined(__CUDACC__) && defined(__CUDA_ARCH__)
+#include <THC/THCAtomics.cuh>
+
+template <typename ScalarT>
+static inline __device__ void atomicAdd(ScalarT* addr, ScalarT value) {
+    gpuAtomicAddNoReturn(addr, value);
+}
+#else
+template <typename ScalarT>
+static inline void atomicAdd(ScalarT* addr, ScalarT value) {
+    *addr += value;
+}
+#endif

third_party/vipe/csrc/utils_ext/utils_bind.cpp

@@ -0,0 +1,25 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> nearestNeighbours(torch::Tensor query, torch::Tensor tree, int knn);
+
+void pybind_utils_ext(py::module &m) {
+    m.def("nearest_neighbours", &nearestNeighbours, "KNN computation (CUDA) distance is squared L2.", py::arg("query"),
+          py::arg("tree"), py::arg("knn"));
+}