30476914782

Committed 29 Jul 2026 05:47PM UTC coverage: 64.787% (+0.5%) from 64.324%

Build # 30476914782

Build Type

Pull #873

github

Committed by

web-flow

Commit Message

Merge 9d147ddd2 into 3eb01880e

Pull Request Pull Request #873: Map fusion migration

Coverage Stats

1388 of 1889 new or added lines in 13 files covered. (73.48%)

195 existing lines in 6 files now uncovered.

44816 of 69174 relevant lines covered (64.79%)

730.43 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.88

/python/docc/benchmarks/rtl.py

"""Read, validate and query daisy RTL instrumentation traces.

The RTL library (``rtl/src/instrumentation.cpp``) emits a Chrome-trace subset
extended with docc-specific metadata. This module loads such a trace, validates
it against the bundled JSON schema (``daisy_trace.schema.json``) and exposes a
typed, queryable :class:`Trace` object.

Two trace shapes exist, selected at capture time by ``__DAISY_INSTRUMENTATION_MODE``:

* per-invocation (``cat == "region,daisy"``) — one event per region entry/exit,
  with ``metrics`` a flat ``{counter: int}`` map.
* aggregated (``cat == "aggregated_region,daisy"``) — one event per region with
  Welford statistics per counter plus a ``runtime`` stat block.

Example::

    from docc.benchmarks import Trace

    trace = Trace.load("daisy_trace.json")
    for region in trace.regions:
        print(region.name, region.runtime_mean_us)

    hot = trace.filter(target_type="CUDA")
"""

from __future__ import annotations

import copy
import json
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Sequence, Union

__all__ = [
    "Trace",
    "TraceRegion",
    "LoopInfo",
    "SourcePos",
    "SourceRange",
    "MetricStat",
    "TraceValidationError",
    "validate",
    "PER_INVOCATION_CAT",
    "AGGREGATED_CAT",
    "STATIC_PREFIX",
]

PER_INVOCATION_CAT = "region,daisy"
AGGREGATED_CAT = "aggregated_region,daisy"
STATIC_PREFIX = "static:::"
_RUNTIME_KEY = "runtime"

_SCHEMA_FILENAME = "daisy_trace.schema.json"


def _find_schema() -> Path:
    """Locate the daisy-trace JSON schema.

    In an installed/built wheel the schema is shipped as package data next to this
    module (CMake installs ``rtl/schema/daisy_trace.schema.json`` into
    ``docc/benchmarks``). When running from a source checkout that install has not
    happened, so fall back to the canonical copy under ``rtl/schema``.
    """
    bundled = Path(__file__).with_name(_SCHEMA_FILENAME)
    if bundled.is_file():
        return bundled

    # rtl.py -> docc/python/docc/benchmarks/rtl.py ; schema -> docc/rtl/schema/...
    repo_schema = (
        Path(__file__).resolve().parents[3] / "rtl" / "schema" / _SCHEMA_FILENAME
    )
    if repo_schema.is_file():
        return repo_schema

    raise FileNotFoundError(
        f"Could not locate {_SCHEMA_FILENAME}; expected it next to {bundled} "
        f"(installed wheel) or at {repo_schema} (source checkout)."
    )


def load_schema() -> dict:
    """Return the bundled daisy-trace JSON schema as a dict."""
    with open(_find_schema(), "r") as f:
        return json.load(f)


class TraceValidationError(ValueError):
    """Raised when a trace does not conform to the daisy-trace schema."""

    def __init__(self, errors: List[str]):
        self.errors = errors
        super().__init__("Trace failed schema validation:\n" + "\n".join(errors))


def validate(instance: Any, schema: Optional[dict] = None) -> List[str]:
    """Validate a parsed trace object against the schema.

    Returns a list of human-readable error strings (empty when valid). Requires
    the optional ``jsonschema`` dependency.
    """
    try:
        from jsonschema import Draft202012Validator
    except ImportError as exc:  # pragma: no cover - exercised only without dep
        raise ImportError(
            "Validating daisy traces requires the 'jsonschema' package. "
            "Install it with `pip install jsonschema`."
        ) from exc

    if schema is None:
        schema = load_schema()

    validator = Draft202012Validator(schema)
    errors: List[str] = []
    for error in sorted(
        validator.iter_errors(instance), key=lambda e: list(e.absolute_path)
    ):
        location = "/".join(str(p) for p in error.absolute_path) or "<root>"
        errors.append(f"{location}: {error.message}")
    return errors


@dataclass(frozen=True)
class SourcePos:
    line: int
    col: int

    @classmethod
    def from_dict(cls, d: dict) -> "SourcePos":
        return cls(line=d["line"], col=d["col"])


@dataclass(frozen=True)
class SourceRange:
    file: str
    begin: SourcePos
    end: SourcePos

    @classmethod
    def from_dict(cls, d: dict) -> "SourceRange":
        return cls(
            file=d["file"],
            begin=SourcePos.from_dict(d["from"]),
            end=SourcePos.from_dict(d["to"]),
        )


@dataclass(frozen=True)
class LoopInfo:
    loopnest_index: int
    num_loops: int
    num_maps: int
    num_fors: int
    num_whiles: int
    max_depth: int
    is_perfectly_nested: bool
    is_perfectly_parallel: bool
    is_elementwise: bool
    has_side_effects: bool

    @classmethod
    def from_dict(cls, d: dict) -> "LoopInfo":
        return cls(
            loopnest_index=d["loopnest_index"],
            num_loops=d["num_loops"],
            num_maps=d["num_maps"],
            num_fors=d["num_fors"],
            num_whiles=d["num_whiles"],
            max_depth=d["max_depth"],
            is_perfectly_nested=d["is_perfectly_nested"],
            is_perfectly_parallel=d["is_perfectly_parallel"],
            is_elementwise=d["is_elementwise"],
            has_side_effects=d["has_side_effects"],
        )


@dataclass(frozen=True)
class MetricStat:
    """Aggregated statistics for a single counter (aggregated traces).

    ``mean``/``variance``/``min``/``max`` may be ``None`` for provided
    (``static:::``) metrics whose value was non-finite at capture time.
    """

    mean: Optional[float]
    variance: Optional[float]
    count: int
    min: Optional[float]
    max: Optional[float]

    @classmethod
    def from_dict(cls, d: dict) -> "MetricStat":
        return cls(
            mean=d["mean"],
            variance=d["variance"],
            count=d["count"],
            min=d["min"],
            max=d["max"],
        )


class TraceRegion:
    """A single region event, wrapping the raw trace dict with typed accessors."""

    def __init__(self, raw: dict):
        self._raw = raw

    # -- Raw / identity ----------------------------------------------------
    @property
    def raw(self) -> dict:
        return self._raw

    @property
    def category(self) -> str:
        return self._raw["cat"]

    @property
    def is_aggregated(self) -> bool:
        return self.category == AGGREGATED_CAT

    @property
    def name(self) -> str:
        return self._raw["name"]

    @property
    def pid(self) -> int:
        return self._raw["pid"]

    @property
    def tid(self) -> int:
        return self._raw["tid"]

    @property
    def ts_us(self) -> Optional[float]:
        return self._raw.get("ts")

    @property
    def dur_us(self) -> Optional[float]:
        return self._raw.get("dur")

    # -- Source metadata ---------------------------------------------------
    @property
    def _args(self) -> dict:
        return self._raw["args"]

    @property
    def function(self) -> str:
        return self._args["function"]

    @property
    def module(self) -> str:
        return self._args["module"]

    @property
    def target_type(self) -> str:
        return self._args.get("target_type", "")

    @cached_property
    def source_ranges(self) -> List[SourceRange]:
        return [SourceRange.from_dict(r) for r in self._args.get("source_ranges", [])]

    # -- docc metadata -----------------------------------------------------
    @property
    def _docc(self) -> dict:
        return self._args["docc"]

    @property
    def sdfg_name(self) -> str:
        return self._docc.get("sdfg_name", "")

    @property
    def sdfg_file(self) -> str:
        return self._docc.get("sdfg_file", "")

    @property
    def arg_capture_path(self) -> str:
        return self._docc.get("arg_capture_path", "")

    @property
    def features_file(self) -> str:
        return self._docc.get("features_file", "")

    @property
    def opt_report_file(self) -> str:
        return self._docc.get("opt_report_file", "")

    @property
    def has_element(self) -> bool:
        """True when SDFG element metadata (id/type/loop_info) is present."""
        return "element_id" in self._docc

    @property
    def element_id(self) -> Optional[int]:
        return self._docc.get("element_id")

    @property
    def element_type(self) -> Optional[str]:
        return self._docc.get("element_type")

    @property
    def loopnest_index(self) -> Optional[int]:
        return self._docc.get("loopnest_index")

    @cached_property
    def loop_info(self) -> Optional[LoopInfo]:
        li = self._docc.get("loop_info")
        return LoopInfo.from_dict(li) if li is not None else None

    # -- Metrics -----------------------------------------------------------
    @property
    def _metrics(self) -> dict:
        return self._args.get("metrics", {})

    @cached_property
    def counters(self) -> Dict[str, int]:
        """PAPI counter values for a per-invocation region (empty if aggregated)."""
        if self.is_aggregated:
            return {}
        return dict(self._metrics)

    @cached_property
    def counter_stats(self) -> Dict[str, MetricStat]:
        """Aggregated PAPI counter statistics (empty for per-invocation regions).

        Excludes the reserved ``runtime`` block and ``static:::`` provided metrics.
        """
        if not self.is_aggregated:
            return {}
        out: Dict[str, MetricStat] = {}
        for key, val in self._metrics.items():
            if key == _RUNTIME_KEY or key.startswith(STATIC_PREFIX):
                continue
            out[key] = MetricStat.from_dict(val)
        return out

    @cached_property
    def static_metrics(self) -> Dict[str, MetricStat]:
        """User-provided (``static:::``) metric statistics, keyed without the prefix."""
        if not self.is_aggregated:
            return {}
        out: Dict[str, MetricStat] = {}
        for key, val in self._metrics.items():
            if key.startswith(STATIC_PREFIX):
                out[key[len(STATIC_PREFIX) :]] = MetricStat.from_dict(val)
        return out

    @cached_property
    def runtime(self) -> Optional[MetricStat]:
        """Aggregated runtime statistics in microseconds (aggregated traces only)."""
        rt = self._metrics.get(_RUNTIME_KEY) if self.is_aggregated else None
        return MetricStat.from_dict(rt) if rt is not None else None

    @property
    def runtime_mean_us(self) -> Optional[float]:
        """Mean runtime in microseconds.

        For aggregated traces this is ``runtime.mean``; for per-invocation traces
        it is this event's own ``dur``.
        """
        if self.is_aggregated:
            return self.runtime.mean if self.runtime is not None else None
        return self.dur_us

    @property
    def runtime_min_us(self) -> Optional[float]:
        """Best (minimum) runtime in microseconds.

        The min over the aggregated samples is the steady-state estimate.
        """
        if self.is_aggregated:
            return self.runtime.min if self.runtime is not None else None
        return self.dur_us

    def metric(self, name: str) -> Optional[Union[int, float]]:
        """Return a representative value for ``name``.

        For per-invocation regions this is the raw counter value; for aggregated
        regions it is the counter/static metric mean. Returns ``None`` if absent.
        """
        if not self.is_aggregated:
            return self.counters.get(name)
        if name in self.counter_stats:
            return self.counter_stats[name].mean
        if name in self.static_metrics:
            return self.static_metrics[name].mean
        return None

    def __repr__(self) -> str:
        return (
            f"TraceRegion(name={self.name!r}, module={self.module!r}, "
            f"target={self.target_type!r}, aggregated={self.is_aggregated})"
        )


class Trace:
    """A parsed, validated daisy trace: a collection of :class:`TraceRegion`."""

    def __init__(self, raw: dict):
        self._raw = raw
        self._regions = [TraceRegion(e) for e in raw.get("traceEvents", [])]

    # -- Construction ------------------------------------------------------
    @classmethod
    def from_dict(cls, data: dict, validate_schema: bool = True) -> "Trace":
        """Build a :class:`Trace` from an already-parsed dict.

        Raises :class:`TraceValidationError` if ``validate_schema`` is set and the
        data does not conform to the schema.
        """
        if validate_schema:
            errors = validate(data)
            if errors:
                raise TraceValidationError(errors)
        return cls(data)

    @classmethod
    def load(cls, path: Union[str, Path], validate_schema: bool = True) -> "Trace":
        """Read a trace JSON file, optionally validate it, and return a :class:`Trace`."""
        with open(path, "r") as f:
            data = json.load(f)
        return cls.from_dict(data, validate_schema=validate_schema)

    def save(self, path: Union[str, Path], indent: Optional[int] = None) -> None:
        """Write this trace back to a JSON file (the daisy-trace format)."""
        with open(path, "w") as f:
            json.dump(self._raw, f, indent=indent)

    # -- Container protocol ------------------------------------------------
    @property
    def regions(self) -> List[TraceRegion]:
        return self._regions

    @property
    def raw(self) -> dict:
        return self._raw

    def __len__(self) -> int:
        return len(self._regions)

    def __iter__(self) -> Iterator[TraceRegion]:
        return iter(self._regions)

    def __getitem__(self, index: int) -> TraceRegion:
        return self._regions[index]

    @property
    def is_aggregated(self) -> bool:
        """True if the trace was captured in aggregate mode.

        Determined from the first region; a valid trace never mixes modes.
        """
        return bool(self._regions) and self._regions[0].is_aggregated

    # -- Queries -----------------------------------------------------------
    def filter(
        self,
        *,
        module: Optional[str] = None,
        function: Optional[str] = None,
        target_type: Optional[str] = None,
        element_type: Optional[str] = None,
        sdfg_name: Optional[str] = None,
    ) -> List[TraceRegion]:
        """Return regions matching all provided (non-``None``) criteria."""
        out = []
        for r in self._regions:
            if module is not None and r.module != module:
                continue
            if function is not None and r.function != function:
                continue
            if target_type is not None and r.target_type != target_type:
                continue
            if element_type is not None and r.element_type != element_type:
                continue
            if sdfg_name is not None and r.sdfg_name != sdfg_name:
                continue
            out.append(r)
        return out

    def by_element_id(self, element_id: int) -> List[TraceRegion]:
        """Return regions carrying the given SDFG ``element_id``."""
        return [r for r in self._regions if r.element_id == element_id]

    def hottest(self, n: Optional[int] = None) -> List[TraceRegion]:
        """Regions sorted by mean runtime (descending); ``n`` limits the count."""
        ranked = sorted(
            (r for r in self._regions if r.runtime_mean_us is not None),
            key=lambda r: r.runtime_mean_us,
            reverse=True,
        )
        return ranked[:n] if n is not None else ranked

    def total_runtime_us(self) -> float:
        """Sum of mean runtimes across all regions (microseconds)."""
        return sum((r.runtime_mean_us or 0.0) for r in self._regions)

    # -- Combining ---------------------------------------------------------
    @staticmethod
    def _region_key(region: "TraceRegion") -> tuple:
        """Stable identity of a region across separate runs of the same workload.

        Matches on source/docc identity and excludes pid/tid/timestamps, which
        differ between runs.
        """
        return (
            region.name,
            region.function,
            region.module,
            region.sdfg_name,
            region.element_id,
        )

    @classmethod
    def combine(
        cls, traces: Sequence["Trace"], validate_schema: bool = False
    ) -> "Trace":
        """Combine multiple traces of the same workload into one.

        Intended for counter multiplexing: when the hardware can only measure a
        subset of counters at once, the workload is run once per counter group,
        producing several traces that each carry a slice of the ``metrics``. This
        takes the first trace as the reference for regions and timesteps (name,
        ts, dur, source/docc metadata, runtime) and unions the per-region
        ``metrics`` maps from the remaining traces into it.

        Regions are matched across traces by :meth:`_region_key` (source/docc
        identity, ignoring pid/tid/timestamps). On a metric-key collision the
        reference wins, so its ``runtime`` and any ``static:::`` values are kept.
        Returns a new :class:`Trace`; the inputs are not mutated.
        """
        traces = list(traces)
        if not traces:
            raise ValueError("Trace.combine() requires at least one trace")
        if len(traces) == 1:
            return cls.from_dict(
                copy.deepcopy(traces[0].raw), validate_schema=validate_schema
            )

        # Index each additional trace's per-region metrics by region key.
        extra_metrics: Dict[tuple, List[dict]] = {}
        for other in traces[1:]:
            for region in other.regions:
                metrics = region.raw.get("args", {}).get("metrics", {})
                extra_metrics.setdefault(cls._region_key(region), []).append(metrics)

        # Deep-copy the reference and union in the other traces' metrics.
        combined_raw = copy.deepcopy(traces[0].raw)
        for event in combined_raw.get("traceEvents", []):
            key = cls._region_key(TraceRegion(event))
            metrics = event.setdefault("args", {}).setdefault("metrics", {})
            for other_metrics in extra_metrics.get(key, []):
                for name, value in other_metrics.items():
                    metrics.setdefault(name, value)

        return cls.from_dict(combined_raw, validate_schema=validate_schema)

1	"""Read, validate and query daisy RTL instrumentation traces.
2
3	The RTL library (``rtl/src/instrumentation.cpp``) emits a Chrome-trace subset
4	extended with docc-specific metadata. This module loads such a trace, validates
5	it against the bundled JSON schema (``daisy_trace.schema.json``) and exposes a
6	typed, queryable :class:`Trace` object.
7
8	Two trace shapes exist, selected at capture time by ``__DAISY_INSTRUMENTATION_MODE``:
9
10	* per-invocation (``cat == "region,daisy"``) — one event per region entry/exit,
11	with ``metrics`` a flat ``{counter: int}`` map.
12	* aggregated (``cat == "aggregated_region,daisy"``) — one event per region with
13	Welford statistics per counter plus a ``runtime`` stat block.
14
15	Example::
16
17	from docc.benchmarks import Trace
18
19	trace = Trace.load("daisy_trace.json")
20	for region in trace.regions:
21	print(region.name, region.runtime_mean_us)
22
23	hot = trace.filter(target_type="CUDA")
24	"""
25
26	from __future__ import annotations	4✔
27
28	import copy	4✔
29	import json	4✔
30	from dataclasses import dataclass	4✔
31	from functools import cached_property	4✔
32	from pathlib import Path	4✔
33	from typing import Any, Dict, Iterator, List, Optional, Sequence, Union	4✔
34
35	__all__ = [	4✔
36	"Trace",
37	"TraceRegion",
38	"LoopInfo",
39	"SourcePos",
40	"SourceRange",
41	"MetricStat",
42	"TraceValidationError",
43	"validate",
44	"PER_INVOCATION_CAT",
45	"AGGREGATED_CAT",
46	"STATIC_PREFIX",
47	]
48
49	PER_INVOCATION_CAT = "region,daisy"	4✔
50	AGGREGATED_CAT = "aggregated_region,daisy"	4✔
51	STATIC_PREFIX = "static:::"	4✔
52	_RUNTIME_KEY = "runtime"	4✔
53
54	_SCHEMA_FILENAME = "daisy_trace.schema.json"	4✔
55
56
57	def _find_schema() -> Path:	4✔
58	"""Locate the daisy-trace JSON schema.
59
60	In an installed/built wheel the schema is shipped as package data next to this
61	module (CMake installs ``rtl/schema/daisy_trace.schema.json`` into
62	``docc/benchmarks``). When running from a source checkout that install has not
63	happened, so fall back to the canonical copy under ``rtl/schema``.
64	"""
65	bundled = Path(__file__).with_name(_SCHEMA_FILENAME)	4✔
66	if bundled.is_file():	4✔
UNCOV 67	return bundled	×
68
69	# rtl.py -> docc/python/docc/benchmarks/rtl.py ; schema -> docc/rtl/schema/...
70	repo_schema = (	4✔
71	Path(__file__).resolve().parents[3] / "rtl" / "schema" / _SCHEMA_FILENAME
72	)
73	if repo_schema.is_file():	4✔
74	return repo_schema	4✔
75
UNCOV 76	raise FileNotFoundError(	×
77	f"Could not locate {_SCHEMA_FILENAME}; expected it next to {bundled} "
78	f"(installed wheel) or at {repo_schema} (source checkout)."
79	)
80
81
82	def load_schema() -> dict:	4✔
83	"""Return the bundled daisy-trace JSON schema as a dict."""
84	with open(_find_schema(), "r") as f:	4✔
85	return json.load(f)	4✔
86
87
88	class TraceValidationError(ValueError):	4✔
89	"""Raised when a trace does not conform to the daisy-trace schema."""
90
91	def __init__(self, errors: List[str]):	4✔
92	self.errors = errors	4✔
93	super().__init__("Trace failed schema validation:\n" + "\n".join(errors))	4✔
94
95
96	def validate(instance: Any, schema: Optional[dict] = None) -> List[str]:	4✔
97	"""Validate a parsed trace object against the schema.
98
99	Returns a list of human-readable error strings (empty when valid). Requires
100	the optional ``jsonschema`` dependency.
101	"""
102	try:	4✔
103	from jsonschema import Draft202012Validator	4✔
104	except ImportError as exc: # pragma: no cover - exercised only without dep
105	raise ImportError(
106	"Validating daisy traces requires the 'jsonschema' package. "
107	"Install it with `pip install jsonschema`."
108	) from exc
109
110	if schema is None:	4✔
111	schema = load_schema()	4✔
112
113	validator = Draft202012Validator(schema)	4✔
114	errors: List[str] = []	4✔
115	for error in sorted(	4✔
116	validator.iter_errors(instance), key=lambda e: list(e.absolute_path)
117	):
118	location = "/".join(str(p) for p in error.absolute_path) or "<root>"	4✔
119	errors.append(f"{location}: {error.message}")	4✔
120	return errors	4✔
121
122
123	@dataclass(frozen=True)	4✔
124	class SourcePos:	4✔
125	line: int	4✔
126	col: int	4✔
127
128	@classmethod	4✔
129	def from_dict(cls, d: dict) -> "SourcePos":	4✔
130	return cls(line=d["line"], col=d["col"])	4✔
131
132
133	@dataclass(frozen=True)	4✔
134	class SourceRange:	4✔
135	file: str	4✔
136	begin: SourcePos	4✔
137	end: SourcePos	4✔
138
139	@classmethod	4✔
140	def from_dict(cls, d: dict) -> "SourceRange":	4✔
141	return cls(	4✔
142	file=d["file"],
143	begin=SourcePos.from_dict(d["from"]),
144	end=SourcePos.from_dict(d["to"]),
145	)
146
147
148	@dataclass(frozen=True)	4✔
149	class LoopInfo:	4✔
150	loopnest_index: int	4✔
151	num_loops: int	4✔
152	num_maps: int	4✔
153	num_fors: int	4✔
154	num_whiles: int	4✔
155	max_depth: int	4✔
156	is_perfectly_nested: bool	4✔
157	is_perfectly_parallel: bool	4✔
158	is_elementwise: bool	4✔
159	has_side_effects: bool	4✔
160
161	@classmethod	4✔
162	def from_dict(cls, d: dict) -> "LoopInfo":	4✔
163	return cls(	4✔
164	loopnest_index=d["loopnest_index"],
165	num_loops=d["num_loops"],
166	num_maps=d["num_maps"],
167	num_fors=d["num_fors"],
168	num_whiles=d["num_whiles"],
169	max_depth=d["max_depth"],
170	is_perfectly_nested=d["is_perfectly_nested"],
171	is_perfectly_parallel=d["is_perfectly_parallel"],
172	is_elementwise=d["is_elementwise"],
173	has_side_effects=d["has_side_effects"],
174	)
175
176
177	@dataclass(frozen=True)	4✔
178	class MetricStat:	4✔
179	"""Aggregated statistics for a single counter (aggregated traces).
180
181	``mean``/``variance``/``min``/``max`` may be ``None`` for provided
182	(``static:::``) metrics whose value was non-finite at capture time.
183	"""
184
185	mean: Optional[float]	4✔
186	variance: Optional[float]	4✔
187	count: int	4✔
188	min: Optional[float]	4✔
189	max: Optional[float]	4✔
190
191	@classmethod	4✔
192	def from_dict(cls, d: dict) -> "MetricStat":	4✔
193	return cls(	4✔
194	mean=d["mean"],
195	variance=d["variance"],
196	count=d["count"],
197	min=d["min"],
198	max=d["max"],
199	)
200
201
202	class TraceRegion:	4✔
203	"""A single region event, wrapping the raw trace dict with typed accessors."""
204
205	def __init__(self, raw: dict):	4✔
206	self._raw = raw	4✔
207
208	# -- Raw / identity ----------------------------------------------------
209	@property	4✔
210	def raw(self) -> dict:	4✔
211	return self._raw	4✔
212
213	@property	4✔
214	def category(self) -> str:	4✔
215	return self._raw["cat"]	4✔
216
217	@property	4✔
218	def is_aggregated(self) -> bool:	4✔
219	return self.category == AGGREGATED_CAT	4✔
220
221	@property	4✔
222	def name(self) -> str:	4✔
223	return self._raw["name"]	4✔
224
225	@property	4✔
226	def pid(self) -> int:	4✔
227	return self._raw["pid"]	4✔
228
229	@property	4✔
230	def tid(self) -> int:	4✔
UNCOV 231	return self._raw["tid"]	×
232
233	@property	4✔
234	def ts_us(self) -> Optional[float]:	4✔
UNCOV 235	return self._raw.get("ts")	×
236
237	@property	4✔
238	def dur_us(self) -> Optional[float]:	4✔
239	return self._raw.get("dur")	4✔
240
241	# -- Source metadata ---------------------------------------------------
242	@property	4✔
243	def _args(self) -> dict:	4✔
244	return self._raw["args"]	4✔
245
246	@property	4✔
247	def function(self) -> str:	4✔
248	return self._args["function"]	4✔
249
250	@property	4✔
251	def module(self) -> str:	4✔
252	return self._args["module"]	4✔
253
254	@property	4✔
255	def target_type(self) -> str:	4✔
256	return self._args.get("target_type", "")	4✔
257
258	@cached_property	4✔
259	def source_ranges(self) -> List[SourceRange]:	4✔
260	return [SourceRange.from_dict(r) for r in self._args.get("source_ranges", [])]	4✔
261
262	# -- docc metadata -----------------------------------------------------
263	@property	4✔
264	def _docc(self) -> dict:	4✔
265	return self._args["docc"]	4✔
266
267	@property	4✔
268	def sdfg_name(self) -> str:	4✔
269	return self._docc.get("sdfg_name", "")	4✔
270
271	@property	4✔
272	def sdfg_file(self) -> str:	4✔
UNCOV 273	return self._docc.get("sdfg_file", "")	×
274
275	@property	4✔
276	def arg_capture_path(self) -> str:	4✔
UNCOV 277	return self._docc.get("arg_capture_path", "")	×
278
279	@property	4✔
280	def features_file(self) -> str:	4✔
UNCOV 281	return self._docc.get("features_file", "")	×
282
283	@property	4✔
284	def opt_report_file(self) -> str:	4✔
UNCOV 285	return self._docc.get("opt_report_file", "")	×
286
287	@property	4✔
288	def has_element(self) -> bool:	4✔
289	"""True when SDFG element metadata (id/type/loop_info) is present."""
290	return "element_id" in self._docc	4✔
291
292	@property	4✔
293	def element_id(self) -> Optional[int]:	4✔
294	return self._docc.get("element_id")	4✔
295
296	@property	4✔
297	def element_type(self) -> Optional[str]:	4✔
298	return self._docc.get("element_type")	4✔
299
300	@property	4✔
301	def loopnest_index(self) -> Optional[int]:	4✔
UNCOV 302	return self._docc.get("loopnest_index")	×
303
304	@cached_property	4✔
305	def loop_info(self) -> Optional[LoopInfo]:	4✔
306	li = self._docc.get("loop_info")	4✔
307	return LoopInfo.from_dict(li) if li is not None else None	4✔
308
309	# -- Metrics -----------------------------------------------------------
310	@property	4✔
311	def _metrics(self) -> dict:	4✔
312	return self._args.get("metrics", {})	4✔
313
314	@cached_property	4✔
315	def counters(self) -> Dict[str, int]:	4✔
316	"""PAPI counter values for a per-invocation region (empty if aggregated)."""
317	if self.is_aggregated:	4✔
UNCOV 318	return {}	×
319	return dict(self._metrics)	4✔
320
321	@cached_property	4✔
322	def counter_stats(self) -> Dict[str, MetricStat]:	4✔
323	"""Aggregated PAPI counter statistics (empty for per-invocation regions).
324
325	Excludes the reserved ``runtime`` block and ``static:::`` provided metrics.
326	"""
327	if not self.is_aggregated:	4✔
328	return {}	4✔
329	out: Dict[str, MetricStat] = {}	4✔
330	for key, val in self._metrics.items():	4✔
331	if key == _RUNTIME_KEY or key.startswith(STATIC_PREFIX):	4✔
332	continue	4✔
333	out[key] = MetricStat.from_dict(val)	4✔
334	return out	4✔
335
336	@cached_property	4✔
337	def static_metrics(self) -> Dict[str, MetricStat]:	4✔
338	"""User-provided (``static:::``) metric statistics, keyed without the prefix."""
339	if not self.is_aggregated:	4✔
UNCOV 340	return {}	×
341	out: Dict[str, MetricStat] = {}	4✔
342	for key, val in self._metrics.items():	4✔
343	if key.startswith(STATIC_PREFIX):	4✔
344	out[key[len(STATIC_PREFIX) :]] = MetricStat.from_dict(val)	4✔
345	return out	4✔
346
347	@cached_property	4✔
348	def runtime(self) -> Optional[MetricStat]:	4✔
349	"""Aggregated runtime statistics in microseconds (aggregated traces only)."""
350	rt = self._metrics.get(_RUNTIME_KEY) if self.is_aggregated else None	4✔
351	return MetricStat.from_dict(rt) if rt is not None else None	4✔
352
353	@property	4✔
354	def runtime_mean_us(self) -> Optional[float]:	4✔
355	"""Mean runtime in microseconds.
356
357	For aggregated traces this is ``runtime.mean``; for per-invocation traces
358	it is this event's own ``dur``.
359	"""
360	if self.is_aggregated:	4✔
361	return self.runtime.mean if self.runtime is not None else None	4✔
362	return self.dur_us	4✔
363
364	@property	4✔
365	def runtime_min_us(self) -> Optional[float]:	4✔
366	"""Best (minimum) runtime in microseconds.
367
368	The min over the aggregated samples is the steady-state estimate.
369	"""
UNCOV 370	if self.is_aggregated:	×
UNCOV 371	return self.runtime.min if self.runtime is not None else None	×
372	return self.dur_us	×
373
374	def metric(self, name: str) -> Optional[Union[int, float]]:	4✔
375	"""Return a representative value for ``name``.
376
377	For per-invocation regions this is the raw counter value; for aggregated
378	regions it is the counter/static metric mean. Returns ``None`` if absent.
379	"""
380	if not self.is_aggregated:	4✔
381	return self.counters.get(name)	4✔
382	if name in self.counter_stats:	4✔
383	return self.counter_stats[name].mean	4✔
384	if name in self.static_metrics:	4✔
385	return self.static_metrics[name].mean	4✔
UNCOV 386	return None	×
387
388	def __repr__(self) -> str:	4✔
UNCOV 389	return (	×
390	f"TraceRegion(name={self.name!r}, module={self.module!r}, "
391	f"target={self.target_type!r}, aggregated={self.is_aggregated})"
392	)
393
394
395	class Trace:	4✔
396	"""A parsed, validated daisy trace: a collection of :class:`TraceRegion`."""
397
398	def __init__(self, raw: dict):	4✔
399	self._raw = raw	4✔
400	self._regions = [TraceRegion(e) for e in raw.get("traceEvents", [])]	4✔
401
402	# -- Construction ------------------------------------------------------
403	@classmethod	4✔
404	def from_dict(cls, data: dict, validate_schema: bool = True) -> "Trace":	4✔
405	"""Build a :class:`Trace` from an already-parsed dict.
406
407	Raises :class:`TraceValidationError` if ``validate_schema`` is set and the
408	data does not conform to the schema.
409	"""
410	if validate_schema:	4✔
411	errors = validate(data)	4✔
412	if errors:	4✔
413	raise TraceValidationError(errors)	4✔
414	return cls(data)	4✔
415
416	@classmethod	4✔
417	def load(cls, path: Union[str, Path], validate_schema: bool = True) -> "Trace":	4✔
418	"""Read a trace JSON file, optionally validate it, and return a :class:`Trace`."""
419	with open(path, "r") as f:	×
UNCOV 420	data = json.load(f)	×
UNCOV 421	return cls.from_dict(data, validate_schema=validate_schema)	×
422
423	def save(self, path: Union[str, Path], indent: Optional[int] = None) -> None:	4✔
424	"""Write this trace back to a JSON file (the daisy-trace format)."""
425	with open(path, "w") as f:	×
UNCOV 426	json.dump(self._raw, f, indent=indent)	×
427
428	# -- Container protocol ------------------------------------------------
429	@property	4✔
430	def regions(self) -> List[TraceRegion]:	4✔
431	return self._regions	4✔
432
433	@property	4✔
434	def raw(self) -> dict:	4✔
435	return self._raw	4✔
436
437	def __len__(self) -> int:	4✔
438	return len(self._regions)	4✔
439
440	def __iter__(self) -> Iterator[TraceRegion]:	4✔
UNCOV 441	return iter(self._regions)	×
442
443	def __getitem__(self, index: int) -> TraceRegion:	4✔
444	return self._regions[index]	4✔
445
446	@property	4✔
447	def is_aggregated(self) -> bool:	4✔
448	"""True if the trace was captured in aggregate mode.
449
450	Determined from the first region; a valid trace never mixes modes.
451	"""
452	return bool(self._regions) and self._regions[0].is_aggregated	4✔
453
454	# -- Queries -----------------------------------------------------------
455	def filter(	4✔
456	self,
457	*,
458	module: Optional[str] = None,
459	function: Optional[str] = None,
460	target_type: Optional[str] = None,
461	element_type: Optional[str] = None,
462	sdfg_name: Optional[str] = None,
463	) -> List[TraceRegion]:
464	"""Return regions matching all provided (non-``None``) criteria."""
465	out = []	4✔
466	for r in self._regions:	4✔
467	if module is not None and r.module != module:	4✔
UNCOV 468	continue	×
469	if function is not None and r.function != function:	4✔
UNCOV 470	continue	×
471	if target_type is not None and r.target_type != target_type:	4✔
472	continue	4✔
473	if element_type is not None and r.element_type != element_type:	4✔
UNCOV 474	continue	×
475	if sdfg_name is not None and r.sdfg_name != sdfg_name:	4✔
UNCOV 476	continue	×
477	out.append(r)	4✔
478	return out	4✔
479
480	def by_element_id(self, element_id: int) -> List[TraceRegion]:	4✔
481	"""Return regions carrying the given SDFG ``element_id``."""
482	return [r for r in self._regions if r.element_id == element_id]	4✔
483
484	def hottest(self, n: Optional[int] = None) -> List[TraceRegion]:	4✔
485	"""Regions sorted by mean runtime (descending); ``n`` limits the count."""
486	ranked = sorted(	4✔
487	(r for r in self._regions if r.runtime_mean_us is not None),
488	key=lambda r: r.runtime_mean_us,
489	reverse=True,
490	)
491	return ranked[:n] if n is not None else ranked	4✔
492
493	def total_runtime_us(self) -> float:	4✔
494	"""Sum of mean runtimes across all regions (microseconds)."""
495	return sum((r.runtime_mean_us or 0.0) for r in self._regions)	4✔
496
497	# -- Combining ---------------------------------------------------------
498	@staticmethod	4✔
499	def _region_key(region: "TraceRegion") -> tuple:	4✔
500	"""Stable identity of a region across separate runs of the same workload.
501
502	Matches on source/docc identity and excludes pid/tid/timestamps, which
503	differ between runs.
504	"""
505	return (	4✔
506	region.name,
507	region.function,
508	region.module,
509	region.sdfg_name,
510	region.element_id,
511	)
512
513	@classmethod	4✔
514	def combine(	4✔
515	cls, traces: Sequence["Trace"], validate_schema: bool = False
516	) -> "Trace":
517	"""Combine multiple traces of the same workload into one.
518
519	Intended for counter multiplexing: when the hardware can only measure a
520	subset of counters at once, the workload is run once per counter group,
521	producing several traces that each carry a slice of the ``metrics``. This
522	takes the first trace as the reference for regions and timesteps (name,
523	ts, dur, source/docc metadata, runtime) and unions the per-region
524	``metrics`` maps from the remaining traces into it.
525
526	Regions are matched across traces by :meth:`_region_key` (source/docc
527	identity, ignoring pid/tid/timestamps). On a metric-key collision the
528	reference wins, so its ``runtime`` and any ``static:::`` values are kept.
529	Returns a new :class:`Trace`; the inputs are not mutated.
530	"""
531	traces = list(traces)	4✔
532	if not traces:	4✔
533	raise ValueError("Trace.combine() requires at least one trace")	4✔
534	if len(traces) == 1:	4✔
535	return cls.from_dict(	4✔
536	copy.deepcopy(traces[0].raw), validate_schema=validate_schema
537	)
538
539	# Index each additional trace's per-region metrics by region key.
540	extra_metrics: Dict[tuple, List[dict]] = {}	4✔
541	for other in traces[1:]:	4✔
542	for region in other.regions:	4✔
543	metrics = region.raw.get("args", {}).get("metrics", {})	4✔
544	extra_metrics.setdefault(cls._region_key(region), []).append(metrics)	4✔
545
546	# Deep-copy the reference and union in the other traces' metrics.
547	combined_raw = copy.deepcopy(traces[0].raw)	4✔
548	for event in combined_raw.get("traceEvents", []):	4✔
549	key = cls._region_key(TraceRegion(event))	4✔
550	metrics = event.setdefault("args", {}).setdefault("metrics", {})	4✔
551	for other_metrics in extra_metrics.get(key, []):	4✔
552	for name, value in other_metrics.items():	4✔
553	metrics.setdefault(name, value)	4✔
554
555	return cls.from_dict(combined_raw, validate_schema=validate_schema)	4✔

daisytuner / docc / 30476914782

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous