[FEAT] add optional CUDA serialization kernels

.gitignore

@@ -1,2 +1,4 @@
 data
 exp
+build
+__pycache__

README.md

@@ -106,6 +106,13 @@ University of Zurich
 
   </details>
 
+- Optional: Install CUDA kernels for faster point cloud serialization in attention layers (automatic fallback to the original PyTorch implementation). Modify the ```all_cuda_archs``` in ```libs/serialization/setup.py``` (see PointROPE above).
+  ```shell
+  cd libs/serialization
+  python setup.py install
+  cd ../..
+  ```
+
 - Additional requirements. The requirements below are optional, and only required for evaluator and PointGroup instance segmentation. 
   * For evaluator:
     ```
@@ -129,6 +136,7 @@ Your_project/
 └── .../
     └── LitePT/
         ├── litept/
+        ├── libs/serialization
         └── libs/pointrope
 ```
 We also provide a ```demo_use.py``` script that illustrates how to use the standalone LitePT with an example point cloud input.

libs/serialization/cuda/bindings.cpp

@@ -0,0 +1,12 @@
+#include <torch/extension.h>
+
+torch::Tensor hilbert_encode_cuda(torch::Tensor coords, int num_bits);
+
+torch::Tensor morton_encode_cuda(torch::Tensor coords);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("hilbert_encode", &hilbert_encode_cuda, "Hilbert encode (CUDA)");
+
+    // morton (or Z-order) encoding
+    m.def("morton_encode", &morton_encode_cuda, "Morton encode (CUDA)");
+}

libs/serialization/cuda/hilbert.cu

@@ -0,0 +1,102 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+__device__ inline uint64_t hilbert3d_single(uint32_t x, uint32_t y, uint32_t z, int num_bits) {
+    uint32_t X = x;
+    uint32_t Y = y;
+    uint32_t Z = z;
+
+    // --- Skilling transform ---
+    uint32_t Q = 1u << (num_bits - 1);
+    while (Q > 1) {
+        uint32_t P = Q - 1;
+
+        // Dimension 0 (X)
+        if (X & Q) {
+            X ^= P; 
+        }
+
+        // Dimension 1 (Y)
+        if (Y & Q) {
+            X ^= P; // Invert X's lower bits
+        } else {
+            // Swap lower bits of X and Y
+            uint32_t t = (X ^ Y) & P;
+            X ^= t;
+            Y ^= t;
+        }
+
+        // Dimension 2 (Z)
+        if (Z & Q) {
+            X ^= P; // Invert X's lower bits
+        } else {
+            // Swap lower bits of X and Z
+            uint32_t t = (X ^ Z) & P;
+            X ^= t;
+            Z ^= t;
+        }
+
+        Q >>= 1;
+    }
+
+    // --- Interleave bits ---
+    // (This produces the Gray-coded Hilbert integer)
+    uint64_t h = 0;
+    for (int i = num_bits - 1; i >= 0; --i) {
+        h <<= 3;
+        h |= ((uint64_t)((X >> i) & 1) << 2) |
+             ((uint64_t)((Y >> i) & 1) << 1) |
+             ((uint64_t)((Z >> i) & 1));
+    }
+
+    // --- Gray to Binary ---
+    // Skilling's outputs a Gray-coded index. We decode it here on the full interleaved integer.
+    h ^= (h >> 1);
+    h ^= (h >> 2);
+    h ^= (h >> 4);
+    h ^= (h >> 8);
+    h ^= (h >> 16);
+    h ^= (h >> 32);
+
+    return h;
+}
+
+
+__global__ void hilbert_kernel(
+    const uint32_t* __restrict__ coords,
+    uint64_t* __restrict__ out,
+    int N,
+    int num_bits
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= N) return;
+
+    uint32_t x = coords[3 * idx + 0];
+    uint32_t y = coords[3 * idx + 1];
+    uint32_t z = coords[3 * idx + 2];
+
+    out[idx] = hilbert3d_single(x, y, z, num_bits);
+}
+
+torch::Tensor hilbert_encode_cuda(torch::Tensor coords, int num_bits) {
+    TORCH_CHECK(coords.is_cuda(), "coords must be CUDA");
+    TORCH_CHECK(coords.size(1) == 3, "coords must be (N, 3)");
+
+    int N = coords.size(0);
+
+    auto coords_u32 = coords.to(torch::kUInt32).contiguous();
+    auto out = torch::empty({N}, torch::dtype(torch::kUInt64).device(coords.device()));
+
+    int threads = 256;
+    int blocks = (N + threads - 1) / threads;
+
+    hilbert_kernel<<<blocks, threads>>>(
+        coords_u32.data_ptr<uint32_t>(),
+        out.data_ptr<uint64_t>(),
+        N,
+        num_bits
+    );
+
+    return out;
+}

libs/serialization/cuda/morton.cu

@@ -0,0 +1,55 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+__device__ inline uint64_t split_by_3(uint32_t x) {
+    uint64_t v = x & 0x1fffff; // up to 21 bits
+
+    v = (v | (v << 32)) & 0x1f00000000ffff;
+    v = (v | (v << 16)) & 0x1f0000ff0000ff;
+    v = (v | (v << 8))  & 0x100f00f00f00f00f;
+    v = (v | (v << 4))  & 0x10c30c30c30c30c3;
+    v = (v | (v << 2))  & 0x1249249249249249;
+
+    return v;
+}
+
+__device__ inline uint64_t morton3d_single(uint32_t x, uint32_t y, uint32_t z) {
+    return (split_by_3(x) << 2) | (split_by_3(y) << 1) | split_by_3(z);
+}
+
+__global__ void morton_kernel(
+    const uint32_t* __restrict__ coords,
+    uint64_t* __restrict__ out,
+    int N
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= N) return;
+
+    uint32_t x = coords[3 * idx + 0];
+    uint32_t y = coords[3 * idx + 1];
+    uint32_t z = coords[3 * idx + 2];
+
+    out[idx] = morton3d_single(x, y, z);
+}
+
+torch::Tensor morton_encode_cuda(torch::Tensor coords) {
+    TORCH_CHECK(coords.is_cuda(), "coords must be CUDA");
+    TORCH_CHECK(coords.size(1) == 3, "coords must be (N,3)");
+
+    int N = coords.size(0);
+
+    auto coords_u32 = coords.to(torch::kUInt32).contiguous();
+    auto out = torch::empty({N}, torch::dtype(torch::kUInt64).device(coords.device()));
+
+    int threads = 256;
+    int blocks = (N + threads - 1) / threads;
+
+    morton_kernel<<<blocks, threads>>>(
+        coords_u32.data_ptr<uint32_t>(),
+        out.data_ptr<uint64_t>(),
+        N
+    );
+
+    return out;
+}

litept/serialization/default.py → libs/serialization/default.py

@@ -4,20 +4,27 @@
 from .hilbert import encode as hilbert_encode_
 from .hilbert import decode as hilbert_decode_
 
-import numpy as np
-from colorhash import ColorHash
+try:
+    import serialization_cuda  # run `python setup.py install`
+    # print("Has CUDA-accelerated serialization.")
+    HAS_CUDA_EXT = True
+except ImportError:
+    # print("No CUDA-accelerated serialization.")
+    HAS_CUDA_EXT = False
 
+# import numpy as np
+# from colorhash import ColorHash
 
-def int_to_plotly_rgb(x):
-    """Convert 1D torch.Tensor of int into plotly-friendly RGB format.
-    This operation is deterministic on the int values.
-    """
-    assert isinstance(x, torch.Tensor)
-    assert x.dim() == 1
-    assert not x.is_floating_point()
-    x = x.cpu().long().numpy()
-    palette = np.array([ColorHash(i).rgb for i in range(x.max() + 1)])
-    return palette[x]
+# def int_to_plotly_rgb(x):
+#     """Convert 1D torch.Tensor of int into plotly-friendly RGB format.
+#     This operation is deterministic on the int values.
+#     """
+#     assert isinstance(x, torch.Tensor)
+#     assert x.dim() == 1
+#     assert not x.is_floating_point()
+#     x = x.cpu().long().numpy()
+#     palette = np.array([ColorHash(i).rgb for i in range(x.max() + 1)])
+#     return palette[x]
 
 @torch.inference_mode()
 def encode(grid_coord, batch=None, depth=16, order="z"):
@@ -59,10 +66,19 @@ def decode(code, depth=16, order="z"):
 
 
 def z_order_encode(grid_coord: torch.Tensor, depth: int = 16):
-    x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long()
-    # we block the support to batch, maintain batched code in Point class
-    code = z_order_encode_(x, y, z, b=None, depth=depth)
-    return code
+    if HAS_CUDA_EXT:
+        code_cuda = serialization_cuda.morton_encode(grid_coord).long()
+        # basic differential testing of the CUDA implementation
+        # x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long()
+        # # we block the support to batch, maintain batched code in Point class
+        # code_torch = z_order_encode_(x, y, z, b=None, depth=depth)
+        # print("morton equal", torch.all(code_cuda == code_torch))
+        return code_cuda
+    else:
+        x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long()
+        # we block the support to batch, maintain batched code in Point class
+        code = z_order_encode_(x, y, z, b=None, depth=depth)
+        return code
 
 
 def z_order_decode(code: torch.Tensor, depth):
@@ -71,8 +87,16 @@ def z_order_decode(code: torch.Tensor, depth):
     return grid_coord
 
 
+
 def hilbert_encode(grid_coord: torch.Tensor, depth: int = 16):
-    return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth)
+    if HAS_CUDA_EXT:
+        code_cuda = serialization_cuda.hilbert_encode(grid_coord.to(torch.uint32).contiguous(), depth).long()
+        # basic differential testing of the CUDA implementation
+        # code_torch = hilbert_encode_(grid_coord, num_dims=3, num_bits=depth)
+        # print("hilbert equal", torch.all(code_cuda == code_torch))
+        return code_cuda
+    else:
+        return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth)
 
 
 def hilbert_decode(code: torch.Tensor, depth: int = 16):

libs/serialization/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+# compile for all possible CUDA architectures
+# all_cuda_archs = cuda.get_gencode_flags().replace('compute=','arch=').split()
+# alternatively, you can list cuda archs that you want, eg:
+# check https://developer.nvidia.com/cuda-gpus to find your arch
+all_cuda_archs = [
+    # '-gencode', 'arch=compute_90,code=sm_90',
+    # '-gencode', 'arch=compute_75,code=sm_75',
+    # '-gencode', 'arch=compute_80,code=sm_80',
+    '-gencode', 'arch=compute_86,code=sm_86'
+]
+
+setup(
+    name='serialization_cuda',
+    ext_modules=[
+        CUDAExtension(
+            name='serialization_cuda',
+            sources=[
+                'cuda/bindings.cpp',
+                'cuda/hilbert.cu',
+                'cuda/morton.cu',
+            ],
+            extra_compile_args={
+                'cxx': ['-O3'],
+                'nvcc': ['-O3']+all_cuda_archs
+            }
+        )
+    ],
+    cmdclass={'build_ext': BuildExtension}
+)

litept/model.py

@@ -16,7 +16,7 @@
 from timm.layers import DropPath
 
 from libs.pointrope import PointROPE
-from .serialization import encode
+from libs.serialization import encode
 
 
 @torch.no_grad()

models/utils/__init__.py

@@ -5,4 +5,3 @@
     batch2offset,
     off_diagonal,
 )
-from .serialization import encode, decode