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 integrate CHAI context with RAJA" 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,96 @@
+..
+    # 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;
+
+  ::RAJA::forall<::RAJA::cuda_exec_async<BLOCK_SIZE>>(::RAJA::TypedRangeSegment<int>(0, N), [=] __device__ (int i) {
+    // Use CHAI data structures in the DEVICE context...
+  });

host-configs/lc/toss_4_x86_64_ib/nvcc_clang.cmake

@@ -0,0 +1,35 @@
+##############################################################################
+# 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
+##############################################################################
+
+# Use gcc std libraries
+set(GCC_VER "13.3.1" CACHE STRING "")
+set(GCC_DIR "/usr/tce/packages/gcc/gcc-${GCC_VER}-magic" CACHE PATH "")
+
+# Use clang toolchain for host code compilers
+set(CLANG_VER "19.1.3" CACHE STRING "")
+set(CLANG_DIR "/usr/tce/packages/clang/clang-${CLANG_VER}-magic" CACHE PATH "")
+
+set(CMAKE_C_COMPILER "${CLANG_DIR}/bin/clang" CACHE PATH "")
+set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} --gcc-toolchain=${GCC_DIR}" CACHE STRING "")
+
+set(CMAKE_CXX_COMPILER "${CLANG_DIR}/bin/clang++" CACHE PATH "")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} --gcc-toolchain=${GCC_DIR}" CACHE STRING "")
+
+# Use nvcc as the device code compiler
+set(ENABLE_CUDA ON CACHE BOOL "")
+set(CUDA_VER "12.9.1" CACHE STRING "")
+set(CUDA_TOOLKIT_ROOT_DIR "/usr/tce/packages/cuda/cuda-${CUDA_VER}" CACHE PATH "")
+set(CMAKE_CUDA_COMPILER "${CUDA_TOOLKIT_ROOT_DIR}/bin/nvcc" CACHE PATH "")
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler=--gcc-toolchain=${GCC_DIR} --expt-relaxed-constexpr" CACHE STRING "")
+set(CMAKE_CUDA_HOST_COMPILER "${CMAKE_CXX_COMPILER}" CACHE PATH "")
+set(CMAKE_CUDA_ARCHITECTURES "90" CACHE STRING "")
+
+# Prevent incorrect implicit libraries from being linked in (if needed)
+set(BLT_CMAKE_IMPLICIT_LINK_DIRECTORIES_EXCLUDE "" CACHE STRING "")
+
+# The header only version of fmt in umpire has issues with nvcc
+set(UMPIRE_FMT_TARGET "fmt::fmt" CACHE STRING "")

host-configs/lc/toss_4_x86_64_ib_cray/amdclang.cmake

@@ -6,8 +6,8 @@
 ##############################################################################
 
 # Set up software versions
-set(ROCM_VERSION "6.2.0" CACHE PATH "")
-set(GCC_VERSION "12.2.1" CACHE PATH "")
+set(ROCM_VERSION "6.4.3" CACHE PATH "")
+set(GCC_VERSION "13.3.1" CACHE PATH "")
 
 # Set up compilers
 set(COMPILER_BASE "/usr/tce/packages/rocmcc/rocmcc-${ROCM_VERSION}-magic" CACHE PATH "")

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