From 3903b1f38f685429e056a5de71e52cb8da2385c7 Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Thu, 30 Apr 2026 12:31:46 -0400 Subject: [PATCH 01/13] token scratch --- doc/plan_histories/token_scratch.md | 231 ++++++++++++++++ src/CMakeLists.txt | 1 + src/utils/cleantypes.hpp | 15 ++ src/utils/token_scratch.hpp | 193 +++++++++++++ tst/unit/CMakeLists.txt | 1 + tst/unit/test_token_scratch.cpp | 404 ++++++++++++++++++++++++++++ 6 files changed, 845 insertions(+) create mode 100644 doc/plan_histories/token_scratch.md create mode 100644 src/utils/token_scratch.hpp create mode 100644 tst/unit/test_token_scratch.cpp diff --git a/doc/plan_histories/token_scratch.md b/doc/plan_histories/token_scratch.md new file mode 100644 index 0000000000000..6ce7c3724cf27 --- /dev/null +++ b/doc/plan_histories/token_scratch.md @@ -0,0 +1,231 @@ +# 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) +}); +``` + +### With TeamPolicy + +```cpp +using TeamPolicy = Kokkos::TeamPolicy; +TokenScratchPool pool(32 * 1024); + +Kokkos::parallel_for("team_kernel", TeamPolicy(n_teams, team_size), + KOKKOS_LAMBDA(const member_type& team) { + // Each team gets its own token + auto scratch = pool.acquire(); + auto shared_data = scratch.template allocate_view(256); + + // Team-parallel work with shared scratch + Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 256), + [&](int i) { shared_data(i) = ...; }); + }); +``` + +## 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/src/CMakeLists.txt b/src/CMakeLists.txt index 5f84e8b33bda3..f05c6e0b90cba 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -312,6 +312,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..8733bb0363921 --- /dev/null +++ b/src/utils/token_scratch.hpp @@ -0,0 +1,193 @@ +//======================================================================================== +// (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 "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 remaining capacity + KOKKOS_INLINE_FUNCTION + std::size_t remaining() const { return capacity_ - offset_; } +}; + +//======================================================================================== +//! \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 = tokens_.acquire(); + 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 84a88a5a9eb9b..6cce9cacbccc2 100644 --- a/tst/unit/CMakeLists.txt +++ b/tst/unit/CMakeLists.txt @@ -44,6 +44,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_unit_integrators.cpp test_upper_bound.cpp diff --git a/tst/unit/test_token_scratch.cpp b/tst/unit/test_token_scratch.cpp new file mode 100644 index 0000000000000..9a972ad9c8a8b --- /dev/null +++ b/tst/unit/test_token_scratch.cpp @@ -0,0 +1,404 @@ +//======================================================================================== +// (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 + +#ifndef CATCH_CONFIG_FAST_COMPILE +#define CATCH_CONFIG_FAST_COMPILE +#include +#endif + +#include + +#include "utils/token_scratch.hpp" + +SCENARIO("TokenScratchPool basic allocation and usage", "[TokenScratch][Basic]") { + GIVEN("A TokenScratchPool with 64KB per token") { + constexpr size_t scratch_bytes = 64 * 1024; // 64KB per token + 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 size_t scratch_bytes = 128 * 1024; + constexpr int n_blocks = 50; + constexpr int ni = 8, nj = 8, nk = 8; + + 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 = 8 * 1024; + 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 supports 4D and higher views", "[TokenScratch][Variadic]") { + GIVEN("A TokenScratchPool with sufficient memory") { + constexpr size_t scratch_bytes = 512 * 1024; + 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 = 256 * 1024; + 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 + } + } + } +} From 0b5a759e6ccd09b9bb296344e9eb458f98bca976 Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Thu, 30 Apr 2026 18:27:58 -0400 Subject: [PATCH 02/13] docs --- doc/sphinx/src/par_for.rst | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index 0dcbada624031..88d83763f5883 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -75,6 +75,27 @@ 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. + }); Cmake Options ------------- @@ -129,4 +150,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. - From 9272d924bb01954e0d51b3204454e87afcad5f9b Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Thu, 30 Apr 2026 18:32:16 -0400 Subject: [PATCH 03/13] token scratch --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index f9103d49edff7..6b4678ddc3863 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 1378]](https://github.com/parthenon-hpc-lab/parthenon/pull/1378) MeshData Swarm Tasks - [[PR 1377]](https://github.com/parthenon-hpc-lab/parthenon/pull/1377) Extend Initialization Hierarchy - [[PR 1332]](https://github.com/parthenon-hpc-lab/parthenon/pull/1332) Add global WatchDog From d7cc55e8e22f17ca3eea936d75b7c1101c39b163 Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 19:31:46 -0400 Subject: [PATCH 04/13] fix token scratch test for cuda by actually allocating hte right amount of memory --- tst/unit/test_token_scratch.cpp | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tst/unit/test_token_scratch.cpp b/tst/unit/test_token_scratch.cpp index 9a972ad9c8a8b..6f3aaf7cd3309 100644 --- a/tst/unit/test_token_scratch.cpp +++ b/tst/unit/test_token_scratch.cpp @@ -26,9 +26,12 @@ #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 = 64 * 1024; // 64KB per token + constexpr size_t scratch_bytes = 100 * bytes_double + 50 * bytes_int; constexpr int n_iterations = 100; parthenon::TokenScratchPool<> pool(scratch_bytes); @@ -90,9 +93,9 @@ SCENARIO("TokenScratchPool handles multi-dimensional views", "[TokenScratch][Mul GIVEN("A TokenScratchPool with 128KB per token") { using ExecSpace = Kokkos::DefaultExecutionSpace; - constexpr size_t scratch_bytes = 128 * 1024; 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); @@ -158,7 +161,7 @@ SCENARIO("TokenScratchPool handles token reuse with many iterations", using ExecSpace = Kokkos::DefaultExecutionSpace; using MemSpace = ExecSpace::memory_space; - constexpr size_t scratch_bytes = 8 * 1024; + constexpr size_t scratch_bytes = 100 * bytes_int; constexpr int n_iterations = 10000; parthenon::TokenScratchPool pool(scratch_bytes); @@ -204,7 +207,7 @@ SCENARIO("TokenScratchPool handles token reuse with many iterations", SCENARIO("TokenScratchPool supports 4D and higher views", "[TokenScratch][Variadic]") { GIVEN("A TokenScratchPool with sufficient memory") { - constexpr size_t scratch_bytes = 512 * 1024; + constexpr size_t scratch_bytes = bytes_double * 5 * 4 * 3 * 2; parthenon::TokenScratchPool<> pool(scratch_bytes); WHEN("We allocate a 4D view") { @@ -309,7 +312,8 @@ SCENARIO("TokenScratchPool supports 4D and higher views", "[TokenScratch][Variad SCENARIO("Variadic allocate_view handles mixed dimensions in single kernel", "[TokenScratch][Variadic][Mixed]") { GIVEN("A TokenScratchPool with sufficient memory") { - constexpr size_t scratch_bytes = 256 * 1024; + 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") { From ca1580bd628c997789bce43cbffac2f1a4723c7c Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 19:38:15 -0400 Subject: [PATCH 05/13] clarifying comment --- doc/sphinx/src/par_for.rst | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index 88d83763f5883..e584c0b4338d4 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -97,6 +97,13 @@ allocations inside the parallel region. // 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. + Cmake Options ------------- From c325a63df851d23a5257dec2635b6a0355a9669a Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 19:39:32 -0400 Subject: [PATCH 06/13] ambiguity --- doc/plan_histories/token_scratch.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/plan_histories/token_scratch.md b/doc/plan_histories/token_scratch.md index 6ce7c3724cf27..a0bc9a4ea25be 100644 --- a/doc/plan_histories/token_scratch.md +++ b/doc/plan_histories/token_scratch.md @@ -22,7 +22,7 @@ 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` +- `TokenScope` (optional) - `UniqueTokenScope::Global` (default) or `UniqueTokenScope::Instance` **Constructor:** ```cpp From 235415a0d23b6144403d36c9a9b3798684e50d5c Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 19:49:00 -0400 Subject: [PATCH 07/13] dont do token scratch in hierarchical parallelism --- doc/plan_histories/token_scratch.md | 18 ------------------ doc/sphinx/src/par_for.rst | 6 ++++++ 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/doc/plan_histories/token_scratch.md b/doc/plan_histories/token_scratch.md index a0bc9a4ea25be..12bfc3f2d5020 100644 --- a/doc/plan_histories/token_scratch.md +++ b/doc/plan_histories/token_scratch.md @@ -95,24 +95,6 @@ Kokkos::parallel_for("multidim", N, KOKKOS_LAMBDA(int i) { }); ``` -### With TeamPolicy - -```cpp -using TeamPolicy = Kokkos::TeamPolicy; -TokenScratchPool pool(32 * 1024); - -Kokkos::parallel_for("team_kernel", TeamPolicy(n_teams, team_size), - KOKKOS_LAMBDA(const member_type& team) { - // Each team gets its own token - auto scratch = pool.acquire(); - auto shared_data = scratch.template allocate_view(256); - - // Team-parallel work with shared scratch - Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 256), - [&](int i) { shared_data(i) = ...; }); - }); -``` - ## Design Considerations ### Memory Layout diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index e584c0b4338d4..e02b093fce56a 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -104,6 +104,12 @@ allocations inside the parallel region. the ``TokenScratchPool`` object is owning and thus **does** fence upon going out of scope. +.. warning:: + + The token scratch infrastructure is not compatible with hierarchical + parallelism. In this case, please use the team scratch provided by + Kokkos. + Cmake Options ------------- From c440cd33e2831f6ab3738a520fa3f4298658c34d Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 19:49:28 -0400 Subject: [PATCH 08/13] softer warning --- doc/sphinx/src/par_for.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index e02b093fce56a..e9e28ec5c38b5 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -106,9 +106,9 @@ allocations inside the parallel region. .. warning:: - The token scratch infrastructure is not compatible with hierarchical - parallelism. In this case, please use the team scratch provided by - Kokkos. + 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 ------------- From c1d1419043d9f87682c7f50bf709c30501b86cf3 Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 20:09:05 -0400 Subject: [PATCH 09/13] add token ID test as suggested by adam --- src/utils/token_scratch.hpp | 4 ++++ tst/unit/test_token_scratch.cpp | 39 +++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp index 8733bb0363921..4742fa006e5eb 100644 --- a/src/utils/token_scratch.hpp +++ b/src/utils/token_scratch.hpp @@ -123,6 +123,10 @@ class ScratchAllocator { 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_; } diff --git a/tst/unit/test_token_scratch.cpp b/tst/unit/test_token_scratch.cpp index 6f3aaf7cd3309..f6294b0191951 100644 --- a/tst/unit/test_token_scratch.cpp +++ b/tst/unit/test_token_scratch.cpp @@ -16,6 +16,7 @@ //! \brief Unit tests for TokenScratchPool #include +#include #ifndef CATCH_CONFIG_FAST_COMPILE #define CATCH_CONFIG_FAST_COMPILE @@ -205,6 +206,44 @@ SCENARIO("TokenScratchPool handles token reuse with many iterations", } } +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 +#endif + { + REQUIRE(actual_h(0) == 0); + } + } + } + } +} + 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; From 92d19b726c830b5191f513fe8d5489e5eade9c28 Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Sat, 2 May 2026 21:50:20 -0400 Subject: [PATCH 10/13] formatting and adempsey comment --- src/utils/token_scratch.hpp | 4 ++-- tst/unit/test_token_scratch.cpp | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp index 4742fa006e5eb..6b9d7170fa41b 100644 --- a/src/utils/token_scratch.hpp +++ b/src/utils/token_scratch.hpp @@ -33,7 +33,7 @@ namespace parthenon { //! released when the allocator goes out of scope. //======================================================================================== template class ScratchAllocator { @@ -154,7 +154,7 @@ class ScratchAllocator { template + Kokkos::Experimental::UniqueTokenScope::Instance> class TokenScratchPool { private: using PoolView = Kokkos::View; diff --git a/tst/unit/test_token_scratch.cpp b/tst/unit/test_token_scratch.cpp index f6294b0191951..5d25003f0803f 100644 --- a/tst/unit/test_token_scratch.cpp +++ b/tst/unit/test_token_scratch.cpp @@ -234,11 +234,12 @@ SCENARIO("TokenScratchPool exposes the expected token id for the execution space if constexpr (std::is_same_v) { REQUIRE(0 <= actual_h(0)); REQUIRE(actual_h(0) < static_cast(pool.size())); - } else -#endif - { + } else { REQUIRE(actual_h(0) == 0); } +#else + REQUIRE(actual_h(0) == 0); +#endif } } } From c508f246254ab41402e474630d8789f24f67280e Mon Sep 17 00:00:00 2001 From: Jonah Miller Date: Mon, 4 May 2026 22:06:15 -0400 Subject: [PATCH 11/13] only acquire a unique token if necessary --- src/utils/token_scratch.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp index 6b9d7170fa41b..d3daf3b16a88a 100644 --- a/src/utils/token_scratch.hpp +++ b/src/utils/token_scratch.hpp @@ -175,7 +175,7 @@ class TokenScratchPool { //! The allocator automatically manages token lifetime via RAII KOKKOS_INLINE_FUNCTION auto acquire() const { - const int token_id = tokens_.acquire(); + const int token_id = bytes_per_token_ > 0 ? tokens_.acquire() : -1; return ScratchAllocator( pool_, &tokens_, token_id, bytes_per_token_); } From f76047bab32b46db321b9ee97c14fcbf7613e38e Mon Sep 17 00:00:00 2001 From: Jonah Maxwell Miller Date: Fri, 31 Jul 2026 10:25:15 -0400 Subject: [PATCH 12/13] lfroberts suggestions --- doc/sphinx/src/par_for.rst | 14 ++++++++++++++ src/utils/token_scratch.hpp | 5 +++++ 2 files changed, 19 insertions(+) diff --git a/doc/sphinx/src/par_for.rst b/doc/sphinx/src/par_for.rst index e9e28ec5c38b5..d666462041db1 100644 --- a/doc/sphinx/src/par_for.rst +++ b/doc/sphinx/src/par_for.rst @@ -97,6 +97,20 @@ allocations inside the parallel region. // 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 diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp index d3daf3b16a88a..a27e30f178745 100644 --- a/src/utils/token_scratch.hpp +++ b/src/utils/token_scratch.hpp @@ -132,6 +132,11 @@ class ScratchAllocator { 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 From 692e8f46f6653ea3b6bb4a3f6b083b91848fcfbc Mon Sep 17 00:00:00 2001 From: Jonah Maxwell Miller Date: Fri, 31 Jul 2026 13:43:26 -0400 Subject: [PATCH 13/13] missing utility include for std::forward --- src/utils/token_scratch.hpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/utils/token_scratch.hpp b/src/utils/token_scratch.hpp index a27e30f178745..debf526ba25a1 100644 --- a/src/utils/token_scratch.hpp +++ b/src/utils/token_scratch.hpp @@ -17,6 +17,7 @@ #include #include +#include #include