Morton/Hilbert fixes (#316)

Cleanup and validation for Morton/Hilbert encoding of ijk values of
gridbatch.

---------

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

Author: blackencino

fvdb/grid_batch.py

@@ -2606,7 +2606,7 @@ def morton(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         else:
             offset = to_Vec3i(offset)
 
-        return self._impl.morton(offset)
+        return JaggedTensor(impl=self._impl.morton(offset))
 
     def morton_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         """
@@ -2628,7 +2628,7 @@ def morton_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         else:
             offset = to_Vec3i(offset)
 
-        return self._impl.morton_zyx(offset)
+        return JaggedTensor(impl=self._impl.morton_zyx(offset))
 
     def hilbert(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         """
@@ -2650,7 +2650,7 @@ def hilbert(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         else:
             offset = to_Vec3i(offset)
 
-        return self._impl.hilbert(offset)
+        return JaggedTensor(impl=self._impl.hilbert(offset))
 
     def hilbert_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         """
@@ -2672,7 +2672,7 @@ def hilbert_zyx(self, offset: NumericMaxRank1 | None = None) -> JaggedTensor:
         else:
             offset = to_Vec3i(offset)
 
-        return self._impl.hilbert_zyx(offset)
+        return JaggedTensor(impl=self._impl.hilbert_zyx(offset))
 
     @property
     def jidx(self) -> torch.Tensor:

src/fvdb/detail/ops/SerializeEncode.cu

@@ -3,186 +3,78 @@
 //
 #include <fvdb/detail/GridBatchImpl.h>
 #include <fvdb/detail/ops/SerializeEncode.h>
-#include <fvdb/detail/utils/AccessorHelpers.cuh>
-#include <fvdb/detail/utils/ForEachCPU.h>
 #include <fvdb/detail/utils/HilbertCode.h>
 #include <fvdb/detail/utils/MortonCode.h>
-#include <fvdb/detail/utils/cuda/ForEachCUDA.cuh>
-#include <fvdb/detail/utils/cuda/ForEachPrivateUse1.cuh>
-
-#include <c10/cuda/CUDAException.h>
+#include <fvdb/detail/utils/SimpleOpHelper.h>
 
 #include <cuda_runtime.h>
 
-#include <vector>
-
 namespace fvdb {
 namespace detail {
 namespace ops {
 
-/// @brief Per-voxel callback which computes the space-filling curve code (Morton or Hilbert) for
-/// each active voxel in a batch of grids
-template <template <typename T, int32_t D> typename TorchAccessor>
-__hostdev__ inline void
-serializeEncodeVoxelCallback(int64_t batchIdx,
-                             int64_t leafIdx,
-                             int64_t voxelIdx,
-                             GridBatchImpl::Accessor gridAccessor,
-                             TorchAccessor<int64_t, 2> outMortonCodes,
-                             const nanovdb::Coord &offset,
-                             int order_type) {
-    const nanovdb::OnIndexGrid *grid = gridAccessor.grid(batchIdx);
-    const typename nanovdb::OnIndexGrid::LeafNodeType &leaf =
-        grid->tree().template getFirstNode<0>()[leafIdx];
-    const int64_t baseOffset = gridAccessor.voxelOffset(batchIdx);
-
-    const nanovdb::Coord &ijk = leaf.offsetToGlobalCoord(voxelIdx);
-    if (leaf.isActive(voxelIdx)) {
-        const int64_t idx = baseOffset + (int64_t)leaf.getValue(voxelIdx) - 1;
+namespace {
 
+template <torch::DeviceType DeviceTag>
+struct Processor : public BaseProcessor<DeviceTag, Processor<DeviceTag>, int64_t> {
+    nanovdb::Coord offset            = nanovdb::Coord{0, 0, 0};
+    SpaceFillingCurveType order_type = SpaceFillingCurveType::ZOrder;
+
+    // Per-voxel callback which computes the space-filling
+    // curve code (Morton or Hilbert) for
+    // each active voxel in a batch of grids
+    __hostdev__ void
+    perActiveVoxel(nanovdb::Coord const &ijk, int64_t const feature_idx, auto out_accessor) const {
         // Apply offset to coordinates
-        int32_t offset_i = offset[0];
-        int32_t offset_j = offset[1];
-        int32_t offset_k = offset[2];
+        auto const i = static_cast<uint32_t>(ijk[0] + offset[0]);
+        auto const j = static_cast<uint32_t>(ijk[1] + offset[1]);
+        auto const k = static_cast<uint32_t>(ijk[2] + offset[2]);
 
         // Compute Morton or Hilbert code with offset to ensure non-negative coordinates
         uint64_t space_filling_code;
-        switch (static_cast<SpaceFillingCurveType>(order_type)) {
-        case SpaceFillingCurveType::ZOrder: // Regular z-order: xyz
-            space_filling_code = utils::morton_with_offset(ijk[0],
-                                                           ijk[1],
-                                                           ijk[2],
-                                                           static_cast<uint32_t>(offset_i),
-                                                           static_cast<uint32_t>(offset_j),
-                                                           static_cast<uint32_t>(offset_k));
+        switch (order_type) {
+        case SpaceFillingCurveType::ZOrder:            // Regular z-order: xyz
+            space_filling_code = utils::morton(i, j, k);
             break;
-        case SpaceFillingCurveType::ZOrderTransposed: // Transposed z-order: zyx
-            space_filling_code = utils::morton_with_offset(ijk[2],
-                                                           ijk[1],
-                                                           ijk[0],
-                                                           static_cast<uint32_t>(offset_k),
-                                                           static_cast<uint32_t>(offset_j),
-                                                           static_cast<uint32_t>(offset_i));
+        case SpaceFillingCurveType::ZOrderTransposed:  // Transposed z-order: zyx
+            space_filling_code = utils::morton(k, j, i);
             break;
-        case SpaceFillingCurveType::Hilbert: // Regular Hilbert curve: xyz
-            space_filling_code = utils::hilbert_with_offset(ijk[0],
-                                                            ijk[1],
-                                                            ijk[2],
-                                                            static_cast<uint32_t>(offset_i),
-                                                            static_cast<uint32_t>(offset_j),
-                                                            static_cast<uint32_t>(offset_k));
+        case SpaceFillingCurveType::Hilbert:           // Regular Hilbert curve: xyz
+            space_filling_code = utils::hilbert(i, j, k);
             break;
         case SpaceFillingCurveType::HilbertTransposed: // Transposed Hilbert curve: zyx
-            space_filling_code = utils::hilbert_with_offset(ijk[2],
-                                                            ijk[1],
-                                                            ijk[0],
-                                                            static_cast<uint32_t>(offset_k),
-                                                            static_cast<uint32_t>(offset_j),
-                                                            static_cast<uint32_t>(offset_i));
+            space_filling_code = utils::hilbert(k, j, i);
             break;
         default:
             // Invalid order type - use assert for device code
             space_filling_code = 0;
             break;
         }
 
-        outMortonCodes[idx][0] = static_cast<int64_t>(space_filling_code);
+        out_accessor[feature_idx] = static_cast<int64_t>(space_filling_code);
     }
-}
+};
 
-/// @brief Get the space-filling curve codes for active voxels in a batch of grids
-/// @param gridBatch The batch of grids
-/// @param outMortonCodes Tensor which will contain the output space-filling curve codes
-/// @param offset Offset to apply to voxel coordinates before encoding
-/// @param order_type Integer representing the order type (0=z, 1=z-trans, 2=hilbert,
-/// 3=hilbert-trans)
-template <torch::DeviceType DeviceTag>
-void
-GetSerializeEncode(const GridBatchImpl &gridBatch,
-                   torch::Tensor &outMortonCodes,
-                   const nanovdb::Coord &offset,
-                   int order_type) {
-    auto outCodesAcc = tensorAccessor<DeviceTag, int64_t, 2>(outMortonCodes);
+} // End anonymous namespace
 
-    if constexpr (DeviceTag == torch::kCUDA) {
-        auto cb = [=] __device__(int64_t batchIdx,
-                                 int64_t leafIdx,
-                                 int64_t voxelIdx,
-                                 int64_t,
-                                 GridBatchImpl::Accessor gridAccessor) {
-            serializeEncodeVoxelCallback<TorchRAcc32>(
-                batchIdx, leafIdx, voxelIdx, gridAccessor, outCodesAcc, offset, order_type);
-        };
-        forEachVoxelCUDA(1024, 1, gridBatch, cb);
-    } else if constexpr (DeviceTag == torch::kPrivateUse1) {
-        auto cb = [=] __device__(int64_t batchIdx,
-                                 int64_t leafIdx,
-                                 int64_t voxelIdx,
-                                 int64_t,
-                                 GridBatchImpl::Accessor gridAccessor) {
-            serializeEncodeVoxelCallback<TorchRAcc32>(
-                batchIdx, leafIdx, voxelIdx, gridAccessor, outCodesAcc, offset, order_type);
-        };
-        forEachVoxelPrivateUse1(1, gridBatch, cb);
-    } else {
-        auto cb = [=](int64_t batchIdx,
-                      int64_t leafIdx,
-                      int64_t voxelIdx,
-                      int64_t,
-                      GridBatchImpl::Accessor gridAccessor) {
-            serializeEncodeVoxelCallback<TorchAcc>(
-                batchIdx, leafIdx, voxelIdx, gridAccessor, outCodesAcc, offset, order_type);
-        };
-        forEachVoxelCPU(1, gridBatch, cb);
-    }
-}
-
-/// @brief Get the space-filling curve codes for active voxels in a batch of grids
-/// @tparam DeviceTag Which device to run on
-/// @param gridBatch The batch of grids to get the space-filling curve codes for
-/// @param order_type The type of space-filling curve to use for encoding
-/// @param offset Offset to apply to voxel coordinates before encoding
-/// @return A JaggedTensor of shape [B, -1, 1] of space-filling curve codes for active voxels
 template <torch::DeviceType DeviceTag>
 JaggedTensor
-SerializeEncode(const GridBatchImpl &gridBatch,
-                SpaceFillingCurveType order_type,
-                const nanovdb::Coord &offset) {
-    gridBatch.checkNonEmptyGrid();
-    auto opts = torch::TensorOptions().dtype(torch::kInt64).device(gridBatch.device());
-    torch::Tensor outMortonCodes = torch::empty({gridBatch.totalVoxels(), 1}, opts);
-
-    // Convert enum to integer for kernel
-    const int order_type_int = static_cast<int>(order_type);
-
-    GetSerializeEncode<DeviceTag>(gridBatch, outMortonCodes, offset, order_type_int);
-
-    return gridBatch.jaggedTensor(outMortonCodes);
-}
-
-template <>
-JaggedTensor
-dispatchSerializeEncode<torch::kCUDA>(const GridBatchImpl &gridBatch,
-                                      SpaceFillingCurveType order_type,
-                                      const nanovdb::Coord &offset) {
-    return SerializeEncode<torch::kCUDA>(gridBatch, order_type, offset);
+dispatchSerializeEncode(GridBatchImpl const &gridBatch,
+                        SpaceFillingCurveType order_type,
+                        nanovdb::Coord const &offset) {
+    Processor<DeviceTag> processor{.offset = offset, .order_type = order_type};
+    return processor.execute(gridBatch);
 }
 
-template <>
-JaggedTensor
-dispatchSerializeEncode<torch::kCPU>(const GridBatchImpl &gridBatch,
-                                     SpaceFillingCurveType order_type,
-                                     const nanovdb::Coord &offset) {
-    return SerializeEncode<torch::kCPU>(gridBatch, order_type, offset);
-}
-
-template <>
-JaggedTensor
-dispatchSerializeEncode<torch::kPrivateUse1>(const GridBatchImpl &gridBatch,
-                                             SpaceFillingCurveType order_type,
-                                             const nanovdb::Coord &offset) {
-    return SerializeEncode<torch::kPrivateUse1>(gridBatch, order_type, offset);
-}
+template JaggedTensor dispatchSerializeEncode<torch::kCUDA>(GridBatchImpl const &,
+                                                            SpaceFillingCurveType,
+                                                            nanovdb::Coord const &);
+template JaggedTensor dispatchSerializeEncode<torch::kCPU>(GridBatchImpl const &,
+                                                           SpaceFillingCurveType,
+                                                           nanovdb::Coord const &);
+template JaggedTensor dispatchSerializeEncode<torch::kPrivateUse1>(GridBatchImpl const &,
+                                                                   SpaceFillingCurveType,
+                                                                   nanovdb::Coord const &);
 
 } // namespace ops
 } // namespace detail

src/fvdb/detail/ops/SerializeEncode.h

@@ -16,10 +16,16 @@ namespace fvdb {
 namespace detail {
 namespace ops {
 
+/// @brief Get the space-filling curve codes for active voxels in a batch of grids
+/// @tparam DeviceTag Which device to run on
+/// @param gridBatch The batch of grids to get the space-filling curve codes for
+/// @param order_type The type of space-filling curve to use for encoding
+/// @param offset Offset to apply to voxel coordinates before encoding
+/// @return A JaggedTensor of shape [B, -1, 1] of space-filling curve codes for active voxels
 template <torch::DeviceType>
-JaggedTensor dispatchSerializeEncode(const GridBatchImpl &gridBatch,
+JaggedTensor dispatchSerializeEncode(GridBatchImpl const &gridBatch,
                                      SpaceFillingCurveType order_type,
-                                     const nanovdb::Coord &offset);
+                                     nanovdb::Coord const &offset);
 
 } // namespace ops
 } // namespace detail

src/fvdb/detail/utils/HilbertCode.h

@@ -97,8 +97,13 @@ gray_stream_to_uint64(uint32_t g0, uint32_t g1, uint32_t g2) {
     return value;
 }
 
+/// @brief Compute Hilbert curve index for 3D coordinates
+/// @param i I coordinate (must be 0 <= i < 2^21)
+/// @param j J coordinate (must be 0 <= j < 2^21)
+/// @param k K coordinate (must be 0 <= k < 2^21)
+/// @return 63-bit Hilbert index
 __hostdev__ inline uint64_t
-hilbert_order21_index(uint32_t i, uint32_t j, uint32_t k) {
+hilbert(uint32_t i, uint32_t j, uint32_t k) {
     // encoder: (i, j, k) -> uint64_t code.
 
     i = i & _HILBERT_MASK;
@@ -117,29 +122,6 @@ hilbert_order21_index(uint32_t i, uint32_t j, uint32_t k) {
     return gray_stream_to_uint64(i, j, k);
 }
 
-/// @brief Compute Hilbert curve index for 3D coordinates with offset handling
-/// @param i I coordinate (can be negative)
-/// @param j J coordinate (can be negative)
-/// @param k K coordinate (can be negative)
-/// @param offset_i Offset to make i non-negative
-/// @param offset_j Offset to make j non-negative
-/// @param offset_k Offset to make k non-negative
-/// @return 63-bit Hilbert index
-__hostdev__ inline uint64_t
-hilbert_with_offset(int32_t const i,
-                    int32_t const j,
-                    int32_t const k,
-                    uint32_t const offset_i,
-                    uint32_t const offset_j,
-                    uint32_t const offset_k) {
-    // Add offsets to ensure non-negative coordinates
-    uint32_t x = static_cast<uint32_t>(i + static_cast<int32_t>(offset_i));
-    uint32_t y = static_cast<uint32_t>(j + static_cast<int32_t>(offset_j));
-    uint32_t z = static_cast<uint32_t>(k + static_cast<int32_t>(offset_k));
-
-    return hilbert_order21_index(x, y, z);
-}
-
 } // namespace utils
 } // namespace detail
 } // namespace fvdb