Implementing morton and hilbert for Grid and GridBatch ijk (#311)

This PR is the work of Hexu Zhao.

Functions have been added to Grid and GridBatch which produce a uint64_t
tensor of the morton or hilbert encodings of the ijk or kji coordinate
orderings.

Signed-off-by: Christopher Horvath <chorvath@nvidia.com>

Author: blackencino

fvdb/_Cpp.pyi

@@ -356,6 +356,10 @@ class GridBatch:
     def inject_to(self, dst_grid: GridBatch, src: JaggedTensor, dst: JaggedTensor) -> None: ...
     def dual_bbox_at(self, arg0: int) -> torch.Tensor: ...
     def dual_grid(self, exclude_border: bool = ...) -> GridBatch: ...
+    def hilbert(self, offset: torch.Tensor) -> JaggedTensor: ...
+    def hilbert_zyx(self, offset: torch.Tensor) -> JaggedTensor: ...
+    def morton(self, offset: torch.Tensor) -> JaggedTensor: ...
+    def morton_zyx(self, offset: torch.Tensor) -> JaggedTensor: ...
     def grid_to_world(self, ijk: JaggedTensor) -> JaggedTensor: ...
     def ijk_to_index(self, ijk: JaggedTensor, cumulative: bool = False) -> JaggedTensor: ...
     def ijk_to_inv_index(self, ijk: JaggedTensor, cumulative: bool = False) -> JaggedTensor: ...

fvdb/grid.py

@@ -2296,6 +2296,83 @@ def ijk(self) -> torch.Tensor:
         """
         return self._impl.ijk.jdata
 
+    def morton(self, offset: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Return Morton codes (Z-order curve) for active voxels in this grid.
+
+        Morton codes use xyz bit interleaving to create a space-filling curve that
+        preserves spatial locality. This is useful for serialization, sorting, and
+        spatial data structures.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the Morton codes for each active voxel.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk, dim=0).values
+        return self._impl.morton(offset).jdata
+
+    def morton_zyx(self, offset: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Return transposed Morton codes (Z-order curve) for active voxels in this grid.
+
+        Transposed Morton codes use zyx bit interleaving to create a space-filling curve.
+        This variant can provide better spatial locality for certain access patterns.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the transposed Morton codes for each active voxel.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk, dim=0).values
+        return self._impl.morton_zyx(offset).jdata
+
+    def hilbert(self, offset: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Return Hilbert curve codes for active voxels in this grid.
+
+        Hilbert curves provide better spatial locality than Morton codes by ensuring
+        that nearby points in 3D space are also nearby in the 1D curve ordering.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the Hilbert codes for each active voxel.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk, dim=0).values
+        return self._impl.hilbert(offset).jdata
+
+    def hilbert_zyx(self, offset: torch.Tensor | None = None) -> torch.Tensor:
+        """
+        Return transposed Hilbert curve codes for active voxels in this grid.
+
+        Transposed Hilbert curves use zyx ordering instead of xyz. This variant can
+        provide better spatial locality for certain access patterns.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the transposed Hilbert codes for each active voxel.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk, dim=0).values
+        return self._impl.hilbert_zyx(offset).jdata
+
     @property
     def num_bytes(self) -> int:
         """

fvdb/grid_batch.py

@@ -2585,6 +2585,95 @@ def ijk(self) -> JaggedTensor:
         """
         return JaggedTensor(impl=self._impl.ijk)
 
+    def morton(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
+        """
+        Return Morton codes (Z-order curve) for active voxels in this grid batch.
+
+        Morton codes use xyz bit interleaving to create a space-filling curve that
+        preserves spatial locality. This is useful for serialization, sorting, and
+        spatial data structures.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the Morton codes for each active voxel in the batch.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk.jdata, dim=0).values
+        else:
+            offset = to_Vec3i(offset)
+
+        return self._impl.morton(offset)
+
+    def morton_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
+        """
+        Return transposed Morton codes (Z-order curve) for active voxels in this grid batch.
+
+        Transposed Morton codes use zyx bit interleaving to create a space-filling curve.
+        This variant can provide better spatial locality for certain access patterns.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the transposed Morton codes for each active voxel in the batch.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk.jdata, dim=0).values
+        else:
+            offset = to_Vec3i(offset)
+
+        return self._impl.morton_zyx(offset)
+
+    def hilbert(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
+        """
+        Return Hilbert curve codes for active voxels in this grid batch.
+
+        Hilbert curves provide better spatial locality than Morton codes by ensuring
+        that nearby points in 3D space are also nearby in the 1D curve ordering.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the Hilbert codes for each active voxel in the batch.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk.jdata, dim=0).values
+        else:
+            offset = to_Vec3i(offset)
+
+        return self._impl.hilbert(offset)
+
+    def hilbert_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
+        """
+        Return transposed Hilbert curve codes for active voxels in this grid batch.
+
+        Transposed Hilbert curves use zyx ordering instead of xyz. This variant can
+        provide better spatial locality for certain access patterns.
+
+        Args:
+            offset: Optional offset to apply to voxel coordinates before encoding.
+                If None, uses the negative minimum coordinate across all voxels.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the transposed Hilbert codes for each active voxel in the batch.
+        """
+        if offset is None:
+            offset = -torch.min(self.ijk.jdata, dim=0).values
+        else:
+            offset = to_Vec3i(offset)
+
+        return self._impl.hilbert_zyx(offset)
+
     @property
     def jidx(self) -> torch.Tensor:
         """

src/CMakeLists.txt

@@ -124,6 +124,7 @@ set(FVDB_CU_FILES
     fvdb/detail/ops/PointsInGrid.cu
     fvdb/detail/ops/IjkToIndex.cu
     fvdb/detail/ops/ActiveGridGoords.cu
+    fvdb/detail/ops/SerializeEncode.cu
     fvdb/detail/ops/IjkToInvIndex.cu
     fvdb/detail/ops/ReadFromDense.cu
     fvdb/detail/ops/NearestIjkForPoints.cu

src/fvdb/GridBatch.cpp

@@ -31,15 +31,20 @@
 #include <fvdb/detail/ops/RayImplicitIntersection.h>
 #include <fvdb/detail/ops/SampleRaysUniform.h>
 #include <fvdb/detail/ops/SegmentsAlongRays.h>
+#include <fvdb/detail/ops/SerializeEncode.h>
 #include <fvdb/detail/ops/VoxelNeighborhood.h>
 #include <fvdb/detail/ops/VoxelsAlongRays.h>
 #include <fvdb/detail/ops/convolution/pack_info/BrickHaloBuffer.h>
 #include <fvdb/detail/ops/convolution/pack_info/ConvolutionKernelMap.h>
 #include <fvdb/detail/ops/convolution/pack_info/IGEMMBitOperations.h>
 #include <fvdb/detail/utils/Utils.h>
+#include <fvdb/detail/utils/nanovdb/TorchNanoConversions.h>
 
 #include <torch/types.h>
 
+#include <tuple>
+#include <vector>
+
 namespace fvdb {
 
 GridBatch::GridBatch() : mImpl() {}
@@ -1080,6 +1085,46 @@ GridBatch::ijk() const {
     });
 }
 
+JaggedTensor
+GridBatch::morton(const torch::Tensor &offset) const {
+    c10::DeviceGuard guard(device());
+    const nanovdb::Coord offsetCoord = tensorToCoord(offset);
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::ZOrder, offsetCoord);
+    });
+}
+
+JaggedTensor
+GridBatch::morton_zyx(const torch::Tensor &offset) const {
+    c10::DeviceGuard guard(device());
+    const nanovdb::Coord offsetCoord = tensorToCoord(offset);
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::ZOrderTransposed, offsetCoord);
+    });
+}
+
+JaggedTensor
+GridBatch::hilbert(const torch::Tensor &offset) const {
+    c10::DeviceGuard guard(device());
+    const nanovdb::Coord offsetCoord = tensorToCoord(offset);
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::Hilbert, offsetCoord);
+    });
+}
+
+JaggedTensor
+GridBatch::hilbert_zyx(const torch::Tensor &offset) const {
+    c10::DeviceGuard guard(device());
+    const nanovdb::Coord offsetCoord = tensorToCoord(offset);
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::HilbertTransposed, offsetCoord);
+    });
+}
+
 std::vector<JaggedTensor>
 GridBatch::viz_edge_network(bool returnVoxelCoordinates) const {
     c10::DeviceGuard guard(device());

src/fvdb/GridBatch.h

@@ -472,6 +472,28 @@ struct GridBatch : torch::CustomClassHolder {
     /// @return A JaggedTensor of voxel coordinates indexed by this grid batch (shape [B, -1, 3])
     JaggedTensor ijk() const;
 
+    /// @brief Return Morton codes (Z-order curve) for active voxels in this grid batch (xyz bit
+    /// interleaving)
+    /// @param offset Offset to apply to voxel coordinates before encoding
+    /// @return A JaggedTensor of Morton codes for active voxels (shape [B, -1, 1])
+    JaggedTensor morton(const torch::Tensor &offset) const;
+
+    /// @brief Return transposed Morton codes (Z-order curve) for active voxels in this grid batch
+    /// (zyx bit interleaving)
+    /// @param offset Offset to apply to voxel coordinates before encoding
+    /// @return A JaggedTensor of transposed Morton codes for active voxels (shape [B, -1, 1])
+    JaggedTensor morton_zyx(const torch::Tensor &offset) const;
+
+    /// @brief Return Hilbert curve codes for active voxels in this grid batch (xyz)
+    /// @param offset Offset to apply to voxel coordinates before encoding
+    /// @return A JaggedTensor of Hilbert codes for active voxels (shape [B, -1, 1])
+    JaggedTensor hilbert(const torch::Tensor &offset) const;
+
+    /// @brief Return transposed Hilbert curve codes for active voxels in this grid batch (zyx)
+    /// @param offset Offset to apply to voxel coordinates before encoding
+    /// @return A JaggedTensor of transposed Hilbert codes for active voxels (shape [B, -1, 1])
+    JaggedTensor hilbert_zyx(const torch::Tensor &offset) const;
+
     /// @brief Find the intersection between a collection of rays and the zero level set of a scalar
     /// field
     ///        at each voxel in the grid batch

src/fvdb/Types.h

@@ -8,6 +8,14 @@
 
 namespace fvdb {
 
+/// @brief Enum class for space-filling curve encoding types
+enum class SpaceFillingCurveType {
+    ZOrder,           ///< Regular z-order curve (xyz)
+    ZOrderTransposed, ///< Transposed z-order curve (zyx)
+    Hilbert,          ///< Regular Hilbert curve (xyz)
+    HilbertTransposed ///< Transposed Hilbert curve (zyx)
+};
+
 // These are union types that can be constructed from nanovdb types, torch tensors, std::vectors,
 // single scalars, etc... They are used to allow the user to pass in a variety of types to the API,
 // and then convert them to the correct type