hpt_allocator/
lib.rs

1//! This crate provides a memory allocator for tensor
2
3#![deny(missing_docs)]
4
5mod allocators;
6mod ptr;
7mod storage;
8pub(crate) mod utils {
9    pub(crate) mod allocate;
10    pub(crate) mod deallocate;
11}
12/// traits for the allocator
13pub mod traits;
14
15pub use crate::allocators::cpu::CACHE;
16#[cfg(feature = "cuda")]
17pub use crate::allocators::cuda::CUDA_CACHE;
18pub use crate::storage::clone_storage;
19pub use storage::cpu::CPU_STORAGE;
20#[cfg(feature = "cuda")]
21pub use storage::cuda::CUDA_STORAGE;
22use traits::Allocator;
23
24/// program will free all the memory before exit
25#[allow(non_snake_case)]
26#[ctor::dtor]
27fn free_pools() {
28    CACHE.lock().unwrap().clear();
29    #[cfg(feature = "cuda")]
30    CUDA_CACHE.lock().unwrap().clear();
31}