Skip to content

Commit

Permalink
Reduce pandas overhead in compilation range name cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
olupton committed Dec 10, 2024
1 parent 0299d5a commit 66715ec
Showing 1 changed file with 7 additions and 7 deletions.
14 changes: 7 additions & 7 deletions .github/container/nsys_jax/nsys_jax/data_loaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -624,19 +624,19 @@ def _load_nvtx_pushpop_trace_single(name: pathlib.Path) -> pd.DataFrame:
# Because the ProgramId and ProgramName ranges provide the same information,
# remove those fields from the compilation range names.
def remove_program_id_and_name(row):
row.Name = (
return (
row.Name.removeprefix("TSL:")
.replace(f",program_id={row.ProgramId}", "")
.replace(f",module={row.ProgramName}", "")
.replace(f":#module={row.ProgramName}#", "")
)
return row

return (
compile_df.drop(columns=["EndMs"])
.astype({"ProgramId": np.int32})
.transform(remove_program_id_and_name, axis="columns")
)
compile_df = compile_df.drop(columns=["EndMs"]).astype({"ProgramId": np.int32})
if len(compile_df):
compile_df["Name"] = compile_df.apply(
remove_program_id_and_name, axis="columns"
)
return compile_df


def _load_nvtx_pushpop_trace(prefix: pathlib.Path, frames: set[str]) -> pd.DataFrame:
Expand Down

0 comments on commit 66715ec

Please sign in to comment.