From 38d52a19efb84ad54a2d30b5cdddb8e70b23cac2 Mon Sep 17 00:00:00 2001 From: Chris Jones Date: Tue, 18 Mar 2025 03:34:17 -0700 Subject: [PATCH] [mosaic_gpu] Force flush all cupti activity, then unsubscribe. With default flushing, it is possible for events to be missed. We should only unsubscribe after we are finished with cupti. PiperOrigin-RevId: 737939327 --- jaxlib/mosaic/gpu/mosaic_gpu_ext.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc b/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc index 2f415912f..4f804c9e2 100644 --- a/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc +++ b/jaxlib/mosaic/gpu/mosaic_gpu_ext.cc @@ -238,11 +238,12 @@ NB_MODULE(_mosaic_gpu_ext, m) { "failed to enable tracking of kernel activity by CUPTI"); }); m.def("_cupti_get_timings", []() { + THROW_IF_CUPTI_ERROR( + cuptiActivityFlushAll(CUPTI_ACTIVITY_FLAG_FLUSH_FORCED), + "failed to flush CUPTI activity buffers"); + THROW_IF_CUPTI_ERROR(cuptiFinalize(), "failed to detach CUPTI"); THROW_IF_CUPTI_ERROR(cuptiUnsubscribe(profiler_state.subscriber), "failed to unsubscribe from CUPTI"); - THROW_IF_CUPTI_ERROR(cuptiActivityFlushAll(CUPTI_ACTIVITY_FLAG_NONE), - "failed to flush CUPTI activity buffers"); - THROW_IF_CUPTI_ERROR(cuptiFinalize(), "failed to detach CUPTI"); return profiler_state.timings; }); }