From a8e572a20251c6a49bbb0faec6f00dc4ade8fec9 Mon Sep 17 00:00:00 2001 From: Francis Williams Date: Wed, 25 Mar 2026 12:51:36 -0700 Subject: [PATCH 1/2] custom ray marcher stub Signed-off-by: Francis Williams --- fvdb/_fvdb_cpp.pyi | 3 + fvdb/grid.py | 11 ++++ src/CMakeLists.txt | 1 + src/fvdb/GridBatch.cpp | 8 +++ src/fvdb/GridBatch.h | 7 +++ src/fvdb/detail/ops/AcousticRaySamples.cu | 70 +++++++++++++++++++++++ src/fvdb/detail/ops/AcousticRaySamples.h | 30 ++++++++++ src/python/GridBatchBinding.cpp | 7 +++ 8 files changed, 137 insertions(+) create mode 100644 src/fvdb/detail/ops/AcousticRaySamples.cu create mode 100644 src/fvdb/detail/ops/AcousticRaySamples.h diff --git a/fvdb/_fvdb_cpp.pyi b/fvdb/_fvdb_cpp.pyi index 99d54e6d..3687afb7 100644 --- a/fvdb/_fvdb_cpp.pyi +++ b/fvdb/_fvdb_cpp.pyi @@ -447,6 +447,9 @@ class GridBatch: def clipped_grid(self, ijk_min: Vec3iBatch, ijk_max: Vec3iBatch) -> GridBatch: ... def coarsened_grid(self, coarsening_factor: Vec3iOrScalar) -> GridBatch: ... def contiguous(self) -> GridBatch: ... + def acoustic_ray_samples( + self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, sound_speeds: JaggedTensor, step_size: float + ) -> JaggedTensor: ... def integrate_tsdf( self, voxel_truncation_distance: float, diff --git a/fvdb/grid.py b/fvdb/grid.py index 6fbfbf47..de537941 100644 --- a/fvdb/grid.py +++ b/fvdb/grid.py @@ -2432,3 +2432,14 @@ def world_to_voxel_matrix(self) -> torch.Tensor: def _gridbatch(self): # Access underlying GridBatchCpp - use sparingly during migration return self._impl + + def acoustic_ray_samples( + self, ray_origins: JaggedTensor, ray_directions: JaggedTensor, sound_speeds: JaggedTensor, step_size: float + ) -> JaggedTensor: + """ + Generate samples along each ray with a given step size where samples are taken + by bending the ray according to Snell's law. + """ + return JaggedTensor( + impl=self._impl.acoustic_ray_samples(ray_origins._impl, ray_directions._impl, sound_speeds._impl, step_size) + ) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5bfa17de..56752bbb 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -75,6 +75,7 @@ set(FVDB_CPP_FILES set(FVDB_CU_FILES fvdb/detail/GridBatchImpl.cu + fvdb/detail/ops/AcousticRaySamples.cu fvdb/detail/ops/ActiveGridGoords.cu fvdb/detail/ops/ActiveVoxelsInBoundsMask.cu fvdb/detail/ops/BuildCoarseGridFromFine.cu diff --git a/src/fvdb/GridBatch.cpp b/src/fvdb/GridBatch.cpp index 2dd5e779..429f39c3 100644 --- a/src/fvdb/GridBatch.cpp +++ b/src/fvdb/GridBatch.cpp @@ -864,6 +864,14 @@ GridBatch::voxels_along_rays(const JaggedTensor &ray_origins, *mImpl, ray_origins, ray_directions, max_vox, eps, return_ijk, cumulative); } +std::vector +GridBatch::acoustic_ray_samples(const JaggedTensor &ray_origins, + const JaggedTensor &ray_directions, + const JaggedTensor &sound_speeds, + float step_size) const { + return fvdb::detail::ops::acousticRaySamples(*mImpl, ray_origins, ray_directions, sound_speeds, step_size); +} + JaggedTensor GridBatch::segments_along_rays(const JaggedTensor &ray_origins, const JaggedTensor &ray_directions, diff --git a/src/fvdb/GridBatch.h b/src/fvdb/GridBatch.h index e79d316c..f14c91ff 100644 --- a/src/fvdb/GridBatch.h +++ b/src/fvdb/GridBatch.h @@ -305,6 +305,13 @@ struct GridBatch : torch::CustomClassHolder { const std::optional &min_coord = std::nullopt, const std::optional &grid_size = std::nullopt) const; + + std::vector acoustic_ray_samples(const JaggedTensor &ray_origins, + const JaggedTensor &ray_directions, + const JaggedTensor &sound_speeds, + float step_size) const; + + /// @brief Convert grid coordinates to world coordinates /// @param ijk A JaggedTensor of grid coordinates with shape [B, -1, 3] (one point set per grid /// in the batch) diff --git a/src/fvdb/detail/ops/AcousticRaySamples.cu b/src/fvdb/detail/ops/AcousticRaySamples.cu new file mode 100644 index 00000000..efc4dd0b --- /dev/null +++ b/src/fvdb/detail/ops/AcousticRaySamples.cu @@ -0,0 +1,70 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#include +#include +#include +#include +#include + +#include +#include + +namespace fvdb { +namespace detail { +namespace ops { +namespace { + +template +__global__ void countAcousticRaySamplesKernel( + BatchGridAccessor batchAccessor, + const JaggedAccessor rayOrigins, + const JaggedAccessor rayDirections, + const JaggedAccessor soundSpeeds, + float stepSize, + JaggedAccessor outCounts) { + + // Kernel parallelizes over rays. + // For each ray, we count the number of samples we will take. + } + +} // anonymous namespace + +template +__global__ void accousticRaySamplesKernel( + BatchGridAccessor batchAccessor, + const JaggedAccessor rayOrigins, + const JaggedAccessor rayDirections, + const JaggedAccessor soundSpeeds, + float stepSize, + JaggedAccessor outSamples) { + + // Kernel parallelizes over rays. + // For each ray, we take a step, query the sound speeds, write a value, and bend the ray according to Snell's law. + } + +} // anonymous namespace + +// Generate samples along each ray with a given step size where samples are taken +// by bending the ray according to Snell's law. +std::vector acousticRaySamples( + const GridBatchImpl &batchHdl, + const JaggedTensor &rayOrigins, + const JaggedTensor &rayDirections, + const JaggedTensor &soundSpeeds, + float stepSize) { + + const JaggedTensor outputSamples = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1), 2}, rayOrigins.options()); + const JaggedTensor outputPositions = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1), 3}, rayOrigins.options()); + const JaggedTensor outputTimes = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1)}, rayOrigins.options()); + + // 1. Count the number of samples we will take for each ray. + + // 2. Generate the samples. + + return {outputSamples, outputPositions, outputTimes}; +} + +} // namespace ops +} // namespace detail +} // namespace fvdb diff --git a/src/fvdb/detail/ops/AcousticRaySamples.h b/src/fvdb/detail/ops/AcousticRaySamples.h new file mode 100644 index 00000000..4ed15a28 --- /dev/null +++ b/src/fvdb/detail/ops/AcousticRaySamples.h @@ -0,0 +1,30 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +#ifndef FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H +#define FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H + +#include +#include + +#include + +#include + +namespace fvdb { +namespace detail { +namespace ops { + +// Generate samples along each ray with a given step size where samples are taken +// by bending the ray according to Snell's law. +std::vector acousticRaySamples(const GridBatchImpl &batchHdl, + const JaggedTensor &rayOrigins, + const JaggedTensor &rayDirections, + const JaggedTensor &soundSpeeds, + float stepSize); + +} // namespace ops +} // namespace detail +} // namespace fvdb + +#endif // FVDB_DETAIL_OPS_ACOUSTICRAYSAMPLES_H diff --git a/src/python/GridBatchBinding.cpp b/src/python/GridBatchBinding.cpp index 9b09c1f5..bc75a739 100644 --- a/src/python/GridBatchBinding.cpp +++ b/src/python/GridBatchBinding.cpp @@ -154,6 +154,13 @@ bind_grid_batch(py::module &m) { .def_property_readonly("dual_bbox", &fvdb::GridBatch::dual_bbox) .def_property_readonly("address", &fvdb::GridBatch::address) + + .def("acoustic_ray_samples", + &fvdb::GridBatch::acoustic_ray_samples, + py::arg("ray_origins"), + py::arg("ray_directions"), + py::arg("sound_speeds"), + py::arg("step_size")) // Read a property for a single grid in the batch .def("voxel_size_at", [](const fvdb::GridBatch &self, int64_t bi) { From fc7e3e9d6ef5df4e9bbe4693a6ef74227fc5eee6 Mon Sep 17 00:00:00 2001 From: Francis Williams Date: Wed, 25 Mar 2026 13:01:49 -0700 Subject: [PATCH 2/2] fix Signed-off-by: Francis Williams --- src/fvdb/GridBatch.cpp | 1 + src/fvdb/detail/ops/AcousticRaySamples.cu | 25 +++++++++++++++-------- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/fvdb/GridBatch.cpp b/src/fvdb/GridBatch.cpp index 429f39c3..cb3e2f97 100644 --- a/src/fvdb/GridBatch.cpp +++ b/src/fvdb/GridBatch.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include diff --git a/src/fvdb/detail/ops/AcousticRaySamples.cu b/src/fvdb/detail/ops/AcousticRaySamples.cu index efc4dd0b..e1af61c8 100644 --- a/src/fvdb/detail/ops/AcousticRaySamples.cu +++ b/src/fvdb/detail/ops/AcousticRaySamples.cu @@ -22,13 +22,11 @@ __global__ void countAcousticRaySamplesKernel( const JaggedAccessor rayDirections, const JaggedAccessor soundSpeeds, float stepSize, - JaggedAccessor outCounts) { + fvdb::TorchRAcc64 outCounts) { // Kernel parallelizes over rays. // For each ray, we count the number of samples we will take. - } - -} // anonymous namespace +} template __global__ void accousticRaySamplesKernel( @@ -41,7 +39,7 @@ __global__ void accousticRaySamplesKernel( // Kernel parallelizes over rays. // For each ray, we take a step, query the sound speeds, write a value, and bend the ray according to Snell's law. - } +} } // anonymous namespace @@ -54,15 +52,24 @@ std::vector acousticRaySamples( const JaggedTensor &soundSpeeds, float stepSize) { - const JaggedTensor outputSamples = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1), 2}, rayOrigins.options()); - const JaggedTensor outputPositions = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1), 3}, rayOrigins.options()); - const JaggedTensor outputTimes = JaggedTensor::empty({rayOrigins.size(0), rayOrigins.size(1)}, rayOrigins.options()); + const auto gridBatchAcc = gridBatchAccessor(batchHdl); + const auto rayOriginsAcc = jaggedAccessor(rayOrigins); + const auto rayDirectionsAcc = jaggedAccessor(rayDirections); + const auto soundSpeedsAcc = jaggedAccessor(soundSpeeds); + + const int64_t N = rayOrigins.jdata().size(0); + auto outSamplesData = torch::empty({N, 2}, rayOrigins.jdata().options()); + auto outPositionsData = torch::empty({N, 3}, rayOrigins.jdata().options()); + auto outTimesData = torch::empty({N}, rayOrigins.jdata().options()); + fvdb::JaggedTensor outSamples = rayOrigins.jagged_like(outSamplesData); + fvdb::JaggedTensor outPositions = rayOrigins.jagged_like(outPositionsData); + fvdb::JaggedTensor outTimes = rayOrigins.jagged_like(outTimesData); // 1. Count the number of samples we will take for each ray. // 2. Generate the samples. - return {outputSamples, outputPositions, outputTimes}; + return {outSamples, outPositions, outTimes}; } } // namespace ops