update with ROCm profiling tests with kernel_details

[cj401-amd]

Motivation

• update with ROCm profiling tests with kernel_details
• upstream PR [ROCm] update to test if there are GPU events when doing profiling on… jax-ml/jax#34135

Technical Details

with rocm-plugin

python3 stack.py develop \
    --rebuild-makefile \
    --fix-bazel-symbols \
    --xla-dir=/work/xla \
    --kernel-jax-dir=/work/jax \
    --rocm-path=${ROCM_PATH}

cd jax_rocm_plugin
make refresh refresh_jaxlib

Test Plan

cd jax;
python3 tests/profiler_test.py

Test Result

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1769609753.588043  150718 session_manager.cc:50] device_tracer_level expects an i value.
W0000 00:00:1769609753.588095  150718 session_manager.cc:50] host_tracer_level expects an i value.
W0000 00:00:1769609753.588100  150718 session_manager.cc:50] python_tracer_level expects an i value.
Starting to trace for 500 ms. Remaining attempt(s): 2
W0128 14:15:53.594123  150715 profiler_session.cc:109] Profiling is late by 5966164 nanoseconds and will start immediately.
W0000 00:00:1769609754.312314  150718 derived_timeline.cc:590] Found derived XLine with clashing display ID: 1. This will cause rendering issues in Trace Viewer.
Dumped profiling information in: /tmp/absl_testing
[       OK ] ProfilerTest.test_remote_profiler
[ RUN      ] ProfilerTest.test_remote_profiler_gcs_path
[  SKIPPED ] ProfilerTest.test_remote_profiler_gcs_path - Profiler takes >30s on Cloud TPUs
[ RUN      ] ProfilerTest.test_rocm_kernel_details_in_trace_json
[       OK ] ProfilerTest.test_rocm_kernel_details_in_trace_json
[ RUN      ] ProfilerTest.test_rocm_profiling
[       OK ] ProfilerTest.test_rocm_profiling
----------------------------------------------------------------------
Ran 22 tests in 74.087s