[ROCm: XLA profiler] Use real HIP stream_id instead of HSA queue handle#46182
Open
cj401-amd wants to merge 1 commit into
Open
[ROCm: XLA profiler] Use real HIP stream_id instead of HSA queue handle#46182cj401-amd wants to merge 1 commit into
cj401-amd wants to merge 1 commit into
Conversation
cj401-amd
requested review from
chsigg,
dimitar-asenov,
magaonka-amd and
thomasjoerg
July 24, 2026 18:13
cj401-amd
force-pushed
the
cj/rocprofiler-sdk-stream-id
branch
from
July 24, 2026 18:49
9039b41 to
0a94bdf
Compare
Fix stream_id tracking in the ROCm profiler so that GPU events are assigned to the correct HIP stream lanes in the trace viewer, matching the behavior of both the rocprofiler-sdk reference tool (tool.cpp) and Kineto (PyTorch). Previously, KernelEvent() used dispatch_info.queue_id.handle (an HSA hardware queue handle) as stream_id, and MemcpyEvent() used kInvalidStreamId (causing memcpy events to be silently dropped from the device timeline). HSA queues and HIP streams are orthogonal — a stream may use different queues over time, and a queue may be shared by multiple streams. The fix uses the rocprofiler-SDK's purpose-built HIP stream tracking: 1. A dedicated hip_stream_ctx_ with ROCPROFILER_CALLBACK_TRACING_HIP_STREAM maintains a thread-local stream stack (absl::InlinedVector<uint64_t, 4>): ROCPROFILER_HIP_STREAM_SET callbacks push/pop the current hipStream_t around every HIP API call. The context is process-lifetime (like utility_context_) so the TLS stack stays warm for the first dispatch after Enable(). 2. An external correlation ID request service on the main tracing context snapshots tls_stream_stack.back() into correlation_id.external.value for every KERNEL_DISPATCH and MEMORY_COPY record. 3. KernelEvent() and MemcpyEvent() read rec.correlation_id.external.value to get the real HIP stream_id. The HSA queue handle is preserved in a new queue_id field on RocmTracerEvent for debugging. 4. PerDeviceCollector::Export() remaps raw stream handle values (64-bit pointers) to dense sequential indices (0, 1, 2, ...) so the trace viewer shows clean "Stream #0", "Stream #1" lane names. Changes: - rocm_tracer.cc: Add thread_local tls_stream_stack (InlinedVector), hip_stream_callback(), stream_external_correlation_callback() with defensive empty check. Create and start hip_stream_ctx_ in InitProfiling(), stop in toolFinalize(). Configure external correlation ID request service on context_ using std::size(kinds). KernelEvent(): stream_id = rec.correlation_id.external.value, queue_id = kinfo.queue_id.handle. MemcpyEvent(): stream_id = rec.correlation_id.external.value. Add VLOG for unexpected HIP stream operations. Add breadcrumb comments at consumption sites. - rocm_tracer.h: Add rocprofiler_context_id_t hip_stream_ctx_ member. - rocm_tracer_utils.h: Add uint64_t queue_id field to RocmTracerEvent. - rocm_collector.cc: Add dense stream remapping in Export().
cj401-amd
force-pushed
the
cj/rocprofiler-sdk-stream-id
branch
from
July 24, 2026 19:21
0a94bdf to
f1b9c72
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
instead of HSA queue handles (kernels) or kInvalidStreamId (memcpy — silently dropped)
ROCPROFILER_CALLBACK_TRACING_HIP_STREAMcallbackwith external correlation ID, matching the rocprofiler-sdk reference tool (tool.cpp)
instead of raw 64-bit pointer values
Motivation
The ROCm profiler had two stream-id bugs:
dispatch_info.queue_id.handle(HSA hardware queue) asstream_id— not the HIP stream the user createdkInvalidStreamId(UINT64_MAX), causing them to be silently dropped from the device timelineApproach
Uses the SDK's purpose-built
ROCPROFILER_CALLBACK_TRACING_HIP_STREAMkind:correlation_id.external.valuefor everyKERNEL_DISPATCHandMEMORY_COPYrecordKernelEvent()andMemcpyEvent()read the real stream from the external correlationqueue_idfield for debuggingChanged files
rocm_tracer.ccrocm_tracer.hhip_stream_ctx_memberrocm_tracer_utils.hqueue_idfield onRocmTracerEventrocm_collector.ccExport()Test plan
jax_matmul.py(JIT bfloat16 matmul, 10 profiled iterations) — trace collected successfully