diff --git a/.github/workflows/codestyle.yml b/.github/workflows/codestyle.yml index b571deae7..5b5235083 100644 --- a/.github/workflows/codestyle.yml +++ b/.github/workflows/codestyle.yml @@ -32,6 +32,11 @@ jobs: extensions: 'h,cpp,cc,cu,cuh' clangFormatVersion: 18 style: file + # Forked nanoVDB headers under nanovdb_overrides/ keep upstream's + # formatting verbatim (apart from documented FVDB FORK deltas) so + # that diff-against-upstream + resync stay tractable. They are + # intentionally exempt from fvdb's clang-format style. + exclude: 'src/fvdb/nanovdb_overrides' include-guards: name: Include guards @@ -41,6 +46,13 @@ jobs: - uses: swahtz/include-guards-check-action@master with: path: 'src/' + # Forked nanoVDB headers keep upstream's include-guard names + # (`NANOVDB_*_HAS_BEEN_INCLUDED`) so the override headers and any + # upstream copy that ends up on the include path share the same + # guard and don't double-include. + # NB: this is a `grep -v -e` pattern matched against paths relative + # to `path:` above, so no `src/` prefix. + ignore: 'fvdb/nanovdb_overrides' check-spdx-identifiers: name: SPDX identifiers diff --git a/CMakeLists.txt b/CMakeLists.txt index 84ecc128f..adc43411b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -169,6 +169,11 @@ set_target_properties(_fvdb_cpp PROPERTIES target_include_directories(_fvdb_cpp PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src ${TORCH_INCLUDE_DIRS} + # Overrides directory must come BEFORE the upstream nanoVDB source so that + # forked headers (e.g. nanovdb/cuda/DeviceBuffer.h) win the include search. + # Mirror the ordering used in src/CMakeLists.txt for the fvdb target. + # See src/fvdb/nanovdb_overrides/README.md. + ${FVDB_NANOVDB_OVERRIDES_DIR} ${nanovdb_SOURCE_DIR}/nanovdb ${NANOVDB_EDITOR_INCLUDE_DIR}) target_link_libraries(_fvdb_cpp PRIVATE diff --git a/pyproject.toml b/pyproject.toml index e7ea4905f..8a4875af2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,8 +11,13 @@ license = { text = "Apache-2.0" } requires-python = ">=3.10" readme = "README.md" dependencies = [ - # Require torch; users will choose the correct CUDA build from PyTorch's index - "torch>=2.8,<2.12", + # Require torch; users will choose the correct CUDA build from PyTorch's index. + # Upper bound is one minor release ahead of the highest torch we currently + # build against in CI; bumping the upper bound here is what `env/{build,test}_ + # requirements.txt` rely on (both are unpinned, so `uv pip install` follows + # the wheel's range). When the build / test pipelines pick up a newer torch + # than this range, bump both. + "torch>=2.8,<2.13", "numpy", ] optional-dependencies = {viewer = ["nanovdb-editor>=0.0.23,<0.2.0"]} diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 1edef1408..03d43a2e3 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -170,6 +170,10 @@ message(STATUS "fvdb: TORCH_INCLUDE_DIRS: ${TORCH_INCLUDE_DIRS}") target_include_directories(fvdb PUBLIC $ + # Overrides directory must come BEFORE the upstream nanoVDB source so that + # forked headers (e.g. nanovdb/cuda/DeviceBuffer.h) win the include search. + # See src/fvdb/nanovdb_overrides/README.md. + $ $ $ $ diff --git a/src/cmake/get_nanovdb.cmake b/src/cmake/get_nanovdb.cmake index ea0c8db41..5de4056d4 100644 --- a/src/cmake/get_nanovdb.cmake +++ b/src/cmake/get_nanovdb.cmake @@ -11,7 +11,24 @@ CPMAddPackage( # NanoVDB is header only, so we don't build it. Instead, we just add the headers # to the include path and create an interface target. +# +# We also prepend an override directory that contains modified copies of a small +# number of upstream nanoVDB headers. The override directory comes first in the +# include search path so that, e.g., `#include ` +# resolves to our forked copy under +# `src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceBuffer.h`. The rest of nanoVDB +# is still picked up from the upstream source tree, so the patch surface stays +# minimal and easy to resync. +# +# See src/fvdb/nanovdb_overrides/README.md for the rationale and the procedure +# for adding new overrides. if(nanovdb_ADDED) + get_filename_component(_fvdb_src_dir "${CMAKE_CURRENT_LIST_DIR}" DIRECTORY) + set(FVDB_NANOVDB_OVERRIDES_DIR + "${_fvdb_src_dir}/fvdb/nanovdb_overrides" + CACHE INTERNAL "Directory with fvdb-local overrides for nanoVDB headers") add_library(nanovdb INTERFACE) - target_include_directories(nanovdb INTERFACE ${nanovdb_SOURCE_DIR}/nanovdb) + target_include_directories(nanovdb INTERFACE + ${FVDB_NANOVDB_OVERRIDES_DIR} + ${nanovdb_SOURCE_DIR}/nanovdb) endif() diff --git a/src/cmake/get_torch.cmake b/src/cmake/get_torch.cmake index 6c7ace171..4a331fafc 100644 --- a/src/cmake/get_torch.cmake +++ b/src/cmake/get_torch.cmake @@ -39,6 +39,24 @@ find_package(Torch REQUIRED PATHS "${TORCH_PACKAGE_DIR}/share/cmake/Torch") # Without this we can't find TH/THC headers set(TORCH_SOURCE_INCLUDE_DIRS ${TORCH_PACKAGE_DIR}/include) +# Conda-forge's `pytorch-gpu` package installs the C++ headers at +# `$CONDA_PREFIX/include/{torch,ATen,c10,caffe2,tensorpipe}/`, and stages +# them into the site-packages tree via symlinks. Recent conda-forge builds +# (e.g. `pytorch-2.10.0-cuda130_mkl_py312_*_304`) have a packaging bug +# where those symlinks land as `torch.c~`, `ATen.c~`, ... (the conda +# file-conflict rename suffix), so `find_package(Torch)`'s +# `/torch/include/<...>` references don't resolve and +# `#include ` fails with "No such file or +# directory" when fvdb builds against a conda-forge torch. +# +# Append the conda env's bare `include/` so the canonical +# `$CONDA_PREFIX/include/torch/...` location is on the include path +# regardless of the symlink state. The IS_DIRECTORY guard keeps this a +# no-op for pip-installed torch (and for any non-conda environment). +if(DEFINED ENV{CONDA_PREFIX} AND IS_DIRECTORY "$ENV{CONDA_PREFIX}/include/torch") + list(APPEND TORCH_INCLUDE_DIRS "$ENV{CONDA_PREFIX}/include") +endif() + if(NOT TORCH_PYTHON_LIBRARY) message(STATUS "Looking for torch_python library...") diff --git a/src/fvdb/nanovdb_overrides/README.md b/src/fvdb/nanovdb_overrides/README.md new file mode 100644 index 000000000..8776cd80f --- /dev/null +++ b/src/fvdb/nanovdb_overrides/README.md @@ -0,0 +1,65 @@ +# fvdb-local overrides for nanoVDB headers + +This directory holds **modified copies of a small number of upstream nanoVDB +headers**. The build prepends this directory to the include search path before +the upstream nanoVDB source tree, so any `#include ` that matches +a file in this tree resolves here. Everything else falls through to upstream. + +The wiring lives in `src/cmake/get_nanovdb.cmake`. + +## Why not just patch the upstream checkout? + +We previously tried `CPM`'s `PATCH_COMMAND` to fix a specific bug (nanoVDB +scratch allocations going through `cudaMallocAsync` instead of PyTorch's +caching allocator, which fragments into two pools and OOMs). Patch-based +approaches are fragile: they silently stop applying if upstream moves the +surrounding lines, and they make it hard to edit nanoVDB during development. + +Forking the exact headers we care about into this tree gives us full edit +access with a narrow, reviewable patch surface and a clean diff-against-upstream +workflow. + +## Layout + +The directory structure **mirrors upstream** starting from the `nanovdb/` +include root: + +``` +nanovdb_overrides/ + nanovdb/ + cuda/ + DeviceBuffer.h # forked: device-handle alloc -> c10::cuda::CUDACachingAllocator + DeviceResource.h # forked: scratch alloc -> c10::cuda::CUDACachingAllocator + tools/ + cuda/ + TopologyBuilder.cuh # forked: opt-in scratch-size trace (FVDB_NANOVDB_TRACE_ALLOCS) + ... # add more as needed +``` + +## Adding a new override + +1. Copy the upstream header from + `build/<...>/_deps/nanovdb-src/nanovdb/nanovdb/` into + `nanovdb_overrides/nanovdb/`, preserving the relative path. +2. Add a short `FVDB FORK:` banner at the top documenting *what* diverges from + upstream and *why*. Keep the rest of the file byte-identical so a future + resync with upstream is a clean 3-way merge. +3. Every non-trivial code change inside the file should be tagged with an + inline `// FVDB FORK:` comment pointing at the banner, so `git blame` and + text searches make the delta obvious. + +## Resyncing with upstream + +When bumping the nanoVDB pin in `get_nanovdb.cmake`: + +1. `diff` each file in this directory against its upstream counterpart. +2. Port the upstream changes over, keeping the `FVDB FORK` deltas. +3. Rebuild + rerun the fvdb test suite. + +## Current overrides + +| File | Reason | +|--------------------------------------------|--------------------------------------------------------------------------------------------------------------| +| `nanovdb/cuda/DeviceBuffer.h` | Route device-handle allocations through PyTorch's caching allocator (avoids dual-pool OOM). | +| `nanovdb/cuda/DeviceResource.h` | Same allocator routing for the per-point scratch buffers used by `PointsToGrid` / `DilateGrid` / `MergeGrids`. | +| `nanovdb/tools/cuda/TopologyBuilder.cuh` | Opt-in `FVDB_NANOVDB_TRACE_ALLOCS` print of tile count + scratch size from `allocateInternalMaskBuffers`. | diff --git a/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceBuffer.h b/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceBuffer.h new file mode 100644 index 000000000..926249b1c --- /dev/null +++ b/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceBuffer.h @@ -0,0 +1,466 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 + +/*! + \file DeviceBuffer.h + + \author Ken Museth + + \date January 8, 2020 + + \brief DeviceBuffer has one pinned host buffer and multiple device CUDA buffers + + \note This file has no device-only kernel functions, + which explains why it's a .h and not .cuh file. + + \note ================================================================ + FVDB FORK: This header overrides the upstream nanoVDB copy of + DeviceBuffer.h. The only functional change is that device-side + allocations go through PyTorch's CUDA caching allocator + (c10::cuda::CUDACachingAllocator) instead of cudaMallocAsync / + cudaFreeAsync. This keeps all transient scratch allocations made + by nanoVDB internals (MergeGrids, TopologyBuilder, DilateGrid, + etc.) inside the same pool that fvdb / PyTorch tensors use, + which prevents the allocator from fragmenting into two + independent pools and avoids OOMs where one pool holds memory + the other cannot see. + + If you need to resync with upstream, diff against + build/.../_deps/nanovdb-src/nanovdb/nanovdb/cuda/DeviceBuffer.h + and port the delta -- only the init() / deviceUpload() / clear() + / operator= sites should differ. + ================================================================ +*/ + +#ifndef NANOVDB_CUDA_DEVICEBUFFER_H_HAS_BEEN_INCLUDED +#define NANOVDB_CUDA_DEVICEBUFFER_H_HAS_BEEN_INCLUDED + +#include +#include // for std::shared_ptr +#include // for BufferTraits +#include // for cudaMalloc/cudaMallocManaged/cudaFree + +// FVDB FORK: route device allocations through the PyTorch caching allocator. +#include +#include + +namespace nanovdb {// ================================================================ + +namespace cuda {// =================================================================== + +// ----------------------------> DeviceBuffer <-------------------------------------- + +/// @brief Simple memory buffer using un-managed pinned host memory when compiled with NVCC. +/// Obviously this class is making explicit used of CUDA so replace it with your own memory +/// allocator if you are not using CUDA. +/// @note While CUDA's pinned host memory allows for asynchronous memory copy between host and device +/// it is significantly slower then cached (un-pinned) memory on the host. +class DeviceBuffer +{ + uint64_t mSize; // total number of bytes managed by this buffer (assumed to be identical for host and device) + void *mCpuData, **mGpuData; // raw pointers to the host and device buffers + int mDeviceCount, mManaged;// if mManaged is non-zero this class is responsible for allocating and freeing memory buffers. Otherwise this is assumed to be handled externally + + /// @brief Initialize buffer + /// @param size byte size of buffer to be initialized + /// @param device id of the device on which to initialize the buffer + /// @note All existing buffers are first cleared + /// @warning size is expected to be non-zero. Use clear() clear buffer! + void init(uint64_t size, int device, cudaStream_t stream); + +public: + + using PtrT = std::shared_ptr; + + /// @brief Default constructor of an empty buffer + DeviceBuffer() : mSize(0), mCpuData(nullptr), mGpuData(nullptr), mDeviceCount(0), mManaged(0){} + + /// @brief Constructor with a specified device and size + /// @param size byte size of buffer to be initialized + /// @param device id of the device on which to initialize the buffer + /// @param stream cuda stream + DeviceBuffer(uint64_t size, int device = cudaCpuDeviceId, cudaStream_t stream = 0) : DeviceBuffer() + { + this->init(size, device, stream); + } + + /// @brief Constructor + /// @param size byte size of buffer to be initialized + /// @param host If true buffer is initialized only on the host/CPU, else on the current device/GPU + /// @param stream optional stream argument (defaults to stream NULL) + DeviceBuffer(uint64_t size, bool host, void* stream) : DeviceBuffer() + { + int device = cudaCpuDeviceId; + if (!host) cudaCheck(cudaGetDevice(&device)); + this->init(size, device, reinterpret_cast(stream)); + } + + /// @brief Constructor for externally managed host and device buffers + /// @param size byte size of the two external buffers + /// @param cpuData host buffer, assumed to NOT be NULL + /// @param gpuData device buffer, assumed to NOT be NULL; + /// @note The device buffer, @c gpuData, will be associated + /// with the current device ID given by cudaGetDevice + DeviceBuffer(uint64_t size, void* cpuData, void* gpuData) + : mSize(size) + , mCpuData(cpuData) + , mManaged(0) + { + cudaCheck(cudaGetDeviceCount(&mDeviceCount)); + mGpuData = new void*[mDeviceCount]();// NULL initialization + NANOVDB_ASSERT(cpuData); + NANOVDB_ASSERT(gpuData); + int device = 0; + cudaCheck(cudaGetDevice(&device)); + mGpuData[device] = gpuData; + } + + /// @brief Constructor for externally managed host and multiple device buffers + /// @param size byte size of the two external buffers + /// @param cpuData host buffer, assumed to NOT be NULL + /// @param list list of device IDs and external device buffers, all assumed to not be NULL + DeviceBuffer(uint64_t size, void* cpuData, std::initializer_list> list) + : mSize(size) + , mCpuData(cpuData) + , mManaged(0) + { + NANOVDB_ASSERT(cpuData); + cudaCheck(cudaGetDeviceCount(&mDeviceCount)); + mGpuData = new void*[mDeviceCount]();// NULL initialization + for (auto &p : list) { + NANOVDB_ASSERT(p.first>=0 && p.firstclear(); }; + + /// @brief Static factory method that return an instance of this buffer + /// @param size byte size of buffer to be initialized + /// @param dummy this argument is currently ignored but required to match the API of the HostBuffer + /// @param host If true buffer is initialized only on the host/CPU, else only on the device/GPU + /// @param stream optional stream argument (defaults to stream NULL) + /// @return An instance of this class using move semantics + static DeviceBuffer create(uint64_t size, const DeviceBuffer* dummy, bool host, void* stream){return DeviceBuffer(size, host, stream);} + + /// @brief Static factory method that returns an instance of this buffer + /// @param size byte size of buffer to be initialized + /// @param dummy this argument is currently ignored but required to match the API of the HostBuffer + /// @param device id of the device on which to initialize the buffer + /// @param stream cuda stream + static DeviceBuffer create(uint64_t size, const DeviceBuffer* dummy = nullptr, int device = cudaCpuDeviceId, cudaStream_t stream = 0){return DeviceBuffer(size, device, stream);} + + /// @brief Static factory method that returns an instance of this buffer that wraps externally managed memory + /// @param size byte size of buffer specified by external memory + /// @param cpuData pointer to externally managed host memory + /// @param gpuData pointer to externally managed device memory + /// @return An instance of this class using move semantics + static DeviceBuffer create(uint64_t size, void* cpuData, void* gpuData) {return DeviceBuffer(size, cpuData, gpuData);} + + /// @brief Static factory method that returns an instance of this buffer that wraps externally managed host and device memory + /// @param size byte size of buffer to be initialized + /// @param cpuData pointer to externally managed host memory + /// @param list list of device IDs and device memory pointers + static DeviceBuffer create(uint64_t size, void* cpuData, std::initializer_list> list) {return DeviceBuffer(size, cpuData, list);} + + /// @brief Static factory method that returns an instance of this buffer constructed from a HostBuffer + /// @param buffer host buffer from which to copy data + /// @param device id of the device on which to initialize the buffer + /// @param stream cuda stream + static DeviceBuffer create(const HostBuffer& buffer, int device = cudaCpuDeviceId, cudaStream_t stream = 0) {return DeviceBuffer(buffer, device, stream);} + + /////////////////////////////////////////////////////////////////////// + + /// @{ + /// @brief Factory methods that create a shared pointer to an DeviceBuffer instance + static PtrT createPtr(uint64_t size, const DeviceBuffer* = nullptr, int device = cudaCpuDeviceId, cudaStream_t stream = 0) {return std::make_shared(size, device, stream);} + static PtrT createPtr(uint64_t size, void* cpuData, void* gpuData) {return std::make_shared(size, cpuData, gpuData);} + static PtrT createPtr(uint64_t size, void* cpuData, std::initializer_list> list) {return std::make_shared(size, cpuData, list);} + static PtrT createPtr(const HostBuffer& buffer, int device = cudaCpuDeviceId, cudaStream_t stream = 0) {return std::make_shared(buffer, device, stream);} + /// @} + + /////////////////////////////////////////////////////////////////////// + + /// @brief Disallow copy assignment operation + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + /// @brief Move copy assignment operation + DeviceBuffer& operator=(DeviceBuffer&& other) noexcept; + + /////////////////////////////////////////////////////////////////////// + + /// @brief Retuns a raw void pointer to the host/CPU buffer managed by this allocator. + /// @warning Note that the pointer can be NULL! + void* data() const { return mCpuData; } + + /// @brief Returns an offset pointer of a specific type from the allocated host memory + /// @tparam T Type of the pointer returned + /// @param count Numbers of elements of @c parameter type T to skip + /// @warning might return NULL + template + T* data(ptrdiff_t count = 0, int device = cudaCpuDeviceId) const + { + NANOVDB_ASSERT(device >= cudaCpuDeviceId && device < mDeviceCount); + void *ptr = device == cudaCpuDeviceId ? mCpuData : mGpuData[device]; + return ptr ? reinterpret_cast(ptr) + count : nullptr; + } + + /// @brief Returns a byte offset void pointer from the allocated host memory + /// @param byteOffset offset of return pointer in units of bytes + /// @warning assumes that this instance is not empty! + void* data(ptrdiff_t byteOffset, int device = cudaCpuDeviceId) const + { + NANOVDB_ASSERT(device >= cudaCpuDeviceId && device < mDeviceCount); + void *ptr = device == cudaCpuDeviceId ? mCpuData : mGpuData[device]; + return ptr ? reinterpret_cast(ptr) + byteOffset : nullptr; + } + + /////////////////////////////////////////////////////////////////////// + + /// @brief Retuns a raw pointer to the specified device/GPU buffer managed by this allocator. + /// @warning Note that the pointer can be NULL! + void* deviceData(int device) const { + NANOVDB_ASSERT(device >= 0 && device < mDeviceCount); + return mGpuData[device]; + } + + /// @brief Retuns a raw pointer to the current device/GPU buffer managed by this allocator. + /// @warning Note that the pointer can be NULL! + void* deviceData() const { + int device = cudaCpuDeviceId; + cudaCheck(cudaGetDevice(&device)); + return this->deviceData(device); + } + + /////////////////////////////////////////////////////////////////////// + + /// @brief Uploads buffer on the host to a specific device. If it doesn't exist it's created first. + /// @param device Device ID that the data is copied to + /// @param stream cuda stream + /// @param sync if false the memory copy is asynchronous. + /// @warning Assumes that the host buffer already exists! + /// @note determine the current device with cudaGetDevice + void deviceUpload(int device = 0, cudaStream_t stream = 0, bool sync = true); + void deviceUpload(int device, void* stream, bool sync){this->deviceUpload(device, cudaStream_t(stream), sync);} + + /// @brief Upload buffer from the host to ALL the existing devices, i.e. CPU -> GPU. + /// If no device buffers exist one is created for the current device (typically 0) + /// and subsequently populated with the host data. + /// @param stream CUDA stream. + /// @param sync if false the memory copy is asynchronous. + /// @warning Assumes that the host buffer already exists! + void deviceUpload(cudaStream_t stream, bool sync); + void deviceUpload(void* stream, bool sync) {this->deviceUpload(cudaStream_t(stream), sync);} + + /////////////////////////////////////////////////////////////////////// + + /// @brief Download data from a specified device to the host. If the host buffer des not exist it will first be allocated + /// @param device device ID to download source data from + /// @param stream cuda stream + /// @param sync if false the memory copy is asynchronous. + /// @warning Assumes that the specifed device buffer already exists! + void deviceDownload(int device = 0, cudaStream_t stream = 0, bool sync = true); + void deviceDownload(int device, void* stream , bool sync) {this->deviceDownload(device, cudaStream_t(stream), sync);} + + /// @brief Download the buffer from the current device to the host, i.e. GPU -> CPU. + /// If the host buffer des not exist it will first be allocated + /// @param stream CUDA stream + /// @param sync if false the memory copy is asynchronous + /// @note If the host/CPU buffer does not exist it is first allocated + /// @warning Assumes that the device/GPU buffer already exists + void deviceDownload(void* stream, bool sync); + + /////////////////////////////////////////////////////////////////////// + + /// @brief Returns the size in bytes of the raw memory buffer managed by this allocator. + uint64_t size() const { return mSize; } + uint64_t capacity() const {return this->size();} + + /// @brief Returns the number of buffers that are not NULL + int bufferCount() const { + int count = mCpuData ? 1 : 0; + for (int i=0; iempty(); } + /// @} + + /// @brief De-allocate all memory managed by this allocator and set all pointers to NULL + void clear(cudaStream_t stream = 0); + void clear(void* stream){this->clear(cudaStream_t(stream));} + +}; // DeviceBuffer class + +// --------------------------> Implementations below <------------------------------------ + +inline DeviceBuffer& DeviceBuffer::operator=(DeviceBuffer&& other) noexcept +{ + if (mManaged) {// first free all the managed data buffers + cudaCheck(cudaFreeHost(mCpuData)); + // FVDB FORK: return device memory to torch's caching allocator instead of cudaFreeAsync. + for (int i=0; i= cudaCpuDeviceId && device < mDeviceCount); + if (device == cudaCpuDeviceId) { + cudaCheck(cudaMallocHost((void**)&mCpuData, size)); // un-managed pinned memory on the host (can be slow to access!). Always 32B aligned + checkPtr(mCpuData, "cuda::DeviceBuffer::init: failed to allocate host buffer"); + } else { + // FVDB FORK: use PyTorch's caching allocator so scratch allocations share + // a pool with tensor memory. We use `raw_alloc_with_stream` (rather than + // plain `raw_alloc`) so torch records this stream against the block and + // defers reuse until work on the stream completes, matching the stream- + // ordered semantics of the original `cudaMallocAsync(..., stream)` call. + // Freeing via `raw_delete` is stream-safe. + if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { + const size_t cutoff = (env[0] == '2') ? 0 : (1ull << 18); // '2' = trace all, else > 256 KB + if (size >= cutoff) { + fprintf(stderr, + "[fvdb/nanovdb] DeviceBuffer::init alloc %12zu bytes (%.3f MB) dev=%d\n", + size, double(size) / 1e6, device); + } + } + c10::cuda::CUDAGuard deviceGuard(device); + mGpuData[device] = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(size, stream); + checkPtr(mGpuData[device], "cuda::DeviceBuffer::init: failed to allocate device buffer"); + } + mSize = size; + mManaged = 1;// i.e. this instance is responsible for allocating and delete memory +} // DeviceBuffer::init + +inline void DeviceBuffer::deviceUpload(int device, cudaStream_t stream, bool sync) +{ + NANOVDB_ASSERT(device >= 0 && device < mDeviceCount);// should be device and not the host + checkPtr(mCpuData, "uninitialized cpu source data"); + if (mGpuData[device] == nullptr) { + if (mManaged==0) throw std::runtime_error("DeviceBuffer::deviceUpload called on externally managed memory that wasn\'t allocated."); + // FVDB FORK: use PyTorch's caching allocator. See init() for rationale. + c10::cuda::CUDAGuard deviceGuard(device); + mGpuData[device] = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(mSize, stream); + } + checkPtr(mGpuData[device], "uninitialized gpu destination data"); + cudaCheck(cudaMemcpyAsync(mGpuData[device], mCpuData, mSize, cudaMemcpyHostToDevice, stream)); + if (sync) cudaCheck(cudaStreamSynchronize(stream)); +} // DeviceBuffer::deviceUpload + +inline void DeviceBuffer::deviceUpload(cudaStream_t stream, bool sync) +{ + int device = 0; + cudaGetDevice(&device); + this->deviceUpload(device, stream, sync); +} // DeviceBuffer::deviceUpload + +inline void DeviceBuffer::deviceDownload(int device, cudaStream_t stream, bool sync) +{ + NANOVDB_ASSERT(device >= 0 && device < mDeviceCount); + checkPtr(mGpuData[device], "uninitialized gpu source data");// no source data on the specified device + if (mCpuData == nullptr) { + if (mManaged==0) throw std::runtime_error("DeviceBuffer::deviceDownload called on uninitialized cpu destination memory that is externally managed."); + cudaCheck(cudaMallocHost((void**)&mCpuData, mSize)); // un-managed pinned memory on the host (can be slow to access!). Always 32B aligned + } + checkPtr(mCpuData, "uninitialized cpu destination data"); + cudaCheck(cudaMemcpyAsync(mCpuData, mGpuData[device], mSize, cudaMemcpyDeviceToHost, stream)); + if (sync) cudaCheck(cudaStreamSynchronize(stream)); +} // DeviceBuffer::deviceDownload + +inline void DeviceBuffer::deviceDownload(void* stream, bool sync) +{ + int device = 0; + cudaCheck(cudaGetDevice(&device)); + this->deviceDownload(device, cudaStream_t(stream), sync); +} // DeviceBuffer::deviceDownload + +inline void DeviceBuffer::clear(cudaStream_t stream) +{ + if (mManaged) {// free all the managed data buffers + cudaCheck(cudaFreeHost(mCpuData)); + // FVDB FORK: return device memory to torch's caching allocator instead of cudaFreeAsync. + (void)stream; + for (int i=0; i +struct BufferTraits +{ + static constexpr bool hasDeviceDual = true; +}; + +}// namespace nanovdb + +#endif // end of NANOVDB_CUDA_DEVICEBUFFER_H_HAS_BEEN_INCLUDED diff --git a/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceResource.h b/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceResource.h new file mode 100644 index 000000000..316ffe153 --- /dev/null +++ b/src/fvdb/nanovdb_overrides/nanovdb/cuda/DeviceResource.h @@ -0,0 +1,93 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 +// +// FVDB FORK: This header overrides the upstream nanoVDB copy of +// DeviceResource.h. The functional change is that device-side +// allocations go through PyTorch's CUDA caching allocator +// (c10::cuda::CUDACachingAllocator) instead of cudaMallocAsync / +// cudaFreeAsync. +// +// Why: `DeviceResource` is the allocator that nanoVDB's +// PointsToGrid / DilateGrid / MergeGrids / TopologyBuilder use for +// their *internal* scratch buffers (O(N_points) sort keys, CUB +// temporary storage, node-count arrays, etc.). Upstream routes these +// through cudaMallocAsync, which on a torch process creates a second, +// independent CUDA memory pool next to torch's caching allocator -- +// so at scale (N ~ 10 M input points) the two pools partition VRAM +// and one of them OOMs even though the aggregate has plenty of free +// memory. Routing nanoVDB scratch through torch's allocator collapses +// the two pools into one, which is exactly what we want. +// +// This is the partner override to nanovdb_overrides/nanovdb/cuda/ +// DeviceBuffer.h -- that file handles the grid-handle-sized +// allocations; this one handles the per-point scratch. Together they +// cover every cudaMallocAsync call site reachable from the nanoVDB +// topology ops we care about. +// +// If you need to resync with upstream, diff against +// build/.../_deps/nanovdb-src/nanovdb/nanovdb/cuda/DeviceResource.h +// and port the delta -- only allocateAsync / deallocateAsync should +// differ. + +#ifndef NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED +#define NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED + +#include +#include + +// FVDB FORK: route device allocations through the PyTorch caching allocator. +#include +#include + +#include +#include + +namespace nanovdb { + +namespace cuda { + +class DeviceResource +{ +public: + // cudaMalloc aligns memory to 256 bytes by default + static constexpr size_t DEFAULT_ALIGNMENT = 256; + + static void* allocateAsync(size_t bytes, size_t /*alignment*/, cudaStream_t stream) { + // FVDB FORK: use PyTorch's caching allocator. + // + // `raw_alloc_with_stream` returns at least 512-byte aligned blocks, + // which is stricter than nanoVDB's DEFAULT_ALIGNMENT of 256, so we + // can safely ignore the alignment parameter here. + // + // We go through `raw_alloc_with_stream` (rather than the plain + // `raw_alloc`) so that torch's allocator records this stream + // against the block and defers reuse until work on the stream + // completes. That matches the stream-ordered semantics the + // original `cudaMallocAsync(..., stream)` call had. + if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { + const size_t cutoff = (env[0] == '2') ? 0 : (1ull << 18); + if (bytes >= cutoff) { + std::fprintf(stderr, + "[fvdb/nanovdb] DeviceResource alloc %12zu bytes (%.3f MB)\n", + bytes, double(bytes) / 1e6); + } + } + void *p = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(bytes, stream); + if (!p) { + throw std::runtime_error("fvdb: DeviceResource::allocateAsync failed"); + } + return p; + } + + static void deallocateAsync(void *p, size_t /*bytes*/, size_t /*alignment*/, cudaStream_t stream) { + (void)stream; + if (p == nullptr) return; + c10::cuda::CUDACachingAllocator::raw_delete(p); + } +}; + +} + +} // namespace nanovdb::cuda + +#endif // end of NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED diff --git a/src/fvdb/nanovdb_overrides/nanovdb/tools/cuda/TopologyBuilder.cuh b/src/fvdb/nanovdb_overrides/nanovdb/tools/cuda/TopologyBuilder.cuh new file mode 100644 index 000000000..b1ac46186 --- /dev/null +++ b/src/fvdb/nanovdb_overrides/nanovdb/tools/cuda/TopologyBuilder.cuh @@ -0,0 +1,583 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 + +/*! + \file nanovdb/tools/cuda/TopologyBuilder.cuh + + \authors Efty Sifakis + + \brief Shared functionality of (mostly morphology) operators that alter the voxel content of grids + + \warning The header file contains cuda device code so be sure + to only include it in .cu files (or other .cuh files) + + \note FVDB FORK: This header is copied from upstream nanoVDB. The only + functional change vs upstream is an env-gated debug print in + `allocateInternalMaskBuffers` that reports the tile count and + per-call scratch size, used to diagnose topology-operator + memory blowup when dilating / merging large scenes. Everything + else is byte-identical with upstream; keep the banner and + matching `FVDB FORK:` tag on any non-trivial future delta so + resync stays easy. See src/fvdb/nanovdb_overrides/README.md. +*/ + +#ifndef NVIDIA_TOOLS_CUDA_TOPOLOGYBUILDER_CUH_HAS_BEEN_INCLUDED +#define NVIDIA_TOOLS_CUDA_TOPOLOGYBUILDER_CUH_HAS_BEEN_INCLUDED + +#include +#include +#include +#include + +// FVDB FORK: needed for the env-gated tileCount trace in allocateInternalMaskBuffers. +#include +#include + +namespace nanovdb { + +namespace tools::cuda { + +template +class TopologyBuilder +{ + static_assert(nanovdb::BuildTraits::is_onindex);// For now, only OnIndexGrids supported + + using GridT = NanoGrid; + using TreeT = NanoTree; + using RootT = NanoRoot; + using UpperT = NanoUpper; + using LowerT = NanoLower; + using LeafT = NanoLeaf; + +public: + + TopologyBuilder(cudaStream_t stream) + { + mData = nanovdb::cuda::DeviceBuffer::create(sizeof(Data)); + } + + struct Data { + void *d_bufferPtr; + uint64_t grid, tree, root, upper, lower, leaf, size;// byte offsets to nodes in buffer + uint32_t nodeCount[3];// 0=leaf,1=lower, 2=upper + uint32_t *d_upperOffsets; + __hostdev__ GridT& getGrid() const {return *util::PtrAdd(d_bufferPtr, grid);} + __hostdev__ TreeT& getTree() const {return *util::PtrAdd(d_bufferPtr, tree);} + __hostdev__ RootT& getRoot() const {return *util::PtrAdd(d_bufferPtr, root);} + __hostdev__ UpperT& getUpper(int i) const {return *(util::PtrAdd(d_bufferPtr, upper)+i);} + __hostdev__ LowerT& getLower(int i) const {return *(util::PtrAdd(d_bufferPtr, lower)+i);} + __hostdev__ LeafT& getLeaf(int i) const {return *(util::PtrAdd(d_bufferPtr, leaf)+i);} + };// Data + + void allocateInternalMaskBuffers(cudaStream_t stream); + + void countNodes(cudaStream_t stream); + + template + BufferT getBuffer(const BufferT &buffer, cudaStream_t stream); + + void processUpperNodes(cudaStream_t stream); + + void processLowerNodes(cudaStream_t stream); + + void processLeafOffsets(cudaStream_t stream); + + void processBBox(cudaStream_t stream); + + void postProcessGridTree(cudaStream_t stream); + + nanovdb::cuda::DeviceBuffer mProcessedRoot; + nanovdb::cuda::DeviceBuffer mUpperMasks; + nanovdb::cuda::DeviceBuffer mLowerMasks; + nanovdb::cuda::DeviceBuffer mUpperOffsets; + nanovdb::cuda::DeviceBuffer mLowerOffsets; + nanovdb::cuda::DeviceBuffer mLeafOffsets; + nanovdb::cuda::DeviceBuffer mVoxelOffsets; + nanovdb::cuda::DeviceBuffer mLowerParents; + nanovdb::cuda::DeviceBuffer mLeafParents; + nanovdb::cuda::DeviceBuffer mData; + CheckMode mChecksum{CheckMode::Disable}; + + auto deviceProcessedRoot() { return static_cast(mProcessedRoot.deviceData()); } + auto hostProcessedRoot() { return static_cast(mProcessedRoot.data()); } + void* deviceUpperMasks() { return mUpperMasks.deviceData(); } + void* deviceLowerMasks() { return mLowerMasks.deviceData(); } + Data* data() { return static_cast(mData.data()); } + Data* deviceData() { return static_cast(mData.deviceData()); } + +private: + static constexpr unsigned int mNumThreads = 128;// for kernels spawned via lambdaKernel (others may specialize) + static unsigned int numBlocks(unsigned int n) {return (n + mNumThreads - 1) / mNumThreads;} + + nanovdb::cuda::TempDevicePool mTempDevicePool; +};// tools::cuda::TopologyBuilder + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +// Define utility macro used to call cub functions that use dynamic temporary storage +#ifndef CALL_CUBS +#ifdef _WIN32 +#define CALL_CUBS(func, ...) \ + cudaCheck(cub::func(nullptr, mTempDevicePool.requestedSize(), __VA_ARGS__, stream)); \ + mTempDevicePool.reallocate(stream); \ + cudaCheck(cub::func(mTempDevicePool.data(), mTempDevicePool.size(), __VA_ARGS__, stream)); +#else// ndef _WIN32 +#define CALL_CUBS(func, args...) \ + cudaCheck(cub::func(nullptr, mTempDevicePool.requestedSize(), args, stream)); \ + mTempDevicePool.reallocate(stream); \ + cudaCheck(cub::func(mTempDevicePool.data(), mTempDevicePool.size(), args, stream)); +#endif// ifdef _WIN32 +#endif// ifndef CALL_CUBS + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +template +void TopologyBuilder::allocateInternalMaskBuffers(cudaStream_t stream) +{ + if (hostProcessedRoot()->tileCount() == 0) return; // Processing empty grid(s); nothing to allocate + + // Allocate (and zero-fill) buffers large enough to hold: + // (a) The serialized masks of all upper nodes, for all tiles in the updated root node, and + // (b) The serialized masks of all densified lower nodes, as if every upper node had a full set of 32^3 lower children + int device = 0; + cudaGetDevice(&device); + uint64_t upperSize = hostProcessedRoot()->tileCount() * sizeof(Mask<5>); + uint64_t lowerSize = hostProcessedRoot()->tileCount() * Mask<5>::SIZE * sizeof(Mask<4>); + // FVDB DEBUG: trace tile counts to diagnose topology blowup. + if (std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { + std::fprintf(stderr, + "[fvdb/nanovdb] TopologyBuilder.allocateInternalMaskBuffers tileCount=%u upper=%.3fMB lower=%.3fMB\n", + hostProcessedRoot()->tileCount(), + double(upperSize)/1e6, double(lowerSize)/1e6); + } + mUpperMasks = nanovdb::cuda::DeviceBuffer::create(upperSize, nullptr, device, stream); + if (mUpperMasks.deviceData() == nullptr) throw std::runtime_error("Failed to allocate upper mask buffer on device"); + cudaCheck(cudaMemsetAsync(mUpperMasks.deviceData(), 0, upperSize, stream)); + mLowerMasks = nanovdb::cuda::DeviceBuffer::create( lowerSize, nullptr, device, stream ); + if (mLowerMasks.deviceData() == nullptr) throw std::runtime_error("Failed to allocate lower mask buffer on device"); + cudaCheck(cudaMemsetAsync(mLowerMasks.deviceData(), 0, lowerSize, stream)); +}// TopologyBuilder::allocateInternalMaskBuffers + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +template +void TopologyBuilder::countNodes(cudaStream_t stream) +{ + auto processedTileCount = hostProcessedRoot()->tileCount(); + if (processedTileCount == 0) { // Processing empty grid(s); zero nodes at all levels + data()->nodeCount[0] = data()->nodeCount[1] = data()->nodeCount[2] = 0; + return; + } + + // Computes prefix sums of (a) non-empty lower nodes, (b) counts of their leaf children, + // and (c) count of the speculatively updated root tiles that have actually been used. + // These are used to reconstruct child offsets for the internal nodes of the updated tree, + // as well as the tile table at the root. + std::size_t size = processedTileCount*Mask<5>::SIZE; + + int device = 0; + cudaGetDevice(&device); + nanovdb::cuda::DeviceBuffer upperCountsBuffer = nanovdb::cuda::DeviceBuffer::create(processedTileCount*sizeof(uint32_t), nullptr, device, stream); + nanovdb::cuda::DeviceBuffer lowerCountsBuffer = nanovdb::cuda::DeviceBuffer::create(size*sizeof(uint32_t), nullptr, device, stream); + nanovdb::cuda::DeviceBuffer leafCountsBuffer = nanovdb::cuda::DeviceBuffer::create(size*sizeof(uint32_t), nullptr, device, stream); + + using CountType = uint32_t (*)[Mask<5>::SIZE]; + auto lowerCounts = reinterpret_cast( lowerCountsBuffer.deviceData() ); + auto leafCounts = reinterpret_cast( leafCountsBuffer.deviceData() ); + + using Op = util::morphology::cuda::EnumerateNodesFunctor; + util::cuda::operatorKernel + <<>> + (deviceUpperMasks(), deviceLowerMasks(), lowerCounts, leafCounts); + + mUpperOffsets = nanovdb::cuda::DeviceBuffer::create((processedTileCount+1)*sizeof(uint32_t), nullptr, device, stream); + mLowerOffsets = nanovdb::cuda::DeviceBuffer::create((size+1)*sizeof(uint32_t), nullptr, device, stream); + mLeafOffsets = nanovdb::cuda::DeviceBuffer::create((size+1)*sizeof(uint32_t), nullptr, device, stream); + + cudaCheck(cudaMemsetAsync(mLowerOffsets.deviceData(), 0, sizeof(uint32_t), stream)); + CALL_CUBS(DeviceScan::InclusiveSum, + static_cast(lowerCountsBuffer.deviceData()), + static_cast(mLowerOffsets.deviceData())+1, + size); + cudaCheck(cudaMemcpyAsync(&data()->nodeCount[1], static_cast(mLowerOffsets.deviceData())+size, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + + cudaCheck(cudaMemsetAsync(mLeafOffsets.deviceData(), 0, sizeof(uint32_t), stream)); + CALL_CUBS(DeviceScan::InclusiveSum, + static_cast(leafCountsBuffer.deviceData()), + static_cast(mLeafOffsets.deviceData())+1, + size); + cudaCheck(cudaMemcpyAsync(&data()->nodeCount[0], static_cast(mLeafOffsets.deviceData())+size, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + + util::cuda::lambdaKernel<<>>( + processedTileCount, + [] __device__(size_t tileID, CountType lowerOffsets, uint32_t* upperCounts) + { upperCounts[tileID] = (lowerOffsets[tileID+1][0] > lowerOffsets[tileID][0]) ? 1 : 0; }, + static_cast(mLowerOffsets.deviceData()), + static_cast(upperCountsBuffer.deviceData())); + + cudaCheck(cudaMemsetAsync( mUpperOffsets.deviceData(), 0, sizeof(uint32_t), stream)); + CALL_CUBS(DeviceScan::InclusiveSum, + static_cast(upperCountsBuffer.deviceData()), + static_cast(mUpperOffsets.deviceData())+1, + processedTileCount); + cudaCheck(cudaMemcpyAsync(&data()->nodeCount[2], static_cast(mUpperOffsets.deviceData())+processedTileCount, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); +}// TopologyBuilder::countNodes + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +template +template +BufferT TopologyBuilder::getBuffer(const BufferT &pool, cudaStream_t stream) +{ + // Allocates a device buffer for the destination grid, once the topology/size of the tree is known + data()->grid = 0;// grid is always stored at the start of the buffer! + data()->tree = GridT::memUsage();// grid ends and tree begins + data()->root = data()->tree + TreeT::memUsage(); // tree ends and root node begins + data()->upper = data()->root + RootT::memUsage(data()->nodeCount[2]);// root node ends and upper internal nodes begin + data()->lower = data()->upper + UpperT::memUsage()*data()->nodeCount[2];// upper internal nodes ends and lower internal nodes begin + data()->leaf = data()->lower + LowerT::memUsage()*data()->nodeCount[1];// lower internal nodes ends and leaf nodes begin + data()->size = data()->leaf + LeafT::DataType::memUsage()*data()->nodeCount[0];// leaf nodes end and blind meta data begins + + int device = 0; + cudaGetDevice(&device); + auto buffer = BufferT::create(data()->size, &pool, device, stream);// only allocate buffer on the device + cudaCheck(cudaMemsetAsync(buffer.deviceData(), 0, data()->size, stream)); + + data()->d_bufferPtr = buffer.deviceData(); + if (data()->d_bufferPtr == nullptr) throw std::runtime_error("Failed to allocate grid buffer on the device"); + if (data()->nodeCount[2] != 0) // Unless the result is an empty grid + data()->d_upperOffsets = static_cast(mUpperOffsets.deviceData()); + mData.deviceUpload(device, stream, false); + + return buffer; +}// TopologyBuilder::getBuffer + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +namespace topology::detail { + +template +struct BuildGridTreeRootFunctor +{ + __device__ + void operator()(size_t, typename TopologyBuilder::Data *d_data) { + + // process Root + auto &root = d_data->getRoot(); + root.mTableSize = d_data->nodeCount[2]; + root.mBackground = NanoRoot::ValueType(0);// background_value + root.mMinimum = root.mMaximum = NanoRoot::ValueType(0); + root.mAverage = root.mStdDevi = NanoRoot::FloatType(0); + root.mBBox = CoordBBox(); // To be further updated after the leaf-level voxel update + + // process Tree + auto &tree = d_data->getTree(); + tree.setRoot(&root); + if (d_data->nodeCount[2]) { + tree.setFirstNode(&d_data->getUpper(0)); + tree.setFirstNode(&d_data->getLower(0)); + tree.setFirstNode(&d_data->getLeaf(0)); + } + else { + tree.template setFirstNode>(nullptr); + tree.template setFirstNode>(nullptr); + tree.template setFirstNode>(nullptr); + } + tree.mNodeCount[2] = d_data->nodeCount[2]; + tree.mNodeCount[1] = d_data->nodeCount[1]; + tree.mNodeCount[0] = d_data->nodeCount[0]; + tree.mVoxelCount = 0; // Actual voxel count (for non-empty grids) will only be known + // once leaf masks have been processed + tree.mTileCount[2] = tree.mTileCount[1] = tree.mTileCount[0] = 0; + + // process Grid + // The GridData header has already been copied from the input; + // reset what is necessary, and assert that others are at the expected values + auto &grid = d_data->getGrid(); + +#ifdef NANOVDB_USE_NEW_MAGIC_NUMBERS + NANOVDB_ASSERT(grid.mMagic == NANOVDB_MAGIC_GRID); +#else + NANOVDB_ASSERT(grid.mMagic == NANOVDB_MAGIC_NUMB); +#endif + grid.mChecksum.disable(); // all 64 bits ON means checksum is disabled + NANOVDB_ASSERT(grid.mVersion == Version()); + NANOVDB_ASSERT(grid.mFlags.isMaskOn(GridFlags::IsBreadthFirst)); + grid.mFlags.initMask({GridFlags::IsBreadthFirst}); // expected flags (HasBBox will be set later if grid is non-empty) + grid.mGridIndex = 0u; // Possibly overwriting input; returned grid has batch size 1 + grid.mGridCount = 1u; // Possibly overwriting input; returned grid has batch size 1 + grid.mGridSize = d_data->size; + // grid.mGridName expected to have been copied verbatim from input + // grid.mMap expected to have been copied verbatim from input + grid.mWorldBBox = Vec3dBBox();// invalid bbox + grid.mVoxelSize = grid.mMap.getVoxelSize(); + NANOVDB_ASSERT(grid.mGridClass == GridClass::IndexGrid); + NANOVDB_ASSERT(grid.mGridType == toGridType()); + grid.mBlindMetadataOffset = d_data->size; // i.e. no blind data, even if the input grid had any + grid.mBlindMetadataCount = 0u; // i.e. no blind data + NANOVDB_ASSERT(grid.mData0 == 0u); // zero padding + grid.mData1 = 1u; // This will be updated (unless this is an empty grid) after voxels have been processed +#ifdef NANOVDB_USE_NEW_MAGIC_NUMBERS + NANOVDB_ASSERT(grid.mData2 == 0u); +#else + NANOVDB_ASSERT(grid.mData2 == NANOVDB_MAGIC_GRID); +#endif + } +}; + +}// namespace topology::detail + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +namespace topology::detail { + +template +struct BuildUpperNodesFunctor +{ + __device__ + void operator()(size_t processedTileID, typename TopologyBuilder::Data *d_data, NanoRoot *d_processedRoot) { + uint32_t tileID = d_data->d_upperOffsets[processedTileID]; + if (tileID != d_data->d_upperOffsets[processedTileID+1]) // if the offsets are the same, this was a speculatively introduced tile which was not necessary + { + auto &root = d_data->getRoot(); + auto &dstUpper = d_data->getUpper(tileID); + auto &processedTile = *d_processedRoot->tile(processedTileID); + root.tile(tileID)->setChild( processedTile.origin(), &dstUpper, &root ); + dstUpper.mBBox = CoordBBox(); // To be further updated after the operation has been applied at leaf-level + // TODO: Is this accurate? Any other flags that should be set? + dstUpper.mFlags = (uint64_t)GridFlags::HasBBox; + } + } +}; + +}// namespace topology::detail + +template +inline void TopologyBuilder::processUpperNodes(cudaStream_t stream) +{ + // Connect all newly allocated upper nodes to their respective tiles + // Also fill in any necessary part of the preamble (in InternalData) of upper nodes + auto processedTileCount = hostProcessedRoot()->tileCount(); + + if (processedTileCount) { // Unless output grid is empty + util::cuda::lambdaKernel<<>>( + processedTileCount, topology::detail::BuildUpperNodesFunctor(), deviceData(), deviceProcessedRoot()); + cudaCheckError(); + } +}// TopologyBuilder::processUpperNodes + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +template +inline void TopologyBuilder::processLowerNodes(cudaStream_t stream) +{ + // Fill out the contents of all newly allocated lower nodes (using the densified upper/lower mask arrays) + // Also fill in the preamble (most of LeafData) for their leaf children + auto processedTileCount = hostProcessedRoot()->tileCount(); + using CountType = uint32_t (*)[Mask<5>::SIZE]; + + if (processedTileCount) { // Unless output grid is empty + int device = 0; + cudaGetDevice(&device); + std::size_t lowerCount = data()->nodeCount[1]; + mLowerParents = nanovdb::cuda::DeviceBuffer::create(lowerCount*sizeof(uint32_t), nullptr, device, stream); + std::size_t leafCount = data()->nodeCount[0]; + mLeafParents = nanovdb::cuda::DeviceBuffer::create(leafCount*sizeof(uint32_t), nullptr, device, stream); + + using Op = util::morphology::cuda::ProcessLowerNodesFunctor; + util::cuda::operatorKernel + <<>>( + deviceUpperMasks(), + deviceLowerMasks(), + static_cast(mUpperOffsets.deviceData()), + static_cast(mLowerOffsets.deviceData()), + static_cast(mLeafOffsets.deviceData()), + static_cast(data()->d_bufferPtr), + static_cast(mLowerParents.deviceData()), + static_cast(mLeafParents.deviceData()) + ); + cudaCheckError(); + } + + mProcessedRoot.clear(stream); + mUpperMasks.clear(stream); + mLowerMasks.clear(stream); + mLowerOffsets.clear(stream); + mLeafOffsets.clear(stream); +}// TopologyBuilder::processLowerNodes + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +namespace topology::detail { + +template +struct UpdateLeafVoxelCountsAndPrefixSumFunctor +{ + __device__ + void operator()(size_t leafID, typename TopologyBuilder::Data *d_data, uint64_t *d_voxelCounts) { + auto &leaf = d_data->getGrid().tree().template getFirstNode<0>()[leafID]; + const uint64_t *w = leaf.mValueMask.words(); + uint64_t prefixSum = 0, sum = util::countOn(*w++); + prefixSum = sum; + for (int n = 9; n < 55; n += 9) {// n=i*9 where i=1,2,..6 + sum += util::countOn(*w++); + prefixSum |= sum << n; }// each pre-fixed sum is encoded in 9 bits + sum += util::countOn(*w); + d_voxelCounts[leafID] = sum; + leaf.mPrefixSum = prefixSum; } +}; + +template +struct UpdateLeafVoxelOffsetsFunctor +{ + __device__ + void operator()(size_t leafID, typename TopologyBuilder::Data *d_data, uint64_t *d_voxelOffsets) { + auto &leaf = d_data->getGrid().tree().template getFirstNode<0>()[leafID]; + leaf.mOffset = d_voxelOffsets[leafID]+1; } +}; + +}// namespace topology::detail + +template +inline void TopologyBuilder::processLeafOffsets(cudaStream_t stream) +{ + int device = 0; + cudaGetDevice(&device); + std::size_t leafCount = data()->nodeCount[0]; + if (leafCount) { // Unless output grid is empty + mVoxelOffsets = nanovdb::cuda::DeviceBuffer::create((leafCount+1)*sizeof(uint64_t), nullptr, device, stream); + cudaCheck(cudaMemsetAsync(mVoxelOffsets.deviceData(), 0, sizeof(uint64_t), stream)); + util::cuda::lambdaKernel<<>>( + leafCount, topology::detail::UpdateLeafVoxelCountsAndPrefixSumFunctor(), deviceData(), static_cast(mVoxelOffsets.deviceData())+1); + CALL_CUBS(DeviceScan::InclusiveSum, + static_cast(mVoxelOffsets.deviceData())+1, + static_cast(mVoxelOffsets.deviceData())+1, + leafCount); + util::cuda::lambdaKernel<<>>( + leafCount, topology::detail::UpdateLeafVoxelOffsetsFunctor(), deviceData(), static_cast(mVoxelOffsets.deviceData())); + } +}// TopologyBuilder::processLeafOffsets + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +// Undefine utility macro for cub functions +#ifdef CALL_CUBS +#undef CALL_CUBS +#endif + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +namespace topology::detail { + +template +struct UpdateAndPropagateLeafBBoxFunctor +{ + __device__ + void operator()(size_t tid, typename TopologyBuilder::Data *d_data, const uint32_t* leafParents) { + auto &lower = d_data->getLower(leafParents[tid]); + auto &leaf = d_data->getLeaf(tid); + leaf.updateBBox(); + lower.mBBox.expandAtomic(leaf.bbox()); + } +}; + +template +struct PropagateLowerBBoxFunctor +{ + __device__ + void operator()(size_t tid, typename TopologyBuilder::Data *d_data, const uint32_t* lowerParents) { + auto &upper = d_data->getUpper(lowerParents[tid]); + auto &lower = d_data->getLower(tid); + upper.mBBox.expandAtomic(lower.bbox()); } +}; + +template +struct PropagateUpperBBoxFunctor +{ + __device__ + void operator()(size_t tid, typename TopologyBuilder::Data *d_data) { + d_data->getRoot().mBBox.expandAtomic(d_data->getUpper(tid).bbox()); + } +}; + +template +struct UpdateRootWorldBBoxFunctor +{ + __device__ + void operator()(size_t tid, typename TopologyBuilder::Data *d_data) { + // TODO: check that the correct semantics are followed in this transformation + auto BBox = d_data->getRoot().mBBox; + BBox.max() += 1; + d_data->getGrid().mFlags.setMaskOn(GridFlags::HasBBox); + d_data->getGrid().mWorldBBox = BBox.transform(d_data->getGrid().data()->mMap); + } +}; + +}// namespace topology::detail + +template +inline void TopologyBuilder::processBBox(cudaStream_t stream) +{ + if (data()->nodeCount[0] == 0) return; // Output grid is empty; retain empty bounding box + + // TODO: Do we need a special case when flags indicates no bounding box? + + // update and propagate bbox from leaf -> lower/parent nodes + util::cuda::lambdaKernel<<nodeCount[0]), mNumThreads, 0, stream>>>( + data()->nodeCount[0], topology::detail::UpdateAndPropagateLeafBBoxFunctor(), deviceData(), static_cast(mLeafParents.deviceData())); + mLeafParents.clear(stream); + cudaCheckError(); + + // propagate bbox from lower -> upper/parent node + util::cuda::lambdaKernel<<nodeCount[1]), mNumThreads, 0, stream>>>( + data()->nodeCount[1], topology::detail::PropagateLowerBBoxFunctor(), deviceData(), static_cast(mLowerParents.deviceData())); + mLowerParents.clear(stream); + cudaCheckError(); + + // propagate bbox from upper -> root/parent node + util::cuda::lambdaKernel<<nodeCount[2]), mNumThreads, 0, stream>>>(data()->nodeCount[2], topology::detail::PropagateUpperBBoxFunctor(), deviceData()); + cudaCheckError(); + + // update the world-bbox in the root node + util::cuda::lambdaKernel<<<1, 1, 0, stream>>>(1, topology::detail::UpdateRootWorldBBoxFunctor(), deviceData()); + cudaCheckError(); +}// TopologyBuilder::processBBox + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +namespace topology::detail { + +template +struct PostProcessGridTreeFunctor +{ + __device__ + void operator()(size_t tid, typename TopologyBuilder::Data *d_data, uint64_t* d_voxelOffsets) { + auto& grid = d_data->getGrid(); + auto& tree = grid.tree(); + auto leafCount = tree.mNodeCount[0]; + tree.mVoxelCount = d_voxelOffsets[leafCount]; + grid.mData1 = tree.mVoxelCount+1; + } +}; + +}// namespace topology::detail + +template +inline void TopologyBuilder::postProcessGridTree(cudaStream_t stream) +{ + // Finish updates to GridData/TreeData and (optionally) update checksum + if (data()->nodeCount[0]) // if grid is empty, the default values are correct + util::cuda::lambdaKernel<<<1, 1, 0, stream>>>(1, topology::detail::PostProcessGridTreeFunctor(), deviceData(), static_cast(mVoxelOffsets.deviceData())); + cudaCheckError(); + mVoxelOffsets.clear(stream); + + tools::cuda::updateChecksum((GridData*)data()->d_bufferPtr, mChecksum, stream); +}// TopologyBuilder::postProcessGridTree + +//------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- + +}// namespace tools::cuda + +}// namespace nanovdb + +#endif // NVIDIA_TOOLS_CUDA_TOPOLOGYBUILDER_CUH_HAS_BEEN_INCLUDED