Update rocprofiler-sdk (v3) along with roctracer (v1) for rocm-jaxlib-v0.6.0

[cj401-amd]

This PR combines rocprofiler-sdk (v3) and roctracer (v1) for rocm-jaxlib-v0.6.0

• Update rocprofiler-sdk integration for improved profiling with rocprofiler_force_configure() and annotations, support both time-based and step-based profiling,
• Keep roctracer(v1) that can be built by providing --bazel_options="--define=xla_rocm_profiler=v1"
• still need to figure out how to add more stats related to kernel, e.g., kernel size, occupancy, DMA copy, etc.

Previously, there was one for v3 only #251.

[i-chaochen]

could you indicate the instructions how to compile it with roctracev1 or rocprofiler-sdk (v3)?

[mmakevic-amd]

I'm not sure I understand this part. How are these hipBLASlt changes related to your profiler work?

[cj401-amd]

HIP_R_8F_E5M2 seems not on ROCm-6.2 when I was testing rocmtracer(v1), so here I was trying to guard it or I should submit a separate PR for this.

[mmakevic-amd]

I think a separate PR might be a good idea. What do you think @i-chaochen ?

[i-chaochen]

I guess it's because hipblaslt doesn't support fp8 on rocm62? Yes, please have a seperate PR just to our local guaranting this TF_ROCM_VERSION <= 60200

[Copilot]

Pull Request Overview

This PR integrates rocprofiler-sdk (v3) along with maintaining roctracer (v1) support for rocm-jaxlib-v0.6.0. The implementation uses compile-time version guards to select between the two profiling systems based on ROCm version, with v3 being used for ROCm >= 6.3.

Key changes:

• Add rocprofiler-sdk (v3) support with improved profiling capabilities including time/step-based profiling
• Maintain backward compatibility with roctracer (v1) for ROCm versions < 6.3
• Implement version-specific compilation guards and API mappings

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
xla/stream_executor/rocm/roctracer_wrapper.h	Adds version guards and rocprofiler-sdk headers/API mappings for v3
xla/stream_executor/rocm/hip_blas_lt.cc	Adds compile guards for FP8 data type constants to prevent undefined symbols
xla/backends/profiler/gpu/rocm_tracer_test.cc	New unit tests for rocm_tracer v3 functionality
xla/backends/profiler/gpu/rocm_tracer.h	Adds v3 RocmTracer class definition with rocprofiler-sdk integration
xla/backends/profiler/gpu/rocm_tracer.cc	Implements v3 tracer with rocprofiler-sdk callbacks and event handling
xla/backends/profiler/gpu/rocm_collector_test.cc	New unit tests for rocm_collector v3 functionality
xla/backends/profiler/gpu/rocm_collector.h	Adds v3 collector interfaces and data structures
xla/backends/profiler/gpu/rocm_collector.cc	Implements v3 collector with event processing and export logic
xla/backends/profiler/gpu/device_tracer_rocm.cc	Updates device tracer to use appropriate version-specific APIs
xla/backends/profiler/gpu/BUILD	Adds rocprofiler-sdk linkage and new test targets

Comments suppressed due to low confidence (4)

xla/backends/profiler/gpu/rocm_tracer.cc:277

• Inconsistent use of optional type - header file uses 'absl::types::optional.h' but implementation uses 'std::optional'. Should be consistent throughout.

      return "Invalid";

xla/backends/profiler/gpu/rocm_tracer.cc:163

• The stream_id type has been changed from int64_t to uint64_t in the header but this change should be consistently applied throughout the codebase to avoid potential issues.

        oss << ", sizeBytes=" << data->args.hipMemcpyDtoD.sizeBytes;

xla/backends/profiler/gpu/rocm_collector.cc:1

• [nitpick] Empty line at the beginning of the file should be removed for consistency with code style.

/* Copyright 2024 The OpenXLA Authors. All Rights Reserved.

[Copilot]

Using compiler-specific attributes like 'attribute((constructor))' reduces portability. Consider using a more portable initialization mechanism.

Suggested change

	void __attribute__((constructor)) init_rocm_lib() {
	rocprofiler_force_configure(xla::profiler::rocprofiler_configure);
	}
	namespace {
	struct RocmLibInitializer {
	RocmLibInitializer() {
	rocprofiler_force_configure(xla::profiler::rocprofiler_configure);
	}
	};
	static RocmLibInitializer rocm_lib_initializer;
	} // namespace

[draganmladjenovic]

Why do you need rocprofiler_force_configure?

[cj401-amd]

We need rocprofiler_force_configure to initialize rocprofiler-sdk's hooks into hip runtime before hipInit. it is called automatically when xla_rocm_plugin.so is loaded. I believe rocprofiler-sdk team is trying to solve this, then we can update our code.

[draganmladjenovic]

Can you explain a bit further. For the look at the source rocprofiler_force_configure cannot work if hipInit was called. So at the point rocprofiler_force_configure is called hipInit was not called yet, so I see no reason for rocprofiler_force_configure to be called.

[i-chaochen]

@cj401-amd I'm confused this PR, what's differences between this and #251 ? why you create separate one?

[mrodden]

Wasn't able to finish my review fully yet but here's what I have so far.

[mrodden]

you shouldn't need this anymore if you are using the macros in roctracer_wrapper.h which will do the dlopen at runtime to load the lib.

If you remove this and it doesn't work then we probably have another issue that needs to be fixed.

[draganmladjenovic]

And you should not link against rocprofiler-sdk if you don't use it. Say on 6.2. If we end up with compile-time switch between the tracers I suggest this be modeled as an proper library and added to dependencies.

[mrodden]

Kinda weird that this has to be under an if. I would think it wouldn't matter since it shouldn't be invoked already due to the rocm-only tag...

[i-chaochen]

this is just to align upstream syntax

[mrodden]

I still think we need to do this switch at runtime, not compile time. It should be possible since you have both symbol sets from roctracer and rocprofiler available from the roctracer_wrapper.h, so you just have to switch on an environment variable or something.

[i-chaochen]

would be great if you can have another PR to make it feasible on those tested workloads.

[draganmladjenovic]

@mrodden But we only have a single rocm version where both rocprofiler-sdk and roctracer coexist?

[mrodden]

I think you still need this?

[pemeliya]

As I see rocm_tracer/collector for rocprofiler v1 and for v3 are two orthogonal implementations. They have very few things in common. If we really insist on keeping rocprofiler v1, perhaps can we split rocm_tracer.h/.cc into rocm_tracer_v1.h/cc and rocm_tracer_v3.h/cc and maybe the same for rocm_collector? This way we can improve code maintainability, and can get rid of major #if / #else' clauses by letting bazel to do the conditional compilation:

e.g. if ROCM_VERSION < 60300 -> compile rocm_tracer/collector_v1
else -> compile rocm_tracer/collector_v3

[draganmladjenovic]

You should include only HIP visible gpus https://github.com/ROCm/rocprofiler-sdk/blob/3a924a0d36a3b4be70ce3ef6e81f7e302e30258e/source/include/rocprofiler-sdk/agent.h#L211C43-L211C61 Otherwise your numbering scheme ignores HIP_VISIBLE_DEVICES.

[draganmladjenovic]

What is the reason we include both CPU and GPU agents in the map?

[cj401-amd]

rocprofiler-sdk retrieves all agents, including both CPUs and GPUs on the system. then gpus are filtered.

[draganmladjenovic]

Where are the gpus filtered?

[cj401-amd]

Where are the gpus filtered?

firstly, filtered out CPU agents and only keep GPU agents https://github.com/ROCm/xla/blob/ci_cj-rocprofv3-v1-rocm-jaxlib-v0.6.0/xla/backends/profiler/gpu/rocm_profiler_sdk.cc#L406

then for writing out events based on GPU ids https://github.com/ROCm/xla/blob/ci_cj-rocprofv3-v1-rocm-jaxlib-v0.6.0/xla/backends/profiler/gpu/rocm_collector.cc#L536

[cj401-amd]

As I see rocm_tracer/collector for rocprofiler v1 and for v3 are two orthogonal implementations. They have very few things in common. If we really insist on keeping rocprofiler v1, perhaps can we split rocm_tracer.h/.cc into rocm_tracer_v1.h/cc and rocm_tracer_v3.h/cc and maybe the same for rocm_collector? This way we can improve code maintainability, and can get rid of major #if / #else' clauses by letting bazel to do the conditional compilation:

e.g. if ROCM_VERSION < 60300 -> compile rocm_tracer/collector_v1 else -> compile rocm_tracer/collector_v3

initially, the idea is to put v1 along v3 for a transition period (v1 is being phased out.) and then remove v1 from the code totally.

[mrodden]

@cj401-amd @i-chaochen Here is a modification of the current changes to switch between them at runtime ae1c5d8

I am not sure my plan of trying to extend and override the one function member of RocmTraceCollector is going to work, but it could be changed to a composition or templated out instead.

[i-chaochen]

@cj401-amd I think you need to rebase your branch to the latest 0.6 to get a green CI pass

[cj401-amd]

As I see rocm_tracer/collector for rocprofiler v1 and for v3 are two orthogonal implementations. They have very few things in common. If we really insist on keeping rocprofiler v1, perhaps can we split rocm_tracer.h/.cc into rocm_tracer_v1.h/cc and rocm_tracer_v3.h/cc and maybe the same for rocm_collector? This way we can improve code maintainability, and can get rid of major #if / #else' clauses by letting bazel to do the conditional compilation:

e.g. if ROCM_VERSION < 60300 -> compile rocm_tracer/collector_v1 else -> compile rocm_tracer/collector_v3

now the backend is split into rocm_tracer_v1 and rocm_profiler_sdk with default to rocm_profiler_sdk (v3), which can be changed to --bazel_options="--define=xla_rocm_profiler=v1" for ROCm < 6.3.

it seems v1 could not be built on ROCm-6.4, and rocprofiler-sdk could not be built on ROCm-6.2.

[cj401-amd]

it seems v1 could not be built on ROCm-6.4, and rocprofiler-sdk could not be built on ROCm-6.2.

[cj401-amd]

@cj401-amd @i-chaochen Here is a modification of the current changes to switch between them at runtime ae1c5d8

I am not sure my plan of trying to extend and override the one function member of RocmTraceCollector is going to work, but it could be changed to a composition or templated out instead.

based on latest exps, v1 seems not be built on ROCm-6.4, and rocprofiler-sdk could not be built on ROCm-6.2.

[i-chaochen]

it's ok to me @pemeliya @ScXfjiang please have a check

[cj401-amd]

I think rocprofiler_force_configure should be removed.

we are going to remove it when rocprofiler-sdk fixes the initialization of its hooks. If we remove it now, rocprofiler-sdk can not trace any GPU events for some workloads due to hipInit is called before rocprofiler_configure, e.g., maxtext LLAMA2-7B workload.

[i-chaochen]

LGTM please squash your commit as one as well, and if @pemeliya or @draganmladjenovic approved we can merge it. Thanks!

[pemeliya]

@cj401-amd , can you also take a look at my two last comments ? I mean we do not really need to access that 'agents_' array to get the device ID. Furthermore, we can experiment with the profiler buffer size to see if increasing it, solves the problem with missing events (as we saw previously)

[pemeliya]

I think we can simplify this - because src_gpu.id.handle seems to be the same as rec.src_agent_id.handle
and same for dst_gpu. So we do not really need to access 'agents_' map here.

[cj401-amd]

The reason I used that way was for checking https://github.com/ROCm/xla/blob/ci_cj-rocprofv3-v1-rocm-jaxlib-v0.6.0/xla/backends/profiler/gpu/rocm_profiler_sdk.cc#L225, otherwise, sometimes we got CPU id, then here we start iterating from cpu_id https://github.com/ROCm/xla/blob/ci_cj-rocprofv3-v1-rocm-jaxlib-v0.6.0/xla/backends/profiler/gpu/rocm_collector.cc#L534.

[cj401-amd]

I think rocprofiler_force_configure should be removed.

Hi @draganmladjenovic, is it possible to help remove the lock now?

[i-chaochen]

I can see there are still 2 failed UTs from CI, is it expecable?

[2025-09-08T21:01:33.346Z] //xla/backends/profiler/gpu:rocm_collector_test_cpu                      FAILED in 3 out of 3 in 3.4s
[2025-09-08T21:01:33.346Z] //xla/backends/profiler/gpu:rocm_collector_test_gpu_amd_any              FAILED in 3 out of 3 in 12.7s

[draganmladjenovic]

Here you go.

[cj401-amd]

I can see there are still 2 failed UTs from CI, is it expecable?

[2025-09-08T21:01:33.346Z] //xla/backends/profiler/gpu:rocm_collector_test_cpu                      FAILED in 3 out of 3 in 3.4s
[2025-09-08T21:01:33.346Z] //xla/backends/profiler/gpu:rocm_collector_test_gpu_amd_any              FAILED in 3 out of 3 in 12.7s

fixed them already.