From 6158599af7f0a07f7f1b15cd5fe7cc98d6333472 Mon Sep 17 00:00:00 2001 From: Max Calman Date: Fri, 21 Nov 2025 14:19:38 -0500 Subject: [PATCH] support privateuse1 kineto activities Signed-off-by: Max Calman --- hta/common/trace_stack_filter.py | 10 ++++++---- hta/common/trace_symbol_table.py | 10 +++++++--- hta/common/types.py | 4 +++- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/hta/common/trace_stack_filter.py b/hta/common/trace_stack_filter.py index 1412cd16..2043a024 100644 --- a/hta/common/trace_stack_filter.py +++ b/hta/common/trace_stack_filter.py @@ -13,18 +13,20 @@ def get_matching_kernels( df_ops: pd.DataFrame, df_both: pd.DataFrame, cat_column: str ) -> pd.DataFrame: - """Get CUDA Kernels launched by operators in df_ops. + """Get Kernels launched by operators in df_ops. Args: df_ops: a DataFrame that contains cpu operators. - df_both: a DataFrame that contains both cpu operators and CUDA kernels. + df_both: a DataFrame that contains both cpu operators and kernels. cat_column: the column name for cat in string type, which can be either `cat` or `s_cat` depending on how the DataFrame symbol columns are encoded/decoded. Returns: - A DataFrame that contains the CUDA kernels launched by ops in df_ops. + A DataFrame that contains the kernels launched by ops in df_ops. """ - df_runtimes = df_ops.loc[df_ops[cat_column].eq("cuda_runtime")] + df_runtimes = df_ops.loc[ + df_ops[cat_column].isin(["cuda_runtime", "privateuse1_runtime"]) + ] df_kernels = df_both.loc[df_both["index_correlation"].isin(df_runtimes["index"])] return df_kernels diff --git a/hta/common/trace_symbol_table.py b/hta/common/trace_symbol_table.py index 6e97de8d..263fac50 100644 --- a/hta/common/trace_symbol_table.py +++ b/hta/common/trace_symbol_table.py @@ -307,14 +307,18 @@ def add_symbols_to_trace_df(self, trace_df: pd.DataFrame, col: str) -> None: def get_operator_or_cuda_runtime_mask(self, df: pd.DataFrame) -> pd.Series: """Returns a boolean mask you can use with pandas dataframes - to filter events that are CUDA runtime events or operators.""" + to filter events that are runtime events or operators.""" cpu_op_id = self.sym_index.get("cpu_op") - cuda_runtime_id = self.sym_index.get("cuda_driver", self.NULL) - cuda_driver_id = self.sym_index.get("cuda_runtime", self.NULL) + cuda_runtime_id = self.sym_index.get("cuda_runtime", self.NULL) + cuda_driver_id = self.sym_index.get("cuda_driver", self.NULL) + privateuse1_runtime_id = self.sym_index.get("privateuse1_runtime", self.NULL) + privateuse1_driver_id = self.sym_index.get("privateuse1_driver", self.NULL) return ( (df["cat"] == cpu_op_id) | (df["cat"] == cuda_runtime_id) | (df["cat"] == cuda_driver_id) + | (df["cat"] == privateuse1_runtime_id) + | (df["cat"] == privateuse1_driver_id) ) def get_runtime_launch_events_mask(self, df: pd.DataFrame) -> pd.Series: diff --git a/hta/common/types.py b/hta/common/types.py index d84e2f63..abf0ebe9 100644 --- a/hta/common/types.py +++ b/hta/common/types.py @@ -147,7 +147,9 @@ def to_grouping_pattern( ) DEVICE_RUNTIME_CATEGORY_PATTERN = GroupingPattern( - re.compile("(cuda_runtime)|(mtia_runtime)"), False, "Device Runtimes" + re.compile("(cuda_runtime)|(mtia_runtime)|(privateuse1_runtime)"), + False, + "Device Runtimes", ) MEMORY_CATEGORY_PATTERN = GroupingPattern(