From 989828b4160d3b9249329ca60cea3b31693e87c8 Mon Sep 17 00:00:00 2001 From: romintomasetti Date: Wed, 21 Aug 2024 14:48:38 +0000 Subject: [PATCH] core(graph): promote `instantiate` to public API --- core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp | 8 ++-- core/src/HIP/Kokkos_HIP_Graph_Impl.hpp | 8 ++-- core/src/Kokkos_Graph.hpp | 5 ++ core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp | 9 ++-- core/src/impl/Kokkos_Default_Graph_Impl.hpp | 9 ++++ core/unit_test/TestGraph.hpp | 53 +++++++++++++++++++-- 6 files changed, 80 insertions(+), 12 deletions(-) diff --git a/core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp b/core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp index 941cc821223..c3043f45f29 100644 --- a/core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp +++ b/core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp @@ -51,7 +51,9 @@ struct GraphImpl { using node_details_t = GraphNodeBackendSpecificDetails; - void _instantiate_graph() { + public: + void instantiate() { + KOKKOS_EXPECTS(!m_graph_exec); constexpr size_t error_log_size = 256; cudaGraphNode_t error_node = nullptr; char error_log[error_log_size]; @@ -60,10 +62,10 @@ struct GraphImpl { ->cuda_graph_instantiate_wrapper(&m_graph_exec, m_graph, &error_node, error_log, error_log_size))); + KOKKOS_ENSURES(m_graph_exec); // TODO @graphs print out errors } - public: using root_node_impl_t = GraphNodeImpl; @@ -160,7 +162,7 @@ struct GraphImpl { void submit() { if (!bool(m_graph_exec)) { - _instantiate_graph(); + instantiate(); } KOKKOS_IMPL_CUDA_SAFE_CALL( (m_execution_space.impl_internal_space_instance() diff --git a/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp b/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp index f73b477c00c..283563c6d65 100644 --- a/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp +++ b/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp @@ -69,15 +69,17 @@ class GraphImpl { template auto create_aggregate_ptr(PredecessorRefs&&...); - private: - void instantiate_graph() { + void instantiate() { + KOKKOS_EXPECTS(!m_graph_exec); constexpr size_t error_log_size = 256; hipGraphNode_t error_node = nullptr; char error_log[error_log_size]; KOKKOS_IMPL_HIP_SAFE_CALL(hipGraphInstantiate( &m_graph_exec, m_graph, &error_node, error_log, error_log_size)); + KOKKOS_ENSURES(m_graph_exec); } + private: Kokkos::HIP m_execution_space; hipGraph_t m_graph = nullptr; hipGraphExec_t m_graph_exec = nullptr; @@ -147,7 +149,7 @@ inline void GraphImpl::add_predecessor( inline void GraphImpl::submit() { if (!m_graph_exec) { - instantiate_graph(); + instantiate(); } KOKKOS_IMPL_HIP_SAFE_CALL( hipGraphLaunch(m_graph_exec, m_execution_space.hip_stream())); diff --git a/core/src/Kokkos_Graph.hpp b/core/src/Kokkos_Graph.hpp index b6f236e8fbe..7dce57eef3f 100644 --- a/core/src/Kokkos_Graph.hpp +++ b/core/src/Kokkos_Graph.hpp @@ -86,6 +86,11 @@ struct [[nodiscard]] Graph { return m_impl_ptr->get_execution_space(); } + void instantiate() { + KOKKOS_EXPECTS(bool(m_impl_ptr)) + (*m_impl_ptr).instantiate(); + } + void submit() const { KOKKOS_EXPECTS(bool(m_impl_ptr)) (*m_impl_ptr).submit(); diff --git a/core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp b/core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp index 4a3ff13b9af..0fac098a50c 100644 --- a/core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp +++ b/core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp @@ -71,9 +71,12 @@ class GraphImpl { template auto create_aggregate_ptr(PredecessorRefs&&...); - private: - void instantiate_graph() { m_graph_exec = m_graph.finalize(); } + void instantiate() { + KOKKOS_EXPECTS(!m_graph_exec.has_value()); + m_graph_exec = m_graph.finalize(); + } + private: Kokkos::SYCL m_execution_space; sycl::ext::oneapi::experimental::command_graph< sycl::ext::oneapi::experimental::graph_state::modifiable> @@ -137,7 +140,7 @@ inline void GraphImpl::add_predecessor( inline void GraphImpl::submit() { if (!m_graph_exec) { - instantiate_graph(); + instantiate(); } m_execution_space.sycl_queue().ext_oneapi_graph(*m_graph_exec); } diff --git a/core/src/impl/Kokkos_Default_Graph_Impl.hpp b/core/src/impl/Kokkos_Default_Graph_Impl.hpp index d47b8715ad5..6023f8e5be7 100644 --- a/core/src/impl/Kokkos_Default_Graph_Impl.hpp +++ b/core/src/impl/Kokkos_Default_Graph_Impl.hpp @@ -136,7 +136,13 @@ struct GraphImpl : private ExecutionSpaceInstanceStorage { return rv; } + void instantiate() { + KOKKOS_EXPECTS(!m_has_been_instantiated); + m_has_been_instantiated = true; + } + void submit() { + if (!m_has_been_instantiated) instantiate(); // This reset is gross, but for the purposes of our simple host // implementation... for (auto& sink : m_sinks) { @@ -147,6 +153,9 @@ struct GraphImpl : private ExecutionSpaceInstanceStorage { } } + private: + bool m_has_been_instantiated = false; + // end required customizations }}}2 //---------------------------------------------------------------------------- }; diff --git a/core/unit_test/TestGraph.hpp b/core/unit_test/TestGraph.hpp index ae0d1fe452e..19860070283 100644 --- a/core/unit_test/TestGraph.hpp +++ b/core/unit_test/TestGraph.hpp @@ -103,7 +103,7 @@ struct TEST_CATEGORY_FIXTURE(graph) : public ::testing::Test { } }; -TEST_F(TEST_CATEGORY_FIXTURE(graph), launch_one) { +TEST_F(TEST_CATEGORY_FIXTURE(graph), submit_once) { auto graph = Kokkos::Experimental::create_graph([&](auto root) { root.then_parallel_for(1, count_functor{count, bugs, 0, 0}); @@ -116,7 +116,7 @@ TEST_F(TEST_CATEGORY_FIXTURE(graph), launch_one) { ASSERT_EQ(0, bugs_host()); } -TEST_F(TEST_CATEGORY_FIXTURE(graph), launch_one_rvalue) { +TEST_F(TEST_CATEGORY_FIXTURE(graph), submit_once_rvalue) { Kokkos::Experimental::create_graph(ex, [&](auto root) { root.then_parallel_for(1, count_functor{count, bugs, 0, 0}); }).submit(); @@ -127,7 +127,54 @@ TEST_F(TEST_CATEGORY_FIXTURE(graph), launch_one_rvalue) { ASSERT_EQ(0, bugs_host()); } -TEST_F(TEST_CATEGORY_FIXTURE(graph), launch_six) { +/// @test Ensure that @c Kokkos::Graph::instantiate works. +/// @note For now, @c Kokkos::Graph::submit will instantiate if needed, +/// so this test is not very strong. +TEST_F(TEST_CATEGORY_FIXTURE(graph), instantiate_and_submit_once) { + auto graph = Kokkos::Experimental::create_graph(ex, [&](auto root) { + root.then_parallel_for(1, count_functor{count, bugs, 0, 0}); + }); + graph.instantiate(); + graph.submit(); + Kokkos::deep_copy(ex, count_host, count); + Kokkos::deep_copy(ex, bugs_host, bugs); + ex.fence(); + ASSERT_EQ(1, count_host()); + ASSERT_EQ(0, bugs_host()); +} + +/// @test Ensure that @c Kokkos::Graph::instantiate can be called only once. +/// This test checks 2 cases: +/// 1. Instantiating after submission is invalid (this also implicitly +/// checks that submission instantiates if need be). +/// 2. Instantiating twice in a row is invalid. +TEST_F(TEST_CATEGORY_FIXTURE(graph), can_instantiate_only_once) { + { + bool checked_assertions = false; + KOKKOS_ASSERT(checked_assertions = true); + if (!checked_assertions) { + GTEST_SKIP() << "Preconditions are not checked."; + } + } + { + auto graph = Kokkos::Experimental::create_graph(ex, [&](auto root) { + root.then_parallel_for(1, count_functor{count, bugs, 0, 0}); + }); + graph.submit(); + ASSERT_DEATH(graph.instantiate(), + "Expected precondition `.*` evaluated false."); + } + { + auto graph = Kokkos::Experimental::create_graph(ex, [&](auto root) { + root.then_parallel_for(1, count_functor{count, bugs, 0, 0}); + }); + graph.instantiate(); + ASSERT_DEATH(graph.instantiate(), + "Expected precondition `.*` evaluated false."); + } +} + +TEST_F(TEST_CATEGORY_FIXTURE(graph), submit_six) { #ifdef KOKKOS_ENABLE_OPENMPTARGET // FIXME_OPENMPTARGET team_size incompatible if (std::is_same_v) GTEST_SKIP() << "skipping since OpenMPTarget can't use team_size 1";