diff --git a/hta/common/trace_parser.py b/hta/common/trace_parser.py index 6e0f470..0caee70 100644 --- a/hta/common/trace_parser.py +++ b/hta/common/trace_parser.py @@ -56,6 +56,27 @@ def __str__(self) -> str: For more details see https://pypi.org/project/ijson/#performance-tips, https://anaconda.org/anaconda/yajl, https://pypi.org/project/yajl/""" +def _is_gpu_type_intel(name_set: dict[str, int]) -> bool: + """ + Check if the GPU type is Intel XPU based on the symbol ID map (name_set). + + Args: + name_set (dict[str, int]): The symbol ID map. + + Returns: + bool: True if the GPU type is Intel XPU, False otherwise. + """ + intel_names: list[str] = [ + "urEnqueueKernelLaunch", + "urEnqueueKernelLaunchExp", + "urEnqueueKernelLaunchCustomExp", + "urEnqueueKernelLaunchWithArgsExp", + "urEnqueueCooperativeKernelLaunchExp", + ] + + return any(name in name_set for name in intel_names) + + def infer_gpu_type( metadata: Optional[MetaData] = None, name_set: dict[str, int] = {} ) -> str: @@ -83,6 +104,8 @@ def infer_gpu_type( return "AMD GPU" if "runFunction - job_prep_and_submit_for_execution" in name_set: return "MTIA" + if _is_gpu_type_intel(name_set): + return "INTEL XPU" return "UNKNOWN GPU" @@ -451,6 +474,20 @@ def _parse_trace_dataframe_ijson( df, local_symbol_table = _compress_df(df, cfg) return meta, df, local_symbol_table +def _remove_syclqueue_column(df: pd.DataFrame): + if "syclqueue" in df.columns: + df.drop(columns=["syclqueue"], inplace=True) + +def _transform_column_syclqueue_to_stream(df: pd.DataFrame): + """ + Remove stream column and rename syclqueue column to stream. + + Args: + df (pd.DataFrame): The trace events DataFrame. + """ + if "syclqueue" in df.columns and "stream" in df.columns: + df.drop(columns=["stream"], inplace=True) + df.rename(columns={"syclqueue": "stream"}, inplace=True) def parse_trace_dataframe( trace_file_path: str, @@ -505,6 +542,13 @@ def parse_trace_dataframe( device_type = infer_gpu_type(meta, local_symbol_table.get_sym_id_map()) meta[str(MetaDataKey.DEVICE_TYPE)] = device_type + # Intel traces do not have stream argument in the traces. + # The equivalent of stream argument in traces from Intel XPU is syclqueue. + if device_type == "INTEL XPU": + _transform_column_syclqueue_to_stream(df) + else: + _remove_syclqueue_column(df) + t_end = time.perf_counter() logger.warning( f"Parsed {trace_file_path} backend={parser_backend} in {(t_end - t_start):.2f} seconds; current PID:{os. getpid()}" diff --git a/hta/configs/event_args_formats/event_args_1.0.0.yaml b/hta/configs/event_args_formats/event_args_1.0.0.yaml index bdc2c76..8c2dcb6 100644 --- a/hta/configs/event_args_formats/event_args_1.0.0.yaml +++ b/hta/configs/event_args_formats/event_args_1.0.0.yaml @@ -264,6 +264,11 @@ AVAILABLE_ARGS: raw_name: batchId value_type: String default_value: "" + xpu::stream: + name: syclqueue + raw_name: sycl queue + value_type: Int + default_value: -1 inference::pre_grad_nodes: name: pre_grad_nodes raw_name: pre_grad_nodes diff --git a/hta/configs/event_args_yaml_parser.py b/hta/configs/event_args_yaml_parser.py index 5ee7c2e..4557692 100644 --- a/hta/configs/event_args_yaml_parser.py +++ b/hta/configs/event_args_yaml_parser.py @@ -42,7 +42,7 @@ ) ARGS_MINIMUM_FUNC: Callable[[Dict[str, AttributeSpec]], List[AttributeSpec]] = ( lambda available_args: [ - available_args[k] for k in ["cuda::stream", "correlation::cpu_gpu"] + available_args[k] for k in ["cuda::stream", "correlation::cpu_gpu", "xpu::stream"] ] ) ARGS_COMPLETE_FUNC: Callable[[Dict[str, AttributeSpec]], List[AttributeSpec]] = (