[rocprofiler-systems] Add support for ompt_callback_thread_begin (#1681)

* Add thread_begin callback

* Make OMPT callbacks that are instant have start_ts = end_ts
This commit is contained in:
Kian Cossettini
2025-11-26 13:38:04 -05:00
zatwierdzone przez GitHub
rodzic bc6f29c04a
commit 76a23eab14
2 zmienionych plików z 38 dodań i 30 usunięć
@@ -836,9 +836,7 @@ metadata_registry::metadata_registry()
overwrite_callback_names({
# if(ROCPROFILER_VERSION >= 600)
{ ROCPROFILER_CALLBACK_TRACING_OMPT,
{ { ROCPROFILER_OMPT_ID_thread_begin, "omp_thread" },
{ ROCPROFILER_OMPT_ID_thread_end, "omp_thread" },
{ ROCPROFILER_OMPT_ID_parallel_begin, "omp_parallel" },
{ { ROCPROFILER_OMPT_ID_parallel_begin, "omp_parallel" },
{ ROCPROFILER_OMPT_ID_parallel_end, "omp_parallel" } } }
# endif
});
@@ -1208,28 +1208,38 @@ tool_tracing_callback(rocprofiler_callback_tracing_record_t record,
};
#if(ROCPROFILER_VERSION >= 600)
// Skip implicit_task associated with an "initial-task-begin" occurrence
// as it is generated by our tool.
// Occurs after our tool initializes OMPT but before the first OpenMP
// region (user code) begins.
// Skip implicit_task associated with an "initial-task-begin" occurrence as
// well as the thread_begin associated with an "initial-thread-begin" occurrence
// as they are generated by our tool.
// The two callbacks occur after our tool initializes OMPT but before the
// first OpenMP region (user code) begins.
// Note: Can occur multiple times (Ex: MPI+OpenMP hybrid)
if(record.kind == ROCPROFILER_CALLBACK_TRACING_OMPT &&
record.operation == ROCPROFILER_OMPT_ID_implicit_task)
if(record.kind == ROCPROFILER_CALLBACK_TRACING_OMPT)
{
auto* payload_data =
static_cast<rocprofiler_callback_tracing_ompt_data_t*>(record.payload);
int flag = payload_data->args.implicit_task.flags;
if(flag & ompt_task_initial) return;
switch(record.operation)
{
case ROCPROFILER_OMPT_ID_implicit_task:
{
int flag = payload_data->args.implicit_task.flags;
if(flag & ompt_task_initial) return; // Skips both the start and end
break;
}
case ROCPROFILER_OMPT_ID_thread_begin:
{
ompt_thread_t thread_type = payload_data->args.thread_begin.thread_type;
if(thread_type == ompt_thread_initial) return;
break;
}
default: break;
}
// TODO: Once finalization issue is fixed, skip the corresponding end
// of the thread_begin callback. Can be identified with:
// - thread_end: The thread_data ptr from the thread_begin callback generated
// by the "initial-thread-begin" needs to match the thread_end's thread_data
// ptr
}
// TODO: Once ompt_callback_thread_begin is supported, we need to skip
// every occurrence of an "initial-thread-begin" for similar reasons.
// This callback is identified with thread_type == ompt_thread_initial.
// Note: If the finalization issue is resolved, the corresponding ends of these
// callbacks (which should be discarded) can be identified with:
// - implicit_task: (flag & ompt_task_initial) && endpoint == ompt_scope_end
// - thread_end: The thread_data ptr from the thread_begin callback generated
// by the "initial-thread-begin" needs to match the thread_end's thread_data ptr
#endif
auto ts = rocprofiler_timestamp_t{};
@@ -1446,10 +1456,7 @@ tool_tracing_callback(rocprofiler_callback_tracing_record_t record,
// Callbacks that are received but that we do not process
static const std::set<rocprofiler_ompt_operation_t> ompt_no_process = {
ROCPROFILER_OMPT_ID_callback_functions, // "Fake" callback
// There is no point in handling ompt_thread_begin events as the
// corresponding ompt_thread_end event will not occur unless
// runtime is finalized earlier
ROCPROFILER_OMPT_ID_thread_begin,
// Not processed as these are received after our tool finalizes
ROCPROFILER_OMPT_ID_thread_end,
};
@@ -1470,6 +1477,10 @@ tool_tracing_callback(rocprofiler_callback_tracing_record_t record,
ompt_tracing_callback_stop(record, user_data, ts, _bt_data);
ompt_pop_parallel_callback(record, ts, _bt_data);
break;
// Unlike parallel callbacks, we cannot receive the corresponding end
// to thread_begin. Set thread_begin as "instant" so the user can
// see callback without it spanning the entire track
case ROCPROFILER_OMPT_ID_thread_begin:
case ROCPROFILER_OMPT_ID_lock_init:
case ROCPROFILER_OMPT_ID_lock_destroy:
// Although this has endpoint arg, treat it as instant event
@@ -1493,12 +1504,11 @@ tool_tracing_callback(rocprofiler_callback_tracing_record_t record,
// These callbacks are considered instant events and should start
// and immediately call stop as no corresponding "end" will be
// received
auto start_ts = ts;
ompt_tracing_callback_start(record, user_data, start_ts);
ROCPROFILER_CALL(
rocprofiler_get_timestamp(&ts)); // Set artificial end ts
ompt_tracing_callback_stop(record, user_data, ts, _bt_data);
ompt_cache_instant_event(record, start_ts, _bt_data);
auto instant_ts = ts;
ompt_tracing_callback_start(record, user_data, instant_ts);
ompt_tracing_callback_stop(record, user_data, instant_ts,
_bt_data);
ompt_cache_instant_event(record, instant_ts, _bt_data);
break;
}
default: