fix: Compile the vis4d kernels here instead of installing them.

pip cannot install vis4d_cuda_ops on a Space. It publishes no wheels, and pip
compiles it in an isolated build environment that resolves its own newest torch
(2.13.0) while the Space caps the runtime torch at 2.11.0, so the extension
imports with `undefined symbol: c10::NotImplementedError`. Setting
PIP_NO_BUILD_ISOLATION as a Space variable does not help, the build does not
see it.

A local vis4d_cuda_ops package now shadows the real one and compiles only the
upstream CPU translation units against the torch the Space runs, which cannot
drift. iou_box3d runs on CPU for the few boxes the track graph compares, and
deformable attention runs the grid-sample path already in mapdet3d, on the GPU.
Verified against the CUDA kernels on a 4090: same 19 tracks over the example
scene, same 5.5 s for 21 frames, and ~10 s to compile on first import.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

README.md

@@ -42,11 +42,14 @@ python app.py
 - `sdk_version` is pinned to a Gradio 5 release because `gradio_rerun==0.26.0`
   is a Gradio 5 custom component. It has to stay on the same version as the
   `rerun-sdk` the model code logs with, so bumping one means bumping all three.
-- `vis4d_cuda_ops` is compiled from source during the build; it provides the
-  multi-scale deformable attention and 3D IoU kernels. The Space needs the
-  variable **`PIP_NO_BUILD_ISOLATION=1`** set in its settings, otherwise pip
-  compiles it against a newer torch than the one the Space runs and the
-  extension fails to import. See the comment at the top of `requirements.txt`.
+- `vis4d_cuda_ops/` shadows the package of the same name, because pip cannot
+  install the real one here: it has no wheels, and pip compiles it in an
+  isolated build environment that resolves its own newest torch while the Space
+  caps the runtime torch lower, so the extension fails to import with
+  `undefined symbol: c10::NotImplementedError`. The local package compiles only
+  the upstream CPU sources, against the torch the Space runs, and routes
+  deformable attention through the grid-sample implementation already in
+  `mapdet3d`. Detections and tracks come out identical; see its docstring.
 - `mapdet3d/` is vendored from the research repository so the Space does not
   depend on a published package.
 - The model weights are pulled from

requirements.txt

@@ -1,13 +1,11 @@
-# NOTE: torch stays unpinned on purpose, and the Space sets the build variable
-# PIP_NO_BUILD_ISOLATION=1. vis4d_cuda_ops declares torch as a build
-# requirement, so with build isolation pip resolves a second, newest torch
-# (2.13.0) to compile the extension against, while the Space caps the runtime
-# torch at 2.11.0. The two ABIs do not match and importing the extension then
-# dies on `undefined symbol: c10::NotImplementedError`. Without isolation the
-# extension compiles against the torch already in the image, and leaving the
-# bound loose here keeps pip from swapping that torch out underneath it.
+# NOTE: torch stays unpinned so it keeps whatever the Space image ships. The
+# kernels in vis4d_cuda_ops/ are compiled against it on first import, so the two
+# cannot drift apart. Do not add vis4d_cuda_ops here: pip builds it in an
+# isolated environment that resolves its own, newest torch, and the extension
+# then fails to import against the older torch the Space runs.
 torch>=2.8.0
 torchvision
+ninja
 
 # Map-Det3D
 numpy
@@ -27,9 +25,6 @@ tqdm
 huggingface_hub
 safetensors
 
-# CUDA ops for multi-scale deformable attention and 3D IoU
-git+https://github.com/SysCV/vis4d_cuda_ops.git
-
 # Demo
 spaces
 rerun-sdk==0.26.0

vis4d_cuda_ops/__init__.py

@@ -0,0 +1,110 @@
+"""CPU build of vis4d_cuda_ops, standing in for the published package.
+
+Map-Det3D imports ``iou_box3d`` and ``ms_deform_attn_forward`` from
+``vis4d_cuda_ops``. Installing that package on a Space does not work: it has no
+wheels, and pip compiles it inside an isolated build environment that resolves
+its own, newest torch, while the Space caps the runtime torch several releases
+lower. The two ABIs disagree and importing the result dies on
+``undefined symbol: c10::NotImplementedError``.
+
+Since ``/home/user/app`` comes first on ``sys.path``, this package answers those
+imports instead. It compiles only the CPU translation units of the upstream
+sources, against the torch the Space actually runs, so no ABI can drift:
+
+* ``iou_box3d`` runs the upstream CPU kernel, on the handful of boxes the track
+  graph compares per frame.
+* ``ms_deform_attn_forward`` runs the grid-sample implementation that ships in
+  ``mapdet3d.op.layer.ms_deform_attn`` and is used whenever the fused kernel is
+  unavailable. It is the same computation and still runs on the GPU, only
+  without the fused kernel.
+
+The upstream sources live under ``src/`` unmodified; see ``LICENSE``.
+"""
+
+from __future__ import annotations
+
+import os
+
+from torch import Tensor
+from torch.utils.cpp_extension import load
+
+_SRC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "src")
+
+_CPU_SOURCES = [
+    os.path.join(_SRC, "vision.cpp"),
+    os.path.join(_SRC, "iou_box3d", "iou_box3d_cpu.cpp"),
+    os.path.join(_SRC, "box_iou_rotated", "box_iou_rotated_cpu.cpp"),
+    os.path.join(_SRC, "nms_rotated", "nms_rotated_cpu.cpp"),
+]
+
+# Compiled on first import and cached in ~/.cache/torch_extensions afterwards.
+_ops = load(
+    name="vis4d_cpu_ops",
+    sources=_CPU_SOURCES,
+    extra_include_paths=[_SRC],
+    extra_cflags=["-O2"],
+    with_cuda=False,
+    verbose=False,
+)
+
+box_iou_rotated = _ops.box_iou_rotated
+nms_rotated = _ops.nms_rotated
+deform_conv_forward = _ops.deform_conv_forward
+deform_conv_backward_input = _ops.deform_conv_backward_input
+deform_conv_backward_filter = _ops.deform_conv_backward_filter
+modulated_deform_conv_forward = _ops.modulated_deform_conv_forward
+modulated_deform_conv_backward = _ops.modulated_deform_conv_backward
+
+
+def iou_box3d(
+    boxes1: Tensor, boxes2: Tensor
+) -> tuple[Tensor, Tensor]:  # pragma: no cover
+    """Intersection volume and IoU of two sets of 3D box corners."""
+    device = boxes1.device
+    volume, iou = _ops.iou_box3d(boxes1.cpu(), boxes2.cpu())
+
+    return volume.to(device), iou.to(device)
+
+
+def ms_deform_attn_forward(
+    value: Tensor,
+    spatial_shapes: Tensor,
+    level_start_index: Tensor,
+    sampling_loc: Tensor,
+    attn_weight: Tensor,
+    im2col_step: int,
+) -> Tensor:  # pragma: no cover
+    """Multi-scale deformable attention, via grid sampling.
+
+    ``level_start_index`` and ``im2col_step`` only exist to drive the fused
+    kernel; the grid-sample path splits ``value`` by ``spatial_shapes`` itself.
+    """
+    # Imported here because mapdet3d imports this module while it is still
+    # setting up its own.
+    from mapdet3d.op.layer.ms_deform_attn import ms_deformable_attention_cpu
+
+    return ms_deformable_attention_cpu(
+        value, spatial_shapes, sampling_loc, attn_weight
+    )
+
+
+def ms_deform_attn_backward(*args: object, **kwargs: object) -> None:
+    """Not built: the demo only ever runs under ``torch.no_grad()``."""
+    raise NotImplementedError(
+        "vis4d_cuda_ops was built without the deformable attention backward "
+        "kernel. Install the upstream package to train."
+    )
+
+
+__all__ = [
+    "box_iou_rotated",
+    "deform_conv_backward_filter",
+    "deform_conv_backward_input",
+    "deform_conv_forward",
+    "iou_box3d",
+    "modulated_deform_conv_backward",
+    "modulated_deform_conv_forward",
+    "ms_deform_attn_backward",
+    "ms_deform_attn_forward",
+    "nms_rotated",
+]

vis4d_cuda_ops/src/box_iou_rotated/box_iou_rotated.h

@@ -0,0 +1,33 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/csrc/box_iou_rotated/box_iou_rotated.h
+#pragma once
+#include <torch/types.h>
+
+at::Tensor box_iou_rotated_cpu(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2);
+
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+at::Tensor box_iou_rotated_cuda(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2);
+#endif
+
+// Interface for Python
+// inline is needed to prevent multiple function definitions when this header is
+// included by different cpps
+inline at::Tensor box_iou_rotated(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  assert(boxes1.device().is_cuda() == boxes2.device().is_cuda());
+  if (boxes1.device().is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    return box_iou_rotated_cuda(boxes1.contiguous(), boxes2.contiguous());
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+
+  return box_iou_rotated_cpu(boxes1.contiguous(), boxes2.contiguous());
+}

vis4d_cuda_ops/src/box_iou_rotated/box_iou_rotated_cpu.cpp

@@ -0,0 +1,37 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/csrc/box_iou_rotated/box_iou_rotated_cpu.cpp
+#include "box_iou_rotated.h"
+#include "box_iou_rotated_utils.h"
+
+template <typename T>
+void box_iou_rotated_cpu_kernel(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2,
+    at::Tensor& ious) {
+  auto num_boxes1 = boxes1.size(0);
+  auto num_boxes2 = boxes2.size(0);
+
+  for (int i = 0; i < num_boxes1; i++) {
+    for (int j = 0; j < num_boxes2; j++) {
+      ious[i * num_boxes2 + j] = single_box_iou_rotated<T>(
+          boxes1[i].data_ptr<T>(), boxes2[j].data_ptr<T>());
+    }
+  }
+}
+
+at::Tensor box_iou_rotated_cpu(
+    // input must be contiguous:
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  auto num_boxes1 = boxes1.size(0);
+  auto num_boxes2 = boxes2.size(0);
+  at::Tensor ious =
+      at::empty({num_boxes1 * num_boxes2}, boxes1.options().dtype(at::kFloat));
+
+  box_iou_rotated_cpu_kernel<float>(boxes1, boxes2, ious);
+
+  // reshape from 1d array to 2d array
+  auto shape = std::vector<int64_t>{num_boxes1, num_boxes2};
+  return ious.reshape(shape);
+}

vis4d_cuda_ops/src/box_iou_rotated/box_iou_rotated_utils.h

@@ -0,0 +1,368 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/csrc/box_iou_rotated/box_iou_rotated_utils.h
+#pragma once
+
+#include <cassert>
+#include <cmath>
+
+#if defined(__CUDACC__) || __HCC__ == 1 || __HIP__ == 1
+// Designates functions callable from the host (CPU) and the device (GPU)
+#define HOST_DEVICE __host__ __device__
+#define HOST_DEVICE_INLINE HOST_DEVICE __forceinline__
+#else
+#include <algorithm>
+#define HOST_DEVICE
+#define HOST_DEVICE_INLINE HOST_DEVICE inline
+#endif
+
+namespace {
+
+template <typename T>
+struct RotatedBox {
+  T x_ctr, y_ctr, w, h, a;
+};
+
+template <typename T>
+struct Point {
+  T x, y;
+  HOST_DEVICE_INLINE Point(const T& px = 0, const T& py = 0) : x(px), y(py) {}
+  HOST_DEVICE_INLINE Point operator+(const Point& p) const {
+    return Point(x + p.x, y + p.y);
+  }
+  HOST_DEVICE_INLINE Point& operator+=(const Point& p) {
+    x += p.x;
+    y += p.y;
+    return *this;
+  }
+  HOST_DEVICE_INLINE Point operator-(const Point& p) const {
+    return Point(x - p.x, y - p.y);
+  }
+  HOST_DEVICE_INLINE Point operator*(const T coeff) const {
+    return Point(x * coeff, y * coeff);
+  }
+};
+
+template <typename T>
+HOST_DEVICE_INLINE T dot_2d(const Point<T>& A, const Point<T>& B) {
+  return A.x * B.x + A.y * B.y;
+}
+
+// R: result type. can be different from input type
+template <typename T, typename R = T>
+HOST_DEVICE_INLINE R cross_2d(const Point<T>& A, const Point<T>& B) {
+  return static_cast<R>(A.x) * static_cast<R>(B.y) -
+      static_cast<R>(B.x) * static_cast<R>(A.y);
+}
+
+template <typename T>
+HOST_DEVICE_INLINE void get_rotated_vertices(
+    const RotatedBox<T>& box,
+    Point<T> (&pts)[4]) {
+  // M_PI / 180. == 0.01745329251
+  double theta = box.a * 0.01745329251;
+  T cosTheta2 = (T)cos(theta) * 0.5f;
+  T sinTheta2 = (T)sin(theta) * 0.5f;
+
+  // y: top --> down; x: left --> right
+  pts[0].x = box.x_ctr + sinTheta2 * box.h + cosTheta2 * box.w;
+  pts[0].y = box.y_ctr + cosTheta2 * box.h - sinTheta2 * box.w;
+  pts[1].x = box.x_ctr - sinTheta2 * box.h + cosTheta2 * box.w;
+  pts[1].y = box.y_ctr - cosTheta2 * box.h - sinTheta2 * box.w;
+  pts[2].x = 2 * box.x_ctr - pts[0].x;
+  pts[2].y = 2 * box.y_ctr - pts[0].y;
+  pts[3].x = 2 * box.x_ctr - pts[1].x;
+  pts[3].y = 2 * box.y_ctr - pts[1].y;
+}
+
+template <typename T>
+HOST_DEVICE_INLINE int get_intersection_points(
+    const Point<T> (&pts1)[4],
+    const Point<T> (&pts2)[4],
+    Point<T> (&intersections)[24]) {
+  // Line vector
+  // A line from p1 to p2 is: p1 + (p2-p1)*t, t=[0,1]
+  Point<T> vec1[4], vec2[4];
+  for (int i = 0; i < 4; i++) {
+    vec1[i] = pts1[(i + 1) % 4] - pts1[i];
+    vec2[i] = pts2[(i + 1) % 4] - pts2[i];
+  }
+
+  // When computing the intersection area, it doesn't hurt if we have
+  // more (duplicated/approximate) intersections/vertices than needed,
+  // while it can cause drastic difference if we miss an intersection/vertex.
+  // Therefore, we add an epsilon to relax the comparisons between
+  // the float point numbers that decide the intersection points.
+  double EPS = 1e-5;
+
+  // Line test - test all line combos for intersection
+  int num = 0; // number of intersections
+  for (int i = 0; i < 4; i++) {
+    for (int j = 0; j < 4; j++) {
+      // Solve for 2x2 Ax=b
+      T det = cross_2d<T>(vec2[j], vec1[i]);
+
+      // This takes care of parallel lines
+      if (fabs(det) <= 1e-14) {
+        continue;
+      }
+
+      auto vec12 = pts2[j] - pts1[i];
+
+      T t1 = cross_2d<T>(vec2[j], vec12) / det;
+      T t2 = cross_2d<T>(vec1[i], vec12) / det;
+
+      if (t1 > -EPS && t1 < 1.0f + EPS && t2 > -EPS && t2 < 1.0f + EPS) {
+        intersections[num++] = pts1[i] + vec1[i] * t1;
+      }
+    }
+  }
+
+  // Check for vertices of rect1 inside rect2
+  {
+    const auto& AB = vec2[0];
+    const auto& DA = vec2[3];
+    auto ABdotAB = dot_2d<T>(AB, AB);
+    auto ADdotAD = dot_2d<T>(DA, DA);
+    for (int i = 0; i < 4; i++) {
+      // assume ABCD is the rectangle, and P is the point to be judged
+      // P is inside ABCD iff. P's projection on AB lies within AB
+      // and P's projection on AD lies within AD
+
+      auto AP = pts1[i] - pts2[0];
+
+      auto APdotAB = dot_2d<T>(AP, AB);
+      auto APdotAD = -dot_2d<T>(AP, DA);
+
+      if ((APdotAB > -EPS) && (APdotAD > -EPS) && (APdotAB < ABdotAB + EPS) &&
+          (APdotAD < ADdotAD + EPS)) {
+        intersections[num++] = pts1[i];
+      }
+    }
+  }
+
+  // Reverse the check - check for vertices of rect2 inside rect1
+  {
+    const auto& AB = vec1[0];
+    const auto& DA = vec1[3];
+    auto ABdotAB = dot_2d<T>(AB, AB);
+    auto ADdotAD = dot_2d<T>(DA, DA);
+    for (int i = 0; i < 4; i++) {
+      auto AP = pts2[i] - pts1[0];
+
+      auto APdotAB = dot_2d<T>(AP, AB);
+      auto APdotAD = -dot_2d<T>(AP, DA);
+
+      if ((APdotAB > -EPS) && (APdotAD > -EPS) && (APdotAB < ABdotAB + EPS) &&
+          (APdotAD < ADdotAD + EPS)) {
+        intersections[num++] = pts2[i];
+      }
+    }
+  }
+
+  return num;
+}
+
+template <typename T>
+HOST_DEVICE_INLINE int convex_hull_graham(
+    const Point<T> (&p)[24],
+    const int& num_in,
+    Point<T> (&q)[24],
+    bool shift_to_zero = false) {
+  assert(num_in >= 2);
+
+  // Step 1:
+  // Find point with minimum y
+  // if more than 1 points have the same minimum y,
+  // pick the one with the minimum x.
+  int t = 0;
+  for (int i = 1; i < num_in; i++) {
+    if (p[i].y < p[t].y || (p[i].y == p[t].y && p[i].x < p[t].x)) {
+      t = i;
+    }
+  }
+  auto& start = p[t]; // starting point
+
+  // Step 2:
+  // Subtract starting point from every points (for sorting in the next step)
+  for (int i = 0; i < num_in; i++) {
+    q[i] = p[i] - start;
+  }
+
+  // Swap the starting point to position 0
+  auto tmp = q[0];
+  q[0] = q[t];
+  q[t] = tmp;
+
+  // Step 3:
+  // Sort point 1 ~ num_in according to their relative cross-product values
+  // (essentially sorting according to angles)
+  // If the angles are the same, sort according to their distance to origin
+  T dist[24];
+#if defined(__CUDACC__) || __HCC__ == 1 || __HIP__ == 1
+  // compute distance to origin before sort, and sort them together with the
+  // points
+  for (int i = 0; i < num_in; i++) {
+    dist[i] = dot_2d<T>(q[i], q[i]);
+  }
+
+  // CUDA version
+  // In the future, we can potentially use thrust
+  // for sorting here to improve speed (though not guaranteed)
+  for (int i = 1; i < num_in - 1; i++) {
+    for (int j = i + 1; j < num_in; j++) {
+      T crossProduct = cross_2d<T>(q[i], q[j]);
+      if ((crossProduct < -1e-6) ||
+          (fabs(crossProduct) < 1e-6 && dist[i] > dist[j])) {
+        auto q_tmp = q[i];
+        q[i] = q[j];
+        q[j] = q_tmp;
+        auto dist_tmp = dist[i];
+        dist[i] = dist[j];
+        dist[j] = dist_tmp;
+      }
+    }
+  }
+#else
+  // CPU version
+  std::sort(
+      q + 1, q + num_in, [](const Point<T>& A, const Point<T>& B) -> bool {
+        T temp = cross_2d<T>(A, B);
+        if (fabs(temp) < 1e-6) {
+          return dot_2d<T>(A, A) < dot_2d<T>(B, B);
+        } else {
+          return temp > 0;
+        }
+      });
+  // compute distance to origin after sort, since the points are now different.
+  for (int i = 0; i < num_in; i++) {
+    dist[i] = dot_2d<T>(q[i], q[i]);
+  }
+#endif
+
+  // Step 4:
+  // Make sure there are at least 2 points (that don't overlap with each other)
+  // in the stack
+  int k; // index of the non-overlapped second point
+  for (k = 1; k < num_in; k++) {
+    if (dist[k] > 1e-8) {
+      break;
+    }
+  }
+  if (k == num_in) {
+    // We reach the end, which means the convex hull is just one point
+    q[0] = p[t];
+    return 1;
+  }
+  q[1] = q[k];
+  int m = 2; // 2 points in the stack
+  // Step 5:
+  // Finally we can start the scanning process.
+  // When a non-convex relationship between the 3 points is found
+  // (either concave shape or duplicated points),
+  // we pop the previous point from the stack
+  // until the 3-point relationship is convex again, or
+  // until the stack only contains two points
+  for (int i = k + 1; i < num_in; i++) {
+    while (m > 1) {
+      auto q1 = q[i] - q[m - 2], q2 = q[m - 1] - q[m - 2];
+      // cross_2d() uses FMA and therefore computes round(round(q1.x*q2.y) -
+      // q2.x*q1.y) So it may not return 0 even when q1==q2. Therefore we
+      // compare round(q1.x*q2.y) and round(q2.x*q1.y) directly. (round means
+      // round to nearest floating point).
+      if (q1.x * q2.y >= q2.x * q1.y)
+        m--;
+      else
+        break;
+    }
+    // Using double also helps, but float can solve the issue for now.
+    // while (m > 1 && cross_2d<T, double>(q[i] - q[m - 2], q[m - 1] - q[m - 2])
+    // >= 0) {
+    //     m--;
+    // }
+    q[m++] = q[i];
+  }
+
+  // Step 6 (Optional):
+  // In general sense we need the original coordinates, so we
+  // need to shift the points back (reverting Step 2)
+  // But if we're only interested in getting the area/perimeter of the shape
+  // We can simply return.
+  if (!shift_to_zero) {
+    for (int i = 0; i < m; i++) {
+      q[i] += start;
+    }
+  }
+
+  return m;
+}
+
+template <typename T>
+HOST_DEVICE_INLINE T polygon_area(const Point<T> (&q)[24], const int& m) {
+  if (m <= 2) {
+    return 0;
+  }
+
+  T area = 0;
+  for (int i = 1; i < m - 1; i++) {
+    area += fabs(cross_2d<T>(q[i] - q[0], q[i + 1] - q[0]));
+  }
+
+  return area / 2.0;
+}
+
+template <typename T>
+HOST_DEVICE_INLINE T rotated_boxes_intersection(
+    const RotatedBox<T>& box1,
+    const RotatedBox<T>& box2) {
+  // There are up to 4 x 4 + 4 + 4 = 24 intersections (including dups) returned
+  // from rotated_rect_intersection_pts
+  Point<T> intersectPts[24], orderedPts[24];
+
+  Point<T> pts1[4];
+  Point<T> pts2[4];
+  get_rotated_vertices<T>(box1, pts1);
+  get_rotated_vertices<T>(box2, pts2);
+
+  int num = get_intersection_points<T>(pts1, pts2, intersectPts);
+
+  if (num <= 2) {
+    return 0.0;
+  }
+
+  // Convex Hull to order the intersection points in clockwise order and find
+  // the contour area.
+  int num_convex = convex_hull_graham<T>(intersectPts, num, orderedPts, true);
+  return polygon_area<T>(orderedPts, num_convex);
+}
+
+} // namespace
+
+template <typename T>
+HOST_DEVICE_INLINE T
+single_box_iou_rotated(T const* const box1_raw, T const* const box2_raw) {
+  // shift center to the middle point to achieve higher precision in result
+  RotatedBox<T> box1, box2;
+  auto center_shift_x = (box1_raw[0] + box2_raw[0]) / 2.0;
+  auto center_shift_y = (box1_raw[1] + box2_raw[1]) / 2.0;
+  box1.x_ctr = box1_raw[0] - center_shift_x;
+  box1.y_ctr = box1_raw[1] - center_shift_y;
+  box1.w = box1_raw[2];
+  box1.h = box1_raw[3];
+  box1.a = box1_raw[4];
+  box2.x_ctr = box2_raw[0] - center_shift_x;
+  box2.y_ctr = box2_raw[1] - center_shift_y;
+  box2.w = box2_raw[2];
+  box2.h = box2_raw[3];
+  box2.a = box2_raw[4];
+
+  T area1 = box1.w * box1.h;
+  T area2 = box2.w * box2.h;
+  if (area1 < 1e-14 || area2 < 1e-14) {
+    return 0.f;
+  }
+
+  T intersection = rotated_boxes_intersection<T>(box1, box2);
+  T iou = intersection / (area1 + area2 - intersection);
+  return iou;
+}

vis4d_cuda_ops/src/deform_conv/deform_conv.h

@@ -0,0 +1,375 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/csrc/deformable/deform_conv.h
+#pragma once
+#include <torch/types.h>
+
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+int deform_conv_forward_cuda(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor offset,
+    at::Tensor output,
+    at::Tensor columns,
+    at::Tensor ones,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    int im2col_step);
+
+int deform_conv_backward_input_cuda(
+    at::Tensor input,
+    at::Tensor offset,
+    at::Tensor gradOutput,
+    at::Tensor gradInput,
+    at::Tensor gradOffset,
+    at::Tensor weight,
+    at::Tensor columns,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    int im2col_step);
+
+int deform_conv_backward_parameters_cuda(
+    at::Tensor input,
+    at::Tensor offset,
+    at::Tensor gradOutput,
+    at::Tensor gradWeight, // at::Tensor gradBias,
+    at::Tensor columns,
+    at::Tensor ones,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    float scale,
+    int im2col_step);
+
+void modulated_deform_conv_cuda_forward(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor bias,
+    at::Tensor ones,
+    at::Tensor offset,
+    at::Tensor mask,
+    at::Tensor output,
+    at::Tensor columns,
+    int kernel_h,
+    int kernel_w,
+    const int stride_h,
+    const int stride_w,
+    const int pad_h,
+    const int pad_w,
+    const int dilation_h,
+    const int dilation_w,
+    const int group,
+    const int deformable_group,
+    const bool with_bias);
+
+void modulated_deform_conv_cuda_backward(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor bias,
+    at::Tensor ones,
+    at::Tensor offset,
+    at::Tensor mask,
+    at::Tensor columns,
+    at::Tensor grad_input,
+    at::Tensor grad_weight,
+    at::Tensor grad_bias,
+    at::Tensor grad_offset,
+    at::Tensor grad_mask,
+    at::Tensor grad_output,
+    int kernel_h,
+    int kernel_w,
+    int stride_h,
+    int stride_w,
+    int pad_h,
+    int pad_w,
+    int dilation_h,
+    int dilation_w,
+    int group,
+    int deformable_group,
+    const bool with_bias);
+
+#endif
+
+inline int deform_conv_forward(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor offset,
+    at::Tensor output,
+    at::Tensor columns,
+    at::Tensor ones,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    int im2col_step) {
+  if (input.is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    TORCH_CHECK(weight.is_cuda(), "weight tensor is not on GPU!");
+    TORCH_CHECK(offset.is_cuda(), "offset tensor is not on GPU!");
+    return deform_conv_forward_cuda(
+        input,
+        weight,
+        offset,
+        output,
+        columns,
+        ones,
+        kW,
+        kH,
+        dW,
+        dH,
+        padW,
+        padH,
+        dilationW,
+        dilationH,
+        group,
+        deformable_group,
+        im2col_step);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+  AT_ERROR("This operator is not implemented on CPU");
+}
+
+inline int deform_conv_backward_input(
+    at::Tensor input,
+    at::Tensor offset,
+    at::Tensor gradOutput,
+    at::Tensor gradInput,
+    at::Tensor gradOffset,
+    at::Tensor weight,
+    at::Tensor columns,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    int im2col_step) {
+  if (gradOutput.is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    TORCH_CHECK(input.is_cuda(), "input tensor is not on GPU!");
+    TORCH_CHECK(weight.is_cuda(), "weight tensor is not on GPU!");
+    TORCH_CHECK(offset.is_cuda(), "offset tensor is not on GPU!");
+    return deform_conv_backward_input_cuda(
+        input,
+        offset,
+        gradOutput,
+        gradInput,
+        gradOffset,
+        weight,
+        columns,
+        kW,
+        kH,
+        dW,
+        dH,
+        padW,
+        padH,
+        dilationW,
+        dilationH,
+        group,
+        deformable_group,
+        im2col_step);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+  AT_ERROR("This operator is not implemented on CPU");
+}
+
+inline int deform_conv_backward_filter(
+    at::Tensor input,
+    at::Tensor offset,
+    at::Tensor gradOutput,
+    at::Tensor gradWeight, // at::Tensor gradBias,
+    at::Tensor columns,
+    at::Tensor ones,
+    int kW,
+    int kH,
+    int dW,
+    int dH,
+    int padW,
+    int padH,
+    int dilationW,
+    int dilationH,
+    int group,
+    int deformable_group,
+    float scale,
+    int im2col_step) {
+  if (gradOutput.is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    TORCH_CHECK(input.is_cuda(), "input tensor is not on GPU!");
+    TORCH_CHECK(offset.is_cuda(), "offset tensor is not on GPU!");
+    return deform_conv_backward_parameters_cuda(
+        input,
+        offset,
+        gradOutput,
+        gradWeight,
+        columns,
+        ones,
+        kW,
+        kH,
+        dW,
+        dH,
+        padW,
+        padH,
+        dilationW,
+        dilationH,
+        group,
+        deformable_group,
+        scale,
+        im2col_step);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+  AT_ERROR("This operator is not implemented on CPU");
+}
+
+inline void modulated_deform_conv_forward(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor bias,
+    at::Tensor ones,
+    at::Tensor offset,
+    at::Tensor mask,
+    at::Tensor output,
+    at::Tensor columns,
+    int kernel_h,
+    int kernel_w,
+    const int stride_h,
+    const int stride_w,
+    const int pad_h,
+    const int pad_w,
+    const int dilation_h,
+    const int dilation_w,
+    const int group,
+    const int deformable_group,
+    const bool with_bias) {
+  if (input.is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    TORCH_CHECK(weight.is_cuda(), "weight tensor is not on GPU!");
+    TORCH_CHECK(bias.is_cuda(), "bias tensor is not on GPU!");
+    TORCH_CHECK(offset.is_cuda(), "offset tensor is not on GPU!");
+    return modulated_deform_conv_cuda_forward(
+        input,
+        weight,
+        bias,
+        ones,
+        offset,
+        mask,
+        output,
+        columns,
+        kernel_h,
+        kernel_w,
+        stride_h,
+        stride_w,
+        pad_h,
+        pad_w,
+        dilation_h,
+        dilation_w,
+        group,
+        deformable_group,
+        with_bias);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+  AT_ERROR("This operator is not implemented on CPU");
+}
+
+inline void modulated_deform_conv_backward(
+    at::Tensor input,
+    at::Tensor weight,
+    at::Tensor bias,
+    at::Tensor ones,
+    at::Tensor offset,
+    at::Tensor mask,
+    at::Tensor columns,
+    at::Tensor grad_input,
+    at::Tensor grad_weight,
+    at::Tensor grad_bias,
+    at::Tensor grad_offset,
+    at::Tensor grad_mask,
+    at::Tensor grad_output,
+    int kernel_h,
+    int kernel_w,
+    int stride_h,
+    int stride_w,
+    int pad_h,
+    int pad_w,
+    int dilation_h,
+    int dilation_w,
+    int group,
+    int deformable_group,
+    const bool with_bias) {
+  if (grad_output.is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    TORCH_CHECK(input.is_cuda(), "input tensor is not on GPU!");
+    TORCH_CHECK(weight.is_cuda(), "weight tensor is not on GPU!");
+    TORCH_CHECK(bias.is_cuda(), "bias tensor is not on GPU!");
+    TORCH_CHECK(offset.is_cuda(), "offset tensor is not on GPU!");
+    return modulated_deform_conv_cuda_backward(
+        input,
+        weight,
+        bias,
+        ones,
+        offset,
+        mask,
+        columns,
+        grad_input,
+        grad_weight,
+        grad_bias,
+        grad_offset,
+        grad_mask,
+        grad_output,
+        kernel_h,
+        kernel_w,
+        stride_h,
+        stride_w,
+        pad_h,
+        pad_w,
+        dilation_h,
+        dilation_w,
+        group,
+        deformable_group,
+        with_bias);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+  AT_ERROR("This operator is not implemented on CPU");
+}

vis4d_cuda_ops/src/iou_box3d/iou_box3d.h

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+#include <torch/extension.h>
+#include <tuple>
+#include "../util/pytorch3d_cutils.h"
+
+// Calculate the intersection volume and IoU metric for two batches of boxes
+//
+// Args:
+//     boxes1: tensor of shape (N, 8, 3) of the coordinates of the 1st boxes
+//     boxes2: tensor of shape (M, 8, 3) of the coordinates of the 2nd boxes
+// Returns:
+//     vol: (N, M) tensor of the volume of the intersecting convex shapes
+//     iou: (N, M) tensor of the intersection over union which is
+//          defined as: `iou = vol / (vol1 + vol2 - vol)`
+
+// CPU implementation
+std::tuple<at::Tensor, at::Tensor> IoUBox3DCpu(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2);
+
+// CUDA implementation
+std::tuple<at::Tensor, at::Tensor> IoUBox3DCuda(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2);
+
+// Implementation which is exposed
+inline std::tuple<at::Tensor, at::Tensor> IoUBox3D(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  if (boxes1.is_cuda() || boxes2.is_cuda()) {
+#ifdef WITH_CUDA
+    CHECK_CUDA(boxes1);
+    CHECK_CUDA(boxes2);
+    return IoUBox3DCuda(boxes1.contiguous(), boxes2.contiguous());
+#else
+    AT_ERROR("Not compiled with GPU support.");
+#endif
+  }
+  return IoUBox3DCpu(boxes1.contiguous(), boxes2.contiguous());
+}

vis4d_cuda_ops/src/iou_box3d/iou_box3d_cpu.cpp

@@ -0,0 +1,122 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <list>
+#include <numeric>
+#include <queue>
+#include <tuple>
+#include "iou_utils.h"
+
+std::tuple<at::Tensor, at::Tensor> IoUBox3DCpu(
+    const at::Tensor& boxes1,
+    const at::Tensor& boxes2) {
+  const int N = boxes1.size(0);
+  const int M = boxes2.size(0);
+  auto float_opts = boxes1.options().dtype(torch::kFloat32);
+  torch::Tensor vols = torch::zeros({N, M}, float_opts);
+  torch::Tensor ious = torch::zeros({N, M}, float_opts);
+
+  // Create tensor accessors
+  auto boxes1_a = boxes1.accessor<float, 3>();
+  auto boxes2_a = boxes2.accessor<float, 3>();
+  auto vols_a = vols.accessor<float, 2>();
+  auto ious_a = ious.accessor<float, 2>();
+
+  // Iterate through the N boxes in boxes1
+  for (int n = 0; n < N; ++n) {
+    const auto& box1 = boxes1_a[n];
+    // Convert to vector of face vertices i.e. effectively (F, 3, 3)
+    // face_verts is a data type defined in iou_utils.h
+    const face_verts box1_tris = GetBoxTris(box1);
+
+    // Calculate the position of the center of the box which is used in
+    // several calculations. This requires a tensor as input.
+    const vec3<float> box1_center = BoxCenter(boxes1[n]);
+
+    // Convert to vector of face vertices i.e. effectively (P, 4, 3)
+    const face_verts box1_planes = GetBoxPlanes(box1);
+
+    // Get Box Volumes
+    const float box1_vol = BoxVolume(box1_tris, box1_center);
+
+    // Iterate through the M boxes in boxes2
+    for (int m = 0; m < M; ++m) {
+      // Repeat above steps for box2
+      // TODO: check if caching these value helps performance.
+      const auto& box2 = boxes2_a[m];
+      const face_verts box2_tris = GetBoxTris(box2);
+      const vec3<float> box2_center = BoxCenter(boxes2[m]);
+      const face_verts box2_planes = GetBoxPlanes(box2);
+      const float box2_vol = BoxVolume(box2_tris, box2_center);
+
+      // Every triangle in one box will be compared to each plane in the other
+      // box. There are 3 possible outcomes:
+      // 1. If the triangle is fully inside, then it will
+      //    remain as is.
+      // 2. If the triagnle it is fully outside, it will be removed.
+      // 3. If the triangle intersects with the (infinite) plane, it
+      //    will be broken into subtriangles such that each subtriangle is full
+      //    inside the plane and part of the intersecting tetrahedron.
+
+      // Tris in Box1 -> Planes in Box2
+      face_verts box1_intersect =
+          BoxIntersections(box1_tris, box2_planes, box2_center);
+      // Tris in Box2 -> Planes in Box1
+      face_verts box2_intersect =
+          BoxIntersections(box2_tris, box1_planes, box1_center);
+
+      // If there are overlapping regions in Box2, remove any coplanar faces
+      if (box2_intersect.size() > 0) {
+        // Identify if any triangles in Box2 are coplanar with Box1
+        std::vector<int> tri2_keep(box2_intersect.size());
+        std::fill(tri2_keep.begin(), tri2_keep.end(), 1);
+        for (int b1 = 0; b1 < box1_intersect.size(); ++b1) {
+          for (int b2 = 0; b2 < box2_intersect.size(); ++b2) {
+            const bool is_coplanar =
+                IsCoplanarTriTri(box1_intersect[b1], box2_intersect[b2]);
+            const float area = FaceArea(box1_intersect[b1]);
+            if ((is_coplanar) && (area > aEpsilon)) {
+              tri2_keep[b2] = 0;
+            }
+          }
+        }
+
+        // Keep only the non coplanar triangles in Box2 - add them to the
+        // Box1 triangles.
+        for (int b2 = 0; b2 < box2_intersect.size(); ++b2) {
+          if (tri2_keep[b2] == 1) {
+            box1_intersect.push_back((box2_intersect[b2]));
+          }
+        }
+      }
+
+      // Initialize the vol and iou to 0.0 in case there are no triangles
+      // in the intersecting shape.
+      float vol = 0.0;
+      float iou = 0.0;
+
+      // If there are triangles in the intersecting shape
+      if (box1_intersect.size() > 0) {
+        // The intersecting shape is a polyhedron made up of the
+        // triangular faces that are all now in box1_intersect.
+        // Calculate the polyhedron center
+        const vec3<float> polyhedron_center = PolyhedronCenter(box1_intersect);
+        // Compute intersecting polyhedron volume
+        vol = BoxVolume(box1_intersect, polyhedron_center);
+        // Compute IoU
+        iou = vol / (box1_vol + box2_vol - vol);
+      }
+      // Save out volume and IoU
+      vols_a[n][m] = vol;
+      ious_a[n][m] = iou;
+    }
+  }
+  return std::make_tuple(vols, ious);
+}

vis4d_cuda_ops/src/iou_box3d/iou_utils.h

@@ -0,0 +1,733 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <ATen/ATen.h>
+#include <assert.h>
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <algorithm>
+#include <list>
+#include <numeric>
+#include <queue>
+#include <tuple>
+#include <type_traits>
+#include "../util/vec3.h"
+
+// dEpsilon: Used in dot products and is used to assess whether two unit vectors
+// are orthogonal (or coplanar). It's an epsilon on cos(θ).
+// With dEpsilon = 0.001, two unit vectors are considered co-planar
+// if their θ = 2.5 deg.
+const auto dEpsilon = 1e-3;
+// aEpsilon: Used once in main function to check for small face areas
+const auto aEpsilon = 1e-4;
+// kEpsilon: Used only for norm(u) = u/max(||u||, kEpsilon)
+const auto kEpsilon = 1e-8;
+
+/*
+_PLANES and _TRIS define the 4- and 3-connectivity
+of the 8 box corners.
+_PLANES gives the quad faces of the 3D box
+_TRIS gives the triangle faces of the 3D box
+*/
+const int NUM_PLANES = 6;
+const int NUM_TRIS = 12;
+const int _PLANES[6][4] = {
+    {0, 1, 2, 3},
+    {3, 2, 6, 7},
+    {0, 1, 5, 4},
+    {0, 3, 7, 4},
+    {1, 5, 6, 2},
+    {4, 5, 6, 7},
+};
+const int _TRIS[12][3] = {
+    {0, 1, 2},
+    {0, 3, 2},
+    {4, 5, 6},
+    {4, 6, 7},
+    {1, 5, 6},
+    {1, 6, 2},
+    {0, 4, 7},
+    {0, 7, 3},
+    {3, 2, 6},
+    {3, 6, 7},
+    {0, 1, 5},
+    {0, 4, 5},
+};
+
+// Create a new data type for representing the
+// verts for each face which can be triangle or plane.
+// This helps make the code more readable.
+using face_verts = std::vector<std::vector<vec3<float>>>;
+
+// Args
+//    box: (8, 3) tensor accessor for the box vertices
+//    plane_idx: index of the plane in the box
+//    vert_idx: index of the vertex in the plane
+//
+// Returns
+//    vec3<T> (x, y, x) vertex coordinates
+//
+template <typename Box>
+inline vec3<float>
+ExtractVertsPlane(const Box& box, const int plane_idx, const int vert_idx) {
+  return vec3<float>(
+      box[_PLANES[plane_idx][vert_idx]][0],
+      box[_PLANES[plane_idx][vert_idx]][1],
+      box[_PLANES[plane_idx][vert_idx]][2]);
+}
+
+// Args
+//    box: (8, 3) tensor accessor for the box vertices
+//    tri_idx: index of the triangle face in the box
+//    vert_idx: index of the vertex in the triangle
+//
+// Returns
+//    vec3<T> (x, y, x) vertex coordinates
+//
+template <typename Box>
+inline vec3<float>
+ExtractVertsTri(const Box& box, const int tri_idx, const int vert_idx) {
+  return vec3<float>(
+      box[_TRIS[tri_idx][vert_idx]][0],
+      box[_TRIS[tri_idx][vert_idx]][1],
+      box[_TRIS[tri_idx][vert_idx]][2]);
+}
+
+// Args
+//    box: (8, 3) tensor accessor for the box vertices
+//
+// Returns
+//    std::vector<std::vector<vec3<T>>> effectively (F, 3, 3)
+//      coordinates of the verts for each face
+//
+template <typename Box>
+inline face_verts GetBoxTris(const Box& box) {
+  face_verts box_tris;
+  for (int t = 0; t < NUM_TRIS; ++t) {
+    vec3<float> v0 = ExtractVertsTri(box, t, 0);
+    vec3<float> v1 = ExtractVertsTri(box, t, 1);
+    vec3<float> v2 = ExtractVertsTri(box, t, 2);
+    box_tris.push_back({v0, v1, v2});
+  }
+  return box_tris;
+}
+
+// Args
+//    box: (8, 3) tensor accessor for the box vertices
+//
+// Returns
+//    std::vector<std::vector<vec3<T>>> effectively (P, 3, 3)
+//      coordinates of the 4 verts for each plane
+//
+template <typename Box>
+inline face_verts GetBoxPlanes(const Box& box) {
+  face_verts box_planes;
+  for (int t = 0; t < NUM_PLANES; ++t) {
+    vec3<float> v0 = ExtractVertsPlane(box, t, 0);
+    vec3<float> v1 = ExtractVertsPlane(box, t, 1);
+    vec3<float> v2 = ExtractVertsPlane(box, t, 2);
+    vec3<float> v3 = ExtractVertsPlane(box, t, 3);
+    box_planes.push_back({v0, v1, v2, v3});
+  }
+  return box_planes;
+}
+
+// The normal of a plane spanned by vectors e0 and e1
+//
+// Args
+//    e0, e1: vec3 vectors defining a plane
+//
+// Returns
+//    vec3: normal of the plane
+//
+inline vec3<float> GetNormal(const vec3<float> e0, const vec3<float> e1) {
+  vec3<float> n = cross(e0, e1);
+  n = n / std::fmaxf(norm(n), kEpsilon);
+  return n;
+}
+
+// The center of a triangle tri
+//
+// Args
+//    tri: vec3 coordinates of the vertices of the triangle
+//
+// Returns
+//    vec3: center of the triangle
+//
+inline vec3<float> TriCenter(const std::vector<vec3<float>>& tri) {
+  // Vertices of the triangle
+  const vec3<float> v0 = tri[0];
+  const vec3<float> v1 = tri[1];
+  const vec3<float> v2 = tri[2];
+
+  return (v0 + v1 + v2) / 3.0f;
+}
+
+// The normal of the triangle defined by vertices (v0, v1, v2)
+// We find the "best" edges connecting the face center to the vertices,
+// such that the cross product between the edges is maximized.
+//
+// Args
+//    tri: vec3 coordinates of the vertices of the face
+//
+// Returns
+//    vec3: normal for the face
+//
+inline vec3<float> TriNormal(const std::vector<vec3<float>>& tri) {
+  // Get center of triangle
+  const vec3<float> ctr = TriCenter(tri);
+
+  // find the "best" normal as cross product of edges from center
+  float max_dist = -1.0f;
+  vec3<float> n = {0.0f, 0.0f, 0.0f};
+  for (int i = 0; i < 2; ++i) {
+    for (int j = i + 1; j < 3; ++j) {
+      const float dist = norm(cross(tri[i] - ctr, tri[j] - ctr));
+      if (dist > max_dist) {
+        n = GetNormal(tri[i] - ctr, tri[j] - ctr);
+      }
+    }
+  }
+  return n;
+}
+
+// The center of a plane
+//
+// Args
+//    plane: vec3 coordinates of the vertices of the plane
+//
+// Returns
+//    vec3: center of the plane
+//
+inline vec3<float> PlaneCenter(const std::vector<vec3<float>>& plane) {
+  // Vertices of the plane
+  const vec3<float> v0 = plane[0];
+  const vec3<float> v1 = plane[1];
+  const vec3<float> v2 = plane[2];
+  const vec3<float> v3 = plane[3];
+
+  return (v0 + v1 + v2 + v3) / 4.0f;
+}
+
+// The normal of a planar face with vertices (v0, v1, v2, v3)
+// We find the "best" edges connecting the face center to the vertices,
+// such that the cross product between the edges is maximized.
+//
+// Args
+//    plane: vec3 coordinates of the vertices of the planar face
+//
+// Returns
+//    vec3: normal of the planar face
+//
+inline vec3<float> PlaneNormal(const std::vector<vec3<float>>& plane) {
+  // Get center of planar face
+  vec3<float> ctr = PlaneCenter(plane);
+
+  // find the "best" normal as cross product of edges from center
+  float max_dist = -1.0f;
+  vec3<float> n = {0.0f, 0.0f, 0.0f};
+  for (int i = 0; i < 3; ++i) {
+    for (int j = i + 1; j < 4; ++j) {
+      const float dist = norm(cross(plane[i] - ctr, plane[j] - ctr));
+      if (dist > max_dist) {
+        n = GetNormal(plane[i] - ctr, plane[j] - ctr);
+      }
+    }
+  }
+  return n;
+}
+
+// The area of the face defined by vertices (v0, v1, v2)
+// Define e0 to be the edge connecting (v1, v0)
+// Define e1 to be the edge connecting (v2, v0)
+// Area is the norm of the cross product of e0, e1 divided by 2.0
+//
+// Args
+//    tri: vec3 coordinates of the vertices of the face
+//
+// Returns
+//    float: area for the face
+//
+inline float FaceArea(const std::vector<vec3<float>>& tri) {
+  // Get verts for face
+  const vec3<float> v0 = tri[0];
+  const vec3<float> v1 = tri[1];
+  const vec3<float> v2 = tri[2];
+  const vec3<float> n = cross(v1 - v0, v2 - v0);
+  return norm(n) / 2.0;
+}
+
+// The normal of a box plane defined by the verts in `plane` such that it
+// points toward the centroid of the box given by `center`.
+//
+// Args
+//    plane: vec3 coordinates of the vertices of the plane
+//    center: vec3 coordinates of the center of the box from
+//        which the plane originated
+//
+// Returns
+//    vec3: normal for the plane such that it points towards
+//      the center of the box
+//
+inline vec3<float> PlaneNormalDirection(
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& center) {
+  // The plane's center & normal
+  const vec3<float> plane_center = PlaneCenter(plane);
+  vec3<float> n = PlaneNormal(plane);
+
+  // We project the center on the plane defined by (v0, v1, v2, v3)
+  // We can write center = plane_center + a * e0 + b * e1 + c * n
+  // We know that <e0, n> = 0 and <e1, n> = 0 and
+  // <a, b> is the dot product between a and b.
+  // This means we can solve for c as:
+  // c = <center - plane_center - a * e0 - b * e1, n>
+  //   = <center - plane_center, n>
+  const float c = dot((center - plane_center), n);
+
+  // If c is negative, then we revert the direction of n such that n
+  // points "inside"
+  if (c < 0.0f) {
+    n = -1.0f * n;
+  }
+
+  return n;
+}
+
+// Calculate the volume of the box by summing the volume of
+// each of the tetrahedrons formed with a triangle face and
+// the box centroid.
+//
+// Args
+//    box_tris: vector of vec3 coordinates of the vertices of each
+//       of the triangles in the box
+//    box_center: vec3 coordinates of the center of the box
+//
+// Returns
+//    float: volume of the box
+//
+inline float BoxVolume(
+    const face_verts& box_tris,
+    const vec3<float>& box_center) {
+  float box_vol = 0.0;
+  // Iterate through each triange, calculate the area of the
+  // tetrahedron formed with the box_center and sum them
+  for (int t = 0; t < box_tris.size(); ++t) {
+    // Subtract the center:
+    const vec3<float> v0 = box_tris[t][0] - box_center;
+    const vec3<float> v1 = box_tris[t][1] - box_center;
+    const vec3<float> v2 = box_tris[t][2] - box_center;
+
+    // Compute the area
+    const float area = dot(v0, cross(v1, v2));
+    const float vol = std::abs(area) / 6.0;
+    box_vol = box_vol + vol;
+  }
+  return box_vol;
+}
+
+// Compute the box center as the mean of the verts
+//
+// Args
+//    box_verts: (8, 3) tensor of the corner vertices of the box
+//
+// Returns
+//    vec3: coordinates of the center of the box
+//
+inline vec3<float> BoxCenter(const at::Tensor& box_verts) {
+  const auto& box_center_t = at::mean(box_verts, 0);
+  const vec3<float> box_center(
+      box_center_t[0].item<float>(),
+      box_center_t[1].item<float>(),
+      box_center_t[2].item<float>());
+  return box_center;
+}
+
+// Compute the polyhedron center as the mean of the face centers
+// of the triangle faces
+//
+// Args
+//    tris: vector of vec3 coordinates of the
+//       vertices of each of the triangles in the polyhedron
+//
+// Returns
+//    vec3: coordinates of the center of the polyhedron
+//
+inline vec3<float> PolyhedronCenter(const face_verts& tris) {
+  float x = 0.0;
+  float y = 0.0;
+  float z = 0.0;
+  const int num_tris = tris.size();
+
+  // Find the center point of each face
+  for (int t = 0; t < num_tris; ++t) {
+    const vec3<float> v0 = tris[t][0];
+    const vec3<float> v1 = tris[t][1];
+    const vec3<float> v2 = tris[t][2];
+    const float x_face = (v0.x + v1.x + v2.x) / 3.0;
+    const float y_face = (v0.y + v1.y + v2.y) / 3.0;
+    const float z_face = (v0.z + v1.z + v2.z) / 3.0;
+    x = x + x_face;
+    y = y + y_face;
+    z = z + z_face;
+  }
+
+  // Take the mean of the centers of all faces
+  x = x / num_tris;
+  y = y / num_tris;
+  z = z / num_tris;
+
+  const vec3<float> center(x, y, z);
+  return center;
+}
+
+// Compute a boolean indicator for whether a point
+// is inside a plane, where inside refers to whether
+// or not the point has a component in the
+// normal direction of the plane.
+//
+// Args
+//    plane: vector of vec3 coordinates of the
+//       vertices of each of the triangles in the box
+//    normal: vec3 of the direction of the plane normal
+//    point: vec3 of the position of the point of interest
+//
+// Returns
+//    bool: whether or not the point is inside the plane
+//
+inline bool IsInside(
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& normal,
+    const vec3<float>& point) {
+  // The center of the plane
+  const vec3<float> plane_ctr = PlaneCenter(plane);
+
+  // Every point p can be written as p = plane_ctr + a e0 + b e1 + c n
+  // Solving for c:
+  // c = (point - plane_ctr - a * e0 - b * e1).dot(n)
+  // We know that <e0, n> = 0 and <e1, n> = 0
+  // So the calculation can be simplified as:
+  const float c = dot((point - plane_ctr), normal);
+  const bool inside = c >= 0.0f;
+  return inside;
+}
+
+// Find the point of intersection between a plane
+// and a line given by the end points (p0, p1)
+//
+// Args
+//    plane: vector of vec3 coordinates of the
+//       vertices of each of the triangles in the box
+//    normal: vec3 of the direction of the plane normal
+//    p0, p1: vec3 of the start and end point of the line
+//
+// Returns
+//    vec3: position of the intersection point
+//
+inline vec3<float> PlaneEdgeIntersection(
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& normal,
+    const vec3<float>& p0,
+    const vec3<float>& p1) {
+  // The center of the plane
+  const vec3<float> plane_ctr = PlaneCenter(plane);
+
+  // The point of intersection can be parametrized
+  // p = p0 + a (p1 - p0) where a in [0, 1]
+  // We want to find a such that p is on plane
+  // <p - ctr, n> = 0
+
+  vec3<float> direc = p1 - p0;
+  direc = direc / std::fmaxf(norm(direc), kEpsilon);
+
+  vec3<float> p = (p1 + p0) / 2.0f;
+
+  if (std::abs(dot(direc, normal)) >= dEpsilon) {
+    const float top = -1.0f * dot(p0 - plane_ctr, normal);
+    const float bot = dot(p1 - p0, normal);
+    const float a = top / bot;
+    p = p0 + a * (p1 - p0);
+  }
+  return p;
+}
+
+// Compute the most distant points between two sets of vertices
+//
+// Args
+//    verts1, verts2: vec3 defining the list of vertices
+//
+// Returns
+//    v1m, v2m: vec3 vectors of the most distant points
+//          in verts1 and verts2 respectively
+//
+inline std::tuple<vec3<float>, vec3<float>> ArgMaxVerts(
+    const std::vector<vec3<float>>& verts1,
+    const std::vector<vec3<float>>& verts2) {
+  vec3<float> v1m = {0.0f, 0.0f, 0.0f};
+  vec3<float> v2m = {0.0f, 0.0f, 0.0f};
+  float maxdist = -1.0f;
+
+  for (const auto& v1 : verts1) {
+    for (const auto& v2 : verts2) {
+      if (norm(v1 - v2) > maxdist) {
+        v1m = v1;
+        v2m = v2;
+        maxdist = norm(v1 - v2);
+      }
+    }
+  }
+  return std::make_tuple(v1m, v2m);
+}
+
+// Compute a boolean indicator for whether or not two faces
+// are coplanar
+//
+// Args
+//    tri1, tri2: std:vector<vec3> of the vertex coordinates of
+//        triangle faces
+//
+// Returns
+//    bool: whether or not the two faces are coplanar
+//
+inline bool IsCoplanarTriTri(
+    const std::vector<vec3<float>>& tri1,
+    const std::vector<vec3<float>>& tri2) {
+  // Get normal for tri 1
+  const vec3<float> n1 = TriNormal(tri1);
+
+  // Get normal for tri 2
+  const vec3<float> n2 = TriNormal(tri2);
+
+  // Check if parallel
+  const bool check1 = std::abs(dot(n1, n2)) > 1 - dEpsilon;
+
+  // Compute most distant points
+  auto argvs = ArgMaxVerts(tri1, tri2);
+  const auto v1m = std::get<0>(argvs);
+  const auto v2m = std::get<1>(argvs);
+
+  vec3<float> n12m = v1m - v2m;
+  n12m = n12m / std::fmaxf(norm(n12m), kEpsilon);
+
+  const bool check2 = (std::abs(dot(n12m, n1)) < dEpsilon) ||
+      (std::abs(dot(n12m, n2)) < dEpsilon);
+
+  return (check1 && check2);
+}
+
+// Compute a boolean indicator for whether or not a triangular and a planar
+// face are coplanar
+//
+// Args
+//    tri, plane: std:vector<vec3> of the vertex coordinates of
+//        triangular face and planar face
+//    normal: the normal direction of the plane pointing "inside"
+//
+// Returns
+//    bool: whether or not the two faces are coplanar
+//
+inline bool IsCoplanarTriPlane(
+    const std::vector<vec3<float>>& tri,
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& normal) {
+  // Get normal for tri
+  const vec3<float> nt = TriNormal(tri);
+
+  // check if parallel
+  const bool check1 = std::abs(dot(nt, normal)) > 1 - dEpsilon;
+
+  // Compute most distant points
+  auto argvs = ArgMaxVerts(tri, plane);
+  const auto v1m = std::get<0>(argvs);
+  const auto v2m = std::get<1>(argvs);
+
+  vec3<float> n12m = v1m - v2m;
+  n12m = n12m / std::fmaxf(norm(n12m), kEpsilon);
+
+  const bool check2 = std::abs(dot(n12m, normal)) < dEpsilon;
+
+  return (check1 && check2);
+}
+
+// Triangle is clipped into a quadrilateral
+// based on the intersection points with the plane.
+// Then the quadrilateral is divided into two triangles.
+//
+// Args
+//    plane: vector of vec3 coordinates of the
+//        vertices of each of the triangles in the box
+//    normal: vec3 of the direction of the plane normal
+//    vout: vec3 of the point in the triangle which is outside
+//       the plane
+//    vin1, vin2: vec3 of the points in the triangle which are
+//        inside the plane
+//
+// Returns
+//    std::vector<std::vector<vec3>>: vector of vertex coordinates
+//      of the new triangle faces
+//
+inline face_verts ClipTriByPlaneOneOut(
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& normal,
+    const vec3<float>& vout,
+    const vec3<float>& vin1,
+    const vec3<float>& vin2) {
+  // point of intersection between plane and (vin1, vout)
+  const vec3<float> pint1 = PlaneEdgeIntersection(plane, normal, vin1, vout);
+  // point of intersection between plane and (vin2, vout)
+  const vec3<float> pint2 = PlaneEdgeIntersection(plane, normal, vin2, vout);
+  const face_verts face_verts = {{vin1, pint1, pint2}, {vin1, pint2, vin2}};
+  return face_verts;
+}
+
+// Triangle is clipped into a smaller triangle based
+// on the intersection points with the plane.
+//
+// Args
+//    plane: vector of vec3 coordinates of the
+//       vertices of each of the triangles in the box
+//    normal: vec3 of the direction of the plane normal
+//    vout1, vout2: vec3 of the points in the triangle which are
+//       outside the plane
+//    vin: vec3 of the point in the triangle which is inside
+//        the plane
+// Returns
+//    std::vector<std::vector<vec3>>: vector of vertex coordinates
+//      of the new triangle face
+//
+inline face_verts ClipTriByPlaneTwoOut(
+    const std::vector<vec3<float>>& plane,
+    const vec3<float>& normal,
+    const vec3<float>& vout1,
+    const vec3<float>& vout2,
+    const vec3<float>& vin) {
+  // point of intersection between plane and (vin, vout1)
+  const vec3<float> pint1 = PlaneEdgeIntersection(plane, normal, vin, vout1);
+  // point of intersection between plane and (vin, vout2)
+  const vec3<float> pint2 = PlaneEdgeIntersection(plane, normal, vin, vout2);
+  const face_verts face_verts = {{vin, pint1, pint2}};
+  return face_verts;
+}
+
+// Clip the triangle faces so that they lie within the
+// plane, creating new triangle faces where necessary.
+//
+// Args
+//    plane: vector of vec3 coordinates of the
+//       vertices of each of the triangles in the box
+//    tri: std:vector<vec3> of the vertex coordinates of the
+//       triangle faces
+//    normal: vec3 of the direction of the plane normal
+//
+// Returns
+//    std::vector<std::vector<vec3>>: vector of vertex coordinates
+//      of the new triangle faces formed after clipping.
+//      All triangles are now "inside" the plane.
+//
+inline face_verts ClipTriByPlane(
+    const std::vector<vec3<float>>& plane,
+    const std::vector<vec3<float>>& tri,
+    const vec3<float>& normal) {
+  // Get Triangle vertices
+  const vec3<float> v0 = tri[0];
+  const vec3<float> v1 = tri[1];
+  const vec3<float> v2 = tri[2];
+
+  // Check coplanar
+  const bool iscoplanar = IsCoplanarTriPlane(tri, plane, normal);
+  if (iscoplanar) {
+    // Return input vertices
+    face_verts tris = {{v0, v1, v2}};
+    return tris;
+  }
+
+  // Check each of the triangle vertices to see if it is inside the plane
+  const bool isin0 = IsInside(plane, normal, v0);
+  const bool isin1 = IsInside(plane, normal, v1);
+  const bool isin2 = IsInside(plane, normal, v2);
+
+  // All in
+  if (isin0 && isin1 && isin2) {
+    // Return input vertices
+    face_verts tris = {{v0, v1, v2}};
+    return tris;
+  }
+
+  face_verts empty_tris = {};
+  // All out
+  if (!isin0 && !isin1 && !isin2) {
+    return empty_tris;
+  }
+
+  // One vert out
+  if (isin0 && isin1 && !isin2) {
+    return ClipTriByPlaneOneOut(plane, normal, v2, v0, v1);
+  }
+  if (isin0 && !isin1 && isin2) {
+    return ClipTriByPlaneOneOut(plane, normal, v1, v0, v2);
+  }
+  if (!isin0 && isin1 && isin2) {
+    return ClipTriByPlaneOneOut(plane, normal, v0, v1, v2);
+  }
+
+  // Two verts out
+  if (isin0 && !isin1 && !isin2) {
+    return ClipTriByPlaneTwoOut(plane, normal, v1, v2, v0);
+  }
+  if (!isin0 && !isin1 && isin2) {
+    return ClipTriByPlaneTwoOut(plane, normal, v0, v1, v2);
+  }
+  if (!isin0 && isin1 && !isin2) {
+    return ClipTriByPlaneTwoOut(plane, normal, v0, v2, v1);
+  }
+
+  // Else return empty (should not be reached)
+  return empty_tris;
+}
+
+// Get the triangles from each box which are part of the
+// intersecting polyhedron by computing the intersection
+// points with each of the planes.
+//
+// Args
+//    tris: vertex coordinates of all the triangle faces
+//       in the box
+//    planes: vertex coordinates of all the planes in the box
+//    center: vec3 coordinates of the center of the box from which
+//        the planes originate
+//
+// Returns
+//    std::vector<std::vector<vec3>>> vector of vertex coordinates
+//      of the new triangle faces formed after clipping.
+//      All triangles are now "inside" the planes.
+//
+inline face_verts BoxIntersections(
+    const face_verts& tris,
+    const face_verts& planes,
+    const vec3<float>& center) {
+  // Create a new vector to avoid modifying in place
+  face_verts out_tris = tris;
+  for (int p = 0; p < NUM_PLANES; ++p) {
+    // Get plane normal direction
+    const vec3<float> n2 = PlaneNormalDirection(planes[p], center);
+    // Iterate through triangles in tris
+    // Create intermediate vector to store the updated tris
+    face_verts tri_verts_updated;
+    for (int t = 0; t < out_tris.size(); ++t) {
+      // Clip tri by plane
+      const face_verts tri_updated = ClipTriByPlane(planes[p], out_tris[t], n2);
+      // Add to the tri_verts_updated output if not empty
+      for (int v = 0; v < tri_updated.size(); ++v) {
+        tri_verts_updated.push_back(tri_updated[v]);
+      }
+    }
+    // Update the tris
+    out_tris = tri_verts_updated;
+  }
+  return out_tris;
+}

vis4d_cuda_ops/src/ms_deform_attn/ms_deform_attn.h

@@ -0,0 +1,60 @@
+/*!
+**************************************************************************************************
+* 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
+
+#ifdef WITH_CUDA
+#include "ms_deform_attn_cuda.h"
+#endif
+
+
+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())
+    {
+#ifdef WITH_CUDA
+        return ms_deform_attn_cuda_forward(
+            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, im2col_step);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    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())
+    {
+#ifdef WITH_CUDA
+        return ms_deform_attn_cuda_backward(
+            value, spatial_shapes, level_start_index, sampling_loc, attn_weight, grad_output, im2col_step);
+#else
+        AT_ERROR("Not compiled with GPU support");
+#endif
+    }
+    AT_ERROR("Not implemented on the CPU");
+}
+

vis4d_cuda_ops/src/ms_deform_attn/ms_deform_attn_cuda.h

@@ -0,0 +1,30 @@
+/*!
+**************************************************************************************************
+* 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>
+
+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);
+

vis4d_cuda_ops/src/nms_rotated/nms_rotated.h

@@ -0,0 +1,37 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/master/detectron2/layers/csrc/nms_rotated/nms_rotated.h
+#pragma once
+#include <torch/types.h>
+
+at::Tensor nms_rotated_cpu(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    const double iou_threshold);
+
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+at::Tensor nms_rotated_cuda(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    const double iou_threshold);
+#endif
+
+// Interface for Python
+// inline is needed to prevent multiple function definitions when this header is
+// included by different cpps
+inline at::Tensor nms_rotated(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    const double iou_threshold) {
+  assert(dets.device().is_cuda() == scores.device().is_cuda());
+  if (dets.device().is_cuda()) {
+#if defined(WITH_CUDA) || defined(WITH_HIP)
+    return nms_rotated_cuda(
+        dets.contiguous(), scores.contiguous(), iou_threshold);
+#else
+    AT_ERROR("Not compiled with GPU support!");
+#endif
+  }
+
+  return nms_rotated_cpu(dets.contiguous(), scores.contiguous(), iou_threshold);
+}

vis4d_cuda_ops/src/nms_rotated/nms_rotated_cpu.cpp

@@ -0,0 +1,73 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+// modified from
+// https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/csrc/nms_rotated/nms_rotated_cpu.cpp
+#include "../box_iou_rotated/box_iou_rotated_utils.h"
+#include "nms_rotated.h"
+
+template <typename scalar_t>
+at::Tensor nms_rotated_cpu_kernel(
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    const double iou_threshold) {
+  // nms_rotated_cpu_kernel is modified from torchvision's nms_cpu_kernel,
+  // however, the code in this function is much shorter because
+  // we delegate the IoU computation for rotated boxes to
+  // the single_box_iou_rotated function in box_iou_rotated_utils.h
+  AT_ASSERTM(dets.device().is_cpu(), "dets must be a CPU tensor");
+  AT_ASSERTM(scores.device().is_cpu(), "scores must be a CPU tensor");
+  AT_ASSERTM(
+      dets.scalar_type() == scores.scalar_type(),
+      "dets should have the same type as scores");
+
+  if (dets.numel() == 0) {
+    return at::empty({0}, dets.options().dtype(at::kLong));
+  }
+
+  auto order_t = std::get<1>(scores.sort(0, /* descending=*/true));
+
+  auto ndets = dets.size(0);
+  at::Tensor suppressed_t = at::zeros({ndets}, dets.options().dtype(at::kByte));
+  at::Tensor keep_t = at::zeros({ndets}, dets.options().dtype(at::kLong));
+
+  auto suppressed = suppressed_t.data_ptr<uint8_t>();
+  auto keep = keep_t.data_ptr<int64_t>();
+  auto order = order_t.data_ptr<int64_t>();
+
+  int64_t num_to_keep = 0;
+
+  for (int64_t _i = 0; _i < ndets; _i++) {
+    auto i = order[_i];
+    if (suppressed[i] == 1) {
+      continue;
+    }
+
+    keep[num_to_keep++] = i;
+
+    for (int64_t _j = _i + 1; _j < ndets; _j++) {
+      auto j = order[_j];
+      if (suppressed[j] == 1) {
+        continue;
+      }
+
+      auto ovr = single_box_iou_rotated<scalar_t>(
+          dets[i].data_ptr<scalar_t>(), dets[j].data_ptr<scalar_t>());
+      if (ovr >= iou_threshold) {
+        suppressed[j] = 1;
+      }
+    }
+  }
+  return keep_t.narrow(/*dim=*/0, /*start=*/0, /*length=*/num_to_keep);
+}
+
+at::Tensor nms_rotated_cpu(
+    // input must be contiguous
+    const at::Tensor& dets,
+    const at::Tensor& scores,
+    const double iou_threshold) {
+  auto result = at::empty({0}, dets.options());
+
+  AT_DISPATCH_FLOATING_TYPES(dets.scalar_type(), "nms_rotated", [&] {
+    result = nms_rotated_cpu_kernel<scalar_t>(dets, scores, iou_threshold);
+  });
+  return result;
+}

vis4d_cuda_ops/src/util/pytorch3d_cutils.h

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+#include <torch/extension.h>
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor.")
+#define CHECK_CONTIGUOUS(x) \
+  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous.")
+#define CHECK_CONTIGUOUS_CUDA(x) \
+  CHECK_CUDA(x);                 \
+  CHECK_CONTIGUOUS(x)

vis4d_cuda_ops/src/util/vec3.h

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+// A fixed-sized vector with basic arithmetic operators useful for
+// representing 3D coordinates.
+// TODO: switch to Eigen if more functionality is needed.
+
+template <
+    typename T,
+    typename = std::enable_if_t<
+        std::is_same<T, double>::value || std::is_same<T, float>::value>>
+struct vec3 {
+  T x, y, z;
+  typedef T scalar_t;
+  vec3(T x, T y, T z) : x(x), y(y), z(z) {}
+};
+
+template <typename T>
+inline vec3<T> operator+(const vec3<T>& a, const vec3<T>& b) {
+  return vec3<T>(a.x + b.x, a.y + b.y, a.z + b.z);
+}
+
+template <typename T>
+inline vec3<T> operator-(const vec3<T>& a, const vec3<T>& b) {
+  return vec3<T>(a.x - b.x, a.y - b.y, a.z - b.z);
+}
+
+template <typename T>
+inline vec3<T> operator/(const vec3<T>& a, const T b) {
+  if (b == 0.0) {
+    AT_ERROR(
+        "denominator in vec3 division is 0"); // prevent divide by 0 errors.
+  }
+  return vec3<T>(a.x / b, a.y / b, a.z / b);
+}
+
+template <typename T>
+inline vec3<T> operator*(const T a, const vec3<T>& b) {
+  return vec3<T>(a * b.x, a * b.y, a * b.z);
+}
+
+template <typename T>
+inline vec3<T> operator*(const vec3<T>& a, const vec3<T>& b) {
+  return vec3<T>(a.x * b.x, a.y * b.y, a.z * b.z);
+}
+
+template <typename T>
+inline T dot(const vec3<T>& a, const vec3<T>& b) {
+  return a.x * b.x + a.y * b.y + a.z * b.z;
+}
+
+template <typename T>
+inline vec3<T> cross(const vec3<T>& a, const vec3<T>& b) {
+  return vec3<T>(
+      a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
+}
+
+template <typename T>
+inline T norm(const vec3<T>& a) {
+  return sqrt(dot(a, a));
+}
+
+template <typename T>
+std::ostream& operator<<(std::ostream& os, const vec3<T>& v) {
+  os << "vec3(" << v.x << ", " << v.y << ", " << v.z << ")";
+  return os;
+}

vis4d_cuda_ops/src/vision.cpp

@@ -0,0 +1,38 @@
+/*!
+**************************************************************************************************
+* Vis4D CUDA Operations
+**************************************************************************************************
+* Modified from https://github.com/fundamentalvision/Deformable-DETR/blob/main/models/ops/src/vision.cpp
+**************************************************************************************************
+*/
+#include <torch/extension.h>
+#include "nms_rotated/nms_rotated.h"
+#include "box_iou_rotated/box_iou_rotated.h"
+#include "ms_deform_attn/ms_deform_attn.h"
+#include "deform_conv/deform_conv.h"
+#include "iou_box3d/iou_box3d.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("nms_rotated", &nms_rotated);
+  m.def("box_iou_rotated", &box_iou_rotated);
+  m.def("deform_conv_forward", &deform_conv_forward, "deform_conv_forward");
+  m.def(
+      "deform_conv_backward_input",
+      &deform_conv_backward_input,
+      "deform_conv_backward_input");
+  m.def(
+      "deform_conv_backward_filter",
+      &deform_conv_backward_filter,
+      "deform_conv_backward_filter");
+  m.def(
+      "modulated_deform_conv_forward",
+      &modulated_deform_conv_forward,
+      "modulated_deform_conv_forward");
+  m.def(
+      "modulated_deform_conv_backward",
+      &modulated_deform_conv_backward,
+      "modulated_deform_conv_backward");
+  m.def("ms_deform_attn_forward", &ms_deform_attn_forward, "ms_deform_attn_forward");
+  m.def("ms_deform_attn_backward", &ms_deform_attn_backward, "ms_deform_attn_backward");
+  m.def("iou_box3d", &IoUBox3D, "iou_box3d");
+}