[hip] Add trace zones to copy/fill/update buffer commands. (iree-org#…

…19544) They were missing in the hip stream case. Signed-off-by: Andrew Woloszyn <[email protected]>
nod-ai · Dec 20, 2024 · 0a0483e · 0a0483e
1 parent 47ccd93
commit 0a0483e
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/runtime/src/iree/hal/drivers/hip/stream_command_buffer.c b/runtime/src/iree/hal/drivers/hip/stream_command_buffer.c
@@ -328,6 +328,9 @@ static iree_status_t iree_hal_hip_stream_command_buffer_fill_buffer(
       iree_hal_buffer_byte_offset(target_ref.buffer) + target_ref.offset;
   hipDeviceptr_t dst = (uint8_t*)target_device_buffer + target_offset;
   size_t num_elements = target_ref.length / pattern_length;
+  IREE_HAL_STREAM_TRACE_ZONE_BEGIN(command_buffer->tracing_context,
+                                   &command_buffer->tracing_event_list,
+                                   IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
 
   switch (pattern_length) {
     case 4: {
@@ -359,7 +362,9 @@ static iree_status_t iree_hal_hip_stream_command_buffer_fill_buffer(
       return iree_make_status(IREE_STATUS_INTERNAL,
                               "unsupported fill pattern length");
   }
-
+  IREE_HAL_STREAM_TRACE_ZONE_END(command_buffer->tracing_context,
+                                 &command_buffer->tracing_event_list,
+                                 IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }
@@ -397,12 +402,17 @@ static iree_status_t iree_hal_hip_stream_command_buffer_update_buffer(
   hipDeviceptr_t dst = (uint8_t*)target_device_buffer +
                        iree_hal_buffer_byte_offset(target_ref.buffer) +
                        target_ref.offset;
+  IREE_HAL_STREAM_TRACE_ZONE_BEGIN(command_buffer->tracing_context,
+                                   &command_buffer->tracing_event_list,
+                                   IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
   IREE_HIP_RETURN_AND_END_ZONE_IF_ERROR(
       z0, command_buffer->hip_symbols,
       hipMemcpyHtoDAsync(dst, (void*)src, target_ref.length,
                          command_buffer->hip_stream),
       "hipMemcpyHtoDAsync");
-
+  IREE_HAL_STREAM_TRACE_ZONE_END(command_buffer->tracing_context,
+                                 &command_buffer->tracing_event_list,
+                                 IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }
@@ -417,6 +427,9 @@ static iree_status_t iree_hal_hip_stream_command_buffer_copy_buffer(
 
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_hal_hip_stream_command_buffer_flush_collectives(command_buffer));
+  IREE_HAL_STREAM_TRACE_ZONE_BEGIN(command_buffer->tracing_context,
+                                   &command_buffer->tracing_event_list,
+                                   IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
 
   hipDeviceptr_t target_device_buffer = iree_hal_hip_buffer_device_pointer(
       iree_hal_buffer_allocated_buffer(target_ref.buffer));
@@ -435,6 +448,10 @@ static iree_status_t iree_hal_hip_stream_command_buffer_copy_buffer(
                      command_buffer->hip_stream),
       "hipMemcpyAsync");
 
+  IREE_HAL_STREAM_TRACE_ZONE_END(command_buffer->tracing_context,
+                                 &command_buffer->tracing_event_list,
+                                 IREE_HAL_STREAM_TRACING_VERBOSITY_FINE);
+
   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }