diff --git a/hta/analyzers/critical_path_analysis.py b/hta/analyzers/critical_path_analysis.py index c5944132..238ccc6f 100644 --- a/hta/analyzers/critical_path_analysis.py +++ b/hta/analyzers/critical_path_analysis.py @@ -1873,8 +1873,14 @@ def critical_path_analysis( f" or (name == {stream_wait_event_id})" ) ) - - clipped_df = trace_df.loc[a.index.union(b.index)].copy() + # Retrieve the kernels initiated by cudaGraphLaunch + if "cudaGraphLaunch" in sym_index: + cudaGraphLaunch_id = sym_index.get("cudaGraphLaunch", -200) + cudaGraphLaunch_event = a[a["name"] == cudaGraphLaunch_id]["correlation"] + graph_launch_kernels = gpu_kernels[gpu_kernels["correlation"].isin(cudaGraphLaunch_event)] + clipped_df = trace_df.loc[a.index.union(b.index).union(graph_launch_kernels.index)].copy() + else: + clipped_df = trace_df.loc[a.index.union(b.index)].copy() logger.info(f"Clipped dataframe has {len(clipped_df)} events")