Morton Order for supporting point transformer v3

fvdb/_Cpp.pyi

@@ -341,6 +341,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 encode_hilbert(self) -> JaggedTensor: ...
+    def encode_hilbert_zyx(self) -> JaggedTensor: ...
+    def encode_morton(self) -> JaggedTensor: ...
+    def encode_morton_zyx(self) -> 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

@@ -1835,6 +1835,124 @@ def has_zero_voxels(self) -> bool:
     def ijk(self) -> torch.Tensor:
         return self._impl.ijk.jdata
 
+    def encode_morton(self) -> 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.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the Morton codes for each active voxel.
+        """
+        return self._impl.encode_morton().jdata
+
+    def encode_morton_zyx(self) -> 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.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the transposed Morton codes for each active voxel.
+        """
+        return self._impl.encode_morton_zyx().jdata
+
+    def encode_hilbert(self) -> 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.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the Hilbert codes for each active voxel.
+        """
+        return self._impl.encode_hilbert().jdata
+
+    def encode_hilbert_zyx(self) -> 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.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the transposed Hilbert codes for each active voxel.
+        """
+        return self._impl.encode_hilbert_zyx().jdata
+
+    def permute(self, curve_codes: torch.Tensor) -> torch.Tensor:
+        """
+        Get permutation indices to sort voxels by spatial order.
+
+        This method takes space-filling curve codes for all active voxels and returns the indices
+        that would sort them according to the curve ordering. This is useful for
+        spatially coherent data access patterns and cache optimization.
+
+        Args:
+            curve_codes (torch.Tensor): Curve codes for each voxel.
+                Shape: (num_voxels, 1).
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the permutation indices. Use these indices to reorder voxel data for spatial coherence.
+        """
+        # Get the curve codes as a flat tensor
+        curve_data = curve_codes.squeeze(-1)  # Shape: [num_voxels]
+
+        # Sort and get indices
+        _, indices = torch.sort(curve_data, dim=0)
+
+        # Return indices as a tensor with shape [num_voxels, 1]
+        return indices.unsqueeze(-1)
+
+    def permutation_morton(self) -> torch.Tensor:
+        """
+        Return permutation indices to sort voxels by Morton curve order.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the permutation indices for Morton (Z-order) curve ordering.
+        """
+        return self.permute(self.encode_morton())
+
+    def permutation_morton_zyx(self) -> torch.Tensor:
+        """
+        Return permutation indices to sort voxels by transposed Morton curve order.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the permutation indices for transposed Morton (Z-order) curve ordering.
+        """
+        return self.permute(self.encode_morton_zyx())
+
+    def permutation_hilbert(self) -> torch.Tensor:
+        """
+        Return permutation indices to sort voxels by Hilbert curve order.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the permutation indices for Hilbert curve ordering.
+        """
+        return self.permute(self.encode_hilbert())
+
+    def permutation_hilbert_zyx(self) -> torch.Tensor:
+        """
+        Return permutation indices to sort voxels by transposed Hilbert curve order.
+
+        Returns:
+            torch.Tensor: A tensor of shape `[num_voxels, 1]` containing
+                the permutation indices for transposed Hilbert curve ordering.
+        """
+        return self.permute(self.encode_hilbert_zyx())
+
     @property
     def num_bytes(self) -> int:
         return self._impl.total_bytes

fvdb/grid_batch.py

@@ -1956,6 +1956,139 @@ def has_zero_grids(self) -> bool:
     def ijk(self) -> JaggedTensor:
         return self._impl.ijk
 
+    def encode_morton(self) -> 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.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the Morton codes for each active voxel in the batch.
+        """
+        return self._impl.encode_morton()
+
+    def encode_morton_zyx(self) -> 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.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the transposed Morton codes for each active voxel in the batch.
+        """
+        return self._impl.encode_morton_zyx()
+
+    def encode_hilbert(self) -> 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.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the Hilbert codes for each active voxel in the batch.
+        """
+        return self._impl.encode_hilbert()
+
+    def encode_hilbert_zyx(self) -> 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.
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the transposed Hilbert codes for each active voxel in the batch.
+        """
+        return self._impl.encode_hilbert_zyx()
+
+    def permute(self, curve_codes: JaggedTensor) -> JaggedTensor:
+        """
+        Get permutation indices to sort voxels by spatial order.
+
+        This method computes space-filling curve codes for all active voxels and returns the indices
+        that would sort them according to the specified ordering. This is useful for
+        spatially coherent data access patterns and cache optimization.
+
+        Args:
+            order_type (str): The type of spatial ordering to use:
+                - "z": Regular Z-order curve (xyz bit interleaving, default)
+                - "z-trans": Transposed Z-order curve (zyx bit interleaving)
+                - "hilbert": Regular Hilbert curve (xyz)
+                - "hilbert-trans": Transposed Hilbert curve (zyx)
+
+        Returns:
+            JaggedTensor: A JaggedTensor of shape `[num_grids, -1, 1]` containing
+                the permutation indices. Use these indices to reorder voxel data for spatial coherence.
+
+        Example:
+            >>> z_indices = grid_batch.permute("z")  # Regular xyz z-order
+            >>> z_trans_indices = grid_batch.permute("z-trans")  # Transposed zyx z-order
+            >>> hilbert_indices = grid_batch.permute("hilbert")  # Regular xyz Hilbert curve
+            >>> hilbert_trans_indices = grid_batch.permute("hilbert-trans")  # Transposed zyx Hilbert curve
+            >>> # Use indices to reorder some voxel data
+            >>> reordered_data = voxel_data.jdata[z_indices.jdata.squeeze(-1)]
+        """
+
+        # Get the curve codes as a flat tensor
+        curve_data = curve_codes.jdata.squeeze(-1)  # Shape: [total_voxels]
+
+        # Create output tensor for permutation indices
+        # permutation_indices = torch.empty_like(curve_data, dtype=torch.long)
+        permutation_indices = torch.empty_like(curve_data, dtype=torch.long)
+
+        # Sort curve codes and get permutation indices for each grid
+        offset = 0
+        for grid_idx in range(self.grid_count):
+            num_voxels = self.num_voxels_at(grid_idx)
+            if num_voxels == 0:
+                continue
+
+            # Extract curve codes for this grid
+            grid_curve_codes = curve_data[offset : offset + num_voxels]
+
+            # Sort and get indices
+            _, indices = torch.sort(grid_curve_codes, dim=0)
+
+            # Store indices with offset
+            permutation_indices[offset : offset + num_voxels] = indices + offset
+
+            offset += num_voxels
+
+        # Return as JaggedTensor with the same structure as the input
+        return self.jagged_like(permutation_indices.unsqueeze(-1))
+
+    def permutation_morton(self) -> JaggedTensor:
+        """
+        Return permutation indices to sort voxels by Morton curve order.
+        """
+        return self.permute(self.encode_morton())
+
+    def permutation_morton_zyx(self) -> JaggedTensor:
+        """
+        Return permutation indices to sort voxels by transposed Morton curve order.
+        """
+        return self.permute(self.encode_morton_zyx())
+
+    def permutation_hilbert(self) -> JaggedTensor:
+        """
+        Return permutation indices to sort voxels by Hilbert curve order.
+        """
+        return self.permute(self.encode_hilbert())
+
+    def permutation_hilbert_zyx(self) -> JaggedTensor:
+        """
+        Return permutation indices to sort voxels by transposed Hilbert curve order.
+        """
+        return self.permute(self.encode_hilbert_zyx())
+
     @property
     def jidx(self) -> torch.Tensor:
         if self.has_zero_grids:

src/CMakeLists.txt

@@ -118,6 +118,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,6 +31,7 @@
 #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>
@@ -40,6 +41,9 @@
 
 #include <torch/types.h>
 
+#include <tuple>
+#include <vector>
+
 namespace fvdb {
 
 GridBatch::GridBatch() : mImpl() {}
@@ -1081,6 +1085,42 @@ GridBatch::ijk() const {
     });
 }
 
+JaggedTensor
+GridBatch::encode_morton() const {
+    c10::DeviceGuard guard(device());
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(*mImpl,
+                                                                     SpaceFillingCurveType::ZOrder);
+    });
+}
+
+JaggedTensor
+GridBatch::encode_morton_zyx() const {
+    c10::DeviceGuard guard(device());
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::ZOrderTransposed);
+    });
+}
+
+JaggedTensor
+GridBatch::encode_hilbert() const {
+    c10::DeviceGuard guard(device());
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::Hilbert);
+    });
+}
+
+JaggedTensor
+GridBatch::encode_hilbert_zyx() const {
+    c10::DeviceGuard guard(device());
+    return FVDB_DISPATCH_KERNEL(this->device(), [&]() {
+        return fvdb::detail::ops::dispatchSerializeEncode<DeviceTag>(
+            *mImpl, SpaceFillingCurveType::HilbertTransposed);
+    });
+}
+
 std::vector<JaggedTensor>
 GridBatch::viz_edge_network(bool returnVoxelCoordinates) const {
     c10::DeviceGuard guard(device());

src/fvdb/GridBatch.h

@@ -472,6 +472,24 @@ 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)
+    /// @return A JaggedTensor of Morton codes for active voxels (shape [B, -1, 1])
+    JaggedTensor encode_morton() const;
+
+    /// @brief Return transposed Morton codes (Z-order curve) for active voxels in this grid batch
+    /// (zyx bit interleaving)
+    /// @return A JaggedTensor of transposed Morton codes for active voxels (shape [B, -1, 1])
+    JaggedTensor encode_morton_zyx() const;
+
+    /// @brief Return Hilbert curve codes for active voxels in this grid batch (xyz)
+    /// @return A JaggedTensor of Hilbert codes for active voxels (shape [B, -1, 1])
+    JaggedTensor encode_hilbert() const;
+
+    /// @brief Return transposed Hilbert curve codes for active voxels in this grid batch (zyx)
+    /// @return A JaggedTensor of transposed Hilbert codes for active voxels (shape [B, -1, 1])
+    JaggedTensor encode_hilbert_zyx() 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