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 cuda_graph() { return m_graph; }
+  cudaGraphExec_t 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 hip_graph() { return m_graph; }
+  hipGraphExec_t hip_graph_exec() { return m_graph_exec; }
+
  private:
   Kokkos::HIP m_execution_space;
   hipGraph_t m_graph          = nullptr;

core/src/HIP/Kokkos_HIP_KernelLaunch.hpp

@@ -25,11 +25,7 @@
 #include <HIP/Kokkos_HIP_Instance.hpp>
 #include <HIP/Kokkos_HIP_Space.hpp>
 
-#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
-#define KOKKOS_IMPL_HIP_GRAPH_ENABLED
-#endif
-
-#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
+#ifdef KOKKOS_HIP_NATIVE_GRAPH
 #include <HIP/Kokkos_HIP_GraphNodeKernel.hpp>
 #include <impl/Kokkos_GraphImpl_fwd.hpp>
 #endif
@@ -384,7 +380,7 @@ struct HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
         driver);
   }
 
-#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
+#ifdef KOKKOS_HIP_NATIVE_GRAPH
   static void create_parallel_launch_graph_node(
       DriverType const &driver, dim3 const &grid, dim3 const &block, int shmem,
       HIPInternal const * /*hip_instance*/) {
@@ -442,7 +438,7 @@ struct HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
         driver_ptr);
   }
 
-#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
+#ifdef KOKKOS_HIP_NATIVE_GRAPH
   static void create_parallel_launch_graph_node(
       DriverType const &driver, dim3 const &grid, dim3 const &block, int shmem,
       HIPInternal const *hip_instance) {
@@ -585,7 +581,7 @@ void hip_parallel_launch(const DriverType &driver, const dim3 &grid,
                          const dim3 &block, const int shmem,
                          const HIPInternal *hip_instance,
                          const bool prefer_shmem) {
-#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
+#ifdef KOKKOS_HIP_NATIVE_GRAPH
   if constexpr (DoGraph) {
     // Graph launch
     using base_t = HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
@@ -628,8 +624,6 @@ void hip_parallel_launch(const DriverType &driver, const dim3 &grid,
 }  // namespace Impl
 }  // namespace Kokkos
 
-#undef KOKKOS_IMPL_HIP_GRAPH_ENABLED
-
 #endif
 
 #endif

core/src/Kokkos_Graph.hpp

@@ -97,6 +97,10 @@ struct [[nodiscard]] Graph {
   }
 
   void submit() const { submit(get_execution_space()); }
+
+  decltype(auto) native_graph();
+
+  decltype(auto) native_graph_exec();
 };
 
 // </editor-fold> end Graph }}}1
@@ -168,6 +172,45 @@ create_graph(Closure&& arg_closure) {
 // </editor-fold> end create_graph }}}1
 //==============================================================================
 
+template <class ExecutionSpace>
+decltype(auto) Graph<ExecutionSpace>::native_graph() {
+  KOKKOS_EXPECTS(bool(m_impl_ptr));
+#if defined(KOKKOS_ENABLE_CUDA)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::Cuda>) {
+    return m_impl_ptr->cuda_graph();
+  }
+#elif defined(KOKKOS_ENABLE_HIP) && defined(KOKKOS_HIP_NATIVE_GRAPH)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::HIP>) {
+    return m_impl_ptr->hip_graph();
+  }
+#elif defined(KOKKOS_ENABLE_SYCL) && defined(SYCL_EXT_ONEAPI_GRAPH)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::SYCL>) {
+    return m_impl_ptr->sycl_graph();
+  }
+#endif
+  Kokkos::abort("Native graph object not available for this execution space");
+}
+
+template <class ExecutionSpace>
+decltype(auto) Graph<ExecutionSpace>::native_graph_exec() {
+  KOKKOS_EXPECTS(bool(m_impl_ptr));
+#if defined(KOKKOS_ENABLE_CUDA)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::Cuda>) {
+    return m_impl_ptr->cuda_graph_exec();
+  }
+#elif defined(KOKKOS_ENABLE_HIP) && defined(KOKKOS_HIP_NATIVE_GRAPH)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::HIP>) {
+    return m_impl_ptr->hip_graph_exec();
+  }
+#elif defined(KOKKOS_ENABLE_SYCL) && defined(SYCL_EXT_ONEAPI_GRAPH)
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::SYCL>) {
+    return m_impl_ptr->sycl_graph_exec();
+  }
+#endif
+  Kokkos::abort(
+      "Native executable graph object not available for this execution space");
+}
+
 }  // end namespace Experimental
 }  // namespace Kokkos
 
@@ -179,8 +222,7 @@ create_graph(Closure&& arg_closure) {
 #include <impl/Kokkos_Default_Graph_Impl.hpp>
 #include <Cuda/Kokkos_Cuda_Graph_Impl.hpp>
 #if defined(KOKKOS_ENABLE_HIP)
-// The implementation of hipGraph in ROCm 5.2 is bugged, so we cannot use it.
-#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
+#if defined(KOKKOS_HIP_NATIVE_GRAPH)
 #include <HIP/Kokkos_HIP_Graph_Impl.hpp>
 #endif
 #endif

core/src/SYCL/Kokkos_SYCL_Graph_Impl.hpp

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

core/src/setup/Kokkos_Setup_HIP.hpp

@@ -43,6 +43,11 @@
 #endif
 // clang-format on
 
+// The implementation of hipGraph in ROCm 5.2 is bugged, so we cannot use it.
+#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
+#define KOKKOS_HIP_NATIVE_GRAPH 1
+#endif
+
 #endif  // #if defined( KOKKOS_ENABLE_HIP )
 
 #endif

core/unit_test/CMakeLists.txt

@@ -841,6 +841,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)
@@ -868,6 +875,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)
@@ -952,6 +966,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

@@ -398,6 +398,16 @@ TEST_F(TEST_CATEGORY_FIXTURE(graph), empty_graph) {
   ex.fence();
 }
 
+// Ensure that an empty graph on the default host execution space
+// can be submitted.
+TEST_F(TEST_CATEGORY_FIXTURE(graph), empty_graph_default_host_exec) {
+  auto graph =
+      Kokkos::Experimental::create_graph(Kokkos::DefaultHostExecutionSpace{});
+  graph.instantiate();
+  graph.submit();
+  graph.get_execution_space().fence();
+}
+
 template <typename ViewType, size_t TargetIndex, size_t NumIndices = 0>
 struct FetchValuesAndContribute {
   static_assert(std::is_same_v<typename ViewType::value_type,

core/unit_test/cuda/TestCuda_InterOp_Graph.cpp

@@ -0,0 +1,125 @@
+//@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 <fstream>
+#include <regex>
+
+#include <TestCuda_Category.hpp>
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Graph.hpp>
+
+#include <gtest/gtest.h>
+
+namespace {
+
+template <typename ViewType>
+struct Increment {
+  ViewType data;
+
+  KOKKOS_FUNCTION
+  void operator()(const int) const { ++data(); }
+};
+
+// 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);
+
+  // Before instantiation, the Cuda graph is valid, but the Cuda executable
+  // graph is still null.
+  cudaGraph_t cuda_graph = graph.native_graph();
+
+  ASSERT_NE(cuda_graph, nullptr);
+  ASSERT_EQ(graph.native_graph_exec(), nullptr);
+
+  // After instantiation, both native objects are valid.
+  graph.instantiate();
+
+  cudaGraphExec_t cuda_graph_exec = graph.native_graph_exec();
+
+  ASSERT_EQ(graph.native_graph(), cuda_graph);
+  ASSERT_NE(cuda_graph_exec, nullptr);
+
+  // Submission should not affect the underlying objects.
+  graph.submit();
+
+  ASSERT_EQ(graph.native_graph(), cuda_graph);
+  ASSERT_EQ(graph.native_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, Increment<view_t>{data});
+
+  graph.instantiate();
+
+  size_t num_nodes = 0;
+
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphGetNodes(graph.native_graph(), nullptr, &num_nodes));
+
+  ASSERT_EQ(num_nodes, 2u);
+
+#if CUDA_VERSION >= 11600
+  const auto dot = std::filesystem::temp_directory_path() / "cuda_graph.dot";
+
+  // Convert path to string then to const char * to make it work on Windows.
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphDebugDotPrint(graph.native_graph(), dot.string().c_str(),
+                             cudaGraphDebugDotFlagsVerbose));
+
+  ASSERT_TRUE(std::filesystem::exists(dot));
+  ASSERT_GT(std::filesystem::file_size(dot), 0u);
+
+  // We could write a check against the full kernel's function signature, but
+  // it would make the test rely too much on internal implementation details.
+  // Therefore, we just look for the functor and policy. Note that the
+  // signature is mangled in the 'dot' output.
+  const std::string expected("[A-Za-z0-9_]+Increment[A-Za-z0-9_]+RangePolicy");
+
+  std::stringstream buffer;
+  buffer << std::ifstream(dot).rdbuf();
+
+  ASSERT_TRUE(std::regex_search(buffer.str(), std::regex(expected)))
+      << "Could not find expected signature regex " << std::quoted(expected)
+      << " in " << dot;
+
+#endif
+
+#if CUDA_VERSION >= 12000
+  unsigned long long flags =
+      Kokkos::Experimental::finite_max_v<unsigned long long>;
+  KOKKOS_IMPL_CUDA_SAFE_CALL(
+      cudaGraphExecGetFlags(graph.native_graph_exec(), &flags));
+
+  ASSERT_EQ(flags, 0u);
+#endif
+}
+
+}  // namespace