graph: allow access to native graph object

The native graph objects should be accessible for advanced users.

cmake/kokkos_arch.cmake

@@ -675,6 +675,8 @@ IF(KOKKOS_ENABLE_SYCL)
       )
     ENDIF()
   ENDIF()
+
+  CHECK_CXX_SYMBOL_EXISTS(SYCL_EXT_ONEAPI_GRAPH "sycl/sycl.hpp" KOKKOS_IMPL_HAVE_SYCL_EXT_ONEAPI_GRAPH)
 ENDIF()
 
 SET(CUDA_ARCH_ALREADY_SPECIFIED "")

core/src/Cuda/Kokkos_Cuda_Graph_Impl.hpp

@@ -199,6 +199,9 @@ struct GraphImpl<Kokkos::Cuda> {
         m_execution_space, _graph_node_kernel_ctor_tag{},
         aggregate_kernel_impl_t{});
   }
+
+  cudaGraph_t get_cuda_graph() { return m_graph; }
+  cudaGraphExec_t get_cuda_graph_exec() { return m_graph_exec; }
 };
 
 }  // end namespace Impl

core/src/HIP/Kokkos_HIP_Graph_Impl.hpp

@@ -79,6 +79,9 @@ class GraphImpl<Kokkos::HIP> {
     KOKKOS_ENSURES(m_graph_exec);
   }
 
+  hipGraph_t get_hip_graph() { return m_graph; }
+  hipGraphExec_t get_hip_graph_exec() { return m_graph_exec; }
+
  private:
   Kokkos::HIP m_execution_space;
   hipGraph_t m_graph          = nullptr;

core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp

@@ -76,6 +76,9 @@ class GraphImpl<Kokkos::SYCL> {
     m_graph_exec = m_graph.finalize();
   }
 
+  auto& get_sycl_graph() { return m_graph; }
+  auto& get_sycl_graph_exec() { return m_graph_exec; }
+
  private:
   Kokkos::SYCL m_execution_space;
   sycl::ext::oneapi::experimental::command_graph<

core/unit_test/CMakeLists.txt

@@ -840,6 +840,13 @@ if(Kokkos_ENABLE_CUDA)
       UnitTestMainInit.cpp
       cuda/TestCuda_InterOp_StreamsMultiGPU.cpp
   )
+
+  KOKKOS_ADD_EXECUTABLE_AND_TEST(
+    CoreUnitTest_CudaInterOpGraph
+    SOURCES
+      UnitTestMainInit.cpp
+      cuda/TestCuda_InterOp_Graph.cpp
+  )
 endif()
 
 if(Kokkos_ENABLE_HIP)
@@ -867,6 +874,13 @@ if(Kokkos_ENABLE_HIP)
       UnitTestMain.cpp
       hip/TestHIP_InterOp_Streams.cpp
   )
+
+  KOKKOS_ADD_EXECUTABLE_AND_TEST(
+    CoreUnitTest_HIPInterOpGraph
+    SOURCES
+      UnitTestMainInit.cpp
+      hip/TestHIP_InterOp_Graph.cpp
+  )
 endif()
 
 if(Kokkos_ENABLE_SYCL)
@@ -951,6 +965,15 @@ if(Kokkos_ENABLE_SYCL)
       UnitTestMainInit.cpp
       sycl/TestSYCL_InterOp_StreamsMultiGPU.cpp
   )
+
+  if(KOKKOS_IMPL_HAVE_SYCL_EXT_ONEAPI_GRAPH)
+    KOKKOS_ADD_EXECUTABLE_AND_TEST(
+      CoreUnitTest_SYCLInterOpGraph
+      SOURCES
+        UnitTestMainInit.cpp
+        sycl/TestSYCL_InterOp_Graph.cpp
+    )
+  endif()
 endif()
 
 SET(DEFAULT_DEVICE_SOURCES

core/unit_test/TestGraph.hpp

@@ -14,12 +14,13 @@
 //
 //@HEADER
 
-#include <Kokkos_Core.hpp>
+#include <TestGraph_helpers.hpp>
+
 #include <Kokkos_Graph.hpp>
 
 #include <gtest/gtest.h>
 
-namespace Test {
+namespace {
 
 template <class ExecSpace, class ValueType>
 struct NoOpReduceFunctor {
@@ -54,20 +55,6 @@ struct CountTestFunctor {
   }
 };
 
-template <class ExecSpace, class T>
-struct SetViewToValueFunctor {
-  using value_type = T;
-  using view_type =
-      Kokkos::View<T, ExecSpace, Kokkos::MemoryTraits<Kokkos::Atomic>>;
-  view_type v;
-  T value;
-
-  template <class... Ts>
-  KOKKOS_FUNCTION void operator()(Ts&&...) const noexcept {
-    v() = value;
-  }
-};
-
 template <class ExecSpace, class T>
 struct SetResultToViewFunctor {
   using value_type = T;
@@ -589,4 +576,4 @@ TEST_F(TEST_CATEGORY_FIXTURE(graph), end_of_submit_control_flow) {
             value_A + 2 * value_B + value_C + value_D + value_E + value_F);
 }
 
-}  // end namespace Test
+}  // namespace

core/unit_test/TestGraph_helpers.hpp

@@ -0,0 +1,35 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#include <Kokkos_Core.hpp>
+
+namespace {
+
+template <class ExecSpace, class T>
+struct SetViewToValueFunctor {
+  using value_type = T;
+  using view_type =
+      Kokkos::View<T, ExecSpace, Kokkos::MemoryTraits<Kokkos::Atomic>>;
+  view_type v;
+  T value;
+
+  template <class... Ts>
+  KOKKOS_FUNCTION void operator()(Ts&&...) const noexcept {
+    v() = value;
+  }
+};
+
+}  // namespace

core/unit_test/cuda/TestCuda_InterOp_Graph.cpp

@@ -0,0 +1,105 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#include <filesystem>
+
+#include <TestCuda_Category.hpp>
+#include <TestGraph_helpers.hpp>
+#include <Kokkos_Graph.hpp>
+
+#include <gtest/gtest.h>
+
+namespace {
+
+// This test checks the promises of Kokkos::Graph against its
+// underlying Cuda native objects.
+TEST(TEST_CATEGORY, graph_promises_on_native_objects) {
+  auto graph = Kokkos::Experimental::create_graph<Kokkos::Cuda>();
+
+  auto root = Kokkos::Impl::GraphAccess::create_root_ref(graph);
+
+  auto graph_ptr_impl =
+      Kokkos::Impl::GraphAccess::get_graph_weak_ptr(root).lock();
+
+  static_assert(std::is_same_v<decltype(graph_ptr_impl)::element_type,
+                               Kokkos::Impl::GraphImpl<Kokkos::Cuda>>);
+
+  // Before instantiation, the Cuda graph is valid, but the Cuda executable
+  // graph is still null.
+  cudaGraph_t cuda_graph = graph_ptr_impl->get_cuda_graph();
+
+  ASSERT_NE(cuda_graph, nullptr);
+  ASSERT_EQ(graph_ptr_impl->get_cuda_graph_exec(), nullptr);
+
+  // After instantiation, both native objects are valid.
+  graph.instantiate();
+
+  cudaGraphExec_t cuda_graph_exec = graph_ptr_impl->get_cuda_graph_exec();
+
+  ASSERT_EQ(graph_ptr_impl->get_cuda_graph(), cuda_graph);
+  ASSERT_NE(cuda_graph_exec, nullptr);
+
+  // Submission should not affect the underlying objects.
+  graph.submit();
+
+  ASSERT_EQ(graph_ptr_impl->get_cuda_graph(), cuda_graph);
+  ASSERT_EQ(graph_ptr_impl->get_cuda_graph_exec(), cuda_graph_exec);
+}
+
+// Use native Cuda graph to generate a DOT representation.
+TEST(TEST_CATEGORY, graph_instantiate_and_debug_dot_print) {
+  using view_t = Kokkos::View<int, Kokkos::Cuda>;
+
+  const Kokkos::Cuda exec{};
+
+  view_t data(Kokkos::view_alloc(exec, "witness"));
+
+  auto graph = Kokkos::Experimental::create_graph(exec);
+
+  auto root = Kokkos::Impl::GraphAccess::create_root_ref(graph);
+
+  root.then_parallel_for(1, SetViewToValueFunctor<Kokkos::Cuda, int>{data, 1});
+
+  auto graph_ptr_impl =
+      Kokkos::Impl::GraphAccess::get_graph_weak_ptr(root).lock();
+
+  graph.instantiate();
+
+  size_t num_nodes = 0;
+
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphGetNodes(graph_ptr_impl->get_cuda_graph(), nullptr, &num_nodes));
+
+  ASSERT_EQ(num_nodes, 2u);
+
+  const auto dot = std::filesystem::temp_directory_path() / "cuda_graph.dot";
+
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphDebugDotPrint(graph_ptr_impl->get_cuda_graph(), dot.c_str(),
+                             cudaGraphDebugDotFlagsVerbose));
+
+  ASSERT_TRUE(std::filesystem::exists(dot));
+  ASSERT_GT(std::filesystem::file_size(dot), 0u);
+
+  unsigned long long flags =
+      Kokkos::Experimental::finite_max_v<unsigned long long>;
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphExecGetFlags(graph_ptr_impl->get_cuda_graph_exec(), &flags));
+
+  ASSERT_EQ(flags, 0u);
+}
+
+}  // namespace

core/unit_test/hip/TestHIP_InterOp_Graph.cpp

@@ -0,0 +1,98 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#include <filesystem>
+
+#include <TestHIP_Category.hpp>
+#include <TestGraph_helpers.hpp>
+#include <Kokkos_Graph.hpp>
+
+#include <gtest/gtest.h>
+
+namespace {
+
+// This test checks the promises of Kokkos::Graph against its
+// underlying HIP native objects.
+TEST(TEST_CATEGORY, graph_promises_on_native_objects) {
+  auto graph = Kokkos::Experimental::create_graph<Kokkos::HIP>();
+
+  auto root = Kokkos::Impl::GraphAccess::create_root_ref(graph);
+
+  auto graph_ptr_impl =
+      Kokkos::Impl::GraphAccess::get_graph_weak_ptr(root).lock();
+
+  static_assert(std::is_same_v<decltype(graph_ptr_impl)::element_type,
+                               Kokkos::Impl::GraphImpl<Kokkos::HIP>>);
+
+  // Before instantiation, the HIP graph is valid, but the HIP executable
+  // graph is still null.
+  hipGraph_t hip_graph = graph_ptr_impl->get_hip_graph();
+
+  ASSERT_NE(hip_graph, nullptr);
+  ASSERT_EQ(graph_ptr_impl->get_hip_graph_exec(), nullptr);
+
+  // After instantiation, both native objects are valid.
+  graph.instantiate();
+
+  hipGraphExec_t hip_graph_exec = graph_ptr_impl->get_hip_graph_exec();
+
+  ASSERT_EQ(graph_ptr_impl->get_hip_graph(), hip_graph);
+  ASSERT_NE(hip_graph_exec, nullptr);
+
+  // Submission should not affect the underlying objects.
+  graph.submit();
+
+  ASSERT_EQ(graph_ptr_impl->get_hip_graph(), hip_graph);
+  ASSERT_EQ(graph_ptr_impl->get_hip_graph_exec(), hip_graph_exec);
+}
+
+// Use native HIP graph to generate a DOT representation.
+TEST(TEST_CATEGORY, graph_instantiate_and_debug_dot_print) {
+  using view_t = Kokkos::View<int, Kokkos::HIP>;
+
+  const Kokkos::HIP exec{};
+
+  view_t data(Kokkos::view_alloc(exec, "witness"));
+
+  auto graph = Kokkos::Experimental::create_graph(exec);
+
+  auto root = Kokkos::Impl::GraphAccess::create_root_ref(graph);
+
+  root.then_parallel_for(1, SetViewToValueFunctor<Kokkos::HIP, int>{data, 1});
+
+  auto graph_ptr_impl =
+      Kokkos::Impl::GraphAccess::get_graph_weak_ptr(root).lock();
+
+  graph.instantiate();
+
+  size_t num_nodes = 0;
+
+  KOKKOS_IMPL_HIP_SAFE_CALL(
+      hipGraphGetNodes(graph_ptr_impl->get_hip_graph(), nullptr, &num_nodes));
+
+  ASSERT_EQ(num_nodes, 2u);
+
+  const auto dot = std::filesystem::temp_directory_path() / "hip_graph.dot";
+
+  KOKKOS_IMPL_HIP_SAFE_CALL(
+      hipGraphDebugDotPrint(graph_ptr_impl->get_hip_graph(), dot.c_str(),
+                            hipGraphDebugDotFlagsVerbose));
+
+  ASSERT_TRUE(std::filesystem::exists(dot));
+  ASSERT_GT(std::filesystem::file_size(dot), 0u);
+}
+
+}  // namespace