[hip][cuda] Merge the tracing implementations. (iree-org#18299)

These were entirely copy-pasted off one another and it's not likely they
will have to diverge in the future.

---------

Signed-off-by: Andrew Woloszyn <andrew.woloszyn@gmail.com>

runtime/src/iree/hal/drivers/cuda/BUILD.bazel

@@ -43,8 +43,6 @@ iree_runtime_cc_library(
         "stream_command_buffer.h",
         "timepoint_pool.c",
         "timepoint_pool.h",
-        "tracing.c",
-        "tracing.h",
     ],
     hdrs = [
         "api.h",
@@ -69,6 +67,7 @@ iree_runtime_cc_library(
         "//runtime/src/iree/hal/utils:memory_file",
         "//runtime/src/iree/hal/utils:resource_set",
         "//runtime/src/iree/hal/utils:semaphore_base",
+        "//runtime/src/iree/hal/utils:stream_tracing",
         "//runtime/src/iree/schemas:cuda_executable_def_c_fbs",
     ],
 )

runtime/src/iree/hal/drivers/cuda/CMakeLists.txt

@@ -44,8 +44,6 @@ iree_cc_library(
     "stream_command_buffer.h"
     "timepoint_pool.c"
     "timepoint_pool.h"
-    "tracing.c"
-    "tracing.h"
   DEPS
     ::dynamic_symbols
     iree::base
@@ -66,6 +64,7 @@ iree_cc_library(
     iree::hal::utils::memory_file
     iree::hal::utils::resource_set
     iree::hal::utils::semaphore_base
+    iree::hal::utils::stream_tracing
     iree::schemas::cuda_executable_def_c_fbs
   PUBLIC
 )

runtime/src/iree/hal/drivers/cuda/cuda_device.c

@@ -26,11 +26,11 @@
 #include "iree/hal/drivers/cuda/pipeline_layout.h"
 #include "iree/hal/drivers/cuda/stream_command_buffer.h"
 #include "iree/hal/drivers/cuda/timepoint_pool.h"
-#include "iree/hal/drivers/cuda/tracing.h"
 #include "iree/hal/utils/deferred_command_buffer.h"
 #include "iree/hal/utils/deferred_work_queue.h"
 #include "iree/hal/utils/file_transfer.h"
 #include "iree/hal/utils/memory_file.h"
+#include "iree/hal/utils/stream_tracing.h"
 
 //===----------------------------------------------------------------------===//
 // iree_hal_cuda_device_t
@@ -62,7 +62,7 @@ typedef struct iree_hal_cuda_device_t {
   // The CUstream used to issue device kernels and allocations.
   CUstream dispatch_cu_stream;
 
-  iree_hal_cuda_tracing_context_t* tracing_context;
+  iree_hal_stream_tracing_context_t* tracing_context;
 
   iree_allocator_t host_allocator;
 
@@ -259,6 +259,108 @@ iree_hal_cuda_deferred_work_queue_device_interface_submit_command_buffer(
   return status;
 }
 
+typedef struct iree_hal_cuda_tracing_device_interface_t {
+  iree_hal_stream_tracing_device_interface_t base;
+  CUdevice cu_device;
+  CUcontext cu_context;
+  CUstream dispatch_cu_stream;
+  iree_allocator_t host_allocator;
+  const iree_hal_cuda_dynamic_symbols_t* cuda_symbols;
+} iree_hal_cuda_tracing_device_interface_t;
+static const iree_hal_stream_tracing_device_interface_vtable_t
+    iree_hal_cuda_tracing_device_interface_vtable_t;
+
+void iree_hal_cuda_tracing_device_interface_destroy(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  iree_allocator_free(device_interface->host_allocator, device_interface);
+}
+
+iree_status_t iree_hal_cuda_tracing_device_interface_synchronize_native_event(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_event_t base_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  return IREE_CURESULT_TO_STATUS(device_interface->cuda_symbols,
+                                 cuEventSynchronize((CUevent)base_event));
+}
+
+iree_status_t iree_hal_cuda_tracing_device_interface_create_native_event(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_event_t* base_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  return IREE_CURESULT_TO_STATUS(
+      device_interface->cuda_symbols,
+      cuEventCreate((CUevent*)base_event, CU_EVENT_DEFAULT));
+}
+
+iree_status_t iree_hal_cuda_tracing_device_interface_query_native_event(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_event_t base_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  return IREE_CURESULT_TO_STATUS(device_interface->cuda_symbols,
+                                 cuEventQuery((CUevent)base_event));
+}
+
+void iree_hal_cuda_tracing_device_interface_event_elapsed_time(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    float* relative_millis, iree_hal_stream_tracing_native_event_t start_event,
+    iree_hal_stream_tracing_native_event_t end_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  IREE_CUDA_IGNORE_ERROR(
+      device_interface->cuda_symbols,
+      cuEventElapsedTime(relative_millis, (CUevent)start_event,
+                         (CUevent)end_event));
+}
+
+void iree_hal_cuda_tracing_device_interface_destroy_native_event(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_event_t base_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  IREE_CUDA_IGNORE_ERROR(device_interface->cuda_symbols,
+                         cuEventDestroy((CUevent)base_event));
+}
+
+iree_status_t iree_hal_cuda_tracing_device_interface_record_native_event(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_event_t base_event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  return IREE_CURESULT_TO_STATUS(
+      device_interface->cuda_symbols,
+      cuEventRecord((CUevent)base_event, device_interface->dispatch_cu_stream));
+}
+
+iree_status_t
+iree_hal_cuda_tracing_device_interface_add_graph_event_record_node(
+    iree_hal_stream_tracing_device_interface_t* base_device_interface,
+    iree_hal_stream_tracing_native_graph_node_t* out_node,
+    iree_hal_stream_tracing_native_graph_t graph,
+    iree_hal_stream_tracing_native_graph_node_t* dependency_nodes,
+    size_t dependency_nodes_count,
+    iree_hal_stream_tracing_native_event_t event) {
+  iree_hal_cuda_tracing_device_interface_t* device_interface =
+      (iree_hal_cuda_tracing_device_interface_t*)base_device_interface;
+
+  return IREE_CURESULT_TO_STATUS(
+      device_interface->cuda_symbols,
+      cuGraphAddEventRecordNode((CUgraphNode*)out_node, (CUgraph)graph,
+                                (CUgraphNode*)dependency_nodes,
+                                dependency_nodes_count, (CUevent)event));
+}
+
 static iree_hal_cuda_device_t* iree_hal_cuda_device_cast(
     iree_hal_device_t* base_value) {
   IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_cuda_device_vtable);
@@ -346,18 +448,36 @@ static iree_status_t iree_hal_cuda_device_create_internal(
 
   // Enable tracing for the (currently only) stream - no-op if disabled.
   if (iree_status_is_ok(status) && device->params.stream_tracing) {
-    if (device->params.stream_tracing >= IREE_HAL_CUDA_TRACING_VERBOSITY_MAX ||
-        device->params.stream_tracing < IREE_HAL_CUDA_TRACING_VERBOSITY_OFF) {
+    if (device->params.stream_tracing >= IREE_HAL_TRACING_VERBOSITY_MAX ||
+        device->params.stream_tracing < IREE_HAL_TRACING_VERBOSITY_OFF) {
       return iree_make_status(
           IREE_STATUS_INVALID_ARGUMENT,
           "invalid stream_tracing argument: expected to be between %d and %d",
-          IREE_HAL_CUDA_TRACING_VERBOSITY_OFF,
-          IREE_HAL_CUDA_TRACING_VERBOSITY_MAX);
+          IREE_HAL_TRACING_VERBOSITY_OFF, IREE_HAL_TRACING_VERBOSITY_MAX);
+    }
+
+    iree_hal_cuda_tracing_device_interface_t* tracing_device_interface = NULL;
+    status = iree_allocator_malloc(
+        host_allocator, sizeof(iree_hal_cuda_tracing_device_interface_t),
+        (void**)&tracing_device_interface);
+
+    if (IREE_UNLIKELY(!iree_status_is_ok(status))) {
+      iree_hal_device_release((iree_hal_device_t*)device);
+      return status;
     }
-    status = iree_hal_cuda_tracing_context_allocate(
-        device->cuda_symbols, device->identifier, dispatch_stream,
-        device->params.stream_tracing, &device->block_pool, host_allocator,
-        &device->tracing_context);
+
+    tracing_device_interface->base.vtable =
+        &iree_hal_cuda_tracing_device_interface_vtable_t;
+    tracing_device_interface->cu_context = context;
+    tracing_device_interface->cu_device = cu_device;
+    tracing_device_interface->dispatch_cu_stream = dispatch_stream;
+    tracing_device_interface->host_allocator = host_allocator;
+    tracing_device_interface->cuda_symbols = cuda_symbols;
+
+    status = iree_hal_stream_tracing_context_allocate(
+        (iree_hal_stream_tracing_device_interface_t*)tracing_device_interface,
+        device->identifier, device->params.stream_tracing, &device->block_pool,
+        host_allocator, &device->tracing_context);
   }
 
   // Memory pool support is conditional.
@@ -505,7 +625,7 @@ static void iree_hal_cuda_device_destroy(iree_hal_device_t* base_device) {
   // Destroy memory pools that hold on to reserved memory.
   iree_hal_cuda_memory_pools_deinitialize(&device->memory_pools);
 
-  iree_hal_cuda_tracing_context_free(device->tracing_context);
+  iree_hal_stream_tracing_context_free(device->tracing_context);
 
   // Destroy various pools for synchronization.
   if (device->timepoint_pool) {
@@ -947,8 +1067,8 @@ static iree_status_t iree_hal_cuda_device_queue_write(
 }
 
 static void iree_hal_cuda_device_collect_tracing_context(void* user_data) {
-  iree_hal_cuda_tracing_context_collect(
-      (iree_hal_cuda_tracing_context_t*)user_data);
+  iree_hal_stream_tracing_context_collect(
+      (iree_hal_stream_tracing_context_t*)user_data);
 }
 
 static iree_status_t iree_hal_cuda_device_queue_execute(
@@ -1074,3 +1194,22 @@ static const iree_hal_deferred_work_queue_device_interface_vtable_t
         .submit_command_buffer =
             iree_hal_cuda_deferred_work_queue_device_interface_submit_command_buffer,
 };
+
+static const iree_hal_stream_tracing_device_interface_vtable_t
+    iree_hal_cuda_tracing_device_interface_vtable_t = {
+        .destroy = iree_hal_cuda_tracing_device_interface_destroy,
+        .synchronize_native_event =
+            iree_hal_cuda_tracing_device_interface_synchronize_native_event,
+        .create_native_event =
+            iree_hal_cuda_tracing_device_interface_create_native_event,
+        .query_native_event =
+            iree_hal_cuda_tracing_device_interface_query_native_event,
+        .event_elapsed_time =
+            iree_hal_cuda_tracing_device_interface_event_elapsed_time,
+        .destroy_native_event =
+            iree_hal_cuda_tracing_device_interface_destroy_native_event,
+        .record_native_event =
+            iree_hal_cuda_tracing_device_interface_record_native_event,
+        .add_graph_event_record_node =
+            iree_hal_cuda_tracing_device_interface_add_graph_event_record_node,
+};