Address reviewer' comments

tcclevenger · Sep 8, 2023 · 035d284 · 035d284
1 parent 9fde326
commit 035d284
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 11 deletions.
diff --git a/core/src/HIP/Kokkos_HIP_GraphNodeKernel.hpp b/core/src/HIP/Kokkos_HIP_GraphNodeKernel.hpp
@@ -27,6 +27,7 @@
 #include <Kokkos_PointerOwnership.hpp>
 
 #include <HIP/Kokkos_HIP_SharedAllocationRecord.hpp>
+#include <HIP/Kokkos_HIP_GraphNode_Impl.hpp>
 
 namespace Kokkos {
 namespace Impl {

diff --git a/core/src/HIP/Kokkos_HIP_GraphNode_Impl.hpp b/core/src/HIP/Kokkos_HIP_GraphNode_Impl.hpp
@@ -51,6 +51,4 @@ struct GraphNodeBackendDetailsBeforeTypeErasure<Kokkos::HIP, Kernel,
 }  // namespace Impl
 }  // namespace Kokkos
 
-#include <HIP/Kokkos_HIP_GraphNodeKernel.hpp>
-
 #endif
diff --git a/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp b/core/src/HIP/Kokkos_HIP_Graph_Impl.hpp
@@ -22,9 +22,9 @@
 #include <Kokkos_Graph_fwd.hpp>
 
 #include <impl/Kokkos_GraphImpl.hpp>
-
 #include <impl/Kokkos_GraphNodeImpl.hpp>
-#include <HIP/Kokkos_HIP_GraphNode_Impl.hpp>
+
+#include <HIP/Kokkos_HIP_GraphNodeKernel.hpp>
 
 namespace Kokkos {
 namespace Impl {

diff --git a/core/src/HIP/Kokkos_HIP_KernelLaunch.hpp b/core/src/HIP/Kokkos_HIP_KernelLaunch.hpp
@@ -26,6 +26,10 @@
 #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
 #include <HIP/Kokkos_HIP_GraphNodeKernel.hpp>
 #include <impl/Kokkos_GraphImpl_fwd.hpp>
 #endif
@@ -380,13 +384,13 @@ struct HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
         driver);
   }
 
-#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
+#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
   static void create_parallel_launch_graph_node(
       DriverType const &driver, dim3 const &grid, dim3 const &block, int shmem,
       HIPInternal const * /*hip_instance*/) {
-    auto const &graph = Impl::get_hip_graph_from_kernel(driver);
+    auto const &graph = get_hip_graph_from_kernel(driver);
     KOKKOS_EXPECTS(graph);
-    auto &graph_node = Impl::get_hip_graph_node_from_kernel(driver);
+    auto &graph_node = get_hip_graph_node_from_kernel(driver);
     // Expect node not yet initialized
     KOKKOS_EXPECTS(!graph_node);
 
@@ -438,13 +442,13 @@ struct HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
         driver_ptr);
   }
 
-#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
+#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
   static void create_parallel_launch_graph_node(
       DriverType const &driver, dim3 const &grid, dim3 const &block, int shmem,
       HIPInternal const *hip_instance) {
-    auto const &graph = Impl::get_hip_graph_from_kernel(driver);
+    auto const &graph = get_hip_graph_from_kernel(driver);
     KOKKOS_EXPECTS(graph);
-    auto &graph_node = Impl::get_hip_graph_node_from_kernel(driver);
+    auto &graph_node = get_hip_graph_node_from_kernel(driver);
     // Expect node not yet initialized
     KOKKOS_EXPECTS(!graph_node);
 
@@ -581,7 +585,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) {
-#if !((HIP_VERSION_MAJOR == 5) && (HIP_VERSION_MINOR == 2))
+#ifdef KOKKOS_IMPL_HIP_GRAPH_ENABLED
   if constexpr (DoGraph) {
     // Graph launch
     using base_t = HIPParallelLaunchKernelInvoker<DriverType, LaunchBounds,
@@ -624,6 +628,8 @@ void hip_parallel_launch(const DriverType &driver, const dim3 &grid,
 }  // namespace Impl
 }  // namespace Kokkos
 
+#undef KOKKOS_IMPL_HIP_GRAPH_ENABLED
+
 #endif
 
 #endif
diff --git a/core/src/Kokkos_Graph.hpp b/core/src/Kokkos_Graph.hpp
@@ -162,6 +162,7 @@ Graph<ExecutionSpace> 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))
 #include <HIP/Kokkos_HIP_Graph_Impl.hpp>
 #endif