Skip to content

Commit

Permalink
core(graph): promote instantiate to public API
Browse files Browse the repository at this point in the history
  • Loading branch information
romintomasetti committed Aug 21, 2024
1 parent f5c527d commit a351610
Show file tree
Hide file tree
Showing 6 changed files with 80 additions and 12 deletions.
8 changes: 5 additions & 3 deletions core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,9 @@ struct GraphImpl<Kokkos::Cuda> {

using node_details_t = GraphNodeBackendSpecificDetails<Kokkos::Cuda>;

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];
Expand All @@ -60,10 +62,10 @@ struct GraphImpl<Kokkos::Cuda> {
->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<Kokkos::Cuda, Kokkos::Experimental::TypeErasedTag,
Kokkos::Experimental::TypeErasedTag>;
Expand Down Expand Up @@ -160,7 +162,7 @@ struct GraphImpl<Kokkos::Cuda> {

void submit() {
if (!bool(m_graph_exec)) {
_instantiate_graph();
instantiate();
}
KOKKOS_IMPL_CUDA_SAFE_CALL(
(m_execution_space.impl_internal_space_instance()
Expand Down
8 changes: 5 additions & 3 deletions core/src/HIP/Kokkos_HIP_Graph_Impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -69,15 +69,17 @@ class GraphImpl<Kokkos::HIP> {
template <class... PredecessorRefs>
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;
Expand Down Expand Up @@ -147,7 +149,7 @@ inline void GraphImpl<Kokkos::HIP>::add_predecessor(

inline void GraphImpl<Kokkos::HIP>::submit() {
if (!m_graph_exec) {
instantiate_graph();
instantiate();
}
KOKKOS_IMPL_HIP_SAFE_CALL(
hipGraphLaunch(m_graph_exec, m_execution_space.hip_stream()));
Expand Down
5 changes: 5 additions & 0 deletions core/src/Kokkos_Graph.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
9 changes: 6 additions & 3 deletions core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,12 @@ class GraphImpl<Kokkos::SYCL> {
template <class... PredecessorRefs>
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>
Expand Down Expand Up @@ -137,7 +140,7 @@ inline void GraphImpl<Kokkos::SYCL>::add_predecessor(

inline void GraphImpl<Kokkos::SYCL>::submit() {
if (!m_graph_exec) {
instantiate_graph();
instantiate();
}
m_execution_space.sycl_queue().ext_oneapi_graph(*m_graph_exec);
}
Expand Down
9 changes: 9 additions & 0 deletions core/src/impl/Kokkos_Default_Graph_Impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,13 @@ struct GraphImpl : private ExecutionSpaceInstanceStorage<ExecutionSpace> {
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) {
Expand All @@ -147,6 +153,9 @@ struct GraphImpl : private ExecutionSpaceInstanceStorage<ExecutionSpace> {
}
}

private:
bool m_has_been_instantiated = false;

// </editor-fold> end required customizations }}}2
//----------------------------------------------------------------------------
};
Expand Down
53 changes: 50 additions & 3 deletions core/unit_test/TestGraph.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<TEST_EXECSPACE>([&](auto root) {
root.then_parallel_for(1, count_functor{count, bugs, 0, 0});
Expand All @@ -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();
Expand All @@ -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());
}

#define GTEST_SKIP_IF_NO_ASSERTION(__msg__) { \
bool checked_assertions = false; \
KOKKOS_ASSERT(checked_assertions = true); \
if (!checked_assertions) { \
GTEST_SKIP() << __msg__; \
}}

/// @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) {
GTEST_SKIP_IF_NO_ASSERTION("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<TEST_EXECSPACE, Kokkos::Experimental::OpenMPTarget>)
GTEST_SKIP() << "skipping since OpenMPTarget can't use team_size 1";
Expand Down

0 comments on commit a351610

Please sign in to comment.