[RFC] ROCm profiler annotation support: ROCTX markers in XProf / TensorBoard

[cj401-amd]

[RFC] ROCm profiler annotation support: ROCTX markers in XProf / TensorBoard

Date: 2026-08-05
Author: cj401-amd <chunyjin@amd.com>
Status: Proposed
Target: openxla/xla main

---

Summary

This RFC proposes upstreaming ROCTX marker support into OpenXLA's profiler
stack, giving ROCm the same annotation capability that CUDA has had via NVTX
since TF 2.x. The change is structured as three layers:

1. Emitter — roctx_utils.cc implements the existing nvtx_utils.h
   interface for ROCm, so every call site that today emits NVTX on CUDA
   automatically emits ROCTX on ROCm with no new API.
2. Listener — rocm_tracer.cc registers ROCPROFILER_CALLBACK_TRACING_MARKER_CORE_API
   to capture roctx ranges and marks as Generic trace events.
3. XProf routing — a new /host:ROCTX XPlane is collected and merged
   into /host:CPU in PostProcessSingleHostXSpace following the identical
   pattern used for NVTX, so TensorBoard's "ROCTX Threads" rows appear in the
   same process lane as CUDA's "NVTX Threads" rows.

The change is ROCm-conditional at build time (Bazel if_rocm), so it has zero
cost on CPU/CUDA/TPU builds. It does not alter any public API, proto schema,
or XPlane format for non-ROCm targets.

---

Motivation

The problem

jax.profiler.TraceAnnotation and tf.profiler.experimental.Trace allow
users to annotate regions of their Python program. On CUDA, the annotation is
emitted both as an AnnotationStack push (attaches the label to kernel events
as a kTfOp stat) and as an NVTX range (captured by CUPTI and shown in
TensorBoard as "NVTX Threads" rows).

On ROCm today, only the AnnotationStack path is active. Users see kTfOp
labels on kernel events but get no visible timeline rows in TensorBoard — the
equivalent of watching a CUDA trace without the green annotation bars.
Third-party ROCm profiling tools (rocprof, Omnitrace, rocprofiler-sdk itself)
use the ROCTX library as the native annotation protocol, so the XLA stack is
the only consumer that does not participate.

Why it belongs in upstream

• nvtx_utils.h is an abstract interface that lives in tsl/profiler/lib/.
  The CUDA implementation is nvtx_utils.cc. Adding roctx_utils.cc as
  the ROCm implementation of the same interface is the natural completion of
  the abstraction the XLA team already designed.
• The NVTX listener path (CuptiTracer) and the corresponding XPlane merge
  in post_process_single_host_xplane.cc established the pattern this change
  follows exactly. A ROCm-only fork would diverge from that pattern
  indefinitely.
• TraceMeWrapper (the Python binding for jax.profiler.TraceAnnotation) was
  extended to call PushAnnotation/PopAnnotation. That change touches
  xla/python/profiler.cc, which is upstream-only code.

---

Design

Emitter: roctx_utils.cc

tsl/profiler/lib/nvtx_utils.h declares a vendor-neutral annotation
interface:

ProfilerDomainHandle DefaultProfilerDomain();
void RangePush(ProfilerDomainHandle, const char* label);
void RangePop(ProfilerDomainHandle);
void NameCurrentThread(const std::string& name);
void NameDevice(int device_id, const std::string& name);
void NameStream(StreamHandle stream, const std::string& name);

The existing nvtx_utils.cc implements this interface for CUDA using
nvtx3/nvToolsExt.h. The new roctx_utils.cc implements the same interface
for ROCm using rocprofiler-sdk-roctx/roctx.h.

The key design decision is a dual-push in RangePush/RangePop:

void RangePush(ProfilerDomainHandle /*domain*/, const char* ascii) {
  // Pipeline A: AnnotationStack → kTfOp on kernel events via correlation ID
  if (AnnotationStack::IsEnabled()) {
    AnnotationStack::PushAnnotation(ascii);
  }
  // Pipeline B: roctx range → MarkerCallback → kNVTXRange in /host:ROCTX
  roctxRangePushA(ascii);
}

This preserves both downstream consumers: the kTfOp stat on kernel events
(which requires AnnotationStack) and the visible timeline row in TensorBoard
(which requires a captured range event). Separating them would require callers
to manage two distinct push/pop sequences, which is error-prone and breaks the
single-call contract expected by scoped_annotation.h.

DefaultProfilerDomain() returns a non-null sentinel. This is required
because scoped_annotation.h uses domain != nullptr to gate the
RangePush/RangePop path vs. the AnnotationStack-only stub path. ROCm
has no domain concept; the sentinel value is never dereferenced or passed to
any roctx API.

Listener: rocm_tracer.cc

RocmTracer::InitProfiling() registers
ROCPROFILER_CALLBACK_TRACING_MARKER_CORE_API alongside the existing HIP API
and kernel dispatch callbacks. The new MarkerCallback handles three
operations:

roctx operation	Handling
roctxRangePushA	Push RoctxFrame{label, start_timestamp} onto a per-thread stack
roctxRangePop	Pop the frame, emit Generic RocmTracerEvent with start/end timestamps and roctx_range field set
roctxMarkA	Emit instantaneous Generic event (start == end)

Per-thread state is protected by roctx_stack_mutex_. Label strings are
interned in roctx_strings_ (an absl::node_hash_set) for pointer stability
and deduplication across multiple pushes of the same label.

Generic events bypass ApiActivityInfoExchange() and are routed directly
to PerDeviceCollector via standalone_events_. This is correct because
ROCTX ranges are host-side events with no GPU activity counterpart; they
carry no correlation ID and require no merge with activity records.

Lock ordering (documented in rocm_tracer.h to prevent future deadlock):

roctx_stack_mutex_ → (released) → roctx_strings_mutex_ → (released) → collector_mutex_

Enable() is the sole exception: it holds collector_mutex_ before acquiring
the roctx mutexes, which is safe because no callback can call collector()
while Enable() holds collector_mutex_.

XPlane routing

A new XPlane constant is added:

// xplane_schema.h / xplane_schema.cc
TF_CONST_INIT extern const absl::string_view kRoctxPlaneName;  // "/host:ROCTX"

PerDeviceCollector::Export() routes Generic events to the /host:ROCTX
plane. In PostProcessSingleHostXSpace, the ROCTX plane is merged into
/host:CPU following the identical pattern used for NVTX:

// Line ID space partition (no collisions after merge):
//   [0,     1<<32)  — host CPU threads
//   [1<<32, 2<<32)  — NVTX-CUPTI plane
//   [2<<32, 3<<32)  — ROCTX marker plane
static constexpr int64_t kRoctxLineIdStart = 2LL << 32;
XPlane* roctx_plane = FindMutablePlaneWithName(space, kRoctxPlaneName);
if (roctx_plane != nullptr) {
  ChangeOccupiedLineIds(roctx_plane, occupied_line_ids, kRoctxLineIdStart);
  MergePlanes({roctx_plane}, host_plane);
  RemovePlanes(space, {roctx_plane});
}

TensorBoard displays the resulting merged lines as "ROCTX Threads" rows within
the process, visually equivalent to CUDA's "NVTX Threads" rows.

TraceMeWrapper changes (xla/python/profiler.cc)

TraceMeWrapper is the Python binding for jax.profiler.TraceAnnotation.
It previously created only a TraceMe CPU event. Two additions:

1. The constructor calls tsl::profiler::PushAnnotation(annotation_name_),
   which routes through roctx_utils.cc::RangePush on ROCm (or
   nvtx_utils.cc::RangePush on CUDA).
2. A ~TraceMeWrapper() destructor calls Stop() to ensure PopAnnotation
   fires even when the object is used outside a Python context manager
   (e.g., t = jax.profiler.TraceAnnotation("name"); ... without with).
   A stopped_ guard makes Stop() idempotent so the context manager's
   __exit__ and the GC-driven destructor cannot double-pop.

Build system

The roctx_utils.cc target is linked via if_rocm_is_configured:

# third_party/tsl/tsl/profiler/lib/BUILD
cc_library(
    name = "nvtx_utils_impl",
    srcs = select({
        "@local_config_rocm//:is_rocm_configured": ["roctx_utils.cc"],
        "//conditions:default": ["nvtx_utils.cc"],
    }),
    ...
    deps = select({
        "@local_config_rocm//:is_rocm_configured": [
            "@local_config_rocm//rocm:rocprofiler_sdk_roctx",
        ],
        "//conditions:default": [...nvtx deps...],
    }),
)

On non-ROCm builds, roctx_utils.cc is never compiled and the
rocprofiler_sdk_roctx target is never pulled in.

---

Alternatives considered

A. Separate roctx_utils.h interface instead of implementing nvtx_utils.h

Rejected. nvtx_utils.h is deliberately abstract and already has exactly
the operations ROCTX needs. A separate interface would duplicate call sites
in scoped_annotation.h and TraceMeWrapper, create a new abstraction that
callers must track, and prevent the automatic "CUDA NVTX → ROCm ROCTX"
equivalence at every existing annotation point.

B. Implement via rocprof system-level tracing only, no XLA listener

Rejected. System-level profilers (rocprof v3, Omnitrace) can capture ROCTX
ranges independently, but XProf/TensorBoard's profiler plugin reads XSpace
protos produced by XLA's own collector. Without the MarkerCallback listener,
ROCTX ranges never enter the XSpace and TensorBoard never shows them —
defeating the goal of parity with the CUDA/TensorBoard experience.

C. Reuse AnnotationStack text as the ROCTX label without roctxRangePushA

Rejected. AnnotationStack is read in the HIP API callback and attached to
kernel events; it does not produce timeline events. Without an actual roctx
push, MarkerCallback never fires, /host:ROCTX remains empty, and the
TensorBoard annotation bars do not appear.

---

Impact assessment

Concern	Assessment
Non-ROCm builds	Zero impact. All new code is behind if_rocm_is_configured.
XSpace proto schema	No new proto fields. kRoctxPlaneName is a string constant; the plane appears only in ROCm profiles.
TensorBoard plugin	No changes required. The plugin already handles merged /host:CPU planes; ROCTX lines appear as additional rows under the existing process lane.
Existing ROCm tests	No regressions. The MarkerCallback path is additive; existing HIP_API and kernel dispatch paths are unchanged.
Lock safety	Lock ordering documented and verified. Enable() is the only site that acquires collector_mutex_ before the roctx mutexes; all callback paths hold only one mutex at a time.
Unmatched pop	MarkerCallback silently drops an unmatched roctxRangePop (empty stack) with a VLOG(2) trace, matching typical defensive handling in annotation stacks.
Null label	A null roctxMarkA message is treated as an empty string; the event is dropped. Verified by test MarkerCallbackNullMessageSafelyIgnored.

---

Testing

All tests run on a CPU-only host (no GPU required for the unit suite).

New unit tests

rocm_tracer_test.cc:

• MarkerCallbackPushPopEmitsRoctxRange — end-to-end: push + pop produces one Generic event with correct label, start, and duration.
• MarkerCallbackMarkEmitsInstantaneousEvent — mark produces event with start == end.
• MarkerCallbackUnmatchedPopIsIgnored — pop with empty stack does not crash or produce a spurious event.
• MarkerCallbackNullMessageSafelyIgnored — null label in roctxMarkA does not crash.
• MarkerEventAppearsInExportedXSpace — exported XSpace contains /host:ROCTX plane with the expected event.
• RealRoctxCallsProduceNvtxRangeInXSpace — integration test against the real librocprofiler-sdk-roctx.so; asserts kNVTXRange stat on exported events.

roctx_utils_test.cc (13 tests):

• Dual-push semantics (both AnnotationStack and roctx paths exercised).
• Nesting correctness (inner label not visible after pop).
• Stub behavior when profiling is disabled.
• GetCurrentRoctxLabel thread-safety and copy-safety.

rocm_tracer_test.cc (existing, extended):

• AnnotationMapStoresRoctxRange — roctx_range field stored alongside annotation text.
• AnnotationMapRoctxRangeEmptyWhenNotProvided — no regression when only annotation text is present.
• AnnotationMapStoresRoctxRangeWhenAnnotationEmpty — standalone ROCTX annotation (no AnnotationStack text) still produces kNVTXRange.

On-GPU validation (ROCm hardware required)

• RealRoctxCallsProduceNvtxRangeInXSpace in rocm_tracer_test.cc links against librocprofiler-sdk-roctx.so and verifies that real roctx calls flow end-to-end through the listener.
• Manual verification: jax.profiler.TraceAnnotation("my_region") produces visible "ROCTX Threads" rows in TensorBoard on an MI300X system.

---

Implementation plan

The implementation is complete and available at
ROCm/xla:cj/rocm-profiler-roctx-markers-with-emitter-pr0 (rebased onto
openxla/xla:main as of 2026-08-05).

The three commits can be submitted as a single PR or as a stack:

Commit	Files	What
1	tsl/profiler/lib/roctx_utils.cc, BUILD.tpl	Emitter: nvtx_utils.h implementation for ROCm
2	rocm_tracer.cc/h, rocm_collector.cc/h, rocm_tracer_utils.cc/h, xplane_schema.cc/h, post_process_single_host_xplane.cc	Listener + XPlane routing
3	xla/python/profiler.cc, xla/python/BUILD	TraceMeWrapper: dual-push for jax.profiler.TraceAnnotation

---

Open questions

1. nvtx_utils_impl select vs. separate target: The current implementation
   uses a select() in the existing nvtx_utils_impl target to swap in
   roctx_utils.cc on ROCm. An alternative is a separate
   roctx_utils_impl target composed at the link layer. Either works;
   preference from the XLA build-system owners is welcome.

2. kNVTXRange stat name on ROCm events: ROCTX range labels are currently
   stored under the kNVTXRange XStat key (matching CUPTI) so TensorBoard's
   existing parser reads them without changes. A new kROCTXRange key would
   be more accurate but would require a TensorBoard plugin update. The current
   choice favors zero-change adoption; feedback welcome.

3. Domain support: ROCTX has no domain concept (unlike NVTX which has
   nvtxDomainCreate). The ProfilerDomainHandle parameter in
   roctx_utils.cc is ignored. If upstream adds domain-specific annotation
   paths that distinguish user vs. framework ranges, ROCm will need a
   different mechanism (e.g., label prefixing).

4. NameStream mapping: The current implementation maps
   NameStream(stream, name) to roctxNameHipStream, which requires a
   hipStream_t. The StreamHandle opaque type passed through nvtx_utils.h
   is cast directly. This works today but couples the abstraction to the
   HIP type; a more principled mapping may be worth discussing.

[dimitar-asenov]

Thank you for the detailed design. Regarding the open questions:

1. Please make a separate roctx_utils_impl target and change the build to switch entire targets, not just C++ sources, based on the platform

2. I think this is fine for now. After the other changes have landed and working it would be good to send a fix to make things more accurate.

3. Fine for now.

4. Fine for now.