Add experimental new design for context handling

cmake/SetupChaiOptions.cmake

@@ -16,6 +16,7 @@ option(CHAI_ENABLE_UM "Use CUDA unified (managed) memory" Off)
 option(CHAI_THIN_GPU_ALLOCATE "Single memory space model" Off)
 option(CHAI_ENABLE_PINNED "Use pinned host memory" Off)
 option(CHAI_ENABLE_RAJA_PLUGIN "Build plugin to set RAJA execution spaces" On)
+option(CHAI_ENABLE_EXPERIMENTAL_RAJA_PLUGIN "Build experimental plugin to set RAJA execution spaces" Off)
 option(CHAI_ENABLE_GPU_ERROR_CHECKING "Enable GPU error checking" On)
 option(CHAI_ENABLE_MANAGED_PTR "Enable managed_ptr" On)
 option(CHAI_DEBUG "Enable Debug Logging." Off)

docs/sphinx/expt/design.rst

@@ -0,0 +1,97 @@
+..
+    # Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and CHAI
+    # project contributors. See the CHAI LICENSE file for details.
+    #
+    # SPDX-License-Identifier: BSD-3-Clause
+
+.. _experimental_design:
+
+*******************
+Experimental Design
+*******************
+
+CHAI provides data structures that implicitly manage coherence across multiple execution contexts.
+
+-------
+Context
+-------
+
+Currently, there are two execution contexts that are handled by CHAI. These are represented in the `Context` enum class.
+The `HOST` enum value represents synchronous execution on a CPU. The `DEVICE` enum value represents asynchronous execution on a GPU.
+Both NVIDIA and AMD GPUs are supported.
+
+--------------
+ContextManager
+--------------
+
+Implicitly managing data coherence requires managing some global state. This is handled by a singleton called `ContextManager`.
+When an application enters an execution context, it uses `ContextManager` to set the current context. `ContextManager` also
+tracks which contexts may need synchronization. CHAI data structures can query `ContextManager` to update data coherence and
+inform `ContextManager` of needed synchronization or synchronization that has been performed.
+
+Note: It is much faster for `ContextManager` to track synchronization than to repeatedly call `cudaDeviceSynchronize()` or `hipDeviceSynchronize()`.
+
+.. code-block:: cpp
+
+  #include "chai/expt/ContextManager.hpp"
+
+  ::chai::expt::ContextManager& contextManager = ::chai::expt::ContextManager::getInstance();
+
+  contextManager.setContext(::chai::expt::Context::HOST);
+  // Use CHAI data structures in the HOST context...
+  contextManager.setContext(::chai::expt::Context::NONE);
+
+  contextManager.setContext(::chai::expt::Context::DEVICE);
+  // Use CHAI data structures in the DEVICE context...
+  contextManager.setContext(::chai::expt::Context::NONE);
+
+------------
+ContextGuard
+------------
+
+It is easy to forget to reset the current context or even to forget the current context
+when writing code. Similar to `std::lock_guard`, CHAI provides `ContextGuard` that sets
+the active context and then resets it upon destruction. This is the recommended approach.
+
+.. code-block:: cpp
+
+  #include "chai/expt/ContextGuard.hpp"
+
+  {
+    ::chai::expt::ContextGuard contextGuard{::chai::expt::Context::HOST};
+    // Use CHAI data structures in the HOST context...
+  }
+
+  {
+    ::chai::expt::ContextGuard contextGuard{::chai::expt::Context::DEVICE};
+    // Use CHAI data structures in the DEVICE context...
+  }
+
+-----------------
+ContextRAJAPlugin
+-----------------
+
+In an application that also uses RAJA, CHAI provides a RAJA plugin, `ContextRAJAPlugin`,
+that implicitly manages the context in calls to RAJA. To enable this plugin, configure with
+`-DCHAI_ENABLE_EXPERIMENTAL_RAJA_PLUGIN=ON` and register the plugin. In the future, registration
+may be handled by CHAI.
+
+.. code-block:: cpp
+
+  #include "chai/expt/ContextRAJAPlugin.hpp"
+  #include "RAJA/RAJA.hpp"
+
+  static ::RAJA::util::PluginRegistry::add<chai::expt::ContextRAJAPlugin> P(
+    "CHAIContextPlugin",
+    "Plugin that integrates CHAI context management with RAJA.");
+
+  ::RAJA::forall<::RAJA::seq_exec>(::RAJA::TypedRangeSegment<int>(0, N), [=] (int i) {
+    // Use CHAI data structures in the HOST context...
+  });
+
+  constexpr int BLOCK_SIZE = 256;
+  constexpr bool ASYNCHRONOUS = true;
+
+  ::RAJA::forall<::RAJA::cuda_exec<BLOCK_SIZE, ASYNCHRONOUS>>(::RAJA::TypedRangeSegment<int>(0, N), [=] __device__ (int i) {
+    // Use CHAI data structures in the DEVICE context...
+  });

src/chai/CMakeLists.txt

@@ -23,30 +23,31 @@ set (chai_headers
   PointerRecord.hpp
   Types.hpp)
 
+set (chai_sources
+  ArrayManager.cpp)
+
+if(CHAI_DISABLE_RM)
+  set(chai_headers
+    ${chai_headers}
+    ManagedArray_thin.inl)
+endif ()
+
 if(CHAI_ENABLE_EXPERIMENTAL)
   set(chai_headers
     ${chai_headers}
+    expt/Context.hpp
+    expt/ContextGuard.hpp
+    expt/ContextManager.hpp
     ManagedSharedPtr.hpp
     SharedPtrCounter.hpp
     SharedPtrManager.hpp
     SharedPtrManager.inl
     SharedPointerRecord.hpp)
-endif()
 
-if(CHAI_DISABLE_RM)
-  set(chai_headers
-    ${chai_headers}
-    ManagedArray_thin.inl)
-endif ()
-
-set (chai_sources
-  ArrayManager.cpp)
-
-if(CHAI_ENABLE_EXPERIMENTAL)
   set (chai_sources
     ${chai_sources}
     SharedPtrManager.cpp)
-endif ()
+endif()
 
 set (chai_depends
   umpire)
@@ -84,6 +85,16 @@ if (CHAI_ENABLE_RAJA_PLUGIN)
   endif ()
 endif ()
 
+if (CHAI_ENABLE_EXPERIMENTAL_RAJA_PLUGIN)
+  set (chai_headers
+    ${chai_headers}
+    expt/ContextRAJAPlugin.hpp)
+
+  set (chai_sources
+    ${chai_sources}
+    expt/ContextRAJAPlugin.cpp)
+endif()
+
 blt_add_library(
   NAME chai
   SOURCES ${chai_sources}

src/chai/expt/Context.hpp

@@ -0,0 +1,24 @@
+//////////////////////////////////////////////////////////////////////////////
+// Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and CHAI
+// project contributors. See the CHAI LICENSE file for details.
+//
+// SPDX-License-Identifier: BSD-3-Clause
+//////////////////////////////////////////////////////////////////////////////
+
+#ifndef CHAI_CONTEXT_HPP
+#define CHAI_CONTEXT_HPP
+
+namespace chai::expt
+{
+  /*!
+   * \brief Execution context identifier.
+   */
+  enum class Context
+  {
+    NONE = 0,   /*!< No context. */
+    HOST = 1,   /*!< Host (CPU) context. */
+    DEVICE = 2  /*!< Device (GPU/accelerator) context. */
+  };  // enum class Context
+}  // namespace chai::expt
+
+#endif  // CHAI_CONTEXT_HPP

src/chai/expt/ContextGuard.hpp

@@ -0,0 +1,49 @@
+//////////////////////////////////////////////////////////////////////////////
+// Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and CHAI
+// project contributors. See the CHAI LICENSE file for details.
+//
+// SPDX-License-Identifier: BSD-3-Clause
+//////////////////////////////////////////////////////////////////////////////
+
+#ifndef CHAI_CONTEXT_GUARD_HPP
+#define CHAI_CONTEXT_GUARD_HPP
+
+#include "chai/expt/Context.hpp"
+#include "chai/expt/ContextManager.hpp"
+
+namespace chai::expt {
+  /*!
+   * \brief RAII guard that temporarily sets the active Context and restores the
+   *        previously active Context upon destruction.
+   */
+  class ContextGuard {
+    public:
+      /*!
+       * \brief Sets the active Context for the lifetime of this guard.
+       * \param context The Context to set as active.
+       */
+      explicit ContextGuard(Context context) {
+        m_context_manager.setContext(context);
+      }
+
+      /*!
+       * \brief Restores the Context that was active when this guard was created.
+       */
+      ~ContextGuard() {
+        m_context_manager.setContext(m_saved_context);
+      }
+
+    private:
+      /*!
+       * \brief Reference to the global ContextManager instance.
+       */
+      ContextManager& m_context_manager{ContextManager::getInstance()};
+
+      /*!
+       * Context that was active at guard construction time.
+       */
+      Context m_saved_context{m_context_manager.getContext()};
+  };  // class ContextGuard
+}  // namespace chai::expt
+
+#endif  // CHAI_CONTEXT_GUARD_HPP

src/chai/expt/ContextManager.hpp

@@ -0,0 +1,143 @@
+//////////////////////////////////////////////////////////////////////////////
+// Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and CHAI
+// project contributors. See the CHAI LICENSE file for details.
+//
+// SPDX-License-Identifier: BSD-3-Clause
+//////////////////////////////////////////////////////////////////////////////
+
+#ifndef CHAI_CONTEXT_MANAGER_HPP
+#define CHAI_CONTEXT_MANAGER_HPP
+
+#include "chai/config.hpp"
+#include "chai/expt/Context.hpp"
+
+#if defined(CHAI_ENABLE_CUDA)
+#include <cuda_runtime.h>
+#elif defined(CHAI_ENABLE_HIP)
+#include <hip/hip_runtime.h>
+#endif
+
+namespace chai::expt {
+  /*!
+   * \brief Singleton class for managing the current context
+   *        and context synchronization across the application.
+   */
+  class ContextManager
+  {
+    public:
+      /*!
+       * \brief Get the singleton instance.
+       */
+      static ContextManager& getInstance()
+      {
+        static ContextManager s_instance;
+        return s_instance;
+      }
+
+      /*!
+       * \brief Disable copy construction.
+       *
+       * ContextManager is a singleton and must not be copied.
+       */
+      ContextManager(const ContextManager&) = delete;
+
+      /*!
+       * \brief Disable copy assignment.
+       *
+       * ContextManager is a singleton and must not be assigned.
+       */
+      ContextManager& operator=(const ContextManager&) = delete;
+
+      /*!
+       * \brief Get the current context.
+       */
+      Context getContext() const
+      {
+        return m_context;
+      }
+
+      /*!
+       * \brief Set the current context.
+       *
+       * Setting the context to DEVICE marks the device as not synchronized.
+       */
+      void setContext(Context context)
+      {
+        m_context = context;
+
+        if (context == Context::DEVICE)
+        {
+          m_device_synchronized = false;
+        }
+      }
+
+      /*!
+       * \brief Synchronize the requested context (no-op if already synchronized).
+       */
+      void synchronize(Context context)
+      {
+        if (context == Context::DEVICE && !m_device_synchronized)
+        {
+#if defined(CHAI_ENABLE_CUDA)
+          cudaDeviceSynchronize();
+#elif defined(CHAI_ENABLE_HIP)
+          hipDeviceSynchronize();
+#endif
+          m_device_synchronized = true;
+        }
+      }
+
+      /*!
+       * \brief Query whether the requested context is synchronized.
+       */
+      bool isSynchronized(Context context) const
+      {
+        return context == Context::DEVICE ? m_device_synchronized : true;
+      }
+
+      /*!
+       * \brief Explicitly set the synchronization state for the requested context.
+       */
+      void setSynchronized(Context context, bool synchronized)
+      {
+        if (context == Context::DEVICE)
+        {
+          m_device_synchronized = synchronized;
+        }
+      }
+
+      /*!
+       * \brief Reset manager state to defaults.
+       */
+      void reset()
+      {
+        m_context = Context::NONE;
+        m_device_synchronized = true;
+      }
+
+    private:
+      /*!
+       * \brief Default constructor.
+       *
+       * Private to enforce singleton access via getInstance().
+       */
+      ContextManager() = default;
+
+      /*!
+       * \brief Current context for the application.
+       *
+       * Defaults to NONE until explicitly set.
+       */
+      Context m_context{Context::NONE};
+
+      /*!
+       * \brief Device synchronization state.
+       *
+       * True if the device context has been synchronized since the last time the
+       * context was set to DEVICE.
+       */
+      bool m_device_synchronized{true};
+  };  // class ContextManager
+}  // namespace chai::expt
+
+#endif  // CHAI_CONTEXT_MANAGER_HPP