Add/restore support for custom resources in PointsToGrid

* Separate data from algorithm

* Add custom resource support

* Add docs

* Add support for cudax::async_resource

* Fix default alignment

* Add cleanup

* Split up DeviceResource into separate file

* Fix file

Signed-off-by: Matthew Cong <[email protected]>

Author: matthewdcong

nanovdb/nanovdb/CMakeLists.txt

@@ -184,9 +184,11 @@ set(NANOVDB_INCLUDE_FILES
 set(NANOVDB_INCLUDE_CUDA_FILES
   cuda/DeviceBuffer.h
   cuda/DeviceMesh.h
+  cuda/DeviceResource.h
   cuda/DeviceStreamMap.h
   cuda/GridHandle.cuh
   cuda/NodeManager.cuh
+  cuda/TempPool.h
   cuda/UnifiedBuffer.h
 )
 

nanovdb/nanovdb/cuda/DeviceResource.h

@@ -0,0 +1,34 @@
+// Copyright Contributors to the OpenVDB Project
+// SPDX-License-Identifier: Apache-2.0
+//
+#ifndef NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED
+#define NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED
+
+#include <cuda_runtime_api.h>
+
+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, cudaStream_t stream) {
+        void* p = nullptr;
+        cudaCheck(util::cuda::mallocAsync(&p, bytes, stream));
+        return p;
+    }
+
+    static void deallocateAsync(void *p, size_t, size_t, cudaStream_t stream) {
+        cudaCheck(util::cuda::freeAsync(p, stream));
+    }
+};
+
+}
+
+} // namespace nanovdb::cuda
+
+#endif // end of NANOVDB_CUDA_DEVICERESOURCE_H_HAS_BEEN_INCLUDED

nanovdb/nanovdb/cuda/TempDevicePool.h

@@ -0,0 +1,55 @@
+// Copyright Contributors to the OpenVDB Project
+// SPDX-License-Identifier: Apache-2.0
+//
+#ifndef NANOVDB_CUDA_TEMPPOOL_H_HAS_BEEN_INCLUDED
+#define NANOVDB_CUDA_TEMPPOOL_H_HAS_BEEN_INCLUDED
+
+#include <nanovdb/cuda/DeviceResource.h>
+
+namespace nanovdb {
+
+namespace cuda {
+
+template <class Resource>
+class TempPool {
+public:
+    TempPool() : mData(nullptr), mSize(0), mRequestedSize(0) {}
+    ~TempPool() {
+        mRequestedSize = 0;
+        Resource::deallocateAsync(mData, mSize, Resource::DEFAULT_ALIGNMENT, nullptr);
+        mData = nullptr;
+        mSize = 0;
+    }
+
+    void* data() {
+        return mData;
+    }
+
+    size_t& size() {
+        return mSize;
+    }
+
+    size_t& requestedSize() {
+        return mRequestedSize;
+    }
+
+    void reallocate(cudaStream_t stream) {
+        if (!mData || mRequestedSize > mSize) {
+            Resource::deallocateAsync(mData, mSize, Resource::DEFAULT_ALIGNMENT, stream);
+            mData = Resource::allocateAsync(mRequestedSize, Resource::DEFAULT_ALIGNMENT, stream);
+            mSize = mRequestedSize;
+        }
+    }
+private:
+    void* mData;
+    size_t mSize;
+    size_t mRequestedSize;
+};
+
+using TempDevicePool = TempPool<DeviceResource>;
+
+}
+
+} // namespace nanovdb::cuda
+
+#endif // end of NANOVDB_CUDA_TEMPPOOL_H_HAS_BEEN_INCLUDED

nanovdb/nanovdb/cuda/TempPool.h

@@ -14,8 +14,8 @@
              to only include it in .cu files (or other .cuh files)
 */
 
-#ifndef NVIDIA_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED
-#define NVIDIA_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED
+#ifndef NANOVDB_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED
+#define NANOVDB_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED
 
 #include <nanovdb/NanoVDB.h>
 #include <nanovdb/cuda/DeviceBuffer.h>
@@ -143,4 +143,4 @@ cudaAddBlindData(const NanoGrid<BuildT> *d_grid,
 
 }// namespace nanovdb
 
-#endif // NVIDIA_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED
+#endif // NANOVDB_TOOLS_CUDA_ADDBLINDDATA_CUH_HAS_BEEN_INCLUDED

nanovdb/nanovdb/tools/cuda/AddBlindData.cuh

@@ -11,12 +11,12 @@
              to only include it in .cu files (or other .cuh files)
 */
 
-#ifndef NVIDIA_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED
-#define NVIDIA_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED
+#ifndef NANOVDB_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED
+#define NANOVDB_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED
 
 #include <nanovdb/GridHandle.h>
 #include <nanovdb/cuda/DeviceMesh.h>
-#include <nanovdb/cuda/TempDevicePool.h>
+#include <nanovdb/cuda/TempPool.h>
 #include <nanovdb/cuda/UnifiedBuffer.h>
 #include <nanovdb/tools/cuda/PointsToGrid.cuh>
 #include <nanovdb/util/cuda/Util.h>
@@ -279,7 +279,7 @@ private:
 
     PointType mPointType;
     std::string mGridName;
-    typename PointsToGrid<BuildT>::Data *mData;
+    PointsToGridData<BuildT> *mData;
     CheckMode mChecksum{CheckMode::Disable};
 
     size_t* mStripeCounts;
@@ -304,7 +304,7 @@ DistributedPointsToGrid<BuildT>::DistributedPointsToGrid(const nanovdb::cuda::De
 {
     mTempDevicePools = new nanovdb::cuda::TempDevicePool[mDeviceMesh.deviceCount()];
 
-    cudaCheck(cudaMallocManaged(&mData, sizeof(typename PointsToGrid<BuildT>::Data)));
+    cudaCheck(cudaMallocManaged(&mData, sizeof(PointsToGridData<BuildT>)));
     mData->flags.initMask({GridFlags::HasBBox, GridFlags::IsBreadthFirst});
     mData->map = map;
 
@@ -1106,4 +1106,4 @@ inline void DistributedPointsToGrid<BuildT>::processBBox()
 
 } // namespace nanovdb
 
-#endif // NVIDIA_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED
+#endif // NANOVDB_TOOLS_CUDA_DISTRIBUTEDPOINTSTOGRID_CUH_HAS_BEEN_INCLUDED

nanovdb/nanovdb/tools/cuda/DistributedPointsToGrid.cuh

@@ -14,8 +14,8 @@
              to only include it in .cu files (or other .cuh files)
 */
 
-#ifndef NVIDIA_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED
-#define NVIDIA_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED
+#ifndef NANOVDB_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED
+#define NANOVDB_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED
 
 #include <nanovdb/NanoVDB.h>
 #include <nanovdb/cuda/DeviceBuffer.h>
@@ -415,4 +415,4 @@ cudaCreateNanoGrid(const NanoGrid<SrcBuildT> *d_srcGrid, const typename BuildToV
 
 }// nanovdb namespace ===================================================================
 
-#endif // NVIDIA_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED
+#endif // NANOVDB_TOOLS_CUDA_INDEXTOGRID_CUH_HAS_BEEN_INCLUDED