.. _program_listing_file_nvcv_types_include_nvcv_alloc_Allocator.hpp:

Program Listing for File Allocator.hpp
======================================

|exhale_lsh| :ref:`Return to documentation for file <file_nvcv_types_include_nvcv_alloc_Allocator.hpp>` (``nvcv_types/include/nvcv/alloc/Allocator.hpp``)

.. |exhale_lsh| unicode:: U+021B0 .. UPWARDS ARROW WITH TIP LEFTWARDS

.. code-block:: cpp

   /*
    * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    * SPDX-License-Identifier: Apache-2.0
    *
    * Licensed under the Apache License, Version 2.0 (the "License");
    * you may not use this file except in compliance with the License.
    * You may obtain a copy of the License at
    *
    * http://www.apache.org/licenses/LICENSE-2.0
    *
    * Unless required by applicable law or agreed to in writing, software
    * distributed under the License is distributed on an "AS IS" BASIS,
    * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    * See the License for the specific language governing permissions and
    * limitations under the License.
    */
   
   #ifndef NVCV_ALLOC_ALLOCATOR_HPP
   #define NVCV_ALLOC_ALLOCATOR_HPP
   
   #include "../CoreResource.hpp"
   #include "../detail/Callback.hpp"
   #include "../detail/CompilerUtils.h"
   #include "../detail/TypeTraits.hpp"
   #include "Allocator.h"
   
   #include <cassert>
   #include <cstddef>
   #include <cstring>
   #include <functional>
   
   namespace nvcv {
   
   // Helper class to explicitly assign
   // address alignments.
   class MemAlignment
   {
   public:
       MemAlignment() = default;
   
       int32_t baseAddr() const
       {
           return m_baseAddrAlignment;
       }
   
       int32_t rowAddr() const
       {
           return m_rowAddrAlignment;
       }
   
       MemAlignment &baseAddr(int32_t alignment)
       {
           m_baseAddrAlignment = alignment;
           return *this;
       }
   
       MemAlignment &rowAddr(int32_t alignment)
       {
           m_rowAddrAlignment = alignment;
           return *this;
       }
   
   private:
       int32_t m_baseAddrAlignment = 0;
       int32_t m_rowAddrAlignment  = 0;
   };
   
   class ResourceAllocator
   {
   public:
       ResourceAllocator() = default;
   
       explicit ResourceAllocator(const NVCVResourceAllocator &alloc)
           : m_data(alloc)
       {
       }
   
       const NVCVResourceAllocator &cdata() const &
       {
           return m_data;
       }
   
       NVCVResourceAllocator cdata() &&
       {
           return m_data;
       }
   
       template<typename Derived>
       Derived cast() const
       {
           static_assert(std::is_base_of<ResourceAllocator, Derived>::value,
                         "The requested type does not inherit from ResourceAllocator");
           static_assert(std::is_constructible<Derived, NVCVResourceAllocator>::value,
                         "The requested type must be constructible from NVCVResourceAllocator");
           return Derived(m_data);
       }
   
       static constexpr bool IsCompatibleKind(NVCVResourceType)
       {
           return true;
       }
   
   protected:
       NVCVResourceAllocator &data() &
       {
           return m_data;
       }
   
       NVCVResourceAllocator m_data{};
   };
   
   class MemAllocator : public ResourceAllocator
   {
   public:
       using ResourceAllocator::ResourceAllocator;
   
       static constexpr int DEFAULT_ALIGN = alignof(std::max_align_t);
   
       void *alloc(int64_t size, int32_t align = DEFAULT_ALIGN)
       {
           return m_data.res.mem.fnAlloc(m_data.ctx, size, align);
       }
   
       void free(void *ptr, int64_t size, int32_t align = DEFAULT_ALIGN) noexcept
       {
           m_data.res.mem.fnFree(m_data.ctx, ptr, size, align);
       }
   
       static constexpr bool IsCompatibleKind(NVCVResourceType resType)
       {
           return resType == NVCV_RESOURCE_MEM_HOST || resType == NVCV_RESOURCE_MEM_HOST_PINNED
               || resType == NVCV_RESOURCE_MEM_CUDA;
       }
   };
   
   namespace detail {
   
   template<NVCVResourceType KIND>
   class MemAllocatorWithKind : public MemAllocator
   {
   public:
       static constexpr NVCVResourceType kResourceType = KIND;
   
       MemAllocatorWithKind() = default;
   
       MemAllocatorWithKind(const NVCVResourceAllocator &data);
   
       static constexpr bool IsCompatibleKind(NVCVResourceType resType)
       {
           return resType == kResourceType;
       }
   };
   
   } // namespace detail
   
   class HostMemAllocator : public detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_HOST>
   {
       using Impl = detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_HOST>;
       using Impl::Impl;
   };
   
   class HostPinnedMemAllocator : public detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_HOST_PINNED>
   {
       using Impl = detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_HOST_PINNED>;
       using Impl::Impl;
   };
   
   class CudaMemAllocator : public detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_CUDA>
   {
       using Impl = detail::MemAllocatorWithKind<NVCV_RESOURCE_MEM_CUDA>;
       using Impl::Impl;
   };
   
   NVCV_IMPL_SHARED_HANDLE(Allocator);
   
   class Allocator : public CoreResource<NVCVAllocatorHandle, Allocator>
   {
   public:
       using CoreResource<NVCVAllocatorHandle, Allocator>::CoreResource;
   
       HostMemAllocator       hostMem() const;
       HostPinnedMemAllocator hostPinnedMem() const;
       CudaMemAllocator       cudaMem() const;
   
       ResourceAllocator get(NVCVResourceType resType) const;
   
       template<typename ResAlloc>
       ResAlloc get() const;
   };
   
   // Custom allocators
   
   template<typename AllocatorType>
   class CustomMemAllocator
   {
   private:
       template<typename Callable>
       struct has_trivial_copy_and_destruction
           : std::integral_constant<bool, std::is_trivially_copyable<Callable>::value
                                              && std::is_trivially_destructible<Callable>::value>
       {
       };
   
       template<typename Callable>
       struct by_value
           : std::integral_constant<bool, has_trivial_copy_and_destruction<Callable>::value
                                              && sizeof(Callable) <= sizeof(void *)
                                              && alignof(Callable) <= alignof(void *)>
       {
       };
   
       template<typename T>
       static constexpr size_t DataSize()
       {
           return std::is_empty<T>::value ? 0 : sizeof(T);
       }
   
   public:
       template<typename AllocFunction, typename FreeFunction,
                typename = detail::EnableIf_t<detail::IsInvocableR<void *, AllocFunction, int64_t, int32_t>::value>,
                typename = detail::EnableIf_t<detail::IsInvocableR<void, FreeFunction, void *, int64_t, int32_t>::value>>
       CustomMemAllocator(AllocFunction &&alloc, FreeFunction &&free);
   
       // TODO(michalz): Add a way of constructing a custom allocator without using lambdas/captures, e.g.
       //                from an object that matches the allocator concept.
   
       CustomMemAllocator(CustomMemAllocator &&other)
       {
           *this = std::move(other);
       }
   
       ~CustomMemAllocator()
       {
           reset();
       }
   
       bool needsCleanup() const noexcept
       {
           return m_data.cleanup != nullptr;
       }
   
       const NVCVResourceAllocator &cdata() const &noexcept
       {
           return m_data;
       }
   
       NVCV_NODISCARD NVCVResourceAllocator release() noexcept
       {
           NVCVResourceAllocator ret = {};
           std::swap(ret, m_data);
           return ret;
       }
   
       void reset(NVCVResourceAllocator &&alloc) noexcept
       {
           reset();
           std::swap(m_data, alloc);
       }
   
       void reset() noexcept
       {
           if (m_data.cleanup)
               m_data.cleanup(m_data.ctx, &m_data);
           m_data = {};
       }
   
       CustomMemAllocator &operator=(CustomMemAllocator &&other) noexcept
       {
           reset(other.release());
           return *this;
       }
   
   private:
       template<typename...>
       friend class CustomAllocator;
   
       template<typename AllocFunction, typename FreeFunction>
       void Construct(AllocFunction &&alloc, FreeFunction &&free, std::true_type);
   
       template<typename AllocFunction, typename FreeFunction>
       void Construct(AllocFunction &&alloc, FreeFunction &&free, std::false_type);
   
       template<typename AllocFunction, typename FreeFunction>
       void ConstructFromDuplicateValues(AllocFunction &&alloc, FreeFunction &&free, std::true_type);
   
   #if __cplusplus < 201703L
       template<typename AllocFunction, typename FreeFunction>
       void ConstructByOneValue(AllocFunction &&, FreeFunction &&, std::true_type, std::false_type)
       {
           assert(!"should never get here");
       }
   #endif
   
       NVCVResourceAllocator m_data{};
   };
   
   using CustomHostMemAllocator       = CustomMemAllocator<HostMemAllocator>;
   using CustomHostPinnedMemAllocator = CustomMemAllocator<HostPinnedMemAllocator>;
   using CustomCudaMemAllocator       = CustomMemAllocator<CudaMemAllocator>;
   
   template<typename... ResourceAllocators>
   class CustomAllocator final : public Allocator
   {
   public:
       explicit CustomAllocator(ResourceAllocators &&...allocators);
   
       ~CustomAllocator()
       {
           preDestroy();
       }
   
   private:
       static constexpr bool kHasReferences
           = detail::Disjunction<detail::IsRefWrapper<detail::RemoveRef_t<ResourceAllocators>>...>::value;
   
       template<bool hasReferences = kHasReferences>
       detail::EnableIf_t<hasReferences> preDestroy()
       {
           if (this->reset() != 0)
               throw std::logic_error(
                   "The allocator context contains references. The handle must not outlive the context.");
       }
   
       template<bool hasReferences = kHasReferences>
       detail::EnableIf_t<!hasReferences> preDestroy() noexcept
       {
       }
   };
   
   template<typename... ResourceAllocators>
   CustomAllocator<ResourceAllocators...> CreateCustomAllocator(ResourceAllocators &&...allocators)
   {
       return CustomAllocator<ResourceAllocators...>{std::move(allocators)...};
   }
   
   } // namespace nvcv
   
   #include "AllocatorImpl.hpp"
   
   #endif // NVCV_ALLOC_ALLOCATOR_HPP