Skip to content

Commit

Permalink
Step2: Add versioning to parserconfig (facebookresearch#188)
Browse files Browse the repository at this point in the history 
Summary:
Pull Request resolved: facebookresearch#188

- Add versioning to parserconfig.
- Hook up yaml parser and the parser config

Reviewed By: briancoutinho

Differential Revision: D64116676

fbshipit-source-id: 9e4b488a2bf622fa3e622735af603b8d79c00924
  • Loading branch information
@niufei8888 @facebook-github-bot
niufei8888 authored and facebook-github-bot committed Oct 11, 2024
1 parent fa9d4dc commit 8e1db0c
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 148 deletions.
207 changes: 60 additions & 147 deletions hta/configs/parser_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,12 @@
from enum import Enum
from typing import Dict, List, Optional, Set

from hta.configs.default_values import AttributeSpec, ValueType
from hta.configs.default_values import AttributeSpec, EventArgs
from hta.configs.event_args_yaml_parser import (
parse_event_args_yaml,
v1_0_0,
YamlVersion,
)


class ParserBackend(str, Enum):
Expand All @@ -27,104 +32,9 @@ class TraceType(str, Enum):
Inference = "inference"


AVAILABLE_ARGS: Dict[str, AttributeSpec] = {
"index::ev_idx": AttributeSpec("ev_idx", "Ev Idx", ValueType.Int, -1),
"index::external_id": AttributeSpec(
"external_id", "External id", ValueType.Int, -1
),
"cpu_op::concrete_inputs": AttributeSpec(
"concrete_inputs", "Concrete Inputs", ValueType.Object, "[]"
),
"cpu_op::fwd_thread": AttributeSpec(
"fwd_thread_id", "Fwd thread id", ValueType.Int, -1
),
"cpu_op::input_dims": AttributeSpec(
"input_dims", "Input Dims", ValueType.Object, "-1"
),
"cpu_op::input_type": AttributeSpec(
"input_type", "Input type", ValueType.Object, "-1"
),
"cpu_op::input_strides": AttributeSpec(
"input_strides",
"Input Strides",
ValueType.Object,
"-1",
),
"cpu_op::sequence_number": AttributeSpec(
"sequence", "Sequence number", ValueType.Int, -1
),
# kernel_backend, for example triton.
"cpu_op::kernel_backend": AttributeSpec(
"kernel_backend", "kernel_backend", ValueType.String, ""
),
"correlation::cbid": AttributeSpec("cbid", "cbid", ValueType.Int, -1),
"correlation::cpu_gpu": AttributeSpec(
"correlation", "correlation", ValueType.Int, -1
),
"sm::blocks": AttributeSpec(
"blocks_per_sm", "blocks per SM", ValueType.Object, "[]"
),
"sm::occupancy": AttributeSpec(
"est_occupancy", "est. achieved occupancy %", ValueType.Int, -1
),
"sm::warps": AttributeSpec("warps_per_sm", "warps per SM", ValueType.Float, 0.0),
"data::bytes": AttributeSpec("bytes", "bytes", ValueType.Int, -1),
"data::bandwidth": AttributeSpec(
"memory_bw_gbps", "memory bandwidth (GB/s)", ValueType.Float, 0.0
),
"cuda::context": AttributeSpec("context", "context", ValueType.Int, -1),
"cuda::device": AttributeSpec("device", "device", ValueType.Int, -1),
"cuda::stream": AttributeSpec("stream", "stream", ValueType.Int, -1),
"kernel::queued": AttributeSpec("queued", "queued", ValueType.Int, -1),
"kernel::shared_memory": AttributeSpec(
"shared_memory", "shared memory", ValueType.Int, -1
),
"threads::block": AttributeSpec("block", "block", ValueType.Object, "[]"),
"threads::grid": AttributeSpec("grid", "grid", ValueType.Object, "[]"),
"threads::registers": AttributeSpec(
"registers_per_thread", "registers per thread", ValueType.Int, -1
),
"cuda_sync::stream": AttributeSpec(
"wait_on_stream", "wait_on_stream", ValueType.Int, -1
),
"cuda_sync::event": AttributeSpec(
"wait_on_cuda_event_record_corr_id",
"wait_on_cuda_event_record_corr_id",
ValueType.Int,
-1,
),
"info::labels": AttributeSpec("labels", "labels", ValueType.String, ""),
"info::name": AttributeSpec("name", "name", ValueType.Int, -1),
"info::op_count": AttributeSpec("op_count", "Op count", ValueType.Int, -1),
"info::sort_index": AttributeSpec("sort_index", "sort_index", ValueType.Int, -1),
"nccl::collective_name": AttributeSpec(
"collective_name", "Collective name", ValueType.String, ""
),
"nccl::in_msg_nelems": AttributeSpec(
"in_msg_nelems", "In msg nelems", ValueType.Int, 0
),
"nccl::out_msg_nelems": AttributeSpec(
"out_msg_nelems", "Out msg nelems", ValueType.Int, 0
),
"nccl::group_size": AttributeSpec("group_size", "Group size", ValueType.Int, 0),
"nccl::dtype": AttributeSpec("msg_dtype", "dtype", ValueType.String, ""),
"nccl::in_split_size": AttributeSpec(
"in_split_size", "In split size", ValueType.Object, "[]"
),
"nccl::out_split_size": AttributeSpec(
"out_split_size", "Out split size", ValueType.Object, "[]"
),
"nccl::process_group_name": AttributeSpec(
"process_group_name", "Process Group Name", ValueType.String, ""
),
"nccl::process_group_desc": AttributeSpec(
"process_group_desc", "Process Group Description", ValueType.String, ""
),
"nccl::process_group_ranks": AttributeSpec(
"process_group_ranks", "Process Group Ranks", ValueType.Object, "[]"
),
"nccl::rank": AttributeSpec("process_rank", "Rank", ValueType.Int, -1),
}
DEFAULT_PARSE_VERSION = v1_0_0
DEFAULT_VER_EVENT_ARGS: EventArgs = parse_event_args_yaml(DEFAULT_PARSE_VERSION)
AVAILABLE_ARGS: Dict[str, AttributeSpec] = DEFAULT_VER_EVENT_ARGS.AVAILABLE_ARGS


class ParserConfig:
Expand All @@ -144,52 +54,24 @@ class ParserConfig:
This class can be extended to support other customizations.
"""

ARGS_INPUT_SHAPE: List[AttributeSpec] = [
AVAILABLE_ARGS[k]
for k in ["cpu_op::input_dims", "cpu_op::input_type", "cpu_op::input_strides"]
]
ARGS_BANDWIDTH: List[AttributeSpec] = [
AVAILABLE_ARGS[k] for k in ["data::bytes", "data::bandwidth"]
]
ARGS_SYNC: List[AttributeSpec] = [
AVAILABLE_ARGS[k] for k in ["cuda_sync::stream", "cuda_sync::event"]
]
ARGS_MINIMUM: List[AttributeSpec] = [
AVAILABLE_ARGS[k] for k in ["cuda::stream", "correlation::cpu_gpu"]
]
ARGS_COMPLETE: List[AttributeSpec] = [
AVAILABLE_ARGS[k] for k in AVAILABLE_ARGS if not k.startswith("info")
]
ARGS_INFO: List[AttributeSpec] = [
AVAILABLE_ARGS[k] for k in ["info::labels", "info::name", "info::sort_index"]
]
ARGS_COMMUNICATION: List[AttributeSpec] = [
AVAILABLE_ARGS[k]
for k in [
"nccl::collective_name",
"nccl::in_msg_nelems",
"nccl::out_msg_nelems",
"nccl::dtype",
"nccl::group_size",
"nccl::rank",
"nccl::in_split_size",
"nccl::out_split_size",
]
]
ARGS_DEFAULT: List[AttributeSpec] = (
ARGS_MINIMUM
+ ARGS_BANDWIDTH
+ ARGS_SYNC
+ ARGS_INPUT_SHAPE
+ [AVAILABLE_ARGS["index::external_id"]]
)
ARGS_INPUT_SHAPE: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_INPUT_SHAPE
ARGS_BANDWIDTH: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_BANDWIDTH
ARGS_SYNC: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_SYNC
ARGS_MINIMUM: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_MINIMUM
ARGS_COMPLETE: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_COMPLETE
ARGS_INFO: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_INFO
ARGS_COMMUNICATION: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_COMMUNICATION
ARGS_DEFAULT: List[AttributeSpec] = DEFAULT_VER_EVENT_ARGS.ARGS_DEFAULT

def __init__(
self,
args: Optional[List[AttributeSpec]] = None,
user_provide_trace_type: Optional[TraceType] = None,
version: YamlVersion = DEFAULT_PARSE_VERSION,
) -> None:
self.args: List[AttributeSpec] = args if args else self.get_default_args()
self.args: List[AttributeSpec] = (
args if args else self.get_default_args(version=version)
)
self.parser_backend: Optional[ParserBackend] = None
self.trace_memory: bool = False
self.user_provide_trace_type: Optional[TraceType] = user_provide_trace_type
Expand All @@ -198,21 +80,31 @@ def __init__(
def get_default_cfg(cls) -> "ParserConfig":
return copy.deepcopy(_DEFAULT_PARSER_CONFIG)

@classmethod
def get_versioned_cfg(cls, version: YamlVersion) -> "ParserConfig":
return copy.deepcopy(ParserConfig(version=version))

@classmethod
def set_default_cfg(cls, cfg: "ParserConfig") -> None:
_DEFAULT_PARSER_CONFIG.set_args(cfg.get_args())

@classmethod
def get_minimum_args(cls) -> List[AttributeSpec]:
return cls.ARGS_MINIMUM.copy()
def get_minimum_args(
cls, version: YamlVersion = DEFAULT_PARSE_VERSION
) -> List[AttributeSpec]:
return parse_event_args_yaml(version).ARGS_MINIMUM.copy()

@classmethod
def get_default_args(cls) -> List[AttributeSpec]:
return cls.ARGS_DEFAULT.copy()
def get_default_args(
cls, version: YamlVersion = DEFAULT_PARSE_VERSION
) -> List[AttributeSpec]:
return parse_event_args_yaml(version).ARGS_DEFAULT.copy()

@classmethod
def get_info_args(cls) -> List[AttributeSpec]:
return cls.ARGS_INFO.copy()
def get_info_args(
cls, version: YamlVersion = DEFAULT_PARSE_VERSION
) -> List[AttributeSpec]:
return parse_event_args_yaml(version).ARGS_INFO.copy()

def set_args(self, args: List[AttributeSpec]) -> None:
if args != self.args:
Expand All @@ -233,8 +125,29 @@ def set_parser_backend(self, parser_backend: ParserBackend) -> None:
self.parser_backend = parser_backend

@staticmethod
def enable_communication_args() -> None:
_DEFAULT_PARSER_CONFIG.add_args(ParserConfig.ARGS_COMMUNICATION)
def enable_communication_args(version: YamlVersion = DEFAULT_PARSE_VERSION) -> None:
_DEFAULT_PARSER_CONFIG.add_args(
parse_event_args_yaml(version).ARGS_COMMUNICATION
)

@staticmethod
def set_global_parser_config_version(version: YamlVersion) -> None:
global _DEFAULT_PARSER_CONFIG
_DEFAULT_PARSER_CONFIG = ParserConfig(version=version)

global AVAILABLE_ARGS
AVAILABLE_ARGS = parse_event_args_yaml(version).AVAILABLE_ARGS

ParserConfig.ARGS_INPUT_SHAPE = parse_event_args_yaml(version).ARGS_INPUT_SHAPE
ParserConfig.ARGS_BANDWIDTH = parse_event_args_yaml(version).ARGS_BANDWIDTH
ParserConfig.ARGS_SYNC = parse_event_args_yaml(version).ARGS_SYNC
ParserConfig.ARGS_MINIMUM = parse_event_args_yaml(version).ARGS_MINIMUM
ParserConfig.ARGS_COMPLETE = parse_event_args_yaml(version).ARGS_COMPLETE
ParserConfig.ARGS_INFO = parse_event_args_yaml(version).ARGS_INFO
ParserConfig.ARGS_COMMUNICATION = parse_event_args_yaml(
version
).ARGS_COMMUNICATION
ParserConfig.ARGS_DEFAULT = parse_event_args_yaml(version).ARGS_DEFAULT


# Define a global ParserConfig variable for internal use. To access this variable,
Expand Down
3 changes: 2 additions & 1 deletion hta/utils/validate_trace.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@

import pandas as pd
from hta.common.trace_file import read_trace
from hta.configs.parser_config import AttributeSpec, ParserConfig, ValueType
from hta.configs.default_values import ValueType
from hta.configs.parser_config import AttributeSpec, ParserConfig


def get_argument_spec(level: str = "default") -> List[AttributeSpec]:
Expand Down

0 comments on commit 8e1db0c

Please sign in to comment.