diff --git a/CHANGELOG.md b/CHANGELOG.md index 54ed64cc9964c..3fcd2e3828a43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ## Current develop ### Added (new features/APIs/variables/...) +- [[PR 1394]](https://github.com/parthenon-hpc-lab/parthenon/pull/1394) Add token scratch utility - [[PR 1431]](https://github.com/parthenon-hpc-lab/parthenon/pull/1431) Add reductions to the loop abstraction - [[PR 1429]](https://github.com/parthenon-hpc-lab/parthenon/pull/1429) Fixes to swarm tensor xdmf - [[PR 1428]](https://github.com/parthenon-hpc-lab/parthenon/pull/1428) Permit larger input file sizes diff --git a/doc/plan_histories/token_scratch.md b/doc/plan_histories/token_scratch.md new file mode 100644 index 0000000000000..12bfc3f2d5020 --- /dev/null +++ b/doc/plan_histories/token_scratch.md @@ -0,0 +1,213 @@ +# Token-Based Scratch Memory System + +## Overview + +The token scratch system provides efficient, thread-safe scratch memory management for Kokkos parallel kernels using `Kokkos::Experimental::UniqueToken`. Each parallel thread acquires a unique token ID and gets access to a pre-allocated scratch buffer that can be carved into multiple typed views. + +## Key Features + +- **Thread-safe**: Uses Kokkos UniqueToken to avoid race conditions +- **Type-safe**: Automatic type deduction and alignment handling +- **Zero runtime overhead**: All inline functions, compile-time resolution +- **Flexible allocation**: Supports 1D, 2D, and 3D views of arbitrary types +- **RAII-based**: Automatic token acquisition and release +- **Debug validation**: Bounds checking in debug builds + +## Components + +### 1. `TokenScratchPool` + +The main pool manager that allocates scratch memory for all tokens. + +**Template Parameters:** +- `MemorySpace` - Kokkos memory space (e.g., `HostSpace`, device memory space) +- `ExecutionSpace` - Kokkos execution space +- `TokenScope` (optional) - `UniqueTokenScope::Global` (default) or `UniqueTokenScope::Instance` + +**Constructor:** +```cpp +explicit TokenScratchPool(size_t bytes_per_token) +``` + +**Methods:** +- `acquire()` - Returns a `ScratchAllocator` for the current thread +- `size()` - Number of tokens in the pool +- `bytes_per_token()` - Bytes allocated per token +- `total_bytes()` - Total pool size + +### 2. `ScratchAllocator` + +Per-token allocator that carves typed views from the token's buffer. + +**Methods:** +- `allocate_view(n)` - Allocate 1D view with n elements +- `allocate_view(n1, n2)` - Allocate 2D view +- `allocate_view(n1, n2, n3)` - Allocate 3D view +- `remaining()` - Get remaining capacity +- `reset()` - Reset allocator to beginning (for reuse) + +### 3. `TypedScratchBundle` + +(Optional) Pre-configured bundle for compile-time known layouts. + +## Usage Examples + +### Basic Usage + +```cpp +using ExecSpace = Kokkos::DefaultExecutionSpace; +using MemSpace = ExecSpace::memory_space; + +// Create pool with 64KB per token +TokenScratchPool pool(64 * 1024); + +Kokkos::parallel_for("my_kernel", N, KOKKOS_LAMBDA(int i) { + // Acquire scratch allocator + auto scratch = pool.acquire(); + + // Allocate multiple typed views + auto doubles = scratch.template allocate_view(100); + auto ints = scratch.template allocate_view(50); + auto flags = scratch.template allocate_view(200); + + // Use the views... + for (int j = 0; j < 100; ++j) { + doubles(j) = compute(i, j); + } + + // Token automatically released when scratch goes out of scope +}); +``` + +### Multi-dimensional Views + +```cpp +TokenScratchPool pool(128 * 1024); + +Kokkos::parallel_for("multidim", N, KOKKOS_LAMBDA(int i) { + auto scratch = pool.acquire(); + + // 2D and 3D scratch arrays + auto work_2d = scratch.template allocate_view(ni, nj); + auto work_3d = scratch.template allocate_view(ni, nj, nk); + + // Access as: work_2d(i, j), work_3d(i, j, k) +}); +``` + +## Design Considerations + +### Memory Layout + +The pool allocates a 2D view: `View` with shape `[num_tokens, bytes_per_token]`. + +- Each row is owned by one token +- Allocations within a token are sequential with proper alignment +- No fragmentation within a token's buffer + +### Token Scope + +Two scope options (specified as a template parameter): + +1. **`UniqueTokenScope::Global`** (default) + - Tokens shared across all kernel invocations + - More token reuse, potentially fewer tokens needed + - May see contention under high concurrency + ```cpp + // Explicit Global scope (or omit for default) + TokenScratchPool + pool(scratch_bytes); + ``` + +2. **`UniqueTokenScope::Instance`** + - Tokens private to each kernel instance + - Better for very high thread counts + - May allocate more tokens + ```cpp + // Use Instance scope for high concurrency + TokenScratchPool + pool(scratch_bytes); + ``` + +### Sizing Guidelines + +Choose `bytes_per_token` based on: + +1. **Per-thread scratch needs**: Sum of all allocations in your kernel +2. **Alignment overhead**: Add ~20% for alignment padding +3. **Safety margin**: Add 10-20% buffer for debug assertions + +Example calculation: +```cpp +// Kernel needs: +// - 100 doubles = 800 bytes +// - 50 ints = 200 bytes +// - 200 bools = 200 bytes +// Total: 1200 bytes +// With alignment (20%): 1440 bytes +// With safety margin (20%): 1728 bytes → round to 2KB +constexpr size_t scratch_bytes = 2 * 1024; +``` + +### Performance Notes + +- **Zero overhead in release builds**: All functions inline, no virtual dispatch +- **Debug validation**: Bounds checking only in debug builds (`#ifndef NDEBUG`) +- **No dynamic allocation**: All memory pre-allocated at pool creation +- **Cache-friendly**: Each token's buffer is contiguous + +## Integration with the rest of Parthenon + +```cpp +// In package initialization or task setup +namespace Package { + +void MyTask(MeshData* md) { + auto pmb = md->GetBlockData(0)->GetBlockPointer(); + const auto &ib = pmb->cellbounds.GetBoundsI(IndexDomain::interior); + const auto &jb = pmb->cellbounds.GetBoundsJ(IndexDomain::interior); + const auto &kb = pmb->cellbounds.GetBoundsK(IndexDomain::interior); + + // Size scratch based on block dimensions + const size_t ni = ib.e - ib.s + 1; + const size_t nj = jb.e - jb.s + 1; + const size_t nk = kb.e - kb.s + 1; + + const size_t scratch_per_k = (ni * nj) * sizeof(Real) * n_vars; + const size_t total_scratch = scratch_per_k * 1.3; // 30% padding + + TokenScratchPool pool(total_scratch); + + // Use in block loop... +} + +} // namespace Package +``` + +## Error Handling + +- **Overflow detection**: In debug builds, `Kokkos::abort()` if allocation exceeds capacity +- **Release builds**: No checks for performance; ensure proper sizing +- **Diagnostic methods**: Use `remaining()` and `current_offset()` to debug sizing + +## Thread Safety + +- **Acquire/release**: Managed by Kokkos UniqueToken (thread-safe) +- **Within token**: No synchronization needed - each thread owns its token +- **Across tokens**: No shared state between tokens + +## Future Extensions + +Possible enhancements: + +1. **Named allocations**: Track allocation names for debugging +2. **High-water mark tracking**: Monitor maximum usage per token +3. **Hierarchical scratch**: Separate team-shared and thread-private levels +4. **Template bundle interface**: Compile-time view configuration + +## References + +- [Kokkos UniqueToken Documentation](https://kokkos.org/kokkos-core-wiki/API/core/UniqueToken.html) +- Parthenon scratch memory patterns diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index 0dcbada624031..d666462041db1 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -75,6 +75,54 @@ When ommitted the ``DEFAULT_LOOP_PATTERN`` is used. flattened into an outer ``TeamThreadRange``. The specializations ``loop_pattern_[tpttr|tptvr|tpttrtvr]_tag`` correspond to ``<1,0>``, ``<0,1>``, ``<1,1>`` respectively. +Token Scratch Utility +--------------------- + +For kernels launched with ``parthenon::par_for`` or directly with Kokkos, Parthenon provides +``parthenon::TokenScratchPool`` in ``utils/token_scratch.hpp`` for temporary per-thread workspace. +The pool pre-allocates a fixed number of bytes for each +``Kokkos::Experimental::UniqueToken``, and ``acquire()`` returns a ``ScratchAllocator`` that can +carve one or more unmanaged typed ``Kokkos::View`` objects from that token-local buffer. This is +useful when a kernel needs short-lived scratch storage without performing repeated dynamic +allocations inside the parallel region. + +.. code:: cpp + + parthenon::TokenScratchPool<> pool(64 * 1024); + + parthenon::par_for( + "token_scratch_example", 0, nblocks - 1, KOKKOS_LAMBDA(const int b) { + auto scratch = pool.acquire(); + auto work = scratch.template allocate_view(ni, nj); + // Use work(...) as temporary storage for this token. + }); + +You can also use the free-floating ``allocate_scratch_view``: + +.. code:: cpp + + parthenon::TokenScratchPool<> pool(64 * 1024); + + parthenon::par_for( + "token_scratch_example", 0, nblocks - 1, KOKKOS_LAMBDA(const int b) { + auto scratch = pool.acquire(); + auto work = parthenon::allocate_scratch_view(ni, nj); + // Use work(...) as temporary storage for this token. + }); + + +.. note:: + + The view created by ``allocate_view`` is unmanaged and thus + non-blocking when destructed. The backing view, which is created by + the ``TokenScratchPool`` object is owning and thus **does** fence + upon going out of scope. + +.. warning:: + + The token scratch infrastructure is not tested for hierarchical + parallelism. In this case, we recommend that you use the team + scratch provided by Kokkos. Cmake Options ------------- @@ -129,4 +177,3 @@ New types can be provided by specializing the ``ProcessLoopBound`` struct in the These structs need to provide a ``GetNumBounds`` method to count the number of start/end bounds contained in the type, as well as a ``GetIndexRanges`` method to fill the ``IndexRange`` bounds used in the parallel dispatch. - diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ee7fcda955b7f..98c373d1b9e5f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -338,6 +338,7 @@ add_library(parthenon utils/sort.hpp utils/string_utils.cpp utils/string_utils.hpp + utils/token_scratch.hpp utils/type_list.hpp utils/unique_id.cpp utils/unique_id.hpp diff --git a/src/utils/cleantypes.hpp b/src/utils/cleantypes.hpp index f2866d17182e4..5622d25fbab18 100644 --- a/src/utils/cleantypes.hpp +++ b/src/utils/cleantypes.hpp @@ -14,6 +14,8 @@ #ifndef UTILS_CLEANTYPES_HPP_ #define UTILS_CLEANTYPES_HPP_ +// This file made with the assistance of generative AI + namespace parthenon { namespace cleantypes { @@ -44,6 +46,19 @@ struct remove_all_pointers { using type = typename remove_all_pointers::type; }; +//! Helper to build pointer types with specified depth +//! E.g., pointer_depth::type = T*** +//! Used for constructing multi-dimensional Kokkos::View data types +template +struct pointer_depth { + using type = typename pointer_depth::type; +}; + +template +struct pointer_depth { + using type = T *; +}; + } // namespace cleantypes } // namespace parthenon diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp new file mode 100644 index 0000000000000..debf526ba25a1 --- /dev/null +++ b/src/utils/token_scratch.hpp @@ -0,0 +1,203 @@ +//======================================================================================== +// (C) (or copyright) 2026. Triad National Security, LLC. All rights reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 for Los +// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC +// for the U.S. Department of Energy/National Nuclear Security Administration. All rights +// in the program are reserved by Triad National Security, LLC, and the U.S. Department +// of Energy/National Nuclear Security Administration. The Government is granted for +// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide +// license in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do so. +//======================================================================================== +#ifndef UTILS_TOKEN_SCRATCH_HPP_ +#define UTILS_TOKEN_SCRATCH_HPP_ + +// This file was created part with generative AI + +#include +#include +#include + +#include + +#include "cleantypes.hpp" + +namespace parthenon { + +//======================================================================================== +//! \class ScratchAllocator +//! \brief RAII wrapper for per-token scratch allocation with type-safe view carving +//! +//! Provides stack-like allocation of typed views from a pre-allocated token-specific +//! buffer. Handles alignment automatically and validates bounds. Token is automatically +//! released when the allocator goes out of scope. +//======================================================================================== +template +class ScratchAllocator { + private: + using PoolView = Kokkos::View; + using TokenType = Kokkos::Experimental::UniqueToken; + + const PoolView &pool_; + TokenType const *tokens_; + int token_id_; + std::size_t capacity_; + std::size_t offset_; + + //! Align offset to satisfy alignment requirements of type T + KOKKOS_INLINE_FUNCTION + static std::size_t AlignOffset(std::size_t offset, std::size_t alignment) { + return (offset + alignment - 1) & ~(alignment - 1); + } + + public: + //! Construct allocator for a specific token + KOKKOS_INLINE_FUNCTION + ScratchAllocator(const PoolView &pool, TokenType const *tokens, int token, + std::size_t capacity) + : pool_(pool), tokens_(tokens), token_id_(token), capacity_(capacity), offset_(0) {} + + // Disable copy to ensure single ownership of token + ScratchAllocator(const ScratchAllocator &) = delete; + ScratchAllocator &operator=(const ScratchAllocator &) = delete; + + // Enable move for flexibility + KOKKOS_INLINE_FUNCTION + ScratchAllocator(ScratchAllocator &&other) + : pool_(other.pool_), tokens_(other.tokens_), token_id_(other.token_id_), + capacity_(other.capacity_), offset_(other.offset_) { + other.token_id_ = -1; // Mark as moved-from + } + + //! Destructor releases the token + KOKKOS_INLINE_FUNCTION + ~ScratchAllocator() { + if (token_id_ >= 0 && tokens_ != nullptr) { + tokens_->release(token_id_); + } + } + + //! Allocate an unmanaged view of type T with arbitrary rank + //! \param dims Variadic pack of dimension sizes + //! \return Unmanaged Kokkos::View with rank = sizeof...(dims) + //! + //! Uses fold expressions to compute total elements and constructs + //! the appropriate View type using pointer_depth helper. + //! + //! Examples: + //! allocate_view(100) -> View + //! allocate_view(10, 20) -> View + //! allocate_view(5, 10, 15) -> View + template + KOKKOS_INLINE_FUNCTION auto allocate_view(Dims... dims) { + static_assert(sizeof...(Dims) > 0, "At least one dimension required"); + static_assert((std::is_convertible_v && ...), + "All dimensions must be convertible to std::size_t"); + + constexpr std::size_t rank = sizeof...(Dims); + using DataType = typename cleantypes::pointer_depth::type; + + // Use fold expression to compute total number of elements + const std::size_t total_elements = (dims * ...); + const std::size_t bytes_needed = total_elements * sizeof(T); + const std::size_t aligned_offset = AlignOffset(offset_, alignof(T)); + +#ifndef NDEBUG + // In debug builds, check for overflow + if (aligned_offset + bytes_needed > capacity_) { + Kokkos::abort("TokenScratch: allocation exceeded capacity"); + } +#endif + + char *base = pool_.data() + token_id_ * capacity_ + aligned_offset; + offset_ = aligned_offset + bytes_needed; + + return Kokkos::View( + reinterpret_cast(base), dims...); + } + + //! Get current offset (useful for debugging) + KOKKOS_INLINE_FUNCTION + std::size_t current_offset() const { return offset_; } + + //! Get the acquired token id (useful for diagnostics and testing) + KOKKOS_INLINE_FUNCTION + int token_id() const { return token_id_; } + + //! Get remaining capacity + KOKKOS_INLINE_FUNCTION + std::size_t remaining() const { return capacity_ - offset_; } +}; + +template +decltype(auto) allocate_scratch_view(scratch_t &scratch, Args &&...args) { + return scratch.template allocate_view(std::forward(args)...); +} + +//======================================================================================== +//! \class TokenScratchPool +//! \brief Manages a pool of per-token scratch memory using Kokkos UniqueToken +//! +//! Allocates a 2D pool where each row corresponds to a unique token's scratch buffer. +//! Threads acquire tokens, get a ScratchAllocator, and carve typed views from their +//! token's buffer. +//! +//! Example: +//! \code +//! TokenScratchPool pool(64*1024); // 64KB per token +//! Kokkos::parallel_for(policy, KOKKOS_LAMBDA(int i) { +//! auto scratch = pool.acquire(); +//! auto doubles = scratch.template allocate_view(100); +//! auto ints = scratch.template allocate_view(50); +//! // Use views... +//! }); +//! \endcode +//======================================================================================== +template +class TokenScratchPool { + private: + using PoolView = Kokkos::View; + using TokenType = Kokkos::Experimental::UniqueToken; + + TokenType tokens_; + PoolView pool_; + std::size_t bytes_per_token_; + + public: + //! Construct pool with specified bytes per token + //! \param bytes_per_token Amount of scratch memory for each unique token + explicit TokenScratchPool(std::size_t bytes_per_token) + : tokens_(), bytes_per_token_(bytes_per_token), + pool_("token_scratch_pool", tokens_.size(), bytes_per_token) {} + + //! Acquire a token and return a scratch allocator for this thread + //! The allocator automatically manages token lifetime via RAII + KOKKOS_INLINE_FUNCTION + auto acquire() const { + const int token_id = bytes_per_token_ > 0 ? tokens_.acquire() : -1; + return ScratchAllocator( + pool_, &tokens_, token_id, bytes_per_token_); + } + + //! Get the number of tokens in the pool + KOKKOS_INLINE_FUNCTION + std::size_t size() const { return tokens_.size(); } + + //! Get bytes per token + KOKKOS_INLINE_FUNCTION + std::size_t bytes_per_token() const { return bytes_per_token_; } + + //! Get total pool size in bytes + KOKKOS_INLINE_FUNCTION + std::size_t total_bytes() const { return tokens_.size() * bytes_per_token_; } +}; +} // namespace parthenon + +#endif // UTILS_TOKEN_SCRATCH_HPP_ diff --git a/tst/unit/CMakeLists.txt b/tst/unit/CMakeLists.txt index 722048446f48a..1ad727c75b20d 100644 --- a/tst/unit/CMakeLists.txt +++ b/tst/unit/CMakeLists.txt @@ -46,6 +46,7 @@ list(APPEND unit_tests_SOURCES test_partitioning.cpp test_scratch_variables.cpp test_state_descriptor.cpp + test_token_scratch.cpp test_swarm.cpp test_swarm_amr_remesh.cpp test_unit_integrators.cpp diff --git a/tst/unit/test_token_scratch.cpp b/tst/unit/test_token_scratch.cpp new file mode 100644 index 0000000000000..5d25003f0803f --- /dev/null +++ b/tst/unit/test_token_scratch.cpp @@ -0,0 +1,448 @@ +//======================================================================================== +// (C) (or copyright) 2026. Triad National Security, LLC. All rights reserved. +// +// This program was produced under U.S. Government contract 89233218CNA000001 for Los +// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC +// for the U.S. Department of Energy/National Nuclear Security Administration. All rights +// in the program are reserved by Triad National Security, LLC, and the U.S. Department +// of Energy/National Nuclear Security Administration. The Government is granted for +// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide +// license in this material to reproduce, prepare derivative works, distribute copies to +// the public, perform publicly and display publicly, and to permit others to do so. +//======================================================================================== +// This file was made in part with generative AI + +//! \file test_token_scratch.cpp +//! \brief Unit tests for TokenScratchPool + +#include +#include + +#ifndef CATCH_CONFIG_FAST_COMPILE +#define CATCH_CONFIG_FAST_COMPILE +#include +#endif + +#include + +#include "utils/token_scratch.hpp" + +constexpr std::size_t bytes_double = 8; +constexpr std::size_t bytes_int = 8; + +SCENARIO("TokenScratchPool basic allocation and usage", "[TokenScratch][Basic]") { + GIVEN("A TokenScratchPool with 64KB per token") { + constexpr size_t scratch_bytes = 100 * bytes_double + 50 * bytes_int; + constexpr int n_iterations = 100; + + parthenon::TokenScratchPool<> pool(scratch_bytes); + + WHEN("We allocate views in parallel iterations") { + Kokkos::View results("results", n_iterations); + + Kokkos::parallel_for( + "test_basic", n_iterations, KOKKOS_LAMBDA(const int i) { + auto scratch = pool.acquire(); + + // Allocate views + auto doubles = scratch.template allocate_view(100); + auto ints = scratch.template allocate_view(50); + + // Initialize + for (int j = 0; j < 100; ++j) { + doubles(j) = static_cast(i + j); + } + for (int j = 0; j < 50; ++j) { + ints(j) = i * j; + } + + // Compute sum + double sum = 0.0; + for (int j = 0; j < 100; ++j) { + sum += doubles(j); + } + for (int j = 0; j < 50; ++j) { + sum += static_cast(ints(j)); + } + results(i) = sum; + }); + + Kokkos::fence(); + + THEN("Results match expected values") { + auto results_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), results); + + std::size_t nwrong = 0; + for (int i = 0; i < n_iterations; ++i) { + double expected = 0.0; + for (int j = 0; j < 100; ++j) { + expected += static_cast(i + j); + } + for (int j = 0; j < 50; ++j) { + expected += static_cast(i * j); + } + nwrong += !(std::abs(results_h(i) - expected) < 1e-10); + } + REQUIRE(nwrong == 0); + } + } + } +} + +SCENARIO("TokenScratchPool handles multi-dimensional views", "[TokenScratch][MultiDim]") { + GIVEN("A TokenScratchPool with 128KB per token") { + using ExecSpace = Kokkos::DefaultExecutionSpace; + + constexpr int n_blocks = 50; + constexpr int ni = 8, nj = 8, nk = 8; + constexpr size_t scratch_bytes = ni * nj * (1 + nk) * bytes_double; + + parthenon::TokenScratchPool pool(scratch_bytes); + + WHEN("We allocate 2D and 3D views in parallel") { + Kokkos::View block_results("block_results", n_blocks); + + Kokkos::parallel_for( + "test_multidim", n_blocks, KOKKOS_LAMBDA(const int b) { + auto scratch = pool.acquire(); + + auto work_2d = scratch.template allocate_view(ni, nj); + auto work_3d = scratch.template allocate_view(ni, nj, nk); + + // Initialize and compute + for (int i = 0; i < ni; ++i) { + for (int j = 0; j < nj; ++j) { + work_2d(i, j) = static_cast(i + j); + for (int k = 0; k < nk; ++k) { + work_3d(i, j, k) = work_2d(i, j) * k; + } + } + } + + double sum = 0.0; + for (int i = 0; i < ni; ++i) { + for (int j = 0; j < nj; ++j) { + for (int k = 0; k < nk; ++k) { + sum += work_3d(i, j, k); + } + } + } + block_results(b) = sum; + }); + + Kokkos::fence(); + + THEN("All blocks produce the same correct result") { + auto results_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), block_results); + + double expected = 0.0; + for (int i = 0; i < ni; ++i) { + for (int j = 0; j < nj; ++j) { + for (int k = 0; k < nk; ++k) { + expected += static_cast(i + j) * k; + } + } + } + + int nwrong = 0; + for (int b = 0; b < n_blocks; ++b) { + nwrong += !(std::abs(results_h(b) - expected) < 1e-10); + } + REQUIRE(nwrong == 0); + } + } + } +} + +SCENARIO("TokenScratchPool handles token reuse with many iterations", + "[TokenScratch][Reuse]") { + GIVEN("A TokenScratchPool with 8KB per token and many iterations") { + using ExecSpace = Kokkos::DefaultExecutionSpace; + using MemSpace = ExecSpace::memory_space; + + constexpr size_t scratch_bytes = 100 * bytes_int; + constexpr int n_iterations = 10000; + + parthenon::TokenScratchPool pool(scratch_bytes); + + WHEN("We run many iterations to force token reuse") { + Kokkos::View counters("counters", n_iterations); + + Kokkos::parallel_for( + "test_reuse", n_iterations, KOKKOS_LAMBDA(const int i) { + auto scratch = pool.acquire(); + auto data = scratch.template allocate_view(100); + + for (int j = 0; j < 100; ++j) { + data(j) = i + j; + } + + int sum = 0; + for (int j = 0; j < 100; ++j) { + sum += data(j); + } + counters(i) = sum; + }); + + Kokkos::fence(); + + THEN("All iterations produce correct results despite token reuse") { + auto counters_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), counters); + + int nwrong = 0; + for (int i = 0; i < n_iterations; ++i) { + int expected = 0; + for (int j = 0; j < 100; ++j) { + expected += i + j; + } + nwrong += (counters_h(i) != expected); + } + REQUIRE(nwrong == 0); + } + } + } +} + +SCENARIO("TokenScratchPool exposes the expected token id for the execution space", + "[TokenScratch][TokenId]") { + GIVEN("A TokenScratchPool using the default execution space") { + using ExecSpace = Kokkos::DefaultExecutionSpace; + + constexpr size_t scratch_bytes = bytes_double; + parthenon::TokenScratchPool pool(scratch_bytes); + + WHEN("We record the token id for a single iteration") { + Kokkos::View actual_ids("actual_ids", 1); + + Kokkos::parallel_for( + "test_token_ids", Kokkos::RangePolicy(0, 1), + KOKKOS_LAMBDA(const int i) { + auto scratch = pool.acquire(); + actual_ids(i) = scratch.token_id(); + }); + + Kokkos::fence(); + + THEN("The observed token ids match the execution-space mapping") { + auto actual_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), actual_ids); + +#ifdef KOKKOS_ENABLE_OPENMPTARGET + if constexpr (std::is_same_v) { + REQUIRE(0 <= actual_h(0)); + REQUIRE(actual_h(0) < static_cast(pool.size())); + } else { + REQUIRE(actual_h(0) == 0); + } +#else + REQUIRE(actual_h(0) == 0); +#endif + } + } + } +} + +SCENARIO("TokenScratchPool supports 4D and higher views", "[TokenScratch][Variadic]") { + GIVEN("A TokenScratchPool with sufficient memory") { + constexpr size_t scratch_bytes = bytes_double * 5 * 4 * 3 * 2; + parthenon::TokenScratchPool<> pool(scratch_bytes); + + WHEN("We allocate a 4D view") { + Kokkos::View result("result", 1); + + Kokkos::parallel_for( + "test_4d", 1, KOKKOS_LAMBDA(const int i) { + auto scratch = pool.acquire(); + + // Allocate a 4D view: 5x4x3x2 = 120 elements + auto view_4d = scratch.template allocate_view(5, 4, 3, 2); + + // Initialize and sum + double sum = 0.0; + for (int i1 = 0; i1 < 5; ++i1) { + for (int i2 = 0; i2 < 4; ++i2) { + for (int i3 = 0; i3 < 3; ++i3) { + for (int i4 = 0; i4 < 2; ++i4) { + view_4d(i1, i2, i3, i4) = static_cast(i1 + i2 + i3 + i4); + sum += view_4d(i1, i2, i3, i4); + } + } + } + } + result(0) = sum; + }); + + Kokkos::fence(); + + THEN("The 4D view works correctly") { + auto result_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), result); + + // Compute expected sum + double expected = 0.0; + for (int i1 = 0; i1 < 5; ++i1) { + for (int i2 = 0; i2 < 4; ++i2) { + for (int i3 = 0; i3 < 3; ++i3) { + for (int i4 = 0; i4 < 2; ++i4) { + expected += static_cast(i1 + i2 + i3 + i4); + } + } + } + } + + REQUIRE(std::abs(result_h(0) - expected) < 1e-10); + } + } + + WHEN("We allocate a 5D view") { + Kokkos::View result("result", 1); + + Kokkos::parallel_for( + "test_5d", 1, KOKKOS_LAMBDA(const int i) { + auto scratch = pool.acquire(); + + // Allocate a 5D view: 4x3x3x2x2 = 144 elements + auto view_5d = scratch.template allocate_view(4, 3, 3, 2, 2); + + // Initialize and sum + int sum = 0; + for (int i1 = 0; i1 < 4; ++i1) { + for (int i2 = 0; i2 < 3; ++i2) { + for (int i3 = 0; i3 < 3; ++i3) { + for (int i4 = 0; i4 < 2; ++i4) { + for (int i5 = 0; i5 < 2; ++i5) { + view_5d(i1, i2, i3, i4, i5) = + i1 * 10000 + i2 * 1000 + i3 * 100 + i4 * 10 + i5; + sum += view_5d(i1, i2, i3, i4, i5); + } + } + } + } + } + result(0) = static_cast(sum); + }); + + Kokkos::fence(); + + THEN("The 5D view works correctly") { + auto result_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), result); + + // Compute expected sum + int expected = 0; + for (int i1 = 0; i1 < 4; ++i1) { + for (int i2 = 0; i2 < 3; ++i2) { + for (int i3 = 0; i3 < 3; ++i3) { + for (int i4 = 0; i4 < 2; ++i4) { + for (int i5 = 0; i5 < 2; ++i5) { + expected += i1 * 10000 + i2 * 1000 + i3 * 100 + i4 * 10 + i5; + } + } + } + } + } + + REQUIRE(std::abs(result_h(0) - static_cast(expected)) < 1e-10); + } + } + } +} + +SCENARIO("Variadic allocate_view handles mixed dimensions in single kernel", + "[TokenScratch][Variadic][Mixed]") { + GIVEN("A TokenScratchPool with sufficient memory") { + constexpr size_t scratch_bytes = + bytes_double * (10 + 5 * 4 + 3 * 3 * 3 + 2 * 2 * 2 * 2 + 2 * 2 * 2 * 2 * 2); + parthenon::TokenScratchPool<> pool(scratch_bytes); + + WHEN("We allocate views of different ranks in the same kernel") { + Kokkos::View results("results", 5); + + Kokkos::parallel_for( + "test_mixed_ranks", 1, KOKKOS_LAMBDA(const int iter) { + auto scratch = pool.acquire(); + + // Allocate views of ranks 1-5 + auto view_1d = scratch.template allocate_view(10); + auto view_2d = scratch.template allocate_view(5, 4); + auto view_3d = scratch.template allocate_view(3, 3, 3); + auto view_4d = scratch.template allocate_view(2, 2, 2, 2); + auto view_5d = scratch.template allocate_view(2, 2, 2, 2, 2); + + // Initialize each view + for (int i = 0; i < 10; ++i) + view_1d(i) = 1.0; + + for (int i = 0; i < 5; ++i) + for (int j = 0; j < 4; ++j) + view_2d(i, j) = 2.0; + + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) + for (int k = 0; k < 3; ++k) + view_3d(i, j, k) = 3.0; + + for (int i = 0; i < 2; ++i) + for (int j = 0; j < 2; ++j) + for (int k = 0; k < 2; ++k) + for (int l = 0; l < 2; ++l) + view_4d(i, j, k, l) = 4.0; + + for (int i = 0; i < 2; ++i) + for (int j = 0; j < 2; ++j) + for (int k = 0; k < 2; ++k) + for (int l = 0; l < 2; ++l) + for (int m = 0; m < 2; ++m) + view_5d(i, j, k, l, m) = 5.0; + + // Compute sums + double sum1 = 0.0, sum2 = 0.0, sum3 = 0.0, sum4 = 0.0, sum5 = 0.0; + + for (int i = 0; i < 10; ++i) + sum1 += view_1d(i); + + for (int i = 0; i < 5; ++i) + for (int j = 0; j < 4; ++j) + sum2 += view_2d(i, j); + + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) + for (int k = 0; k < 3; ++k) + sum3 += view_3d(i, j, k); + + for (int i = 0; i < 2; ++i) + for (int j = 0; j < 2; ++j) + for (int k = 0; k < 2; ++k) + for (int l = 0; l < 2; ++l) + sum4 += view_4d(i, j, k, l); + + for (int i = 0; i < 2; ++i) + for (int j = 0; j < 2; ++j) + for (int k = 0; k < 2; ++k) + for (int l = 0; l < 2; ++l) + for (int m = 0; m < 2; ++m) + sum5 += view_5d(i, j, k, l, m); + + results(0) = sum1; + results(1) = sum2; + results(2) = sum3; + results(3) = sum4; + results(4) = sum5; + }); + + Kokkos::fence(); + + THEN("All views work correctly") { + auto results_h = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), results); + + REQUIRE(std::abs(results_h(0) - 10.0) < 1e-10); // 10 * 1.0 + REQUIRE(std::abs(results_h(1) - 40.0) < 1e-10); // 20 * 2.0 + REQUIRE(std::abs(results_h(2) - 81.0) < 1e-10); // 27 * 3.0 + REQUIRE(std::abs(results_h(3) - 64.0) < 1e-10); // 16 * 4.0 + REQUIRE(std::abs(results_h(4) - 160.0) < 1e-10); // 32 * 5.0 + } + } + } +}