Last active
October 8, 2019 13:39
-
-
Save jrhemstad/36df9016cc5d1caff5ae36d9ba91fa51 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| namespace { | |
| device_memory_resource* cuda_resource() { | |
| static cuda_memory_resource resource{}; | |
| return &resource; | |
| std::atomic<device_memory_resource*>& get_default() { | |
| static std::atomic<device_memory_resource*> res{detail::cuda_resource()}; | |
| return res; | |
| } | |
| } // namespace | |
| device_memory_resource* get_default_resource() { return get_default().load(); } | |
| device_memory_resource* set_default_resource( | |
| device_memory_resource* new_resource) { | |
| new_resource = | |
| (new_resource == nullptr) ? detail::cuda_resource() : new_resource; | |
| return get_default().exchange(new_resource); | |
| } | |
| enum class resource { DEFAULT, POOL, MANAGED, MANAGED_POOL }; | |
| /**---------------------------------------------------------------------------* | |
| * @brief Maintains unique global state for lifetime of library. | |
| *---------------------------------------------------------------------------**/ | |
| class state_singleton { | |
| public: | |
| state_singleton& Instance() { | |
| static state_singleton instance; | |
| return instance; | |
| } | |
| /**---------------------------------------------------------------------------* | |
| * @brief Set the default resource to the requested resource. | |
| * | |
| * Does nothing if `initialize()` was already called without then calling | |
| * `finalize()`. | |
| * | |
| * @param mode The requested resource | |
| *---------------------------------------------------------------------------**/ | |
| void initialize(resource mode) { | |
| if (not initialized) { | |
| initialized = true; | |
| switch (mode) { | |
| case resource::DEFAULT: | |
| set_default_resource(nullptr); | |
| case resource::POOL: | |
| set_default_resource(pool.reset(new pool_memory_resource{}).get()); | |
| case resource::MANAGED: | |
| set_default_resource(managed.reset(new managed_memory_resource{}).get()); | |
| case resource::MANAGED_POOL: | |
| set_default_resource(managed_pool.reset(new managed_pool_memory_resource{}).get()); | |
| } | |
| } | |
| } | |
| /**---------------------------------------------------------------------------* | |
| * @brief Frees memory held by all resources. | |
| * | |
| * Does nothing if `initialize()` was not already called. | |
| * | |
| *---------------------------------------------------------------------------**/ | |
| void finalize() { | |
| if (initiazed) { | |
| pool.reset(); | |
| managed.reset(); | |
| managed_pool.reset(); | |
| set_default_resource(nullptr); | |
| } | |
| } | |
| private: | |
| bool initialized{false}; | |
| std::unique_ptr<pool_memory_resource> pool{}; | |
| std::unique_ptr<managed_memory_resource> managed{}; | |
| std::unique_ptr<managed_pool_memory_resource> managed_pool{}; | |
| }; |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment