|
| 1 | +// Copyright Contributors to the OpenVDB Project |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | +// |
| 4 | +#ifndef FVDB_TORCHRESOURCE_H |
| 5 | +#define FVDB_TORCHRESOURCE_H |
| 6 | + |
| 7 | +#include <nanovdb/cuda/DeviceResource.h> |
| 8 | + |
| 9 | +#include <c10/cuda/CUDACachingAllocator.h> |
| 10 | + |
| 11 | +#include <cstdio> |
| 12 | +#include <cstdlib> |
| 13 | +#include <stdexcept> |
| 14 | + |
| 15 | +namespace fvdb { |
| 16 | + |
| 17 | +/// @brief NanoVDB stream-ordered memory resource backed by PyTorch's currently |
| 18 | +/// active CUDA allocator. |
| 19 | +/// |
| 20 | +/// c10::cuda::CUDACachingAllocator is a namespace, not a concrete |
| 21 | +/// allocator: its free functions raw_alloc_with_stream / raw_delete |
| 22 | +/// dispatch through CUDACachingAllocator::get(), the runtime-swappable |
| 23 | +/// c10::cuda::CUDAAllocator* Torch itself allocates tensors from. This |
| 24 | +/// resource therefore follows whatever allocator the user has installed — |
| 25 | +/// the native caching allocator (including PYTORCH_CUDA_ALLOC_CONF knobs), |
| 26 | +/// the cudaMallocAsync backend (PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync), |
| 27 | +/// or a user-provided allocator installed via |
| 28 | +/// torch.cuda.memory.change_current_allocator(CUDAPluggableAllocator(...)). |
| 29 | +/// |
| 30 | +/// Passed as the ResourceT template parameter of NanoVDB's CUDA builders |
| 31 | +/// (PointsToGrid / DilateGrid / MergeGrids / PruneGrid / RefineGrid / |
| 32 | +/// CoarsenGrid) — always via the fvdb::BuilderResource alias |
| 33 | +/// (BuilderResource.h), never named directly at call sites — it routes |
| 34 | +/// their internal device scratch — O(N-points) sort |
| 35 | +/// keys, CUB temp storage, topology mask buffers — through the same pool |
| 36 | +/// that fvdb / PyTorch tensors use. Without this, nanoVDB's default |
| 37 | +/// DeviceResource allocates from a second cudaMallocAsync pool that |
| 38 | +/// partitions VRAM against torch's pool, and large workloads (e.g. |
| 39 | +/// multi-frame TSDF integration) OOM even when the GPU has free memory in |
| 40 | +/// aggregate. |
| 41 | +/// |
| 42 | +/// The resource is stateless, so builders can bind the shared instance |
| 43 | +/// returned by nanovdb::cuda::default_resource<TorchResource>() — naming |
| 44 | +/// the template parameter at a call site is sufficient, no instance needs |
| 45 | +/// to be threaded through. |
| 46 | +/// |
| 47 | +/// Set FVDB_NANOVDB_TRACE_ALLOCS=1 in the environment to trace allocations |
| 48 | +/// of 256 KiB and larger to stderr (a value starting with '2' traces every |
| 49 | +/// allocation). Useful for diagnosing topology-op memory blowup on large |
| 50 | +/// scenes. |
| 51 | +struct TorchResource : nanovdb::cuda::SyncFromAsync<TorchResource> { |
| 52 | + /// Alignment guaranteed by every allocation. Torch's native caching |
| 53 | + /// allocator returns blocks aligned to at least 512 bytes and the |
| 54 | + /// cudaMallocAsync backend to at least 256, so advertising nanoVDB's |
| 55 | + /// conventional 256 (matching cuda::DeviceResource) is satisfied and the |
| 56 | + /// alignment parameter below can be ignored. A pluggable allocator wrapping |
| 57 | + /// any cudaMalloc-family call satisfies 256 as well. |
| 58 | + static constexpr size_t DEFAULT_ALIGNMENT = 256; |
| 59 | + |
| 60 | + /// @brief Stream-ordered allocation from torch's active CUDA allocator. |
| 61 | + /// @note raw_alloc_with_stream records @p stream against the block so torch |
| 62 | + /// defers reuse until work on it completes, matching the stream-ordered |
| 63 | + /// semantics of the cudaMallocAsync call it replaces. Allocation |
| 64 | + /// happens on the current device, like cudaMallocAsync. The call |
| 65 | + /// dispatches to CUDACachingAllocator::get(), so a swapped-in backend |
| 66 | + /// or pluggable allocator is honored. |
| 67 | + void * |
| 68 | + allocate_async(size_t bytes, size_t /*alignment*/, cudaStream_t stream) { |
| 69 | + if (const char *env = std::getenv("FVDB_NANOVDB_TRACE_ALLOCS")) { |
| 70 | + const size_t cutoff = |
| 71 | + (env[0] == '2') ? 0 : (1ull << 18); // '2' = trace all, else >= 256 KiB |
| 72 | + if (bytes >= cutoff) { |
| 73 | + std::fprintf(stderr, |
| 74 | + "[fvdb/nanovdb] TorchResource alloc %12zu bytes (%.3f MB)\n", |
| 75 | + bytes, |
| 76 | + double(bytes) / 1e6); |
| 77 | + } |
| 78 | + } |
| 79 | + void *p = c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(bytes, stream); |
| 80 | + if (!p) { |
| 81 | + throw std::runtime_error("fvdb: TorchResource::allocate_async failed"); |
| 82 | + } |
| 83 | + return p; |
| 84 | + } |
| 85 | + |
| 86 | + /// @brief Free through torch's active CUDA allocator. |
| 87 | + /// @note The stream argument is deliberately ignored: raw_delete relies on |
| 88 | + /// the stream recorded at allocation time — the native backend's |
| 89 | + /// per-stream event tracking, or the alloc-time stream Torch hands a |
| 90 | + /// pluggable allocator's free function — so the free is safe without |
| 91 | + /// ordering on the caller's stream. This is the same contract Torch's |
| 92 | + /// own tensor frees rely on. |
| 93 | + void |
| 94 | + deallocate_async(void *p, size_t /*bytes*/, size_t /*alignment*/, cudaStream_t /*stream*/) { |
| 95 | + if (p == nullptr) { |
| 96 | + return; |
| 97 | + } |
| 98 | + c10::cuda::CUDACachingAllocator::raw_delete(p); |
| 99 | + } |
| 100 | +}; |
| 101 | + |
| 102 | +static_assert(nanovdb::cuda::is_async_resource<TorchResource>::value, |
| 103 | + "TorchResource must model nanoVDB's stream-ordered AsyncResource concept"); |
| 104 | + |
| 105 | +} // namespace fvdb |
| 106 | + |
| 107 | +#endif // FVDB_TORCHRESOURCE_H |
0 commit comments