diff --git a/core/src/Cuda/Kokkos_Cuda_Event.hpp b/core/src/Cuda/Kokkos_Cuda_Event.hpp new file mode 100644 index 00000000000..b3665088648 --- /dev/null +++ b/core/src/Cuda/Kokkos_Cuda_Event.hpp @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project + +#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE +#include +static_assert(false, + "Including non-public Kokkos header files is not allowed."); +#endif + +#ifndef KOKKOS_CUDA_EVENT_HPP +#define KOKKOS_CUDA_EVENT_HPP + +#include +#if defined(KOKKOS_ENABLE_CUDA) + +#include + +#include +#include + +#include +#include + +namespace Kokkos { +namespace Impl { + +template <> +struct EventResource { + std::string label = "unknown"; + cudaEvent_t m_event = nullptr; + int m_cudaDev = -1; + + explicit EventResource(const std::string& label, + const Kokkos::Cuda& exec_space) + : m_label(label), m_cudaDev(exec_space.cuda_device()) { + KOKKOS_IMPL_CUDA_SAFE_CALL(cudaSetDevice(m_cudaDev)); + KOKKOS_IMPL_CUDA_SAFE_CALL( + cudaEventCreateWithFlags(&m_event, cudaEventDisableTiming)); + } + + ~EventResource() { + if (m_event != nullptr) { + KOKKOS_IMPL_CUDA_SAFE_CALL(cudaEventDestroy(m_event)); + } + } + + EventResource(const EventResource&) = delete; + EventResource& operator=(const EventResource&) = delete; +}; + +} // namespace Impl + +namespace Experimental { + +//============================================================================ +// CUDA specialization — native cudaEvent_t implementation +//============================================================================ + +/// CUDA specialization of Event. +/// +/// Copyable: copies share the underlying cudaEvent_t via reference +/// counting. The last copy standing destroys the event. Re-recording +/// through any copy affects all copies (same semantics as sharing a +/// raw cudaEvent_t). +template <> +class Event { + public: + Event(const std::string& label) + : m_handle(std::make_shared>( + label, Kokkos::Cuda())) {} + + Event(const std::string& label, const Kokkos::Cuda& exec_space) + : m_handle(std::make_shared>( + label, exec_space)) { + record(exec_space); + } + + void record(const Kokkos::Cuda& exec_space) { + KOKKOS_IMPL_CUDA_SAFE_CALL( + cudaEventRecord(m_handle->m_event, exec_space.cuda_stream())); + } + + void fence() const { + KOKKOS_IMPL_CUDA_SAFE_CALL(cudaEventSynchronize(m_handle->m_event)); + } + + bool is_complete() const { + cudaError_t err = cudaEventQuery(m_handle->m_event); + if (err == cudaSuccess) return true; + if (err == cudaErrorNotReady) return false; + KOKKOS_IMPL_CUDA_SAFE_CALL(err); + return false; + } + + const std::string& label() const { return m_handle->m_label; } + cudaEvent_t cuda_event() const noexcept { return m_handle->m_event; } + + private: + std::shared_ptr> m_handle; +}; + +/// CUDA: insert a stream wait for the recorded event (non-blocking on host). +inline void space_depends_on(const Kokkos::Cuda& exec_space, + const Event& event) { + KOKKOS_IMPL_CUDA_SAFE_CALL( + cudaStreamWaitEvent(exec_space.cuda_stream(), event.cuda_event(), 0)); +} + +} // namespace Experimental +} // namespace Kokkos + +#endif // KOKKOS_ENABLE_CUDA +#endif // KOKKOS_CUDA_EVENT_HPP diff --git a/core/src/Kokkos_Core.cppm b/core/src/Kokkos_Core.cppm index 794c39a8be0..4685ad15cc7 100644 --- a/core/src/Kokkos_Core.cppm +++ b/core/src/Kokkos_Core.cppm @@ -52,7 +52,9 @@ export { #endif } // namespace Experimental namespace Experimental { + using ::Kokkos::Experimental::Event; using ::Kokkos::Experimental::partition_space; + using ::Kokkos::Experimental::space_depends_on; } // namespace Experimental using ::Kokkos::AnonymousSpace; using ::Kokkos::DefaultExecutionSpace; diff --git a/core/src/Kokkos_Core.hpp b/core/src/Kokkos_Core.hpp index ff9ce95b1d4..78b32da52bc 100644 --- a/core/src/Kokkos_Core.hpp +++ b/core/src/Kokkos_Core.hpp @@ -61,6 +61,7 @@ //---------------------------------------------------------------------------- //---------------------------------------------------------------------------- +#include #include #include // Including this in Kokkos_Parallel_Reduce.hpp led to a circular dependency diff --git a/core/src/Kokkos_Event.hpp b/core/src/Kokkos_Event.hpp new file mode 100644 index 00000000000..3dd0a50fdec --- /dev/null +++ b/core/src/Kokkos_Event.hpp @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project + +/// \file Kokkos_Event.hpp +/// \brief Experimental event API for fine-grained stream dependencies. +/// +/// Events capture a point in an execution space's asynchronous timeline. +/// They enable cross-stream dependencies without a full fence, and +/// selective host synchronisation. +/// +/// API: +/// - space_depends_on(exec_space, event) — GPU-side dependency +/// (non-blocking on host) +/// - event.fence() — host-side blocking synchronisation +/// - event.is_complete() — non-blocking query +/// +/// Currently only the CUDA backend provides a native implementation. +/// For other backends the fallback records a fence on record() and +/// space_depends_on / fence / is_complete are no-ops or trivially satisfied. + +#ifndef KOKKOS_EVENT_HPP +#define KOKKOS_EVENT_HPP +#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE +#define KOKKOS_IMPL_PUBLIC_INCLUDE +#define KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_EVENT +#endif + +#include +#include +#include + +namespace Kokkos { +namespace Impl { +template +struct EventResource { + EventResource(const std::string& label_, + const Kokkos::View& flag_, + const ExecutionSpace& exec_) + : label(label_), flag(flag_), exec(exec_) {} + std::string label; + Kokkos::View flag; + ExecutionSpace exec; +}; +} // namespace Impl + +namespace Experimental { + +//============================================================================ +// Backend-agnostic Event — fallback for non-native-event backends +//============================================================================ + +/// Portable fallback event for backends without native event support. +/// +/// On record(), a fence is issued so that subsequent space_depends_on() +/// and fence() are trivially satisfied. This preserves correctness at +/// the cost of synchronisation -- the same trade-off existing Kokkos +/// code already makes. +/// +/// Backends that provide a native implementation (e.g. CUDA) specialize +/// this template in backend-specific headers. + +// forward declare the class and the friend function space_depends_on +// so that we can make namespace qualified call work +template +struct Event; + +// Device-side dependency: the given execution space waits until the event +// has occured. +template +void space_depends_on(const Exec& exec_space, const Event& event); + +template +struct Event { + using execution_space = Exec; + + private: + using resource_t = Kokkos::Impl::EventResource; + using handle_t = std::shared_ptr; + using flag_t = Kokkos::View; + + public: + Event(const std::string& label_) + : m_handle(std::make_shared( + label_, flag_t(std::string("Kokkos::Event::flag:" + label_)), + execution_space())) { + m_handle->flag() = 1; + }; + + Event(const std::string& label_, const execution_space& exec_space) + : m_handle(std::make_shared( + label_, + Kokkos::View( + std::string("Kokkos::Event::flag:") + label_), + execution_space())) { + record(exec_space); + }; + + // Create an event at the current spot in the execution space queue + void record(const execution_space& exec_space) { + m_handle->flag() = 0; + m_handle->exec = execution_space(); + auto flag = m_handle->flag; + Kokkos::parallel_for( + std::string("Kokkos::Event::record:" + m_handle->label), 1, + KOKKOS_LAMBDA(int) { flag() = 1; }); + } + + // Wait untile the even occurs + void fence() const { + while (m_handle->flag() != 1) std::this_thread::yield(); + } + + // Check whether the even has occured + bool is_complete() const { return m_handle->flag() == 1; } + + const std::string& label() const { return m_handle->label; } + + // Enqueue a dependency on the event in an execution space instance + friend void space_depends_on( + const execution_space& exec_space, const Event& event); + + private: + handle_t m_handle; +}; + +template +void space_depends_on(const Exec& exec_space, const Event& event) { + // Only need to wait if its not the same execution space instance + // Otherwise any work issues to + if (exec_space != event.m_handle->exec) event.fence(); +} +} // namespace Experimental +} // namespace Kokkos + +#ifdef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_EVENT +#undef KOKKOS_IMPL_PUBLIC_INCLUDE +#undef KOKKOS_IMPL_PUBLIC_INCLUDE_NOTDEFINED_EVENT +#endif +#endif // KOKKOS_EVENT_HPP diff --git a/core/src/decl/Kokkos_Declare_CUDA.hpp b/core/src/decl/Kokkos_Declare_CUDA.hpp index b7e0bf7b405..b4d35ca64a4 100644 --- a/core/src/decl/Kokkos_Declare_CUDA.hpp +++ b/core/src/decl/Kokkos_Declare_CUDA.hpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include diff --git a/core/unit_test/CMakeLists.txt b/core/unit_test/CMakeLists.txt index b128d32e462..46710fc1e0b 100644 --- a/core/unit_test/CMakeLists.txt +++ b/core/unit_test/CMakeLists.txt @@ -125,6 +125,7 @@ foreach(Tag Threads;Serial;OpenMP;Cuda;HPX;OpenACC;HIP;SYCL) DeepCopy_Assignment DeepCopy_Narrowing DeepCopyAlignment + Event ExecSpacePartitioning ExecSpaceThreadSafety ExecutionSpace @@ -603,6 +604,7 @@ if(Kokkos_ENABLE_CUDA) kokkos_add_executable_and_test( CoreUnitTest_CudaInterOpGraphMultiGPU SOURCES UnitTestMainInit.cpp cuda/TestCuda_InterOp_GraphMultiGPU.cpp ) + endif() if(Kokkos_ENABLE_HIP) diff --git a/core/unit_test/TestEvent.hpp b/core/unit_test/TestEvent.hpp new file mode 100644 index 00000000000..20d0d9e9f24 --- /dev/null +++ b/core/unit_test/TestEvent.hpp @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project + +#include +#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES +import kokkos.core; +#else +#include +#endif + +#include + +#include + +namespace Test { + +template +struct FillFunctor { + ViewType data; + template + KOKKOS_FUNCTION void operator()(T i) const { + data(i) = i + 1; + } +}; + +template +struct ProduceFunctor { + ViewType data; + template + KOKKOS_FUNCTION void operator()(T i) const { + data(i) = i * 2; + } +}; + +template +struct ConsumeFunctor { + ViewType data; + ResultType result; + template + KOKKOS_FUNCTION void operator()(T i) const { + Kokkos::atomic_add(&result(), + static_cast(data(i))); + } +}; + +// ============================================================================ +// Portable tests -- run for every enabled backend via TEST_EXECSPACE +// ============================================================================ + +TEST(TEST_CATEGORY, event_record_and_wait) { + using exec_space = TEST_EXECSPACE; + using memory_space = typename exec_space::memory_space; + using view_type = Kokkos::View; + + exec_space space; + constexpr int N = 1000; + + view_type data(Kokkos::view_alloc(space, "data", Kokkos::WithoutInitializing), + N); + + Kokkos::parallel_for("fill", Kokkos::RangePolicy(space, 0, N), + FillFunctor{data}); + + Kokkos::Experimental::Event evt("test_event"); + evt.record(space); + evt.fence(); + + auto h_data = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), data); + for (int i = 0; i < N; ++i) { + ASSERT_EQ(h_data(i), i + 1); + } +} + +TEST(TEST_CATEGORY, event_is_complete) { + using exec_space = TEST_EXECSPACE; + + exec_space space; + + Kokkos::Experimental::Event evt("test_event"); + evt.record(space); + evt.fence(); + + ASSERT_TRUE(evt.is_complete()); +} + +TEST(TEST_CATEGORY, event_space_depends_on) { + using exec_space = TEST_EXECSPACE; + using memory_space = typename exec_space::memory_space; + using view_type = Kokkos::View; + using result_type = Kokkos::View; + + exec_space space_a; + exec_space space_b; + constexpr int N = 10000; + + view_type data( + Kokkos::view_alloc(space_a, "data", Kokkos::WithoutInitializing), N); + + Kokkos::parallel_for("produce", + Kokkos::RangePolicy(space_a, 0, N), + ProduceFunctor{data}); + + Kokkos::Experimental::Event evt("test_event"); + evt.record(space_a); + Kokkos::Experimental::space_depends_on(space_b, evt); + + result_type result(Kokkos::view_alloc(space_b, "result")); + + Kokkos::parallel_for("consume", + Kokkos::RangePolicy(space_b, 0, N), + ConsumeFunctor{data, result}); + + space_b.fence(); + + auto h_result = + Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), result); + + int64_t expected = 0; + for (int i = 0; i < N; ++i) expected += i * 2; + ASSERT_EQ(h_result(), expected); +} + +TEST(TEST_CATEGORY, event_move_semantics) { + using exec_space = TEST_EXECSPACE; + + exec_space space; + + Kokkos::Experimental::Event evt1("test_event"); + evt1.record(space); + + Kokkos::Experimental::Event evt2(std::move(evt1)); + + evt2.fence(); + ASSERT_TRUE(evt2.is_complete()); +} + +TEST(TEST_CATEGORY, event_copy_semantics) { + using exec_space = TEST_EXECSPACE; + using memory_space = typename exec_space::memory_space; + using view_type = Kokkos::View; + + exec_space space_a; + exec_space space_b; + constexpr int N = 1000; + + view_type data( + Kokkos::view_alloc(space_a, "data", Kokkos::WithoutInitializing), N); + + Kokkos::parallel_for("fill", Kokkos::RangePolicy(space_a, 0, N), + FillFunctor{data}); + + Kokkos::Experimental::Event evt("test_event", space_a); + auto evt_copy = evt; + + Kokkos::Experimental::space_depends_on(space_b, evt); + evt_copy.fence(); + + ASSERT_TRUE(evt_copy.is_complete()); + ASSERT_TRUE(evt.is_complete()); + + auto h_data = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), data); + for (int i = 0; i < N; ++i) { + ASSERT_EQ(h_data(i), i + 1); + } +} + +} // namespace Test