From d2f87078732f866ea1cb0c6ea944242abec4c4a9 Mon Sep 17 00:00:00 2001 From: Alan Dayton Date: Tue, 10 Feb 2026 12:47:25 -0800 Subject: [PATCH 1/3] Add UnifiedArrayManager --- src/chai/CMakeLists.txt | 1 + src/chai/expt/UnifiedArrayManager.hpp | 151 ++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) create mode 100644 src/chai/expt/UnifiedArrayManager.hpp diff --git a/src/chai/CMakeLists.txt b/src/chai/CMakeLists.txt index 2cd8ceee..87a3a452 100644 --- a/src/chai/CMakeLists.txt +++ b/src/chai/CMakeLists.txt @@ -42,6 +42,7 @@ if(CHAI_ENABLE_EXPERIMENTAL) expt/HostSharedPointer.hpp expt/ManagedArrayPointer.hpp expt/ManagedArraySharedPointer.hpp + expt/UnifiedArrayManager.hpp ManagedSharedPtr.hpp SharedPtrCounter.hpp SharedPtrManager.hpp diff --git a/src/chai/expt/UnifiedArrayManager.hpp b/src/chai/expt/UnifiedArrayManager.hpp new file mode 100644 index 00000000..b6c3f7f0 --- /dev/null +++ b/src/chai/expt/UnifiedArrayManager.hpp @@ -0,0 +1,151 @@ +////////////////////////////////////////////////////////////////////////////// +// Copyright (c) 2016-25, Lawrence Livermore National Security, LLC and CHAI +// project contributors. See the CHAI LICENSE file for details. +// +// SPDX-License-Identifier: BSD-3-Clause +////////////////////////////////////////////////////////////////////////////// +#ifndef CHAI_UNIFIED_ARRAY_MANAGER_HPP +#define CHAI_UNIFIED_ARRAY_MANAGER_HPP + +#include "umpire/ResourceManager.hpp" +#include "umpire/TypedAllocator.hpp" +#include +#include + +namespace chai::expt +{ + /*! + * \brief This class manages a host array. It is designed for use with + * ManagedArrayPointer. + * + * \tparam ElementType The type of elements contained in this array. + * + * \note UnifiedArrayManager performs value initialization of each array element. + * That is to say, numeric types will be initialized to zero and nontrivial + * types will be default constructed. In the future, this behavior may + * change to default initialization for performance reasons, such that + * numeric types will be left in an indeterminate state and nontrivial + * types will be default constructed. + * + * \note UnifiedArrayManager does not rely on ContextManager, so it will behave + * differently than other array managers. The major difference is that + * when used with ManagedArrayPointer, the ManagedArrayPointer does not + * need the update method called or to be copy constructed before it can + * be used on the host. Since it will not respect the current Context, + * be extra careful to avoid using it on the device. + */ + template + class UnifiedArrayManager { + private: + /*! + * \brief Allocator used by the managed host storage. + */ + using AllocatorType = ::umpire::TypedAllocator; + + /*! + * \brief Underlying contiguous host storage type for managed elements. + */ + using StorageType = std::vector; + + public: + /*! + * \brief Default-constructs a UnifiedArrayManager with zero elements + * and a default allocator for host memory allocations. + */ + UnifiedArrayManager() = default; + + /*! + * \brief Constructs a UnifiedArrayManager with zero elements + * and \p allocator for host memory allocations. + * + * \param allocator Allocator used for host memory allocations. + */ + explicit UnifiedArrayManager(const umpire::Allocator& allocator) + : m_storage{StorageType(AllocatorType(allocator))} + { + } + + /*! + * \brief Constructs a UnifiedArrayManager with \p size elements + * using the default allocator for host memory allocations. + * + * \param size Number of elements to allocate. + */ + explicit UnifiedArrayManager(std::size_t size) + : m_storage{StorageType(size, AllocatorType(::umpire::ResourceManager::getInstance().getAllocator("HOST")))} + { + } + + /*! + * \brief Constructs a UnifiedArrayManager with \p size elements + * using \p allocator for host memory allocations. + * + * \param size Number of elements to allocate. + * \param allocator Allocator used for host memory allocations. + */ + UnifiedArrayManager(std::size_t size, + const umpire::Allocator& allocator) + : m_storage{StorageType(size, AllocatorType(allocator))} + { + } + + /*! + * \brief Resizes the managed storage to \p new_size elements. + * + * \param new_size New number of elements. + */ + void resize(std::size_t new_size) + { + m_storage.resize(new_size); + } + + /*! + * \brief Returns the number of elements currently managed. + * + * \return Number of elements in the managed storage. + */ + std::size_t size() const + { + return m_storage.size(); + } + + /*! + * \brief Returns a pointer to the underlying contiguous storage. + * + * \return Pointer to the first element, or nullptr if the storage is empty. + */ + ElementType* data(bool touch) + { + ContextManager& contextManager = ContextManager::getInstance(); + Context context = contextManager.getContext(); + + if (context != m_modified) + { + contextManager.synchronize(m_modified); + } + + if (touch) + { + m_modified = context; + } + + return m_storage.empty() ? nullptr : m_storage.data(); + } + + private: + /*! + * \brief Underlying host storage for the managed elements. + */ + StorageType m_storage{AllocatorType(::umpire::ResourceManager::getInstance().getAllocator("UM"))}; + + /*! + * \brief Context in which the managed storage was most recently modified. + * + * \note Used to determine when synchronization is required before accessing + * the underlying storage from the current context. + */ + Context m_modified{Context::NONE}; + }; // class UnifiedArrayManager +} // namespace chai::expt + +#endif // CHAI_UNIFIED_ARRAY_MANAGER_HPP From 7cc351139115fff98c9e2d5370de349f56aebb81 Mon Sep 17 00:00:00 2001 From: Alan Dayton Date: Tue, 10 Feb 2026 12:52:52 -0800 Subject: [PATCH 2/3] Clean up documentation --- src/chai/expt/UnifiedArrayManager.hpp | 37 +++++++++++---------------- 1 file changed, 15 insertions(+), 22 deletions(-) diff --git a/src/chai/expt/UnifiedArrayManager.hpp b/src/chai/expt/UnifiedArrayManager.hpp index b6c3f7f0..6f2e0740 100644 --- a/src/chai/expt/UnifiedArrayManager.hpp +++ b/src/chai/expt/UnifiedArrayManager.hpp @@ -15,50 +15,43 @@ namespace chai::expt { /*! - * \brief This class manages a host array. It is designed for use with - * ManagedArrayPointer. + * \brief This class manages a unified memory array. It is designed for use + * with ManagedArrayPointer. * * \tparam ElementType The type of elements contained in this array. * * \note UnifiedArrayManager performs value initialization of each array element. * That is to say, numeric types will be initialized to zero and nontrivial - * types will be default constructed. In the future, this behavior may - * change to default initialization for performance reasons, such that - * numeric types will be left in an indeterminate state and nontrivial - * types will be default constructed. - * - * \note UnifiedArrayManager does not rely on ContextManager, so it will behave - * differently than other array managers. The major difference is that - * when used with ManagedArrayPointer, the ManagedArrayPointer does not - * need the update method called or to be copy constructed before it can - * be used on the host. Since it will not respect the current Context, - * be extra careful to avoid using it on the device. + * types will be default constructed. This initialization occurs on the host. + * In the future, this behavior may change to default initialization for + * performance reasons, such that numeric types will be left in an + * indeterminate state and nontrivial types will be default constructed. */ template class UnifiedArrayManager { private: /*! - * \brief Allocator used by the managed host storage. + * \brief Allocator used by the managed unified memory storage. */ using AllocatorType = ::umpire::TypedAllocator; /*! - * \brief Underlying contiguous host storage type for managed elements. + * \brief Underlying contiguous unified memory storage type for managed elements. */ using StorageType = std::vector; public: /*! * \brief Default-constructs a UnifiedArrayManager with zero elements - * and a default allocator for host memory allocations. + * and a default allocator for unified memory allocations. */ UnifiedArrayManager() = default; /*! * \brief Constructs a UnifiedArrayManager with zero elements - * and \p allocator for host memory allocations. + * and \p allocator for unified memory allocations. * - * \param allocator Allocator used for host memory allocations. + * \param allocator Allocator used for unified memory allocations. */ explicit UnifiedArrayManager(const umpire::Allocator& allocator) : m_storage{StorageType(AllocatorType(allocator))} @@ -67,7 +60,7 @@ namespace chai::expt /*! * \brief Constructs a UnifiedArrayManager with \p size elements - * using the default allocator for host memory allocations. + * using the default allocator for unified memory allocations. * * \param size Number of elements to allocate. */ @@ -78,10 +71,10 @@ namespace chai::expt /*! * \brief Constructs a UnifiedArrayManager with \p size elements - * using \p allocator for host memory allocations. + * using \p allocator for unified memory allocations. * * \param size Number of elements to allocate. - * \param allocator Allocator used for host memory allocations. + * \param allocator Allocator used for unified memory allocations. */ UnifiedArrayManager(std::size_t size, const umpire::Allocator& allocator) @@ -134,7 +127,7 @@ namespace chai::expt private: /*! - * \brief Underlying host storage for the managed elements. + * \brief Underlying unified memory storage for the managed elements. */ StorageType m_storage{AllocatorType(::umpire::ResourceManager::getInstance().getAllocator("UM"))}; From 4736f8f0a0dfb0cc543ca92d35cc022ab5df4603 Mon Sep 17 00:00:00 2001 From: Alan Dayton Date: Tue, 10 Feb 2026 15:44:58 -0800 Subject: [PATCH 3/3] Use UM allocator for UnifiedArrayManager --- src/chai/expt/UnifiedArrayManager.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/chai/expt/UnifiedArrayManager.hpp b/src/chai/expt/UnifiedArrayManager.hpp index 6f2e0740..fa7333fa 100644 --- a/src/chai/expt/UnifiedArrayManager.hpp +++ b/src/chai/expt/UnifiedArrayManager.hpp @@ -65,7 +65,7 @@ namespace chai::expt * \param size Number of elements to allocate. */ explicit UnifiedArrayManager(std::size_t size) - : m_storage{StorageType(size, AllocatorType(::umpire::ResourceManager::getInstance().getAllocator("HOST")))} + : m_storage{StorageType(size, AllocatorType(::umpire::ResourceManager::getInstance().getAllocator("UM")))} { }