Cleaning up the iree/base/tracing.h header a bit. (iree-org#14089)

There's more work to generalize out some tracy-isms we may not want to
enforce on other implementations. GPU and mutex tracing is still behind
the tracy wall today.

experimental/cuda2/cuda_allocator.c

@@ -360,7 +360,7 @@ static iree_status_t iree_hal_cuda2_allocator_allocate_buffer(
   void* host_ptr = NULL;
   CUdeviceptr device_ptr = 0;
   IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree_hal_cuda2_buffer_allocate");
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, allocation_size);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, allocation_size);
   if (iree_all_bits_set(compat_params.type,
                         IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) {
     if (iree_all_bits_set(compat_params.type,

experimental/cuda2/memory_pools.c

@@ -214,7 +214,7 @@ iree_status_t iree_hal_cuda2_memory_pools_alloca(
     iree_device_size_t allocation_size,
     iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (int64_t)allocation_size);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)allocation_size);
 
   iree_hal_buffer_params_canonicalize(&params);
 
@@ -272,7 +272,7 @@ iree_status_t iree_hal_cuda2_memory_pools_dealloca(
     iree_hal_cuda2_memory_pools_t* pools, CUstream stream,
     iree_hal_buffer_t* buffer) {
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(
       z0, (int64_t)iree_hal_buffer_allocation_size(buffer));
 
   // Only process the request if the buffer came from an async pool.

experimental/webgpu/simple_allocator.c

@@ -201,7 +201,7 @@ static iree_status_t iree_hal_webgpu_simple_allocator_allocate_buffer(
   // mappedAtCreation and populating it before unmapping.
   if (has_initial_data) {
     IREE_TRACE_ZONE_BEGIN(z1);
-    IREE_TRACE_ZONE_APPEND_VALUE(z1, (uint64_t)initial_data.data_length);
+    IREE_TRACE_ZONE_APPEND_VALUE_I64(z1, (uint64_t)initial_data.data_length);
     void* mapped_ptr =
         wgpuBufferGetMappedRange(buffer_handle, 0, initial_data.data_length);
     memcpy(mapped_ptr, initial_data.data, initial_data.data_length);

runtime/bindings/tflite/interpreter.c

@@ -393,7 +393,7 @@ TFL_CAPI_EXPORT extern TfLiteInterpreter* TfLiteInterpreterCreate(
       _TfLiteInterpreterCreate(model, optional_options, &interpreter);
   if (iree_status_is_ok(iree_status_consume_code(status))) {
     IREE_TRACE_ZONE_APPEND_TEXT(z0, "num_threads=", strlen("num_threads="));
-    IREE_TRACE_ZONE_APPEND_VALUE(z0, interpreter->options.num_threads);
+    IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, interpreter->options.num_threads);
   } else {
     IREE_TRACE_MESSAGE(ERROR, "failed interpreter creation");
     TfLiteInterpreterDelete(interpreter);
@@ -563,7 +563,7 @@ TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterAllocateTensors(
     total_input_size +=
         iree_hal_buffer_byte_length(interpreter->input_tensors[i].buffer);
   }
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, total_input_size);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, total_input_size);
 #endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
 
   IREE_TRACE_ZONE_END(z0);
@@ -608,7 +608,7 @@ TFL_CAPI_EXPORT extern TfLiteStatus TfLiteInterpreterInvoke(
     total_output_size +=
         iree_hal_buffer_byte_length(interpreter->output_tensors[i].buffer);
   }
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, total_output_size);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, total_output_size);
 #endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
 
   IREE_TRACE_ZONE_END(z0);

runtime/bindings/tflite/options.c

@@ -43,7 +43,7 @@ TFL_CAPI_EXPORT extern void TfLiteInterpreterOptionsDelete(
 TFL_CAPI_EXPORT extern void TfLiteInterpreterOptionsSetNumThreads(
     TfLiteInterpreterOptions* options, int32_t num_threads) {
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, num_threads);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, num_threads);
   options->num_threads = num_threads;
   IREE_TRACE_ZONE_END(z0);
 }

runtime/bindings/tflite/tensor.c

@@ -245,7 +245,8 @@ TFL_CAPI_EXPORT extern TfLiteStatus TfLiteTensorCopyFromBuffer(
     return kTfLiteApplicationError;
   }
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, tensor->buffer_mapping.contents.data_length);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0,
+                                   tensor->buffer_mapping.contents.data_length);
 
   // NOTE: we could use a iree_hal_buffer_map_write here but we already
   // have the buffer mapped. If we knew the user would never use
@@ -264,7 +265,7 @@ TFL_CAPI_EXPORT extern TfLiteStatus TfLiteTensorCopyToBuffer(
     return kTfLiteApplicationError;
   }
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(
       z0, output_tensor->buffer_mapping.contents.data_length);
 
   // NOTE: as with above we should use an iree_hal_buffer_map_read here.

runtime/src/iree/base/internal/wait_handle_inproc.c

@@ -118,7 +118,7 @@ iree_status_t iree_wait_set_allocate(iree_host_size_t capacity,
   }
 
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (int64_t)capacity);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)capacity);
   *out_set = NULL;
 
   iree_wait_set_t* set = NULL;

runtime/src/iree/base/loop.c

@@ -44,9 +44,9 @@ IREE_API_EXPORT iree_status_t iree_loop_dispatch(
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "null loop");
   }
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)workgroup_count_xyz[0]);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)workgroup_count_xyz[1]);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)workgroup_count_xyz[2]);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)workgroup_count_xyz[0]);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)workgroup_count_xyz[1]);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)workgroup_count_xyz[2]);
 
   const iree_loop_dispatch_params_t params = {
       .callback =
@@ -159,7 +159,7 @@ IREE_API_EXPORT iree_status_t iree_loop_wait_any(
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "null loop");
   }
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)count);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)count);
   iree_status_t status =
       iree_loop_wait_multi(IREE_LOOP_COMMAND_WAIT_ANY, loop, count,
                            wait_sources, timeout, callback, user_data);
@@ -174,7 +174,7 @@ IREE_API_EXPORT iree_status_t iree_loop_wait_all(
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "null loop");
   }
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)count);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)count);
   iree_status_t status =
       iree_loop_wait_multi(IREE_LOOP_COMMAND_WAIT_ALL, loop, count,
                            wait_sources, timeout, callback, user_data);

runtime/src/iree/base/loop_sync.c

@@ -576,7 +576,7 @@ static void iree_loop_wait_list_handle_wake(iree_loop_wait_list_t* wait_list,
   int woken_tasks = 0;
 
   (void)woken_tasks;
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, woken_tasks);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, woken_tasks);
   IREE_TRACE_ZONE_END(z0);
 }
 
@@ -647,7 +647,7 @@ static iree_status_t iree_loop_wait_list_commit(
 
   // Real system wait.
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(z0, (int64_t)wait_list->count);
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)wait_list->count);
 
   // Enter the system wait API.
   iree_wait_handle_t wake_handle = iree_wait_handle_immediate();

runtime/src/iree/base/time.c

@@ -169,7 +169,7 @@ bool iree_wait_until(iree_time_t deadline_ns) {
   if (deadline_ns == IREE_TIME_INFINITE_PAST) return true;
 
   IREE_TRACE_ZONE_BEGIN(z0);
-  IREE_TRACE_ZONE_APPEND_VALUE(
+  IREE_TRACE_ZONE_APPEND_VALUE_I64(
       z0, (uint64_t)iree_absolute_deadline_to_timeout_ns(deadline_ns));
 
   // NOTE: we want to use sleep APIs with absolute times as that makes retrying

runtime/src/iree/base/tracing.h

@@ -157,58 +157,145 @@ enum {
 };
 
 #if !(IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION)
+
+// Evaluates the expression code only if tracing is enabled.
+//
+// Example:
+//  struct {
+//    IREE_TRACE(uint32_t trace_only_value);
+//  } my_object;
+//  IREE_TRACE(my_object.trace_only_value = 5);
+#define IREE_TRACE(expr)
+
+// Sets an application-specific payload that will be stored in the trace.
+// This can be used to fingerprint traces to particular versions and denote
+// compilation options or configuration. The given string value will be copied.
 #define IREE_TRACE_SET_APP_INFO(value, value_length)
+
+// Sets the current thread name to the given string value.
+// This will only set the thread name as it appears in the tracing backend and
+// not set the OS thread name as it would appear in a debugger.
+// The C-string |name| will be copied and does not need to be a literal.
 #define IREE_TRACE_SET_THREAD_NAME(name)
-#define IREE_TRACE(expr)
+
+// Notifies the tracing implementation that the app is about to exit.
+// This allows implementations to flush their buffers. |exit_code| may be
+// provided to indicate the application exit code if available.
+#define IREE_TRACE_APP_EXIT(exit_code)
+
+// Enters a fiber context.
+// |fiber| must be a unique pointer and remain live for the process lifetime.
 #define IREE_TRACE_FIBER_ENTER(fiber)
+// Exits a fiber context.
 #define IREE_TRACE_FIBER_LEAVE()
+
+// Begins a new zone with the parent function name.
 #define IREE_TRACE_ZONE_BEGIN(zone_id) \
   iree_zone_id_t zone_id = 0;          \
   (void)zone_id;
+// Begins a new zone with the given compile-time |name_literal|.
+// The literal must be static const and will be embedded in the trace buffer by
+// reference.
 #define IREE_TRACE_ZONE_BEGIN_NAMED(zone_id, name_literal) \
   IREE_TRACE_ZONE_BEGIN(zone_id)
+// Begins a new zone with the given runtime dynamic string |name|.
+// The |name| string will be copied into the trace buffer and does not require
+// a NUL terminator.
 #define IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(zone_id, name, name_length) \
   IREE_TRACE_ZONE_BEGIN(zone_id)
+// Begins an externally defined zone with a dynamic source location.
+// The |file_name|, |function_name|, and optional |name| strings will be copied
+// into the trace buffer and do not need to persist or have a NUL terminator.
 #define IREE_TRACE_ZONE_BEGIN_EXTERNAL(                        \
     zone_id, file_name, file_name_length, line, function_name, \
     function_name_length, name, name_length)                   \
   IREE_TRACE_ZONE_BEGIN(zone_id)
+
+// Ends the current zone. Must be passed the |zone_id| from the _BEGIN.
+#define IREE_TRACE_ZONE_END(zone_id)
+
+// Ends the current zone before returning on a failure.
+// Sugar for IREE_TRACE_ZONE_END + IREE_RETURN_IF_ERROR.
+#define IREE_RETURN_AND_END_ZONE_IF_ERROR(zone_id, ...) \
+  IREE_RETURN_IF_ERROR(__VA_ARGS__)
+
+// Sets the dynamic color of the zone to an XXBBGGRR value.
 #define IREE_TRACE_ZONE_SET_COLOR(zone_id, color_xrgb)
-#define IREE_TRACE_ZONE_APPEND_VALUE(zone_id, value)
+
+// Appends an int64_t value to the parent zone. May be called multiple times.
+#define IREE_TRACE_ZONE_APPEND_VALUE_I64(zone_id, value)
+
+// Appends a string value to the parent zone. May be called multiple times.
+// The provided NUL-terminated C string or string view will be copied into the
+// trace buffer.
 #define IREE_TRACE_ZONE_APPEND_TEXT(zone_id, ...)
 #define IREE_TRACE_ZONE_APPEND_TEXT_CSTRING(zone_id, value)
 #define IREE_TRACE_ZONE_APPEND_TEXT_STRING_VIEW(zone_id, value, value_length)
-#define IREE_TRACE_ZONE_END(zone_id)
-#define IREE_RETURN_AND_END_ZONE_IF_ERROR(zone_id, ...) \
-  IREE_RETURN_IF_ERROR(__VA_ARGS__)
+
+// Configures the named plot with an IREE_TRACING_PLOT_TYPE_* representation.
 #define IREE_TRACE_SET_PLOT_TYPE(name_literal, plot_type, step, fill, color)
+// Plots a value in the named plot group as an int64_t.
 #define IREE_TRACE_PLOT_VALUE_I64(name_literal, value)
+// Plots a value in the named plot group as a single-precision float.
 #define IREE_TRACE_PLOT_VALUE_F32(name_literal, value)
+// Plots a value in the named plot group as a double-precision float.
 #define IREE_TRACE_PLOT_VALUE_F64(name_literal, value)
+
+// Demarcates an advancement of the top-level unnamed frame group.
 #define IREE_TRACE_FRAME_MARK()
+// Demarcates an advancement of a named frame group.
 #define IREE_TRACE_FRAME_MARK_NAMED(name_literal)
+// Begins a discontinuous frame in a named frame group.
+// Must be properly matched with a IREE_TRACE_FRAME_MARK_NAMED_END.
 #define IREE_TRACE_FRAME_MARK_BEGIN_NAMED(name_literal)
+// Ends a discontinuous frame in a named frame group.
 #define IREE_TRACE_FRAME_MARK_END_NAMED(name_literal)
+
+// Logs a message at the given logging level to the trace.
+// The message text must be a compile-time string literal.
 #define IREE_TRACE_MESSAGE(level, value_literal)
+// Logs a message with the given color to the trace.
+// Standard colors are defined as IREE_TRACING_MESSAGE_LEVEL_* values.
+// The message text must be a compile-time string literal.
 #define IREE_TRACE_MESSAGE_COLORED(color, value_literal)
+// Logs a dynamically-allocated message at the given logging level to the trace.
+// The string |value| will be copied into the trace buffer.
 #define IREE_TRACE_MESSAGE_DYNAMIC(level, value, value_length)
+// Logs a dynamically-allocated message with the given color to the trace.
+// Standard colors are defined as IREE_TRACING_MESSAGE_LEVEL_* values.
+// The string |value| will be copied into the trace buffer.
 #define IREE_TRACE_MESSAGE_DYNAMIC_COLORED(color, value, value_length)
+
 #endif  // IREE_TRACING_FEATURE_INSTRUMENTATION
 
 //===----------------------------------------------------------------------===//
 // Allocation tracking macros (C/C++)
 //===----------------------------------------------------------------------===//
-//
-// IREE_TRACE_ALLOC: records an malloc.
-// IREE_TRACE_FREE: records a free.
-//
-// NOTE: realloc must be recorded as a FREE/ALLOC pair.
 
 #if !(IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_ALLOCATION_TRACKING)
+
+// Traces a new memory allocation with host |ptr| and the given |size|.
+// A balanced IREE_TRACE_FREE on the same |ptr| is required for proper memory
+// tracking. Allocations will be attributed to their parent zone.
+// Reallocations must be recorded as an IREE_TRACE_ALLOC/IREE_TRACE_FREE pair.
 #define IREE_TRACE_ALLOC(ptr, size)
+
+// Traces a free of an existing allocation traced with IREE_TRACE_ALLOC.
 #define IREE_TRACE_FREE(ptr)
-#define IREE_TRACE_ALLOC_NAMED(name, ptr, size)
-#define IREE_TRACE_FREE_NAMED(name, ptr)
+
+// Traces a new memory allocation in a named memory pool.
+// The |name_literal| pointer must be static const and be identical to the
+// one passed to a balanced IREE_TRACE_FREE_NAMED. This is best accomplished
+// with a compilation-unit scoped `static const char* MY_POOL_ID = "foo";` as
+// in debug modes compilers may not deduplicate local literals.
+// Reallocations must be recorded as an
+// IREE_TRACE_ALLOC_NAMED/IREE_TRACE_FREE_NAMED pair.
+#define IREE_TRACE_ALLOC_NAMED(name_literal, ptr, size)
+
+// Traces a free of an existing allocation traced with IREE_TRACE_ALLOC_NAMED.
+// Note that the |name_literal| must be the same pointer as passed to the alloc.
+#define IREE_TRACE_FREE_NAMED(name_literal, ptr)
+
 #endif  // IREE_TRACING_FEATURE_ALLOCATION_TRACKING
 
 //===----------------------------------------------------------------------===//
@@ -218,8 +305,23 @@ enum {
 #ifdef __cplusplus
 
 #if !(IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION)
+
+// Automatically instruments the calling scope as if calling
+// IREE_TRACE_ZONE_BEGIN/IREE_TRACE_ZONE_END.
+// The scope's zone ID can be retrieved with IREE_TRACE_SCOPE_ID for use with
+// other tracing macros.
 #define IREE_TRACE_SCOPE()
+
+// Automatically instruments the calling scope as if calling
+// IREE_TRACE_ZONE_BEGIN_NAMED/IREE_TRACE_ZONE_END.
+// The scope's zone ID can be retrieved with IREE_TRACE_SCOPE_ID for use with
+// other tracing macros.
 #define IREE_TRACE_SCOPE_NAMED(name_literal)
+
+// The zone ID of the current scope as set by IREE_TRACE_SCOPE.
+// This can be passed to other tracing macros like IREE_TRACE_ZONE_APPEND_TEXT.
+#define IREE_TRACE_SCOPE_ID
+
 #endif  // IREE_TRACING_FEATURE_INSTRUMENTATION
 
 #endif  // __cplusplus