mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request from h-vetinari:packed

Browse Source 
PiperOrigin-RevId: 733696873
This commit is contained in:
jax authors 2025-03-05 05:46:01 -08:00
commit a13b3cedad
1 changed files with 2 additions and 1 deletions

3
jaxlib/mosaic/gpu/mosaic_gpu_ext.cc

@ -182,8 +182,9 @@ void callback_complete(CUcontext context, uint32_t streamId,
// Convert integer nanoseconds to floating point milliseconds to match
// the interface of the events-based profiler.
double duration_ms = (kernel->end - kernel->start) / 1e6;
const char* kernel_name = kernel->name;
profiler_state.timings.push_back(
std::make_tuple(kernel->name, duration_ms));
std::make_tuple(kernel_name, duration_ms));
}
} else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) {
// no more records available