NanoVDB: SyncFromAsync, ResourceRef, and resource seams for MeshToGrid and TempPool (CUDA)

nanovdb/nanovdb/cuda/DeviceResource.h

@@ -145,6 +145,106 @@ struct is_resource<R, std::void_t<
     decltype(std::declval<R&>().deallocate(std::declval<void*>(), size_t{0}, size_t{0}))>>
     : std::true_type {};
 
+/// @brief CRTP base supplying the synchronous half of the resource concept in
+///        terms of the stream-ordered half, so a custom stream-ordered resource
+///        only has to write allocate_async and deallocate_async.
+/// @tparam Derived the resource deriving from this base
+/// @details A stream-ordered resource must also model the synchronous concept
+///          (is_async_resource implies is_resource), which means writing four
+///          methods where two would do. The synchronous pair is not a bare
+///          delegate: memory from allocate must be usable immediately on any
+///          stream, so the null-stream allocation has to be synchronized before
+///          it is returned. Omitting that synchronization yields memory that
+///          satisfies the concept but is not actually synchronous -- a race
+///          rather than a compile error -- so it lives here rather than being
+///          rewritten per resource.
+/// @code
+/// struct MyResource : nanovdb::cuda::SyncFromAsync<MyResource> {
+///     static constexpr size_t DEFAULT_ALIGNMENT = 256;
+///     void* allocate_async(size_t bytes, size_t alignment, cudaStream_t stream);
+///     void  deallocate_async(void* p, size_t bytes, size_t alignment, cudaStream_t stream);
+/// };
+/// @endcode
+template <class Derived>
+struct SyncFromAsync
+{
+    /// @brief Allocates @c bytes usable on any stream when this returns.
+    /// @param bytes number of bytes to allocate
+    /// @param alignment requested alignment
+    /// @note Every call synchronizes the null stream; on hot paths prefer the
+    ///       stream-ordered pair.
+    void* allocate(size_t bytes, size_t alignment)
+    {
+        void* p = static_cast<Derived&>(*this).allocate_async(bytes, alignment, cudaStream_t{0});
+        cudaCheck(cudaStreamSynchronize(cudaStream_t{0}));
+        return p;
+    }
+
+    /// @brief Frees @c p on the null stream.
+    /// @param p pointer previously returned by allocate
+    /// @param bytes size passed to the matching allocate
+    /// @param alignment alignment passed to the matching allocate
+    /// @note No synchronization here: the synchronous concept's contract is
+    ///       that the memory is already quiescent when deallocate is called.
+    void deallocate(void* p, size_t bytes, size_t alignment)
+    {
+        static_cast<Derived&>(*this).deallocate_async(p, bytes, alignment, cudaStream_t{0});
+    }
+};
+
+/// @brief Non-owning reference to a memory resource that is itself a resource:
+///        copying the ref shares the underlying resource rather than copying it.
+/// @tparam R the referenced resource type
+/// @details Types that hold their resource by value -- cuda::Buffer, matching
+///          cuda::buffer -- select their ownership semantics by what is placed
+///          in that slot: a concrete resource is owned as a copy, while a
+///          ResourceRef borrows. This is the same division cuda::mr draws
+///          between any_resource (owning) and resource_ref (borrowing), and the
+///          same shape as std::pmr::polymorphic_allocator over memory_resource*.
+///          Use it when a resource is stateful or long-lived and a container
+///          must allocate through *that* instance rather than a copy of it.
+/// @warning The referenced resource must outlive every use of this ref and of
+///          all copies of it, including any container holding one.
+template <class R>
+struct ResourceRef
+{
+    static_assert(is_async_resource<R>::value || is_resource<R>::value,
+                  "ResourceRef requires R to model the AsyncResource or the Resource concept");
+
+    static constexpr size_t DEFAULT_ALIGNMENT = R::DEFAULT_ALIGNMENT;
+
+    /// @brief Constructs a ref borrowing @c resource.
+    /// @param resource resource to allocate from; must outlive this ref
+    ResourceRef(R& resource) : mResource(&resource) {}
+
+    /// @{
+    /// @brief Stream-ordered pair, present only when @c R models AsyncResource,
+    ///        so a ref over a synchronous resource does not misreport its tier.
+    template<class S = R, std::enable_if_t<is_async_resource<S>::value, int> = 0>
+    void* allocate_async(size_t bytes, size_t alignment, cudaStream_t stream)
+    {
+        return mResource->allocate_async(bytes, alignment, stream);
+    }
+    template<class S = R, std::enable_if_t<is_async_resource<S>::value, int> = 0>
+    void deallocate_async(void* p, size_t bytes, size_t alignment, cudaStream_t stream)
+    {
+        mResource->deallocate_async(p, bytes, alignment, stream);
+    }
+    /// @}
+
+    /// @brief Synchronous pair, forwarding to the referenced resource.
+    void* allocate(size_t bytes, size_t alignment) { return mResource->allocate(bytes, alignment); }
+    void  deallocate(void* p, size_t bytes, size_t alignment) { mResource->deallocate(p, bytes, alignment); }
+
+    /// @brief Two refs compare equal iff they reference the same resource, i.e.
+    ///        memory allocated through one may be deallocated through the other.
+    friend bool operator==(ResourceRef lhs, ResourceRef rhs) { return lhs.mResource == rhs.mResource; }
+    friend bool operator!=(ResourceRef lhs, ResourceRef rhs) { return lhs.mResource != rhs.mResource; }
+
+private:
+    R* mResource;
+};// ResourceRef<R>
+
 }
 
 } // namespace nanovdb::cuda

nanovdb/nanovdb/cuda/TempPool.h

@@ -10,6 +10,7 @@
 #ifndef NANOVDB_CUDA_TEMPPOOL_H_HAS_BEEN_INCLUDED
 #define NANOVDB_CUDA_TEMPPOOL_H_HAS_BEEN_INCLUDED
 
+#include <nanovdb/cuda/Buffer.h>
 #include <nanovdb/cuda/DeviceResource.h>
 
 #include <cstddef>
@@ -21,31 +22,33 @@ namespace cuda {
 
 template <class Resource>
 class TempPool {
+    static_assert(is_async_resource<Resource>::value,
+                  "TempPool allocates stream-ordered scratch and requires an AsyncResource");
+    // The buffer borrows the pool's resource through a ResourceRef rather than
+    // copying it, preserving the pool's contract that all traffic reaches the
+    // caller's resource instance (which may be stateful).
+    using BufferT = Buffer<std::byte, ResourceRef<Resource>>;
 public:
 
     /// @brief Default c-tor of an empty memory pool that uses the default
     ///        instance of @c Resource for all allocations.
-    TempPool() : mResource(&default_resource<Resource>()), mData(nullptr), mSize(0), mRequestedSize(0), mStream(nullptr) {}
+    TempPool() : TempPool(default_resource<Resource>()) {}
 
     /// @brief C-tor of an empty memory pool that routes all allocations through
     ///        the supplied @c Resource instance.
     /// @param resource resource instance to allocate from; must outlive this pool.
-    explicit TempPool(Resource& resource) : mResource(&resource), mData(nullptr), mSize(0), mRequestedSize(0), mStream(nullptr) {}
-
-    /// @brief Destructor. Frees the managed memory on the stream of the most
-    ///        recent reallocate(), so the stream-ordered free is ordered after
-    ///        the work that used the memory (rather than on the null stream).
-    ~TempPool() {
-        mRequestedSize = 0;
-        mResource->deallocate_async(mData, mSize, Resource::DEFAULT_ALIGNMENT, mStream);
-        mData = nullptr;
-        mSize = 0;
+    explicit TempPool(Resource& resource)
+        : mResource(&resource)
+        , mBuffer(cudaStream_t{0}, ResourceRef<Resource>(resource), 0, noInit)
+    {
     }
 
     /// @brief Returns a non-const void pointer to the data managed by this instance.
-    void* data() {return mData;}
+    void* data() {return mBuffer.data();}
 
     /// @brief Returns a non-const reference to the actual size of the data managed by this instance.
+    /// @note Returned by reference because cub's two-pass API takes the storage
+    ///       size as a size_t&, so this cannot forward Buffer::size() by value.
     size_t& size() {return mSize;}
 
     /// @brief Returns a non-const reference to the requested size of the data managed by this instance.
@@ -54,25 +57,27 @@ class TempPool {
 
     /// @brief Returns the stream that the managed memory was last (re)allocated on,
     ///        i.e. the stream this pool will free on at destruction.
-    cudaStream_t stream() const {return mStream;}
+    cudaStream_t stream() const {return mBuffer.stream();}
 
     /// @brief Re-allocation of the data managed by this instance. Only has affect if the pool in empty or
     ///        the requested memory is larger than the existing size.
     /// @param stream cuda stream used for asynchronous de-allocation and allocation.
+    /// @note Scratch is discarded, never resized: preserving a prefix of
+    ///       temporary storage would be a wasted copy.
     void reallocate(cudaStream_t stream) {
-        if (!mData || mRequestedSize > mSize) {
-            mResource->deallocate_async(mData, mSize, Resource::DEFAULT_ALIGNMENT, stream);
-            mData = mResource->allocate_async(mRequestedSize, Resource::DEFAULT_ALIGNMENT, stream);
-            mSize = mRequestedSize;
+        if (mBuffer.empty() || mRequestedSize > mSize) {
+            mBuffer.destroy(stream);// free the outgrown block on this stream
+            mBuffer = BufferT(stream, ResourceRef<Resource>(*mResource), mRequestedSize, noInit);
+            mSize = mBuffer.size();
+        } else {
+            mBuffer.set_stream(stream);// retained so the d-tor frees on the most-recently-used stream
         }
-        mStream = stream;// retained so the destructor frees on the most-recently-used stream
     }
 private:
-    Resource    *mResource;
-    void        *mData;
-    size_t       mSize;
-    size_t       mRequestedSize;
-    cudaStream_t mStream;
+    Resource *mResource;// non-owning; must outlive this pool and its buffer
+    BufferT   mBuffer;
+    size_t    mSize{0};
+    size_t    mRequestedSize{0};
 };// TempPool<Resource> class
 
 using TempDevicePool = TempPool<DeviceResource>;