core(graph): promote instantiate to public API

core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp

@@ -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];
@@ -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>;
@@ -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()

core/src/HIP/Kokkos_HIP_Graph_Impl.hpp

@@ -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;
@@ -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()));

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();

core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp

@@ -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>
@@ -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);
 }

core/src/impl/Kokkos_Default_Graph_Impl.hpp

@@ -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) {
@@ -147,6 +153,9 @@ struct GraphImpl : private ExecutionSpaceInstanceStorage<ExecutionSpace> {
     }
   }
 
+ private:
+  bool m_has_been_instantiated = false;
+
   // </editor-fold> end required customizations }}}2
   //----------------------------------------------------------------------------
 };

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<TEST_EXECSPACE>([&](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) {
+// Ensure that Kokkos::Graph::instantiate works.
+// For now, 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());
+}
+
+// Ensure that 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<TEST_EXECSPACE, Kokkos::Experimental::OpenMPTarget>)
     GTEST_SKIP() << "skipping since OpenMPTarget can't use team_size 1";