From f294db31bea86d08b8d875d218f24c65221dca76 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Sun, 22 Sep 2024 06:32:39 -0400 Subject: [PATCH] TA::host_allocator is serializable, so that btas::Tensor can be used as a tile again --- src/CMakeLists.txt | 2 - src/TiledArray/device/allocators.h | 138 ---------------------------- src/TiledArray/device/um_storage.cu | 2 +- src/TiledArray/device/um_storage.h | 2 +- src/TiledArray/external/device.h | 15 ++- src/TiledArray/external/umpire.h | 83 ++++++++++++++++- src/TiledArray/fwd.h | 32 ++++--- src/TiledArray/host/allocator.h | 78 ---------------- src/TiledArray/host/env.h | 10 ++ src/TiledArray/tensor/tensor.h | 3 +- 10 files changed, 127 insertions(+), 238 deletions(-) delete mode 100644 src/TiledArray/device/allocators.h delete mode 100644 src/TiledArray/host/allocator.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c426d1ffbe..3d6b94ea9a 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -134,7 +134,6 @@ TiledArray/external/btas.h TiledArray/external/madness.h TiledArray/external/umpire.h TiledArray/host/env.h -TiledArray/host/allocator.h TiledArray/math/blas.h TiledArray/math/gemm_helper.h TiledArray/math/outer.h @@ -223,7 +222,6 @@ if(CUDA_FOUND OR HIP_FOUND) TiledArray/device/kernel/thrust/reduce_kernel.h TiledArray/device/platform.h TiledArray/device/thrust.h - TiledArray/device/allocators.h TiledArray/device/um_storage.h) if(CUDA_FOUND) list(APPEND TILEDARRAY_HEADER_FILES diff --git a/src/TiledArray/device/allocators.h b/src/TiledArray/device/allocators.h deleted file mode 100644 index 2bda79e768..0000000000 --- a/src/TiledArray/device/allocators.h +++ /dev/null @@ -1,138 +0,0 @@ -/* - * This file is a part of TiledArray. - * Copyright (C) 2018 Virginia Tech - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU General Public License as published by - * the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - * - * You should have received a copy of the GNU General Public License - * along with this program. If not, see . - * - * Eduard Valeyev - * Department of Chemistry, Virginia Tech - * Jan 31, 2018 - * - */ - -#ifndef TILEDARRAY_DEVICE_ALLOCATORS_H___INCLUDED -#define TILEDARRAY_DEVICE_ALLOCATORS_H___INCLUDED - -#include - -#ifdef TILEDARRAY_HAS_DEVICE - -#include -#include - -#include - -#include -#include - -namespace TiledArray { - -template -class umpire_based_allocator - : public umpire_based_allocator_impl { - public: - using base_type = umpire_based_allocator_impl; - using typename base_type::const_pointer; - using typename base_type::const_reference; - using typename base_type::pointer; - using typename base_type::reference; - using typename base_type::value_type; - - umpire_based_allocator() noexcept : base_type(&UmpireAllocatorAccessor{}()) {} - - template - umpire_based_allocator( - const umpire_based_allocator& - rhs) noexcept - : base_type( - static_cast&>( - rhs)) {} - - template - friend bool operator==( - const umpire_based_allocator& - lhs, - const umpire_based_allocator& - rhs) noexcept; -}; // class umpire_based_allocator - -template -bool operator==( - const umpire_based_allocator& lhs, - const umpire_based_allocator& - rhs) noexcept { - return lhs.umpire_allocator() == rhs.umpire_allocator(); -} - -template -bool operator!=( - const umpire_based_allocator& lhs, - const umpire_based_allocator& - rhs) noexcept { - return !(lhs == rhs); -} - -namespace detail { - -struct get_um_allocator { - umpire::Allocator& operator()() { - return deviceEnv::instance()->um_allocator(); - } -}; - -struct get_pinned_allocator { - umpire::Allocator& operator()() { - return deviceEnv::instance()->pinned_allocator(); - } -}; - -} // namespace detail - -} // namespace TiledArray - -namespace madness { -namespace archive { - -template -struct ArchiveLoadImpl> { - static inline void load( - const Archive& ar, - TiledArray::umpire_based_allocator& allocator) { - allocator = TiledArray::umpire_based_allocator{}; - } -}; - -template -struct ArchiveStoreImpl> { - static inline void store( - const Archive& ar, - const TiledArray::umpire_based_allocator< - T, StaticLock, UmpireAllocatorAccessor>& allocator) {} -}; - -} // namespace archive -} // namespace madness - -#endif // TILEDARRAY_HAS_DEVICE - -#endif // TILEDARRAY_DEVICE_ALLOCATORS_H___INCLUDED diff --git a/src/TiledArray/device/um_storage.cu b/src/TiledArray/device/um_storage.cu index cc3a1aae55..8879c246f8 100644 --- a/src/TiledArray/device/um_storage.cu +++ b/src/TiledArray/device/um_storage.cu @@ -22,7 +22,7 @@ */ -#include +#include #include #ifdef TILEDARRAY_HAS_CUDA diff --git a/src/TiledArray/device/um_storage.h b/src/TiledArray/device/um_storage.h index d151a3c316..d91c032312 100644 --- a/src/TiledArray/device/um_storage.h +++ b/src/TiledArray/device/um_storage.h @@ -24,7 +24,7 @@ #ifndef TILEDARRAY_DEVICE_UM_VECTOR_H__INCLUDED #define TILEDARRAY_DEVICE_UM_VECTOR_H__INCLUDED -#include +#include #ifdef TILEDARRAY_HAS_DEVICE diff --git a/src/TiledArray/external/device.h b/src/TiledArray/external/device.h index 38bcbbc745..597643b225 100644 --- a/src/TiledArray/external/device.h +++ b/src/TiledArray/external/device.h @@ -798,9 +798,22 @@ class Env { static std::unique_ptr instance_{nullptr}; return instance_; } -}; +}; // class Env namespace detail { + +struct get_um_allocator { + umpire::Allocator& operator()() { + return deviceEnv::instance()->um_allocator(); + } +}; + +struct get_pinned_allocator { + umpire::Allocator& operator()() { + return deviceEnv::instance()->pinned_allocator(); + } +}; + // in a madness device task point to its local optional stream to use by // madness_task_stream_opt; set to nullptr after task callable finished inline std::optional*& madness_task_stream_opt_ptr_accessor() { diff --git a/src/TiledArray/external/umpire.h b/src/TiledArray/external/umpire.h index e8d0d48632..ac23a60260 100644 --- a/src/TiledArray/external/umpire.h +++ b/src/TiledArray/external/umpire.h @@ -156,6 +156,54 @@ bool operator!=( return !(lhs == rhs); } +template +class umpire_based_allocator + : public umpire_based_allocator_impl { + public: + using base_type = umpire_based_allocator_impl; + using typename base_type::const_pointer; + using typename base_type::const_reference; + using typename base_type::pointer; + using typename base_type::reference; + using typename base_type::value_type; + + umpire_based_allocator() noexcept : base_type(&UmpireAllocatorAccessor{}()) {} + + template + umpire_based_allocator( + const umpire_based_allocator& + rhs) noexcept + : base_type( + static_cast&>( + rhs)) {} + + template + friend bool operator==( + const umpire_based_allocator& + lhs, + const umpire_based_allocator& + rhs) noexcept; +}; // class umpire_based_allocator + +template +bool operator==( + const umpire_based_allocator& lhs, + const umpire_based_allocator& + rhs) noexcept { + return lhs.umpire_allocator() == rhs.umpire_allocator(); +} + +template +bool operator!=( + const umpire_based_allocator& lhs, + const umpire_based_allocator& + rhs) noexcept { + return !(lhs == rhs); +} + /// see /// https://stackoverflow.com/questions/21028299/is-this-behavior-of-vectorresizesize-type-n-under-c11-and-boost-container/21028912#21028912 template @@ -202,7 +250,7 @@ struct ArchiveLoadImpl& allocator) { std::string allocator_name; - ar& allocator_name; + ar & allocator_name; allocator = TiledArray::umpire_based_allocator_impl( umpire::ResourceManager::getInstance().getAllocator(allocator_name)); } @@ -214,7 +262,7 @@ struct ArchiveStoreImpl< static inline void store( const Archive& ar, const TiledArray::umpire_based_allocator_impl& allocator) { - ar& allocator.umpire_allocator()->getName(); + ar & allocator.umpire_allocator()->getName(); } }; @@ -224,7 +272,7 @@ struct ArchiveLoadImpl> { TiledArray::default_init_allocator& allocator) { if constexpr (!std::allocator_traits::is_always_equal::value) { A base_allocator; - ar& base_allocator; + ar & base_allocator; allocator = TiledArray::default_init_allocator(base_allocator); } } @@ -244,4 +292,33 @@ struct ArchiveStoreImpl> { } // namespace archive } // namespace madness +namespace madness { +namespace archive { + +template +struct ArchiveLoadImpl> { + static inline void load( + const Archive& ar, + TiledArray::umpire_based_allocator& allocator) { + allocator = TiledArray::umpire_based_allocator{}; + } +}; + +template +struct ArchiveStoreImpl> { + static inline void store( + const Archive& ar, + const TiledArray::umpire_based_allocator< + T, StaticLock, UmpireAllocatorAccessor>& allocator) {} +}; + +} // namespace archive +} // namespace madness + #endif // TILEDARRAY_EXTERNAL_UMPIRE_H___INCLUDED diff --git a/src/TiledArray/fwd.h b/src/TiledArray/fwd.h index 97d91a9a00..6127db32f3 100644 --- a/src/TiledArray/fwd.h +++ b/src/TiledArray/fwd.h @@ -36,12 +36,27 @@ class aligned_allocator; // fwddecl host_allocator namespace TiledArray { -template -class host_allocator_impl; -template +namespace detail { +struct get_host_allocator; +struct NullLock; +template +class MutexLock; +} // namespace detail + +template +class umpire_based_allocator; + +template > class default_init_allocator; + +class hostEnv; + +/// pooled thread-safe host memory allocator template -using host_allocator = default_init_allocator>; +using host_allocator = + default_init_allocator, + detail::get_host_allocator>>; } // namespace TiledArray namespace madness { @@ -87,18 +102,9 @@ class Env; } using deviceEnv = device::Env; -template -class umpire_based_allocator; - -template > -class default_init_allocator; - namespace detail { struct get_um_allocator; struct get_pinned_allocator; -struct NullLock; -template -class MutexLock; } // namespace detail /// pooled thread-safe unified memory (UM) allocator for device computing diff --git a/src/TiledArray/host/allocator.h b/src/TiledArray/host/allocator.h deleted file mode 100644 index a22613fb38..0000000000 --- a/src/TiledArray/host/allocator.h +++ /dev/null @@ -1,78 +0,0 @@ -/* - * This file is a part of TiledArray. - * Copyright (C) 2021 Virginia Tech - * - * This program is free software: you can redistribute it and/or modify - * it under the terms of the GNU General Public License as published by - * the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. - * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU General Public License for more details. - * - * You should have received a copy of the GNU General Public License - * along with this program. If not, see . - * - * Eduard Valeyev - * Department of Chemistry, Virginia Tech - * Jan 31, 2018 - * - */ - -#ifndef TILEDARRAY_HOST_ALLOCATOR_H___INCLUDED -#define TILEDARRAY_HOST_ALLOCATOR_H___INCLUDED - -#include - -#include -#include - -#include - -#include -#include - -namespace TiledArray { - -/// pooled, thread-safe allocator for host memory -template -class host_allocator_impl - : public umpire_based_allocator_impl> { - public: - using base_type = umpire_based_allocator_impl>; - using typename base_type::const_pointer; - using typename base_type::const_reference; - using typename base_type::pointer; - using typename base_type::reference; - using typename base_type::value_type; - - host_allocator_impl() noexcept - : base_type(&hostEnv::instance()->host_allocator()) {} - - template - host_allocator_impl(const host_allocator_impl& rhs) noexcept - : base_type(static_cast>&>(rhs)) {} - - template - friend bool operator==(const host_allocator_impl& lhs, - const host_allocator_impl& rhs) noexcept; -}; // class host_allocator_impl - -template -bool operator==(const host_allocator_impl& lhs, - const host_allocator_impl& rhs) noexcept { - return lhs.umpire_allocator() == rhs.umpire_allocator(); -} - -template -bool operator!=(const host_allocator_impl& lhs, - const host_allocator_impl& rhs) noexcept { - return !(lhs == rhs); -} - -} // namespace TiledArray - -#endif // TILEDARRAY_HOST_ALLOCATOR_H___INCLUDED diff --git a/src/TiledArray/host/env.h b/src/TiledArray/host/env.h index 1b3c4f277f..be1de5369c 100644 --- a/src/TiledArray/host/env.h +++ b/src/TiledArray/host/env.h @@ -148,6 +148,16 @@ class hostEnv { } }; +namespace detail { + +struct get_host_allocator { + umpire::Allocator& operator()() { + return hostEnv::instance()->host_allocator(); + } +}; + +} // namespace detail + } // namespace TiledArray #endif // TILEDARRAY_HOST_ENV_H__INCLUDED diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 12479ef53c..171dac2eea 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -22,7 +22,8 @@ #include "TiledArray/config.h" -#include "TiledArray/host/allocator.h" +#include "TiledArray/external/umpire.h" +#include "TiledArray/host/env.h" #include "TiledArray/math/blas.h" #include "TiledArray/math/gemm_helper.h"