diff --git a/.coveragerc b/.coveragerc index c4bb7d51c..23f815467 100644 --- a/.coveragerc +++ b/.coveragerc @@ -6,6 +6,8 @@ omit = TraceLens/PerfModel/benchmarking/* TraceLens/PerfModel/origami_helper.py TraceLens/PerfModel/run_perf_model.py + # pattern_finder.py is slated for replacement by an external package; omit until then + TraceLens/Agent/Analysis/semantic_analyses/pattern_finder.py [paths] source = diff --git a/TraceLens/Agent/Analysis/semantic_analyses/__init__.py b/TraceLens/Agent/Analysis/semantic_analyses/__init__.py new file mode 100644 index 000000000..35a14a24b --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""TraceLens Agent Semantic Comparison Scripts""" diff --git a/TraceLens/Agent/Analysis/semantic_analyses/_helpers.py b/TraceLens/Agent/Analysis/semantic_analyses/_helpers.py new file mode 100644 index 000000000..8a90b37fd --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/_helpers.py @@ -0,0 +1,75 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared helper functions for the semantic_analyses scripts.""" + +import gzip +import json + + +def build_rle(kernel_indices, cls_by_idx): + """Run-length encode kernel indices by perf_category. + + Returns list of (perf_category, count, [kernel_indices], [kernel_types]). + """ + if not kernel_indices: + return [] + + def _cls(idx): + c = cls_by_idx.get(idx, {}) + return c.get("perf_category", "Others"), c.get("kernel_type", "Unknown") + + first_cat, first_kt = _cls(kernel_indices[0]) + groups = [] + cur_cat = first_cat + cur_indices = [kernel_indices[0]] + cur_types = [first_kt] + + for idx in kernel_indices[1:]: + cat, kt = _cls(idx) + if cat == cur_cat: + cur_indices.append(idx) + cur_types.append(kt) + else: + groups.append((cur_cat, len(cur_indices), cur_indices[:], cur_types[:])) + cur_cat = cat + cur_indices = [idx] + cur_types = [kt] + + groups.append((cur_cat, len(cur_indices), cur_indices[:], cur_types[:])) + return groups + + +def detect_period(rle_groups): + """Find the shortest repeating period in an RLE group sequence. + + Returns the period length (number of RLE groups per super-cycle). + """ + cats = [g[0] for g in rle_groups] + n = len(cats) + if n < 6: + return n + + for p in range(3, n // 2 + 1): + prefix = cats[:p] + matches = sum(1 for i in range(p, n) if cats[i] == prefix[i % p]) + total = n - p + if total > 0 and matches / total > 0.85 and total >= 2 * p: + return p + + return n + + +def load_json(path): + """Load a JSON file, transparently decompressing .gz input.""" + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt") as f: + return json.load(f) + + +def load_labels(path): + """Load a semantic_labels.json file.""" + return load_json(path) diff --git a/TraceLens/Agent/Analysis/semantic_analyses/annotation_metadata.py b/TraceLens/Agent/Analysis/semantic_analyses/annotation_metadata.py new file mode 100644 index 000000000..805ef370f --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/annotation_metadata.py @@ -0,0 +1,199 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Multi-source metadata gathering for roofline analysis. + +Gathers metadata from annotation, filename (isl/osl/conc), trace Input Dims, +and user input. Runs sanity checks and merges into a unified dict. +""" + +import logging +import re +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + +# Filename regex: isl1024, osl8, conc4, tp1 +FILENAME_ISL_RE = re.compile(r"isl(\d+)", re.I) +FILENAME_OSL_RE = re.compile(r"osl(\d+)", re.I) +FILENAME_CONC_RE = re.compile(r"conc(\d+)", re.I) +FILENAME_TP_RE = re.compile(r"tp(\d+)", re.I) + + +def parse_filename_metadata(filepath: str) -> Dict[str, Any]: + """ + Parse isl, osl, conc, tp from trace filename. + Example: mi355_tp1_isl1024_osl8_conc4_opt_asm64x256.pt.trace.json.gz + """ + basename = filepath.split("/")[-1] if "/" in filepath else filepath + result = {} + m = FILENAME_ISL_RE.search(basename) + if m: + result["isl"] = int(m.group(1)) + m = FILENAME_OSL_RE.search(basename) + if m: + result["osl"] = int(m.group(1)) + m = FILENAME_CONC_RE.search(basename) + if m: + result["conc"] = int(m.group(1)) + m = FILENAME_TP_RE.search(basename) + if m: + result["tp"] = int(m.group(1)) + if "isl" in result and "conc" in result: + result["num_tokens_prefill"] = result["isl"] * result["conc"] + if "conc" in result: + result["num_tokens_decode"] = result["conc"] + return result + + +def parse_trace_input_dims(events: List[dict]) -> Dict[str, Any]: + """ + Aggregate N_Q, N_KV from kernel Input Dims (for attention kernels). + Returns block-level aggregates when available. + For multi-layer models, each layer has one attention kernel with same N_Q/N_KV. + """ + nq_values = [] + nkv_values = [] + for e in events: + if e.get("cat") != "kernel": + continue + name = e.get("name", "") + if "unified_attention" not in name and "attention" not in name.lower(): + continue + dims = e.get("args", {}).get("Input Dims") + if not dims or len(dims) < 2: + continue + try: + q_shape = dims[0] + k_shape = dims[1] + if isinstance(q_shape, (list, tuple)) and len(q_shape) >= 3: + nq = q_shape[0] if len(q_shape) == 3 else q_shape[-3] + nq_values.append(nq) + if isinstance(k_shape, (list, tuple)) and len(k_shape) >= 3: + nkv = k_shape[-3] if len(k_shape) >= 3 else k_shape[0] + nkv_values.append(nkv) + except (IndexError, TypeError): + continue + result = {} + if nq_values: + result["trace_avg_nq"] = sum(nq_values) // len(nq_values) + result["trace_attention_kernel_count"] = len(nq_values) + if nkv_values: + result["trace_avg_nkv"] = sum(nkv_values) // len(nkv_values) + return result + + +def run_sanity_checks( + annotation_meta: Dict[str, Any], + filename_meta: Dict[str, Any], + trace_meta: Dict[str, Any], + user_meta: Dict[str, Any], +) -> List[str]: + """Cross-validate metadata from different sources. Returns list of warning messages.""" + warnings = [] + batch_ann = annotation_meta.get("batch_size") + batch_file = None + if "isl" in filename_meta and "conc" in filename_meta: + batch_file = filename_meta["isl"] * filename_meta["conc"] + if batch_ann is not None and batch_file is not None: + if abs(batch_ann - batch_file) > max(1, 0.1 * batch_ann): + warnings.append( + f"Batch size mismatch: annotation={batch_ann}, filename(isl*conc)={batch_file}" + ) + num_tokens_user = user_meta.get("num_tokens") + if num_tokens_user is not None and batch_ann is not None: + if abs(num_tokens_user - batch_ann) > max(1, 0.1 * batch_ann): + warnings.append( + f"num_tokens mismatch: user={num_tokens_user}, annotation batch_size={batch_ann}" + ) + ctx_ann = annotation_meta.get("context_sum") + ctx_trace = trace_meta.get("trace_avg_nq") + if ctx_ann is not None and ctx_trace is not None: + if abs(ctx_ann - ctx_trace) > max(1, 0.2 * ctx_ann): + warnings.append( + f"Context sum mismatch: annotation={ctx_ann}, trace avg N_Q={ctx_trace}" + ) + return warnings + + +def merge_metadata( + annotation_meta: Optional[Dict[str, Any]] = None, + filename_meta: Optional[Dict[str, Any]] = None, + trace_meta: Optional[Dict[str, Any]] = None, + user_meta: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """ + Merge metadata from all sources. Prefer annotation for per-iteration, + filename for run config, trace for per-kernel accuracy, user for overrides. + """ + merged = {} + annotation_meta = annotation_meta or {} + filename_meta = filename_meta or {} + trace_meta = trace_meta or {} + user_meta = user_meta or {} + merged["annotation"] = annotation_meta + merged["filename"] = filename_meta + merged["trace"] = trace_meta + merged["user"] = user_meta + num_tokens = ( + user_meta.get("num_tokens") + or annotation_meta.get("batch_size") + or filename_meta.get("num_tokens_prefill") + or filename_meta.get("num_tokens_decode") + ) + context_length = ( + user_meta.get("context_length") + or annotation_meta.get("context_sum") + or annotation_meta.get("generation_sum") + or filename_meta.get("isl") + or num_tokens + ) + merged["num_tokens"] = num_tokens + merged["context_length"] = context_length + merged["batch_size"] = annotation_meta.get("batch_size") or num_tokens + merged["context_sum"] = annotation_meta.get("context_sum") + merged["generation_sum"] = annotation_meta.get("generation_sum") + warnings = run_sanity_checks(annotation_meta, filename_meta, trace_meta, user_meta) + for w in warnings: + logger.warning("Metadata sanity check: %s", w) + merged["_warnings"] = warnings + return merged + + +def gather_metadata( + trace_path: str, + events: Optional[List[dict]] = None, + annotation_meta: Optional[Dict[str, Any]] = None, + num_tokens: Optional[int] = None, + context_length: Optional[int] = None, +) -> Dict[str, Any]: + """ + Gather metadata from all available sources. + + Args: + trace_path: Path to trace file (for filename parsing) + events: Trace events (for Input Dims); if None, trace Input Dims are skipped + annotation_meta: Pre-parsed annotation metadata (e.g., from vllm_trace_split) + num_tokens: User-provided num_tokens + context_length: User-provided context_length + + Returns: + Merged metadata dict with num_tokens, context_length, batch_size, etc. + """ + filename_meta = parse_filename_metadata(trace_path) + trace_meta = parse_trace_input_dims(events or []) + user_meta = {} + if num_tokens is not None: + user_meta["num_tokens"] = num_tokens + if context_length is not None: + user_meta["context_length"] = context_length + return merge_metadata( + annotation_meta=annotation_meta, + filename_meta=filename_meta, + trace_meta=trace_meta, + user_meta=user_meta, + ) diff --git a/TraceLens/Agent/Analysis/semantic_analyses/build_semantic_labels.py b/TraceLens/Agent/Analysis/semantic_analyses/build_semantic_labels.py new file mode 100644 index 000000000..02320726c --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/build_semantic_labels.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Build semantic_labels.json deterministically from breakdown artifacts. + +Combines extracted.json + classified.json + pattern.json +into the labeled kernel format consumed by the comparison pipeline. + +Each kernel gets a positionally-indexed ``semantic_block`` name (e.g. +``GEMM_0``, ``GEMM_1``, ``Normalization_0``) derived from its position +within the repeating layer cycle. A global counter ensures block names +are unique across regions (pre-layer, body, post-layer, secondary). +The harmonization agent later renames these to descriptive labels +(e.g. ``QKV Projection``). + +Per-kernel output fields: + index, name, dur, gpu_op_uid, perf_category, semantic_block, region, + nn_module, cpu_op, input_dims, layer. + +Usage: + python build_semantic_labels.py extracted.json classified.json pattern.json \ + [-o semantic_labels.json] +""" + +import argparse +import json +import sys + +from _helpers import build_rle, detect_period + + +def _build_cycle_names(body_rle, period, cat_counter): + """Build positionally-indexed block names for one layer cycle. + + Uses *cat_counter* (mutated in-place) so indices are globally unique + across regions. Returns a list of length ``period`` like + ``["GEMM_2", "Normalization_1", "GEMM_3", "SDPA_0", ...]``. + """ + if period <= 0: + return [] + cycle_names = [] + for cat, _count, _indices, _types in body_rle[:period]: + idx = cat_counter.get(cat, 0) + cat_counter[cat] = idx + 1 + cycle_names.append(f"{cat}_{idx}") + return cycle_names + + +def _build_region_block_names(index_set, cls_by_idx, cat_counter): + """Build indexed block names for a region of kernels. + + Groups consecutive kernels by perf_category and assigns globally + unique ``Category_N`` names using *cat_counter* (mutated in-place). + Returns a mapping ``{kernel_index: "Category_N"}``. + """ + if not index_set: + return {} + sorted_indices = sorted(index_set) + groups = [] + cur_cat = cls_by_idx.get(sorted_indices[0], {}).get("perf_category", "Others") + cur_group = [sorted_indices[0]] + for idx in sorted_indices[1:]: + cat = cls_by_idx.get(idx, {}).get("perf_category", "Others") + if cat == cur_cat: + cur_group.append(idx) + else: + groups.append((cur_cat, cur_group)) + cur_cat = cat + cur_group = [idx] + groups.append((cur_cat, cur_group)) + + result = {} + for cat, indices in groups: + n = cat_counter.get(cat, 0) + cat_counter[cat] = n + 1 + block_name = f"{cat}_{n}" + for idx in indices: + result[idx] = block_name + return result + + +def build_labels(extracted, classified, pattern): + """Build semantic_labels.json from breakdown artifacts.""" + kernels = extracted["kernels"] + total_kernels = len(kernels) + + cls_by_idx = {c["index"]: c for c in classified["classified_kernels"]} + + preamble_set = set(pattern.get("preamble_indices", [])) + epilogue_set = set(pattern.get("epilogue_indices", [])) + secondary_set = set(pattern.get("secondary_stream_indices", [])) + body_indices = [ + i + for i in range(total_kernels) + if i not in preamble_set and i not in epilogue_set and i not in secondary_set + ] + + body_rle = build_rle(body_indices, cls_by_idx) + period = detect_period(body_rle) + num_layers = len(body_rle) // period if period > 0 else 0 + + cat_counter = {} + + preamble_blocks = _build_region_block_names(preamble_set, cls_by_idx, cat_counter) + cycle_names = _build_cycle_names(body_rle, period, cat_counter) + epilogue_blocks = _build_region_block_names(epilogue_set, cls_by_idx, cat_counter) + secondary_blocks = _build_region_block_names(secondary_set, cls_by_idx, cat_counter) + + body_index_to_rle_group = {} + for g_idx, (_cat, _count, indices, _types) in enumerate(body_rle): + for idx in indices: + body_index_to_rle_group[idx] = g_idx + + labeled_kernels = [] + for i in range(total_kernels): + k = kernels[i] + c = cls_by_idx.get(i, {}) + + entry = { + "index": i, + "name": k["name"], + "dur": k["dur"], + # Raw-index UID stamped by extract_trace_data.py; aligns with + # the perf report's kernel_details gpu_op_uid. + "gpu_op_uid": k.get("gpu_op_uid"), + "perf_category": c.get("perf_category", "Others"), + # nn_module / cpu_op / input_dims need cpu_op ancestry from a + # trace tree, unavailable for graph-mode / no-capture traces. + # Left empty; downstream falls back to semantic_block. + "nn_module": "", + "cpu_op": "", + "input_dims": [], + } + + if i in body_index_to_rle_group: + g_idx = body_index_to_rle_group[i] + entry["region"] = "body" + entry["layer"] = g_idx // period if period > 0 else 0 + entry["semantic_block"] = cycle_names[g_idx % period] + elif i in preamble_blocks: + entry["region"] = "pre" + entry["layer"] = None + entry["semantic_block"] = preamble_blocks[i] + elif i in epilogue_blocks: + entry["region"] = "post" + entry["layer"] = None + entry["semantic_block"] = epilogue_blocks[i] + elif i in secondary_blocks: + entry["region"] = "secondary" + entry["layer"] = None + entry["semantic_block"] = secondary_blocks[i] + else: # pragma: no cover - defensive: every index is body/pre/post/secondary + entry["region"] = "body" + entry["layer"] = None + entry["semantic_block"] = c.get("perf_category", "Others") + "_0" + + labeled_kernels.append(entry) + + result = { + "source_file": extracted.get("source_file", ""), + "total_kernel_time_us": round( + extracted.get("metadata", {}).get("total_kernel_time_us", 0), 2 + ), + "model_info": { + "num_layers": num_layers, + "period": period, + "graph_mode": extracted.get("metadata", {}).get("is_graph_mode", False), + }, + "labeled_kernels": labeled_kernels, + } + return result + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser( + description="Build semantic_labels.json deterministically from breakdown artifacts" + ) + parser.add_argument("extracted_json", help="Path to extracted.json") + parser.add_argument("classified_json", help="Path to classified.json") + parser.add_argument("pattern_json", help="Path to pattern.json") + parser.add_argument("-o", "--output", help="Output JSON path (default: stdout)") + args = parser.parse_args() + + with open(args.extracted_json) as f: + extracted = json.load(f) + with open(args.classified_json) as f: + classified = json.load(f) + with open(args.pattern_json) as f: + pattern = json.load(f) + + result = build_labels(extracted, classified, pattern) + + output = json.dumps(result, indent=2) + if args.output: + with open(args.output, "w") as f: + f.write(output) + n = len(result["labeled_kernels"]) + info = result["model_info"] + print( + f"Wrote {args.output} ({n} kernels, {info['num_layers']} layers, " + f"period {info['period']})", + file=sys.stderr, + ) + else: + print(output) + + +if __name__ == "__main__": + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/extract_trace_data.py b/TraceLens/Agent/Analysis/semantic_analyses/extract_trace_data.py new file mode 100644 index 000000000..ef9b5d632 --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/extract_trace_data.py @@ -0,0 +1,345 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Step 1+3: Load a Chrome trace JSON and extract structured data. + +Outputs a JSON with: + - ordered kernel list (name, duration, timestamp, gpu_op_uid) + - metadata (categories, graph mode detection, total kernel time) + +vLLM traces with annotation iterations are auto-detected and split into +per-region subdirectories. Pass ``--no-split`` to force single-file mode. + +Usage: + python extract_trace_data.py -o output_dir/ +""" + +import argparse +import json +import logging +import os +import sys +from collections import Counter + +from trace_split_adapter import split_vllm_trace, get_steady_state_key +from annotation_metadata import gather_metadata +from _helpers import load_json + +from TraceLens import GPUEventAnalyser + +logger = logging.getLogger(__name__) + + +def load_trace(path_or_data): + """Load trace from path (str) or use dict directly.""" + if isinstance(path_or_data, dict): + data = path_or_data + else: + data = load_json(path_or_data) + events = data.get("traceEvents", []) + by_cat = {} + for e in events: + if not isinstance(e, dict): + continue + cat = e.get("cat", "unknown") + by_cat.setdefault(cat, []).append(e) + for cat in by_cat: + by_cat[cat].sort(key=lambda e: e.get("ts", 0)) + return data, by_cat + + +def _stamp_raw_uid(data): + """Tag each event with its position in the raw traceEvents array. + + This mirrors the UID scheme TraceToTree/TreePerfAnalyzer assign + (enumerate() over the full, unfiltered traceEvents list), so a kernel's + ``_gpu_op_uid`` here lines up with the same kernel's UID in a perf report + built from the same trace file -- without ever building a tree. Only + valid for a single, non-split trace file; do not call this on a + per-region slice produced by vLLM trace splitting, since a region's + traceEvents subset does not share the original file's indexing. + """ + for i, e in enumerate(data.get("traceEvents", [])): + if isinstance(e, dict): + e["_gpu_op_uid"] = i + + +def get_stream_id(event): + """Get stream from args['stream'] or tid (fallback for magic-trace style).""" + stream = event.get("args", {}).get("stream") + if stream is not None: + try: + return int(stream) + except (TypeError, ValueError): + pass + tid = event.get("tid") + if tid is not None: + try: + return int(tid) + except (TypeError, ValueError): + pass + return None + + +def filter_to_primary_stream(by_cat): + """If multiple streams in kernel events, keep only primary (most kernels). + + Skips filtering when secondary streams carry significant compute (>5% + of total kernel time), since some runtimes schedule MoE / communication + kernels on secondary CUDA streams. + """ + kernels = by_cat.get("kernel", []) + if not kernels: + return + stream_counts = Counter( + get_stream_id(k) for k in kernels if get_stream_id(k) is not None + ) + if len(stream_counts) <= 1: + return + total_time = sum(k.get("dur", 0) for k in kernels) + primary = max(stream_counts, key=stream_counts.get) + secondary_time = sum( + k.get("dur", 0) + for k in kernels + if get_stream_id(k) != primary and get_stream_id(k) is not None + ) + if total_time > 0 and secondary_time / total_time > 0.05: + logger.info( + "Keeping all %d streams: secondary streams have %.1f%% of kernel time", + len(stream_counts), + 100 * secondary_time / total_time, + ) + return + by_cat["kernel"] = [k for k in kernels if get_stream_id(k) == primary] + + +def extract_kernel_sequence(by_cat): + kernels = by_cat.get("kernel", []) + memcpy = by_cat.get("gpu_memcpy", []) + combined = sorted(kernels + memcpy, key=lambda e: e["ts"]) + return [ + { + "name": k["name"], + "cat": k.get("cat", "kernel"), + "dur": k["dur"], + "ts": k["ts"], + "args": k.get("args", {}), + "stream_id": get_stream_id(k), + # Single-trace path: _stamp_raw_uid() sets "_gpu_op_uid" in-memory. + # Split path: split_inference_trace_annotation.py (invoked with + # --emit-gpu-op-uid) already persisted "gpu_op_uid" onto the event + # before writing the per-region file, so it survives the subprocess + # boundary as literal JSON content. + "gpu_op_uid": k.get("_gpu_op_uid", k.get("gpu_op_uid")), + } + for k in combined + ] + + +def detect_graph_mode(by_cat): + rt = by_cat.get("cuda_runtime", []) + graph_launches = [e for e in rt if "GraphLaunch" in e.get("name", "")] + return len(graph_launches) > 0, graph_launches + + +def run_assertions(data, by_cat, kernels, is_graph_mode, strict=True): + errors = [] + + if "traceEvents" not in data: + errors.append("A1.1 FAIL: Missing traceEvents key") + + required_cats = {"kernel", "cpu_op"} if strict else {"kernel"} + missing = required_cats - set(by_cat.keys()) + if missing: + errors.append(f"A1.2 FAIL: Missing categories: {missing}") + + if len(kernels) == 0: + errors.append("A1.3 FAIL: No GPU kernels found") + + for i, k in enumerate(kernels): + if k["dur"] <= 0: + errors.append( + f"A3.2 FAIL: Kernel {i} ({k['name'][:50]}) has non-positive duration {k['dur']}" + ) + break + + timestamps = [k["ts"] for k in kernels] + for i in range(1, len(timestamps)): + if timestamps[i] < timestamps[i - 1]: + errors.append(f"A3.1 FAIL: Kernel timestamps not monotonic at index {i}") + break + + total_time = sum(k["dur"] for k in kernels) + if total_time <= 0: + errors.append("A1.5 FAIL: Zero total kernel time") + + return errors + + +def compute_gpu_timeline_metrics(events): # pragma: no cover + """ + Run GPUEventAnalyser on events and return gpu_timeline dict for metadata. + Returns None on failure. + """ + try: + # GPUEventAnalyser needs a unique UID per event for overlap + # computation. Raw trace events don't carry one, so assign a contiguous + # 0..N-1. If an event already has a "UID", our assumption is broken and + # uniqueness is no longer guaranteed -- fail loudly rather than proceed. + for i, e in enumerate(events): + if "UID" not in e: + e["UID"] = i + else: + raise ValueError( + "Event unexpectedly already carries a 'UID'; " + "cannot guarantee unique ids for GPUEventAnalyser" + ) + analyzer = GPUEventAnalyser(events) + metrics = analyzer.compute_metrics() + total = metrics.get("total_time", 0) + busy = metrics.get("busy_time", 0) + idle = metrics.get("idle_time", 0) + if total <= 0: + return None + return { + "busy_time_us": metrics.get("busy_time", 0), + "idle_time_us": metrics.get("idle_time", 0), + "total_time_us": total, + "computation_time_us": metrics.get("computation_time", 0), + "exposed_comm_time_us": metrics.get("exposed_comm_time", 0), + "exposed_memcpy_time_us": metrics.get("exposed_memcpy_time", 0), + "idle_pct": 100 * idle / total, + "busy_pct": 100 * busy / total, + } + except Exception as e: + logger.warning("GPUEventAnalyser failed: %s", e) + return None + + +def extract_and_build_result(data, by_cat, source_file, region_metadata=None): + """Build extraction result dict.""" + kernels = extract_kernel_sequence(by_cat) + is_graph_mode, graph_launches = detect_graph_mode(by_cat) + total_kernel_time = sum(k["dur"] for k in kernels) + categories_found = sorted(by_cat.keys()) + result = { + "source_file": source_file, + "metadata": { + "total_kernels": len(kernels), + "total_kernel_time_us": round(total_kernel_time, 2), + "is_graph_mode": is_graph_mode, + "graph_launch_count": len(graph_launches), + "categories": categories_found, + }, + "kernels": kernels, + } + if region_metadata: + result["region_metadata"] = region_metadata + return result, kernels + + +def _write_split_regions(split_result, trace_path, output_dir): # pragma: no cover + """Write per-region extracted.json + metadata.json files.""" + os.makedirs(output_dir, exist_ok=True) + for trace_dict, region_meta in split_result: + key = get_steady_state_key(region_meta) + region_dir = os.path.join(output_dir, key) + os.makedirs(region_dir, exist_ok=True) + merged_meta = gather_metadata( + trace_path, + trace_dict.get("traceEvents", []), + annotation_meta=region_meta, + ) + gpu_timeline = compute_gpu_timeline_metrics(trace_dict.get("traceEvents", [])) + if gpu_timeline: + merged_meta["gpu_timeline"] = gpu_timeline + region_meta = {**region_meta, "gpu_timeline": gpu_timeline} + data, by_cat = load_trace(trace_dict) + filter_to_primary_stream(by_cat) + kernels_tmp = extract_kernel_sequence(by_cat) + is_graph_tmp, _ = detect_graph_mode(by_cat) + errors = run_assertions(data, by_cat, kernels_tmp, is_graph_tmp, strict=False) + if errors: + print(f"Skipping {key}: {'; '.join(errors)}", file=sys.stderr) + continue + result, kernels = extract_and_build_result( + data, by_cat, trace_path, region_metadata=region_meta + ) + extracted_path = os.path.join(region_dir, "extracted.json") + meta_path = os.path.join(region_dir, "metadata.json") + with open(extracted_path, "w") as f: + json.dump(result, f, indent=2) + with open(meta_path, "w") as f: + json.dump(merged_meta, f, indent=2) + print(f"Wrote {extracted_path} ({len(kernels)} kernels)", file=sys.stderr) + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser( + description="Extract structured data from a Chrome trace JSON" + ) + parser.add_argument("trace", help="Path to trace JSON file") + parser.add_argument( + "-o", + "--output", + help="Output directory. Single-trace writes extracted.json inside " + "it; multi-region (vLLM) writes per-region subdirectories.", + ) + parser.add_argument( + "--split-vllm", + action="store_true", + help="(deprecated, now auto-detected) kept for backward compatibility", + ) + parser.add_argument( + "--no-split", + action="store_true", + help="Skip annotation auto-detection; always extract as one trace", + ) + args = parser.parse_args() + + # --- Auto-detect vLLM annotation regions (unless suppressed) ----------- + if not args.no_split: + split_result = split_vllm_trace(args.trace) + if split_result: + output_dir = args.output or "." + _write_split_regions(split_result, args.trace, output_dir) + return + + # --- Single-trace extraction ------------------------------------------- + data, by_cat = load_trace(args.trace) + _stamp_raw_uid(data) + kernels = extract_kernel_sequence(by_cat) + is_graph_mode, graph_launches = detect_graph_mode(by_cat) + + errors = run_assertions(data, by_cat, kernels, is_graph_mode) + if errors: + for e in errors: + print(e, file=sys.stderr) + sys.exit(1) + + result, _ = extract_and_build_result(data, by_cat, args.trace) + output = json.dumps(result, indent=2) + + if args.output: + out_path = args.output + if os.path.isdir(out_path) or out_path.endswith("/"): + os.makedirs(out_path, exist_ok=True) + out_path = os.path.join(out_path, "extracted.json") + with open(out_path, "w") as f: + f.write(output) + print( + f"Wrote {out_path} ({len(kernels)} kernels, " + f"{sum(k['dur'] for k in kernels):.1f}us total)", + file=sys.stderr, + ) + else: + print(output) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/generate_semantic_diff.py b/TraceLens/Agent/Analysis/semantic_analyses/generate_semantic_diff.py new file mode 100644 index 000000000..a28baf86a --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/generate_semantic_diff.py @@ -0,0 +1,444 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Generate TraceDiff-compatible output from two semantic breakdowns. + +Takes two semantic_labels.json files (from the semantic breakdown pipeline) +and produces output files identical in schema to TraceDiff's +print_tracediff_report_files(), enabling the existing comparative mode +(tracelens_diff_analyzer.py) to consume graph-mode trace comparisons. + +Instead of the Wagner-Fischer DP tree-merge algorithm, kernels are matched +across traces by their semantic_block label. + +Input: + - Two semantic_labels.json files (one per trace) + +Output (in output directory): + - diff_stats.csv (includes per-kernel ``gpu_op_uid`` -- the TraceTree UID -- + and per-LCA ``busy_time``, so the output is consumable by the perf-report + comparison enrichment in tracediff_comparison_extension.py) + - diff_stats_unique_args_summary.csv + - cpu_op_map.json + - cpu_op_map_trace1.json + - cpu_op_map_trace2.json + - merged_tree_output.txt + +Usage: + python generate_semantic_diff.py \\ + trace_a/semantic_labels.json trace_b/semantic_labels.json \\ + --name-a MI355 --name-b B200 \\ + -o output_dir/ +""" + +import argparse +import json +import os +import sys +from collections import OrderedDict + +import pandas as pd + +from _helpers import load_labels + +# --------------------------------------------------------------------------- +# diff_stats.csv generation +# --------------------------------------------------------------------------- + + +def build_diff_stats(labeled_a, labeled_b): + """Build a list of row dicts matching TraceDiff's diff_stats.csv schema. + + Each kernel becomes one row. Kernels are grouped by semantic_block, + which is the functional label (e.g. "QKV Projection") when available. + + The kernel's ``nn_module`` field is used for the nn_module columns. + Falls back to perf_category. + + Returns (rows, block_id_map) where block_id_map is {semantic_block: int}. + """ + all_blocks_ordered = list( + OrderedDict.fromkeys( + [k["semantic_block"] for k in labeled_a] + + [k["semantic_block"] for k in labeled_b] + ) + ) + block_id_map = {block: idx for idx, block in enumerate(all_blocks_ordered)} + + def _nn_module(k): + if k.get("nn_module"): + return k["nn_module"] + return k.get("perf_category", "Others") + + def _format_dims(dims_list): + if not dims_list: + return "" + parts = [] + for d in dims_list: + if isinstance(d, (list, tuple)): + parts.append(str(tuple(d))) + else: + parts.append(str(d)) + return ", ".join(parts) + + rows = [] + for source_tag, kernels in [ + ("trace1", labeled_a), + ("trace2", labeled_b), + ]: + for k in kernels: + block = k["semantic_block"] + nn_mod = _nn_module(k) + cpu_op = k.get("cpu_op", "") or block + rows.append( + { + "name": k["name"], + "cpu_op_name": cpu_op, + "source": source_tag, + "Input Dims": _format_dims(k.get("input_dims", [])), + "Input Strides": "", + "Input type": "", + "Concrete Inputs": "", + "kernel_time": k["dur"], + "lowest_common_ancestor_name": block, + "lowest_common_ancestor_id": block_id_map[block], + "nn_module_stack": nn_mod, + "nn_module_parent": nn_mod, + "gpu_op_uid": k.get("gpu_op_uid"), + } + ) + + return rows, block_id_map + + +# --------------------------------------------------------------------------- +# diff_stats_unique_args_summary.csv generation +# --------------------------------------------------------------------------- + + +def build_unique_args_summary(diff_stats_df): + """Aggregate diff_stats rows by all non-metric columns. + + Mirrors TraceDiff.get_df_diff_stats_unique_args() with sum aggregation + on kernel_time. + """ + metric_columns = ["kernel_time"] + # gpu_op_uid is unique per kernel (would explode groups) and busy_time is + # a derived per-LCA aggregate, so neither belongs in the grouping key. + excluded_cols = {"lowest_common_ancestor_id", "gpu_op_uid", "busy_time"} + grouping_cols = [ + c + for c in diff_stats_df.columns + if c not in metric_columns and c not in excluded_cols + ] + + agg_dict = {mcol: ["sum", "mean"] for mcol in metric_columns} + for col in grouping_cols: + agg_dict[col] = "first" + + try: + df_agg = diff_stats_df.groupby(grouping_cols, dropna=False).agg(agg_dict) + df_agg["operation_count"] = diff_stats_df.groupby( + grouping_cols, dropna=False + ).size() + except TypeError: + str_cols = [f"{col}_str_repr" for col in grouping_cols] + df_temp = diff_stats_df.copy() + for col, str_col in zip(grouping_cols, str_cols): + df_temp[str_col] = df_temp[col].astype(str) + df_agg = df_temp.groupby(str_cols, dropna=False).agg(agg_dict) + df_agg["operation_count"] = df_temp.groupby(str_cols, dropna=False).size() + + df_agg.columns = ["_".join(col).strip() for col in df_agg.columns.values] + df_agg = df_agg.reset_index(drop=True) + + rename_map = {} + for col in grouping_cols: + col_first = f"{col}_first" + if col_first in df_agg.columns: + rename_map[col_first] = col + df_agg = df_agg.rename(columns=rename_map) + + primary_cols = grouping_cols + metric_cols = [] + for metric in metric_columns: + for agg in ["sum", "mean"]: + col_name = f"{metric}_{agg}" + if col_name in df_agg.columns: + metric_cols.append(col_name) + metric_cols = list(dict.fromkeys(metric_cols)) + other_cols = [ + col for col in df_agg.columns if col not in primary_cols + metric_cols + ] + df_agg = df_agg[primary_cols + metric_cols + other_cols] + + df_agg = df_agg.rename(columns={"operation_count_": "operation_count"}) + if "operation_count" in df_agg.columns: + cols = list(df_agg.columns) + cols.remove("operation_count") + cols.insert(1, "operation_count") + df_agg = df_agg[cols] + + sort_col = "kernel_time_sum" + if sort_col in df_agg.columns: + df_agg = df_agg.sort_values(by=sort_col, ascending=False, ignore_index=True) + + return df_agg + + +# --------------------------------------------------------------------------- +# cpu_op_map JSON generation +# --------------------------------------------------------------------------- + + +def build_cpu_op_maps(diff_stats_df): + """Build cpu_op_map dicts analogous to TraceDiff.get_cpu_op_to_kernels_json(). + + Returns (cpu_op_map, cpu_op_map_trace1_df, cpu_op_map_trace2_df). + """ + cpu_op_map = {} + for cpu_op in diff_stats_df["cpu_op_name"].unique(): + cpu_op_map[cpu_op] = {} + sub = diff_stats_df[diff_stats_df["cpu_op_name"] == cpu_op] + for source, group in sub.groupby("source"): + cpu_op_map[cpu_op][source] = { + "kernels": sorted(list(group["name"].unique())), + "nn_module_parents": sorted(list(group["nn_module_parent"].unique())), + } + + cpu_op_map_trace1 = ( + diff_stats_df[diff_stats_df["source"] == "trace1"] + .groupby("cpu_op_name") + .agg({"name": lambda x: sorted(set(x))}) + .sort_index() + ) + cpu_op_map_trace2 = ( + diff_stats_df[diff_stats_df["source"] == "trace2"] + .groupby("cpu_op_name") + .agg({"name": lambda x: sorted(set(x))}) + .sort_index() + ) + + return cpu_op_map, cpu_op_map_trace1, cpu_op_map_trace2 + + +# --------------------------------------------------------------------------- +# merged_tree_output.txt generation +# --------------------------------------------------------------------------- + + +def build_merged_tree_text(block_id_map, labeled_a, labeled_b, name_a, name_b): + """Build a text tree representation mimicking TraceDiff's merged tree. + + Structure: + Root + └── + ├── (combined / trace1-only / trace2-only) + │ ├── kernel_name_a [trace1] + │ └── kernel_name_b [trace2] + """ + blocks_a = OrderedDict() + for k in labeled_a: + b = k["semantic_block"] + blocks_a.setdefault(b, []).append(k["name"]) + blocks_b = OrderedDict() + for k in labeled_b: + b = k["semantic_block"] + blocks_b.setdefault(b, []).append(k["name"]) + + block_nn_module = {} + for k in labeled_a + labeled_b: + b = k["semantic_block"] + if b not in block_nn_module: + if k.get("nn_module"): + block_nn_module[b] = k["nn_module"] + else: + block_nn_module[b] = k.get("perf_category", "Others") + + all_blocks = list(block_id_map.keys()) + + groups = OrderedDict() + for block in all_blocks: + g = block_nn_module.get(block, "Others") + groups.setdefault(g, []).append(block) + + lines = [] + lines.append(f"└── Root ({name_a} vs {name_b})") + + group_list = list(groups.items()) + for gi, (group_name, group_blocks) in enumerate(group_list): + is_last_group = gi == len(group_list) - 1 + g_connector = "└── " if is_last_group else "├── " + g_prefix = " " if is_last_group else "│ " + lines.append(f" {g_connector}{group_name}") + + for bi, block in enumerate(group_blocks): + is_last_block = bi == len(group_blocks) - 1 + b_connector = "└── " if is_last_block else "├── " + b_prefix = " " if is_last_block else "│ " + + in_a = block in blocks_a + in_b = block in blocks_b + + if in_a and in_b: + label = block + elif in_a: + label = f">> trace1: {block}" + else: + label = f"<< trace2: {block}" + + lines.append(f" {g_prefix}{b_connector}{label}") + + kernels_a = blocks_a.get(block, []) + kernels_b = blocks_b.get(block, []) + kernel_names_a = sorted(set(kernels_a)) + kernel_names_b = sorted(set(kernels_b)) + + all_kernel_entries = [(kn, "trace1") for kn in kernel_names_a] + [ + (kn, "trace2") for kn in kernel_names_b + ] + + for ki, (kn, src) in enumerate(all_kernel_entries): + is_last_kernel = ki == len(all_kernel_entries) - 1 + k_connector = "└── " if is_last_kernel else "├── " + if in_a and in_b: + k_line = kn + elif src == "trace1": + k_line = f">> {src}: {kn}" + else: + k_line = f"<< {src}: {kn}" + lines.append(f" {g_prefix}{b_prefix}{k_connector}{k_line}") + + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Output writer +# --------------------------------------------------------------------------- + + +def write_outputs( # pragma: no cover + output_dir, + diff_stats_df, + summary_df, + cpu_op_map, + cpu_op_map_trace1, + cpu_op_map_trace2, + merged_tree_text, +): + """Write all output files to output_dir.""" + os.makedirs(output_dir, exist_ok=True) + + diff_stats_df.to_csv(os.path.join(output_dir, "diff_stats.csv"), index=False) + + summary_df.to_csv( + os.path.join(output_dir, "diff_stats_unique_args_summary.csv"), index=False + ) + + with open(os.path.join(output_dir, "cpu_op_map.json"), "w") as f: + json.dump(cpu_op_map, f, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "cpu_op_map_trace1.json"), "w") as f: + json.dump(cpu_op_map_trace1.to_dict()["name"], f, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "cpu_op_map_trace2.json"), "w") as f: + json.dump(cpu_op_map_trace2.to_dict()["name"], f, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "merged_tree_output.txt"), "w") as f: + f.write(merged_tree_text + "\n") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser( + description="Generate TraceDiff-compatible output from two semantic breakdowns" + ) + parser.add_argument("trace_a", help="Path to trace A semantic_labels.json") + parser.add_argument("trace_b", help="Path to trace B semantic_labels.json") + parser.add_argument("--name-a", default="trace_a", help="Short name for trace A") + parser.add_argument("--name-b", default="trace_b", help="Short name for trace B") + parser.add_argument( + "-o", + "--output", + default="semantic_diff_output", + help="Output directory (default: semantic_diff_output)", + ) + args = parser.parse_args() + + data_a = load_labels(args.trace_a) + data_b = load_labels(args.trace_b) + + labeled_a = data_a["labeled_kernels"] + labeled_b = data_b["labeled_kernels"] + + print(f"Trace A ({args.name_a}): {len(labeled_a)} kernels", file=sys.stderr) + print(f"Trace B ({args.name_b}): {len(labeled_b)} kernels", file=sys.stderr) + + rows, block_id_map = build_diff_stats(labeled_a, labeled_b) + diff_stats_df = pd.DataFrame(rows) + + # Per-LCA block total kernel time, broadcast to each row (mirrors + # TraceDiff's busy_time, which is identical for every row in an LCA + # group). semantic_labels.json has no timestamps, so sum-of-durations + # is used as the busy-time proxy (exact for non-overlapping kernels). + diff_stats_df["busy_time"] = ( + diff_stats_df.groupby(["source", "lowest_common_ancestor_id"])["kernel_time"] + .transform("sum") + .round(3) + ) + + summary_df = build_unique_args_summary(diff_stats_df) + + cpu_op_map, cpu_op_map_trace1, cpu_op_map_trace2 = build_cpu_op_maps(diff_stats_df) + + merged_tree_text = build_merged_tree_text( + block_id_map, labeled_a, labeled_b, args.name_a, args.name_b + ) + + write_outputs( + args.output, + diff_stats_df, + summary_df, + cpu_op_map, + cpu_op_map_trace1, + cpu_op_map_trace2, + merged_tree_text, + ) + + blocks_a = set(k["semantic_block"] for k in labeled_a) + blocks_b = set(k["semantic_block"] for k in labeled_b) + matched = blocks_a & blocks_b + only_a = blocks_a - blocks_b + only_b = blocks_b - blocks_a + + print(f"\nSemantic blocks: {len(block_id_map)} total", file=sys.stderr) + print(f" Matched (in both traces): {len(matched)}", file=sys.stderr) + if only_a: + print(f" Only in {args.name_a}: {sorted(only_a)}", file=sys.stderr) + if only_b: + print(f" Only in {args.name_b}: {sorted(only_b)}", file=sys.stderr) + + print(f"\nOutput written to: {args.output}/", file=sys.stderr) + print(f" diff_stats.csv ({len(diff_stats_df)} rows)", file=sys.stderr) + print( + f" diff_stats_unique_args_summary.csv ({len(summary_df)} rows)", + file=sys.stderr, + ) + print( + f" cpu_op_map.json, cpu_op_map_trace1.json, cpu_op_map_trace2.json", + file=sys.stderr, + ) + print(f" merged_tree_output.txt", file=sys.stderr) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/kernel_coherence.py b/TraceLens/Agent/Analysis/semantic_analyses/kernel_coherence.py new file mode 100644 index 000000000..a717ad28d --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/kernel_coherence.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Second-pass kernel-name coherence (LLM-assisted). + +After the first-pass name-first unification (``kernel_unification.py``), the +comparison still has **one-sided** buckets: kernels whose unified name appears +in only one trace (e.g. vendor GEMM families -- platform-A ``_*`` vs +platform-B ``_*`` -- that could not be paired by name alone). + +This pass uses the first-pass **shared** buckets as cross-trace-stable +positional anchors. For each one-sided bucket it derives the *neighbor context* +(nearest shared symbols to its left and right in the run-length-collapsed kernel +sequence) and lets an LLM: + + * pair a one-sided bucket in trace A with a one-sided bucket in trace B that + occupies the **same** neighbor context (e.g. the GEMM between ``add_rmsnorm`` + and ``rotary_embedding`` is the QKV projection on both traces), assigning + both a new shared name; and + * split a single name that occurs in **different** contexts into different + buckets (context-dependent granularity). + +Two subcommands: + + prepare-context Build the LLM packet: condensed sequences, shared vs + one-sided symbol sets, and per one-sided symbol the unique + neighbor contexts with evidence (perf_category, kernels in + the run, duration, sample input_dims, raw name). Emits a flat + ``context_catalog`` with stable ids. + + apply Apply the LLM decisions (``context_renames`` + + ``fallback_remap_a`` / ``fallback_remap_b``) to recompute each + kernel's ``semantic_block`` in place, emit an audit CSV, and + warn about any residual one-sided condensed symbols. + +Usage: + python kernel_coherence.py prepare-context \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a platform_a --name-b platform_b \ + --neighbor-radius 1 \ + -o kernel_coherence_context.json + + python kernel_coherence.py apply \ + --context kernel_coherence_context.json \ + --decisions kernel_coherence_decisions.json \ + [--audit-csv-a per_kernel_final_a.csv] \ + [--audit-csv-b per_kernel_final_b.csv] +""" + +import argparse +import csv +import json +import sys +from collections import OrderedDict, defaultdict + +from kernel_runlength import ( + collapse_consecutive, + run_index_per_kernel, + shared_neighbor_windows_skip_non_shared, +) +from _helpers import load_json + +DEFAULT_RADIUS = 1 +DEFAULT_TOP_KERNELS = 5 + + +def _load(path): + return load_json(path) + + +def _dims_repr(dims, limit=160): + if not dims: + return "" + s = json.dumps(dims, separators=(",", ":")) + return s[:limit] + "..." if len(s) > limit else s + + +# =========================================================================== +# prepare-context +# =========================================================================== + + +def _run_to_kernel_indices(run_per_kernel): + """Map run index -> list of kernel positions in that run.""" + run_to_k = defaultdict(list) + for ki, rj in enumerate(run_per_kernel): + run_to_k[rj].append(ki) + return run_to_k + + +def _symbol_evidence(kernels, indices, top_kernels): + """Aggregate evidence for the kernels of one run (or symbol).""" + by_name_dur = defaultdict(float) + by_name_cnt = defaultdict(int) + cats = set() + sample_dims = "" + for ki in indices: + k = kernels[ki] + nm = k.get("name", "") + by_name_dur[nm] += k.get("dur", 0.0) or 0.0 + by_name_cnt[nm] += 1 + pc = k.get("perf_category") + if pc: + cats.add(pc) + if not sample_dims: + sample_dims = _dims_repr(k.get("input_dims")) + ranked = sorted(by_name_dur.items(), key=lambda x: -x[1])[:top_kernels] + top = [ + {"kernel_name": n, "total_us": round(d, 3), "kernel_count": by_name_cnt[n]} + for n, d in ranked + ] + return sorted(cats), sample_dims, top + + +def _collect_contexts( + name, kernels, condensed, run_per_kernel, shared, problematic, radius, top_kernels +): + """Per one-sided symbol, the unique (left,right) shared-neighbor contexts.""" + run_to_k = _run_to_kernel_indices(run_per_kernel) + out = OrderedDict() + ctr = 0 + for sym in sorted(problematic): + seen = set() + contexts = [] + for j, c in enumerate(condensed): + if c != sym: + continue + left, right = shared_neighbor_windows_skip_non_shared( + condensed, j, shared, radius + ) + key = (tuple(left), tuple(right)) + if key in seen: + continue + seen.add(key) + indices = run_to_k.get(j, []) + cats, dims, top = _symbol_evidence(kernels, indices, top_kernels) + contexts.append( + OrderedDict( + id=f"{name}:{ctr}", + first_pass_block=sym, + left_window=list(left), + right_window=list(right), + kernels_in_run=len(indices), + perf_categories=cats, + sample_input_dims=dims, + top_kernel_names_by_dur=top, + ) + ) + ctr += 1 + out[sym] = {"contexts": contexts, "context_count": len(contexts)} + return out + + +def cmd_prepare_context(args): # pragma: no cover + labels_a = _load(args.labels_a) + labels_b = _load(args.labels_b) + kernels_a = labels_a["labeled_kernels"] + kernels_b = labels_b["labeled_kernels"] + + seq_a = [k.get("semantic_block", "") for k in kernels_a] + seq_b = [k.get("semantic_block", "") for k in kernels_b] + cond_a = collapse_consecutive(seq_a) + cond_b = collapse_consecutive(seq_b) + set_a, set_b = set(cond_a), set(cond_b) + shared = set_a & set_b + prob_a = set_a - set_b + prob_b = set_b - set_a + + run_a = run_index_per_kernel(seq_a) + run_b = run_index_per_kernel(seq_b) + + detail_a = _collect_contexts( + args.name_a, + kernels_a, + cond_a, + run_a, + shared, + prob_a, + args.neighbor_radius, + args.top_kernels, + ) + detail_b = _collect_contexts( + args.name_b, + kernels_b, + cond_b, + run_b, + shared, + prob_b, + args.neighbor_radius, + args.top_kernels, + ) + + catalog = [] + for detail, wl in ((detail_a, args.name_a), (detail_b, args.name_b)): + for sym, pack in detail.items(): + for c in pack["contexts"]: + catalog.append( + OrderedDict( + id=c["id"], + workload=wl, + first_pass_block=sym, + left_window=c["left_window"], + right_window=c["right_window"], + ) + ) + + out = OrderedDict() + out["name_a"] = args.name_a + out["name_b"] = args.name_b + out["definitions"] = { + "condensed_sequence": "Kernel-order semantic_block values with consecutive " + "duplicates removed (one symbol per contiguous run).", + "shared_block": "semantic_block present in the condensed sequence of both traces.", + "one_sided_block": "semantic_block present in only one trace's condensed sequence.", + "neighbor_context": "left_window/right_window: nearest `neighbor_radius` shared " + "symbols on each side of the run, skipping non-shared symbols.", + } + out["hyperparameters"] = { + "neighbor_radius": args.neighbor_radius, + "top_kernels": args.top_kernels, + } + out["inputs"] = {"labels_a": args.labels_a, "labels_b": args.labels_b} + out["condensed_sequence_a"] = cond_a + out["condensed_sequence_b"] = cond_b + out["shared_blocks"] = sorted(shared) + out["one_sided_in_a"] = sorted(prob_a) + out["one_sided_in_b"] = sorted(prob_b) + out["one_sided_details_a"] = detail_a + out["one_sided_details_b"] = detail_b + out["llm_task"] = ( + "Re-label one-sided blocks so the comparison has no one-sided condensed " + "symbols. Pair a one-sided block in A with a one-sided block in B that has " + "the SAME (left_window,right_window) and perf_category by giving both the " + "same new shared name. Split a block that appears in different contexts via " + "distinct context ids. See the kernel-coherence agent for the output schema." + ) + out["llm_output_schema"] = { + "context_renames": "{context_id -> final_block}", + "fallback_remap_a": "{first_pass_block -> final_block} for trace A", + "fallback_remap_b": "{first_pass_block -> final_block} for trace B", + "notes": "optional string", + } + out["context_catalog"] = catalog + + with open(args.output, "w") as f: + json.dump(out, f, indent=2) + + nca = sum(v["context_count"] for v in detail_a.values()) + ncb = sum(v["context_count"] for v in detail_b.values()) + print( + f"Wrote {args.output}: shared={len(shared)} " + f"one_sided_a={len(prob_a)} one_sided_b={len(prob_b)} " + f"contexts_a={nca} contexts_b={ncb}", + file=sys.stderr, + ) + + +# =========================================================================== +# apply +# =========================================================================== + + +def _context_lookup(catalog): + """(workload, symbol, left, right) -> context_id.""" + out = {} + for row in catalog: + out[ + ( + row["workload"], + row["first_pass_block"], + tuple(row.get("left_window") or []), + tuple(row.get("right_window") or []), + ) + ] = row["id"] + return out + + +def _final_blocks( + workload, kernels, shared, problematic, radius, lookup, context_renames, fallback +): + """Compute the final semantic_block for each kernel; return (finals, audit).""" + seq = [k.get("semantic_block", "") for k in kernels] + cond = collapse_consecutive(seq) + run_per_kernel = run_index_per_kernel(seq) + + finals = [] + audit = [] + for i, k in enumerate(kernels): + fp = seq[i] + j = run_per_kernel[i] + sym = cond[j] + cid = "" + if sym not in problematic: + final = fp + else: + left, right = shared_neighbor_windows_skip_non_shared( + cond, j, shared, radius + ) + cid = lookup.get((workload, sym, tuple(left), tuple(right)), "") + if cid and cid in context_renames: + final = context_renames[cid] + elif sym in fallback: + final = fallback[sym] + else: + final = fp + finals.append(final) + audit.append( + { + "kernel_index": k.get("index", i), + "name": k.get("name", ""), + "first_pass_block": fp, + "final_block": final, + "context_id": cid, + } + ) + return finals, audit + + +def _write_audit(path, rows): # pragma: no cover + with open(path, "w", newline="") as f: + w = csv.DictWriter( + f, + fieldnames=[ + "kernel_index", + "name", + "first_pass_block", + "final_block", + "context_id", + ], + ) + w.writeheader() + w.writerows(rows) + + +def _residual_one_sided(labels_a, labels_b): + """Return one-sided condensed symbols remaining after apply.""" + ca = set( + collapse_consecutive( + [k.get("semantic_block", "") for k in labels_a["labeled_kernels"]] + ) + ) + cb = set( + collapse_consecutive( + [k.get("semantic_block", "") for k in labels_b["labeled_kernels"]] + ) + ) + return sorted(ca - cb), sorted(cb - ca) + + +def cmd_apply(args): # pragma: no cover + ctx = _load(args.context) + dec = _load(args.decisions) + name_a = ctx["name_a"] + name_b = ctx["name_b"] + radius = int(ctx.get("hyperparameters", {}).get("neighbor_radius", DEFAULT_RADIUS)) + labels_a_path = ctx["inputs"]["labels_a"] + labels_b_path = ctx["inputs"]["labels_b"] + + context_renames = { + str(k): str(v) for k, v in (dec.get("context_renames") or {}).items() + } + fb_a = {str(k): str(v) for k, v in (dec.get("fallback_remap_a") or {}).items()} + fb_b = {str(k): str(v) for k, v in (dec.get("fallback_remap_b") or {}).items()} + + labels_a = _load(labels_a_path) + labels_b = _load(labels_b_path) + ka = labels_a["labeled_kernels"] + kb = labels_b["labeled_kernels"] + + cond_a = collapse_consecutive([k.get("semantic_block", "") for k in ka]) + cond_b = collapse_consecutive([k.get("semantic_block", "") for k in kb]) + shared = set(cond_a) & set(cond_b) + prob_a = set(cond_a) - set(cond_b) + prob_b = set(cond_b) - set(cond_a) + lookup = _context_lookup(ctx.get("context_catalog", [])) + + finals_a, audit_a = _final_blocks( + name_a, ka, shared, prob_a, radius, lookup, context_renames, fb_a + ) + finals_b, audit_b = _final_blocks( + name_b, kb, shared, prob_b, radius, lookup, context_renames, fb_b + ) + + changed_a = 0 + for k, fin in zip(ka, finals_a): + if k.get("semantic_block") != fin: + changed_a += 1 + k["semantic_block"] = fin + changed_b = 0 + for k, fin in zip(kb, finals_b): + if k.get("semantic_block") != fin: + changed_b += 1 + k["semantic_block"] = fin + + with open(labels_a_path, "w") as f: + json.dump(labels_a, f, indent=2) + with open(labels_b_path, "w") as f: + json.dump(labels_b, f, indent=2) + + if args.audit_csv_a: + _write_audit(args.audit_csv_a, audit_a) + if args.audit_csv_b: + _write_audit(args.audit_csv_b, audit_b) + + blocks_a = set(k["semantic_block"] for k in ka) + blocks_b = set(k["semantic_block"] for k in kb) + res_a, res_b = _residual_one_sided(labels_a, labels_b) + + print( + f"Applied: {name_a} {changed_a} kernels relabeled, " + f"{name_b} {changed_b} kernels relabeled.", + file=sys.stderr, + ) + print( + f"Shared blocks now: {len(blocks_a & blocks_b)} " + f"({name_a}-only {len(blocks_a - blocks_b)}, " + f"{name_b}-only {len(blocks_b - blocks_a)}).", + file=sys.stderr, + ) + if res_a or res_b: + print( + f"WARNING: condensed one-sided symbols remain -- " + f"{name_a}: {res_a} {name_b}: {res_b}", + file=sys.stderr, + ) + else: + print("No one-sided condensed symbols remain.", file=sys.stderr) + + +# =========================================================================== +# CLI +# =========================================================================== + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser( + description="Second-pass kernel-name coherence (prepare-context / apply)" + ) + sub = parser.add_subparsers(dest="command", required=True) + + p_prep = sub.add_parser("prepare-context", help="Build the coherence LLM context") + p_prep.add_argument("--labels-a", required=True) + p_prep.add_argument("--labels-b", required=True) + p_prep.add_argument("--name-a", default="trace_a") + p_prep.add_argument("--name-b", default="trace_b") + p_prep.add_argument( + "--neighbor-radius", + type=int, + default=DEFAULT_RADIUS, + help=f"Shared symbols to collect on each side (default {DEFAULT_RADIUS})", + ) + p_prep.add_argument( + "--top-kernels", + type=int, + default=DEFAULT_TOP_KERNELS, + help=f"Top kernel names by duration per context (default {DEFAULT_TOP_KERNELS})", + ) + p_prep.add_argument("-o", "--output", required=True) + p_prep.set_defaults(func=cmd_prepare_context) + + p_apply = sub.add_parser( + "apply", help="Apply coherence decisions to labels in place" + ) + p_apply.add_argument( + "--context", required=True, help="kernel_coherence_context.json" + ) + p_apply.add_argument("--decisions", required=True, help="LLM decisions JSON") + p_apply.add_argument("--audit-csv-a", help="Per-kernel audit CSV for trace A") + p_apply.add_argument("--audit-csv-b", help="Per-kernel audit CSV for trace B") + p_apply.set_defaults(func=cmd_apply) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/kernel_runlength.py b/TraceLens/Agent/Analysis/semantic_analyses/kernel_runlength.py new file mode 100644 index 000000000..250541f49 --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/kernel_runlength.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Run-length helpers over per-kernel ``semantic_block`` sequences. + +Shared by the kernel-name coherence (second) pass. Operates purely on the +``semantic_block`` field of ``semantic_labels.json`` -- no ``nn_module`` / +tree_context is required, so it works on graph-mode traces. + +Terminology: + sequence -- kernel-order list of semantic_block values (one per kernel). + condensed -- the sequence with consecutive duplicates collapsed + (``A A B D D`` -> ``A B D``); one symbol per contiguous run. + shared -- a symbol present in the condensed sets of BOTH traces. + one-sided -- a symbol present in only one trace's condensed set. +""" + +from _helpers import load_json + + +def load_sequence(labels_path): + """Return the kernel-order list of semantic_block values for a labels file.""" + data = load_json(labels_path) + return [k.get("semantic_block", "") for k in data.get("labeled_kernels", [])] + + +def collapse_consecutive(seq): + """Collapse consecutive duplicates: ``A A B D D`` -> ``A B D``.""" + if not seq: + return [] + out = [seq[0]] + for s in seq[1:]: + if s != out[-1]: + out.append(s) + return out + + +def run_index_per_kernel(seq): + """Map each kernel position to the index of its run in ``collapse_consecutive(seq)``.""" + if not seq: + return [] + out = [] + run_idx = 0 + prev = None + for s in seq: + if prev is not None and s != prev: + run_idx += 1 + out.append(run_idx) + prev = s + return out + + +def shared_neighbor_windows_skip_non_shared(condensed, center_j, shared, radius): + """Nearest ``radius`` *shared* symbols on each side of ``condensed[center_j]``. + + Walks outward from the center, skipping any symbol not in ``shared`` + (including other one-sided symbols), so the windows are always expressed + in cross-trace-stable anchors. Nearest-to-center first on each side. + Returns ``(left_window, right_window)`` as lists. + """ + left = [] + i = center_j - 1 + while i >= 0 and len(left) < radius: + if condensed[i] in shared: + left.append(condensed[i]) + i -= 1 + + right = [] + i = center_j + 1 + while i < len(condensed) and len(right) < radius: + if condensed[i] in shared: + right.append(condensed[i]) + i += 1 + + return left, right diff --git a/TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py b/TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py new file mode 100644 index 000000000..c3111b3dc --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py @@ -0,0 +1,550 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Name-first cross-trace kernel-name unification (LLM-assisted). + +Graph-mode traces collapse the CPU->GPU call stack under ``hip/cudaGraphLaunch``, +so ``nn_module`` / ``cpu_op`` context is unavailable and the block-alignment +harmonization cannot work. The only reliable cross-trace signal is the raw GPU +**kernel name**. This module unifies kernel names across two traces so that the +downstream comparison (which matches on the per-kernel ``semantic_block`` field) +can pair equivalent kernels. + +The approach is *name-first*: an LLM inspects the unique kernel names from both +traces and writes a conservative map of names it is **certain** are equivalent +(e.g. ``moe_attn_vllm`` and ``sglang_moe_attention`` -> ``moe_attn``). Names +that are already identical need no entry -- they unify by default. The goal is +to establish clear anchors, not to resolve every ambiguity. + +Three subcommands: + + prepare-context Aggregate the unique kernel names from both traces (with + per-name stats) into a compact LLM packet. When the + combined unique-name count exceeds ``--threshold`` it emits + a representative *sample* plus ``needs_stem_preprocessing: + true`` instead of the full lists (see apply-stem-rules). + + apply-stem-rules Apply an LLM-authored ``stem_rules.json`` (custom regexes + + collapse/preserve/drop actions) to the full unique-name set, + emitting a ``raw_to_stem`` map and a reduced, + stem-level context. Prints resulting cardinality so the + model can iterate. + + apply-map Apply the LLM's ``kernel_unification_map.json`` back onto + both ``semantic_labels.json`` files, writing the unified + name into ``semantic_block`` (default = raw name / stem when + the map has no entry). + +Usage: + python kernel_unification.py prepare-context \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a --name-b \ + -o kernel_unification_context.json + + python kernel_unification.py apply-stem-rules \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a --name-b \ + --rules stem_rules.json \ + --raw-to-stem raw_to_stem.json \ + -o kernel_unification_context.json + + python kernel_unification.py apply-map \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a --name-b \ + --map kernel_unification_map.json \ + [--raw-to-stem raw_to_stem.json] +""" + +import argparse +import json +import re +import sys +from collections import OrderedDict + +from _helpers import load_json + +DEFAULT_THRESHOLD = 5000 +DEFAULT_SAMPLE_SIZE = 300 + + +# =========================================================================== +# shared helpers +# =========================================================================== + + +def _load(path): + return load_json(path) + + +def _dims_repr(dims, limit=160): + """Compact, length-capped string form of an input_dims value.""" + if not dims: + return "" + s = json.dumps(dims, separators=(",", ":")) + if len(s) > limit: + s = s[:limit] + "..." + return s + + +def aggregate_names(labels, key_fn=None): + """Aggregate a labels file's kernels by name (or by ``key_fn(name)``). + + Returns an OrderedDict ``{key: entry}`` where entry carries count, + total duration, the set of perf_categories seen, and one sample + input_dims. Order is by descending total duration then key, so the + most impactful kernels come first for the LLM. + """ + acc = {} + for k in labels.get("labeled_kernels", []): + name = k.get("name", "") + key = key_fn(name) if key_fn else name + if key is None: + continue # dropped by stem rules -> excluded from the unification set + entry = acc.get(key) + if entry is None: + entry = acc[key] = { + "name": key, + "kernel_count": 0, + "total_dur_us": 0.0, + "perf_categories": set(), + "sample_input_dims": "", + "sample_raw_names": set(), + } + entry["kernel_count"] += 1 + entry["total_dur_us"] += k.get("dur", 0.0) or 0.0 + pc = k.get("perf_category") + if pc: + entry["perf_categories"].add(pc) + if not entry["sample_input_dims"]: + entry["sample_input_dims"] = _dims_repr(k.get("input_dims")) + if key_fn and len(entry["sample_raw_names"]) < 5 and name != key: + entry["sample_raw_names"].add(name) + + ordered = sorted(acc.values(), key=lambda e: (-e["total_dur_us"], e["name"])) + out = OrderedDict() + for e in ordered: + e["total_dur_us"] = round(e["total_dur_us"], 3) + e["perf_categories"] = sorted(e["perf_categories"]) + e["sample_raw_names"] = sorted(e["sample_raw_names"]) + if not e["sample_raw_names"]: + del e["sample_raw_names"] + out[e["name"]] = e + return out + + +def _entry_list(agg, keys): + """Materialize entries for *keys* preserving *agg* ordering.""" + keyset = set(keys) + return [e for name, e in agg.items() if name in keyset] + + +def _build_context(agg_a, agg_b, name_a, name_b, level, extra=None): + """Assemble the unification-context dict from two aggregations.""" + set_a = set(agg_a) + set_b = set(agg_b) + only_a = [n for n in agg_a if n not in set_b] + only_b = [n for n in agg_b if n not in set_a] + in_both = sorted(set_a & set_b) + + ctx = OrderedDict() + ctx["name_a"] = name_a + ctx["name_b"] = name_b + ctx["key_level"] = level # "raw_name" or "stem" + ctx["summary"] = { + f"unique_{name_a}": len(set_a), + f"unique_{name_b}": len(set_b), + "combined_unique": len(set_a | set_b), + "in_both": len(in_both), + f"only_in_{name_a}": len(only_a), + f"only_in_{name_b}": len(only_b), + } + ctx[f"only_in_{name_a}"] = _entry_list(agg_a, only_a) + ctx[f"only_in_{name_b}"] = _entry_list(agg_b, only_b) + ctx["in_both"] = in_both + if extra: + ctx.update(extra) + return ctx + + +# =========================================================================== +# prepare-context +# =========================================================================== + + +def _sample_names(agg_a, agg_b, name_a, name_b, sample_size): + """Deterministic representative sample across both traces. + + Interleaves the two per-trace aggregations (already ordered by impact) + and takes evenly-spaced picks so the sample spans high- and low-impact + kernels rather than only the top-N. + """ + + def _spaced(agg, budget): + items = list(agg.values()) + if len(items) <= budget: + return items + step = len(items) / float(budget) + return [items[int(i * step)] for i in range(budget)] + + half = max(1, sample_size // 2) + sample = [] + for e in _spaced(agg_a, half): + row = dict(e) + row["trace"] = name_a + sample.append(row) + for e in _spaced(agg_b, sample_size - half): + row = dict(e) + row["trace"] = name_b + sample.append(row) + return sample + + +def cmd_prepare_context(args): # pragma: no cover + labels_a = _load(args.labels_a) + labels_b = _load(args.labels_b) + agg_a = aggregate_names(labels_a) + agg_b = aggregate_names(labels_b) + + combined = len(set(agg_a) | set(agg_b)) + needs_stem = combined > args.threshold + + if needs_stem: + ctx = OrderedDict() + ctx["name_a"] = args.name_a + ctx["name_b"] = args.name_b + ctx["key_level"] = "raw_name" + ctx["summary"] = { + f"unique_{args.name_a}": len(agg_a), + f"unique_{args.name_b}": len(agg_b), + "combined_unique": combined, + } + ctx["needs_stem_preprocessing"] = True + ctx["threshold"] = args.threshold + ctx["instructions"] = ( + "Combined unique kernel-name count exceeds the threshold and will " + "likely overwhelm the context. Inspect the sample below, author a " + "stem_rules.json (see the kernel-stem-preprocessing agent), and run " + "'kernel_unification.py apply-stem-rules' to reduce the name set " + "before unification." + ) + ctx["sample"] = _sample_names( + agg_a, agg_b, args.name_a, args.name_b, args.sample_size + ) + else: + ctx = _build_context( + agg_a, + agg_b, + args.name_a, + args.name_b, + "raw_name", + extra={"needs_stem_preprocessing": False}, + ) + + with open(args.output, "w") as f: + json.dump(ctx, f, indent=2) + + if needs_stem: + print( + f"Wrote {args.output}: {combined} combined unique names > " + f"threshold {args.threshold} -> STEM PREPROCESSING NEEDED " + f"({len(ctx['sample'])} sampled names emitted).", + file=sys.stderr, + ) + else: + print( + f"Wrote {args.output}: {combined} combined unique names " + f"({ctx['summary']['in_both']} shared, " + f"{ctx['summary'][f'only_in_{args.name_a}']} {args.name_a}-only, " + f"{ctx['summary'][f'only_in_{args.name_b}']} {args.name_b}-only).", + file=sys.stderr, + ) + + +# =========================================================================== +# apply-stem-rules +# =========================================================================== + + +def _compile_rules(rules): + """Compile stem_rules entries; validate shape and actions.""" + compiled = [] + for i, r in enumerate(rules): + action = r.get("action", "collapse") + if action not in ("collapse", "preserve", "drop"): + raise SystemExit( + f"rule {i}: invalid action {action!r} " + f"(expected collapse/preserve/drop)" + ) + pattern = r.get("pattern", "") + try: + rx = re.compile(pattern) + except re.error as e: + raise SystemExit(f"rule {i}: bad regex {pattern!r}: {e}") + compiled.append( + { + "regex": rx, + "action": action, + "replacement": r.get("replacement", ""), + "note": r.get("note", ""), + } + ) + return compiled + + +def stem_for(name, compiled): + """Return (stem, action) for *name* using the first matching rule. + + Unmatched names default to (name, "preserve"). + """ + for r in compiled: + if r["regex"].search(name): + if r["action"] == "collapse": + try: + return r["regex"].sub(r["replacement"], name), "collapse" + except re.error as e: + # A malformed replacement template (e.g. a backreference with + # no matching capture group) is only detected by re at + # substitution time. Don't let one bad rule crash the whole + # run: warn once and fall back to preserving the name. + if not r.get("_warned"): + print( + f"[kernel_unification] ignoring stem rule with bad " + f"replacement {r['replacement']!r} for pattern " + f"{r['regex'].pattern!r}: {e}", + file=sys.stderr, + ) + r["_warned"] = True + return name, "preserve" + if r["action"] == "preserve": + return name, "preserve" + return name, "drop" + return name, "preserve" + + +def cmd_apply_stem_rules(args): # pragma: no cover + labels_a = _load(args.labels_a) + labels_b = _load(args.labels_b) + rules_doc = _load(args.rules) + rules = rules_doc["rules"] if isinstance(rules_doc, dict) else rules_doc + compiled = _compile_rules(rules) + + agg_a_raw = aggregate_names(labels_a) + agg_b_raw = aggregate_names(labels_b) + all_names = set(agg_a_raw) | set(agg_b_raw) + + raw_to_stem = {} + dropped = set() + action_counts = {"collapse": 0, "preserve": 0, "drop": 0} + for name in all_names: + stem, action = stem_for(name, compiled) + action_counts[action] += 1 + if action == "drop": + dropped.add(name) + raw_to_stem[name] = name # falls back to raw; excluded from context + else: + raw_to_stem[name] = stem + + if args.raw_to_stem: + with open(args.raw_to_stem, "w") as f: + json.dump(raw_to_stem, f, indent=2) + + def _stem_key(name): + if name in dropped: + return None # exclude dropped kernels from the stem aggregation + return raw_to_stem.get(name, name) + + agg_a = aggregate_names(labels_a, key_fn=lambda n: _stem_key(n)) + agg_b = aggregate_names(labels_b, key_fn=lambda n: _stem_key(n)) + agg_a.pop(None, None) + agg_b.pop(None, None) + + combined_stems = len(set(agg_a) | set(agg_b)) + ctx = _build_context( + agg_a, + agg_b, + args.name_a, + args.name_b, + "stem", + extra={ + "needs_stem_preprocessing": False, + "stem_preprocessing_applied": True, + "raw_unique_before": len(all_names), + "stem_unique_after": combined_stems, + "action_counts": action_counts, + "dropped_count": len(dropped), + }, + ) + + with open(args.output, "w") as f: + json.dump(ctx, f, indent=2) + + status = "OK" if combined_stems <= args.threshold else "STILL ABOVE THRESHOLD" + print( + f"Stem rules: {len(all_names)} raw -> {combined_stems} stems " + f"(collapse {action_counts['collapse']}, preserve " + f"{action_counts['preserve']}, drop {action_counts['drop']}). " + f"threshold {args.threshold}: {status}.", + file=sys.stderr, + ) + if combined_stems > args.threshold: + print( + " Reduce further: broaden collapse rules or drop more low-value " + "families, then re-run apply-stem-rules.", + file=sys.stderr, + ) + + +# =========================================================================== +# apply-map +# =========================================================================== + + +def _load_map_side(map_doc, side, name): + """Extract one trace's {key: unified} map from the LLM map document. + + Accepts either ``map_`` / ``map_`` keys or a nested + ``{trace_a: {map: {...}}}`` shape. + """ + for cand in (f"map_{side}", f"map_{name}", side, name): + if cand in map_doc: + val = map_doc[cand] + if isinstance(val, dict) and "map" in val: + return val["map"] + if isinstance(val, dict): + return val + return {} + + +def _apply_side(labels, unified_map, raw_to_stem): + """Write semantic_block = unified name for every kernel; return stats.""" + n_mapped = 0 + n_stemmed = 0 + for k in labels.get("labeled_kernels", []): + raw = k.get("name", "") + base = raw + if raw_to_stem is not None: + base = raw_to_stem.get(raw, raw) + if base != raw: + n_stemmed += 1 + unified = unified_map.get(base, base) + if base in unified_map: + n_mapped += 1 + k["semantic_block"] = unified + return { + "kernels": len(labels.get("labeled_kernels", [])), + "mapped": n_mapped, + "stemmed": n_stemmed, + } + + +def cmd_apply_map(args): # pragma: no cover + labels_a = _load(args.labels_a) + labels_b = _load(args.labels_b) + map_doc = _load(args.map) + raw_to_stem = _load(args.raw_to_stem) if args.raw_to_stem else None + + map_a = _load_map_side(map_doc, "a", args.name_a) + map_b = _load_map_side(map_doc, "b", args.name_b) + + stats_a = _apply_side(labels_a, map_a, raw_to_stem) + stats_b = _apply_side(labels_b, map_b, raw_to_stem) + + with open(args.labels_a, "w") as f: + json.dump(labels_a, f, indent=2) + with open(args.labels_b, "w") as f: + json.dump(labels_b, f, indent=2) + + blocks_a = set(k["semantic_block"] for k in labels_a["labeled_kernels"]) + blocks_b = set(k["semantic_block"] for k in labels_b["labeled_kernels"]) + shared = blocks_a & blocks_b + + print(f"Applied to {args.name_a}: {stats_a}", file=sys.stderr) + print(f"Applied to {args.name_b}: {stats_b}", file=sys.stderr) + print( + f"Unified vocabulary: {len(shared)} shared blocks, " + f"{len(blocks_a - shared)} {args.name_a}-only, " + f"{len(blocks_b - shared)} {args.name_b}-only.", + file=sys.stderr, + ) + + +# =========================================================================== +# CLI +# =========================================================================== + + +def _add_common(p): # pragma: no cover + p.add_argument("--labels-a", required=True, help="Trace A semantic_labels.json") + p.add_argument("--labels-b", required=True, help="Trace B semantic_labels.json") + p.add_argument("--name-a", default="trace_a", help="Short name for trace A") + p.add_argument("--name-b", default="trace_b", help="Short name for trace B") + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser( + description="Name-first cross-trace kernel-name unification " + "(prepare-context / apply-stem-rules / apply-map)" + ) + sub = parser.add_subparsers(dest="command", required=True) + + p_prep = sub.add_parser("prepare-context", help="Build the LLM unification context") + _add_common(p_prep) + p_prep.add_argument("-o", "--output", required=True, help="Output context JSON") + p_prep.add_argument( + "--threshold", + type=int, + default=DEFAULT_THRESHOLD, + help=f"Combined-unique count above which stem preprocessing is flagged " + f"(default {DEFAULT_THRESHOLD})", + ) + p_prep.add_argument( + "--sample-size", + type=int, + default=DEFAULT_SAMPLE_SIZE, + help=f"Names to sample when over threshold (default {DEFAULT_SAMPLE_SIZE})", + ) + p_prep.set_defaults(func=cmd_prepare_context) + + p_stem = sub.add_parser( + "apply-stem-rules", + help="Apply LLM-authored stem_rules.json and emit a reduced context", + ) + _add_common(p_stem) + p_stem.add_argument("--rules", required=True, help="stem_rules.json from the LLM") + p_stem.add_argument( + "--raw-to-stem", help="Output path for the raw-name -> stem map" + ) + p_stem.add_argument("-o", "--output", required=True, help="Reduced context JSON") + p_stem.add_argument( + "--threshold", + type=int, + default=DEFAULT_THRESHOLD, + help=f"Target combined-stem count (default {DEFAULT_THRESHOLD})", + ) + p_stem.set_defaults(func=cmd_apply_stem_rules) + + p_apply = sub.add_parser( + "apply-map", + help="Apply kernel_unification_map.json onto both label files", + ) + _add_common(p_apply) + p_apply.add_argument("--map", required=True, help="kernel_unification_map.json") + p_apply.add_argument( + "--raw-to-stem", help="raw-name -> stem map (if stem preprocessing was used)" + ) + p_apply.set_defaults(func=cmd_apply_map) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/match_and_compare.py b/TraceLens/Agent/Analysis/semantic_analyses/match_and_compare.py new file mode 100644 index 000000000..956c3eae3 --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/match_and_compare.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Match two semantic breakdowns and compute comparison stats. + +Takes two semantic_labels.json files (one per trace) that have already been +individually broken down with matching semantic_block vocabularies. + +Input: two semantic_labels.json files +Output: comparison.csv + +Usage: + python match_and_compare.py \ + --name-a --name-b \ + [-o comparison.csv] + + For multi-region (per steady-state), use --regions-dir-a and --regions-dir-b: + python match_and_compare.py --regions-dir-a output/ --regions-dir-b output/ \ + --name-a --name-b -o comparison.csv + + This matches regions by subdir name (e.g. prefill_only_3072) and compares + only corresponding regions (apples-to-apples). +""" + +import argparse +import csv +import json +import os +import sys +from collections import OrderedDict + +from _helpers import load_labels, load_json + + +def aggregate(labeled_kernels): + """Aggregate labeled kernels by semantic_block. + + Carries forward perf_category and nn_module from kernel data. + """ + blocks = OrderedDict() + for k in labeled_kernels: + block = k["semantic_block"] + if block not in blocks: + blocks[block] = { + "names": set(), + "durs": [], + "count": 0, + "perf_category": k.get("perf_category"), + "nn_module": k.get("nn_module"), + } + b = blocks[block] + b["names"].add(k["name"]) + b["durs"].append(k["dur"]) + b["count"] += 1 + return blocks + + +def load_metadata(path): # pragma: no cover + """Load metadata.json and return gpu_timeline dict if present.""" + if not path or not os.path.exists(path): + return None + try: + meta = load_json(path) + return meta.get("gpu_timeline") + except (json.JSONDecodeError, OSError): + return None + + +def build_comparison( + agg_a, + agg_b, + total_a, + total_b, + name_a, + name_b, + region=None, + gpu_timeline_a=None, + gpu_timeline_b=None, +): + """Build comparison rows for all semantic blocks present in either trace.""" + all_blocks = list(OrderedDict.fromkeys(list(agg_a.keys()) + list(agg_b.keys()))) + has_gpu_timeline = gpu_timeline_a is not None and gpu_timeline_b is not None + + rows = [] + for i, block in enumerate(all_blocks): + a = agg_a.get(block, {"names": set(), "durs": [], "count": 0}) + b = agg_b.get(block, {"names": set(), "durs": [], "count": 0}) + + a_total = sum(a["durs"]) + b_total = sum(b["durs"]) + a_avg = a_total / a["count"] if a["count"] else 0 + b_avg = b_total / b["count"] if b["count"] else 0 + a_pct = 100 * a_total / total_a if total_a > 0 else 0 + b_pct = 100 * b_total / total_b if total_b > 0 else 0 + ratio = a_total / b_total if b_total > 0 else float("inf") + gap = a_total - b_total + + row = OrderedDict() + if region is not None: + row["region"] = region + row["semantic_block"] = block + kernel_names_a_str = " | ".join(sorted(a["names"])) + kernel_names_b_str = " | ".join(sorted(b["names"])) + row["perf_category"] = ( + a.get("perf_category") or b.get("perf_category") or "Others" + ) + row["nn_module"] = a.get("nn_module") or b.get("nn_module") or "" + row["algorithm_order"] = i + 1 + row[f"{name_a}_kernel_names"] = kernel_names_a_str + row[f"{name_a}_kernel_count"] = a["count"] + row[f"{name_a}_total_us"] = round(a_total, 2) + row[f"{name_a}_avg_us"] = round(a_avg, 2) + row[f"{name_a}_pct"] = round(a_pct, 1) + row[f"{name_b}_kernel_names"] = kernel_names_b_str + row[f"{name_b}_kernel_count"] = b["count"] + row[f"{name_b}_total_us"] = round(b_total, 2) + row[f"{name_b}_avg_us"] = round(b_avg, 2) + row[f"{name_b}_pct"] = round(b_pct, 1) + row[f"{name_a}_vs_{name_b}_ratio"] = ( + round(ratio, 3) if ratio != float("inf") else "inf" + ) + row[f"{name_a}_minus_{name_b}_us"] = round(gap, 2) + + if has_gpu_timeline: + row[f"{name_a}_busy_ms"] = round( + gpu_timeline_a.get("busy_time_us", 0) / 1000, 2 + ) + row[f"{name_a}_idle_pct"] = gpu_timeline_a.get("idle_pct") + row[f"{name_b}_busy_ms"] = round( + gpu_timeline_b.get("busy_time_us", 0) / 1000, 2 + ) + row[f"{name_b}_idle_pct"] = gpu_timeline_b.get("idle_pct") + + rows.append(row) + return rows + + +def run_assertions(rows, labeled_a, labeled_b, total_a, total_b, name_a, name_b): + errors = [] + + a_count = sum(r[f"{name_a}_kernel_count"] for r in rows) + if a_count != len(labeled_a): + errors.append( + f"A6.1 FAIL: {name_a} kernel count mismatch: {a_count} matched vs {len(labeled_a)} total" + ) + + b_count = sum(r[f"{name_b}_kernel_count"] for r in rows) + if b_count != len(labeled_b): + errors.append( + f"A6.2 FAIL: {name_b} kernel count mismatch: {b_count} matched vs {len(labeled_b)} total" + ) + + a_time = sum(r[f"{name_a}_total_us"] for r in rows) + if abs(a_time - total_a) > 1.0: + errors.append( + f"A6.3 FAIL: {name_a} time mismatch: {a_time:.1f} vs {total_a:.1f}" + ) + b_time = sum(r[f"{name_b}_total_us"] for r in rows) + if abs(b_time - total_b) > 1.0: + errors.append( + f"A6.3 FAIL: {name_b} time mismatch: {b_time:.1f} vs {total_b:.1f}" + ) + + a_pct = sum(r[f"{name_a}_pct"] for r in rows) + if abs(a_pct - 100.0) > 2.0: + errors.append(f"A7.2 FAIL: {name_a} percentages sum to {a_pct:.1f}%") + b_pct = sum(r[f"{name_b}_pct"] for r in rows) + if abs(b_pct - 100.0) > 2.0: + errors.append(f"A7.2 FAIL: {name_b} percentages sum to {b_pct:.1f}%") + + for r in rows: + a_t = r[f"{name_a}_total_us"] + b_t = r[f"{name_b}_total_us"] + if b_t > 0: + expected_ratio = round(a_t / b_t, 3) + actual_ratio = r[f"{name_a}_vs_{name_b}_ratio"] + if actual_ratio != "inf" and abs(expected_ratio - actual_ratio) > 0.1: + errors.append( + f"A7.5 FAIL: {r['semantic_block']}: ratio mismatch " + f"{expected_ratio} vs {actual_ratio}" + ) + + return errors + + +def main(): # pragma: no cover + parser = argparse.ArgumentParser(description="Compare two semantic breakdowns") + parser.add_argument( + "labels_a", nargs="?", help="Path to trace A semantic_labels.json" + ) + parser.add_argument( + "labels_b", nargs="?", help="Path to trace B semantic_labels.json" + ) + parser.add_argument( + "--regions-dir-a", + help="Dir with per-region subdirs (e.g. prefill_only_3072/semantic_labels.json)", + ) + parser.add_argument( + "--regions-dir-b", help="Dir with per-region subdirs for trace B" + ) + parser.add_argument("--name-a", default="trace_a", help="Short name for trace A") + parser.add_argument("--name-b", default="trace_b", help="Short name for trace B") + parser.add_argument( + "--region", + help="Region descriptor (e.g. prefill_only_3072) for multi-section reports", + ) + parser.add_argument( + "-o", "--output", default="comparison.csv", help="Output CSV path" + ) + args = parser.parse_args() + + all_rows = [] + fieldnames = None + + if args.regions_dir_a and args.regions_dir_b: + regions_a = { + d + for d in os.listdir(args.regions_dir_a) + if os.path.isdir(os.path.join(args.regions_dir_a, d)) + } + regions_b = { + d + for d in os.listdir(args.regions_dir_b) + if os.path.isdir(os.path.join(args.regions_dir_b, d)) + } + common_regions = sorted(regions_a & regions_b) + for region in common_regions: + labels_a = os.path.join(args.regions_dir_a, region, "semantic_labels.json") + labels_b = os.path.join(args.regions_dir_b, region, "semantic_labels.json") + metadata_a_path = os.path.join(args.regions_dir_a, region, "metadata.json") + metadata_b_path = os.path.join(args.regions_dir_b, region, "metadata.json") + if not os.path.exists(labels_a) or not os.path.exists(labels_b): + continue + data_a = load_labels(labels_a) + data_b = load_labels(labels_b) + labeled_a = data_a["labeled_kernels"] + labeled_b = data_b["labeled_kernels"] + total_a = data_a.get( + "total_kernel_time_us", sum(k["dur"] for k in labeled_a) + ) + total_b = data_b.get( + "total_kernel_time_us", sum(k["dur"] for k in labeled_b) + ) + gpu_timeline_a = load_metadata(metadata_a_path) + gpu_timeline_b = load_metadata(metadata_b_path) + agg_a = aggregate(labeled_a) + agg_b = aggregate(labeled_b) + rows = build_comparison( + agg_a, + agg_b, + total_a, + total_b, + args.name_a, + args.name_b, + region=region, + gpu_timeline_a=gpu_timeline_a, + gpu_timeline_b=gpu_timeline_b, + ) + all_rows.extend(rows) + if fieldnames is None and rows: + fieldnames = list(rows[0].keys()) + if not all_rows: + print("No matching regions found", file=sys.stderr) + sys.exit(1) + else: + if not args.labels_a or not args.labels_b: + parser.error( + "Provide labels_a and labels_b, or --regions-dir-a and --regions-dir-b" + ) + data_a = load_labels(args.labels_a) + data_b = load_labels(args.labels_b) + labeled_a = data_a["labeled_kernels"] + labeled_b = data_b["labeled_kernels"] + total_a = data_a.get("total_kernel_time_us", sum(k["dur"] for k in labeled_a)) + total_b = data_b.get("total_kernel_time_us", sum(k["dur"] for k in labeled_b)) + agg_a = aggregate(labeled_a) + agg_b = aggregate(labeled_b) + all_rows = build_comparison( + agg_a, + agg_b, + total_a, + total_b, + args.name_a, + args.name_b, + region=args.region, + ) + fieldnames = list(all_rows[0].keys()) if all_rows else [] + if all_rows: + errors = run_assertions( + all_rows, + labeled_a, + labeled_b, + total_a, + total_b, + args.name_a, + args.name_b, + ) + for e in errors: + print(e, file=sys.stderr) + if any("FAIL" in e for e in errors): + sys.exit(1) + + if fieldnames and all_rows: + with open(args.output, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(all_rows) + print(f"Wrote {args.output} ({len(all_rows)} rows)", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/pattern_finder.py b/TraceLens/Agent/Analysis/semantic_analyses/pattern_finder.py new file mode 100644 index 000000000..85a259a14 --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/pattern_finder.py @@ -0,0 +1,946 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Discover repeating kernel patterns in GPU traces using loop detection. + +Handles multi-stream traces by detecting the primary GPU stream (most +kernels) and running pattern discovery on that stream only. All output +indices refer to the **original** kernel list so downstream tools do not +need to know about the stream split. + +The loop-detection engine (data classes + iterative pattern-finding +algorithm with extension, splitting, and rotation) is included in-file. + +Usage: + python pattern_finder.py [-o pattern.json] +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from collections import Counter +from dataclasses import dataclass +from typing import Dict, List, Optional, Set, Tuple, Union + +logger = logging.getLogger(__name__) + + +# =========================================================================== +# Loop-detection engine (data classes + iterative pattern-finding algorithm) +# =========================================================================== + +# --------------------------------------------------------------------------- +# Data classes +# --------------------------------------------------------------------------- + + +@dataclass(slots=True) +class Operation: + name: str + ts: float + dur: float + + +@dataclass(slots=True) +class GPUOperation(Operation): + external_id: Optional[int] = None + correlation: Optional[int] = None + cat: str = "" + stream: Optional[int] = None + + @property + def end_ts(self) -> float: + return self.ts + self.dur + + +# --------------------------------------------------------------------------- +# Patterns – shared pattern vocabulary +# --------------------------------------------------------------------------- + + +class Patterns: + __slots__ = ("patterns", "pattern_operation_names", "operation_name_to_idx") + + def __init__( + self, + patterns: Optional[List[List[int]]] = None, + pattern_operation_names: Optional[Dict[int, str]] = None, + operation_name_to_idx: Optional[Dict[str, int]] = None, + ): + self.patterns: List[List[int]] = patterns if patterns is not None else [] + self.pattern_operation_names: Dict[int, str] = ( + pattern_operation_names if pattern_operation_names is not None else {} + ) + self.operation_name_to_idx: Dict[str, int] = ( + operation_name_to_idx if operation_name_to_idx is not None else {} + ) + if self.pattern_operation_names and not self.operation_name_to_idx: + self.operation_name_to_idx = { + name: idx for idx, name in self.pattern_operation_names.items() + } + elif self.operation_name_to_idx and not self.pattern_operation_names: + self.pattern_operation_names = { + idx: name for name, idx in self.operation_name_to_idx.items() + } + + def get_pattern_name(self, pattern_idx: int) -> str: + if pattern_idx < 0 or pattern_idx >= len(self.patterns): + return "?" + if pattern_idx < 26: + return chr(ord("A") + pattern_idx) + offset = pattern_idx - 26 + letter = chr(ord("A") + (offset % 26)) + suffix = (offset // 26) + 1 + return f"{letter}{suffix}" + + def _validate_pattern_indices_mapped(self) -> None: + if not self.patterns: + return + for op_name, op_idx in self.operation_name_to_idx.items(): + if op_idx not in self.pattern_operation_names: + self.pattern_operation_names[op_idx] = op_name + used_indices = {op_idx for pattern in self.patterns for op_idx in pattern} + missing = sorted( + idx for idx in used_indices if idx not in self.pattern_operation_names + ) + if missing: + raise ValueError( + "All pattern indices must map to operation names. " + f"Missing mappings for indices: {missing}" + ) + + +# --------------------------------------------------------------------------- +# LoopStructures – per-stream detection results +# --------------------------------------------------------------------------- + + +class LoopStructures: + + def __init__( + self, + operations: List[Operation], + _patterns_obj: Optional[Patterns] = None, + stream_id: Optional[int] = None, + ): + self.operations = operations + self.stream_id = stream_id + self.sequences: List[List[List[int]]] = [] + self._patterns_obj: Patterns + + if _patterns_obj is not None: + self._patterns_obj = _patterns_obj + else: + self._patterns_obj = Patterns() + if operations: + unique_names: List[str] = [] + seen: set = set() + for op in operations: + if op.name not in seen: + unique_names.append(op.name) + seen.add(op.name) + self._patterns_obj.pattern_operation_names = { + i: name for i, name in enumerate(unique_names) + } + self._patterns_obj.operation_name_to_idx = { + name: idx + for idx, name in self._patterns_obj.pattern_operation_names.items() + } + + if self.operations: + for op in self.operations: + if op.name not in self._patterns_obj.operation_name_to_idx: + new_idx = len(self._patterns_obj.operation_name_to_idx) + self._patterns_obj.operation_name_to_idx[op.name] = new_idx + self._patterns_obj.pattern_operation_names[new_idx] = op.name + + self._patterns_obj._validate_pattern_indices_mapped() + + # -- Property wrappers --------------------------------------------------- + + @property + def patterns(self) -> List[List[int]]: + return self._patterns_obj.patterns + + @property + def pattern_operation_names(self) -> Dict[int, str]: + return self._patterns_obj.pattern_operation_names + + @property + def operation_name_to_idx(self) -> Dict[str, int]: + return self._patterns_obj.operation_name_to_idx + + def get_pattern_name(self, pattern_idx: int) -> str: + return self._patterns_obj.get_pattern_name(pattern_idx) + + # -- Coverage helpers ---------------------------------------------------- + + def get_covered_indices( + self, + pattern_indices: Optional[Union[int, List[int]]] = None, + ) -> set: + covered: set = set() + if pattern_indices is None: + patterns_to_include = range(len(self.patterns)) + elif isinstance(pattern_indices, int): + patterns_to_include = [pattern_indices] + else: + patterns_to_include = pattern_indices + for idx in patterns_to_include: + if idx < len(self.sequences): + for seq in self.sequences[idx]: + covered.update(seq) + return covered + + def _get_pattern_coverage(self, pattern_idx: Optional[int] = None) -> float: + if not self.operations: + return 0.0 + covered_indices = self.get_covered_indices(pattern_idx) + return len(covered_indices) / len(self.operations) + + def get_pattern_coverage(self) -> float: + return self._get_pattern_coverage() + + +# =========================================================================== +# Algorithm – iterative pattern finding with extension, splitting, rotation +# =========================================================================== + + +def get_operation_frequencies(loop_structures: LoopStructures) -> Counter: + covered = loop_structures.get_covered_indices() + unused_op_indices = [ + loop_structures.operation_name_to_idx[op.name] + for i, op in enumerate(loop_structures.operations) + if i not in covered + ] + return Counter(unused_op_indices) + + +# --------------------------------------------------------------------------- +# Pattern extension +# --------------------------------------------------------------------------- + + +def extend_pattern( + loop_structures: LoopStructures, + pattern_idx: int, + min_loop_count: int, + min_pattern_length: int, + snippets: bool = False, +) -> int: + covered = loop_structures.get_covered_indices( + [i for i in range(len(loop_structures.patterns)) if i != pattern_idx] + ) + best_agree = 0 + + while True: + sequences = loop_structures.sequences[pattern_idx] + names_before = [ + ( + loop_structures.operations[seq[0] - 1].name + if not (seq[0] == 0 or seq[0] - 1 in covered) + else None + ) + for i, seq in enumerate(sequences) + ] + names_after = [ + ( + loop_structures.operations[seq[-1] + 1].name + if not ( + seq[-1] == len(loop_structures.operations) - 1 + or seq[-1] + 1 in covered + ) + else None + ) + for i, seq in enumerate(sequences) + ] + + extendable_before = sum(1 for name in names_before if name is not None) + extendable_after = sum(1 for name in names_after if name is not None) + if extendable_before < min_loop_count and extendable_after < min_loop_count: + best_agree = max(extendable_before, extendable_after) + break + + name_counts_before = Counter( + n for n in names_before if n is not None + ).most_common(1) + name_counts_after = Counter( + n for n in names_after if n is not None + ).most_common(1) + + kernel_before, agree_before = ( + name_counts_before[0] if name_counts_before else (None, 0) + ) + kernel_after, agree_after = ( + name_counts_after[0] if name_counts_after else (None, 0) + ) + best_agree = max(agree_before, agree_after) + + if agree_before < min_loop_count and agree_after < min_loop_count: + break + + current_seq_count = len(sequences) + if ( + snippets + and len(loop_structures.patterns[pattern_idx]) >= min_pattern_length + and best_agree < current_seq_count + ): + break + + extend_before = agree_before >= agree_after and agree_before >= min_loop_count + kernel_name = kernel_before if extend_before else kernel_after + op_idx = loop_structures.operation_name_to_idx[kernel_name] + + if extend_before: + loop_structures.patterns[pattern_idx].insert(0, op_idx) + else: + loop_structures.patterns[pattern_idx].append(op_idx) + + new_sequences = [] + for i, seq in enumerate(sequences): + if extend_before: + if names_before[i] == kernel_name: + new_start = seq[0] - 1 + if new_sequences and new_sequences[-1][-1] >= new_start: + continue + seq.insert(0, new_start) + new_sequences.append(seq) + else: + if names_after[i] == kernel_name: + if new_sequences and new_sequences[-1][-1] >= seq[0]: + continue + seq.append(seq[-1] + 1) + new_sequences.append(seq) + + loop_structures.sequences[pattern_idx] = new_sequences + + if len(new_sequences) >= min_pattern_length and any( + new_sequences[i - 1][-1] + 1 == new_sequences[i][0] + and new_sequences[i][-1] + 1 == new_sequences[i + 1][0] + for i in range(1, len(new_sequences) - 1) + ): + break + + return best_agree + + +# --------------------------------------------------------------------------- +# Pattern splitting +# --------------------------------------------------------------------------- + + +def _analyze_pattern_for_splits( + loop_structures: LoopStructures, + pattern_idx: int, + min_pattern_length: int, + split_pattern_length_limit: int, +) -> List[List[int]]: + pattern = loop_structures.patterns[pattern_idx] + if len(pattern) <= split_pattern_length_limit: + return [] + if not loop_structures.sequences[pattern_idx]: + return [] + + mini_operations = [ + GPUOperation( + name=loop_structures.pattern_operation_names[op_idx], + ts=0.0, + dur=0.0, + ) + for op_idx in pattern + ] + + mini_ls = find_kernel_loops_single_stream( + mini_operations, + min_loop_count=2, + min_pattern_length=min_pattern_length, + split_pattern_length_limit=split_pattern_length_limit, + ) + + if not mini_ls.patterns: + return [] + + sub_patterns: List[List[int]] = [] + for sp_idx, mini_seqs in enumerate(mini_ls.sequences): + all_consecutive = len(mini_seqs) >= 2 and all( + mini_seqs[i][0] == mini_seqs[i - 1][-1] + 1 + for i in range(1, len(mini_seqs)) + ) + if not all_consecutive: + continue + global_pat = [ + loop_structures.operation_name_to_idx[mini_ls.pattern_operation_names[mi]] + for mi in mini_ls.patterns[sp_idx] + ] + sub_patterns.append(global_pat) + + return sub_patterns + + +def _apply_pattern_split( + loop_structures: LoopStructures, + pattern_idx: int, + sub_patterns: List[List[int]], +) -> None: + loop_structures.patterns.pop(pattern_idx) + loop_structures.sequences.pop(pattern_idx) + + affected_indices: Set[int] = set() + for _sp_idx, pat in enumerate(sub_patterns): + existing_idx = next( + (i for i, p in enumerate(loop_structures.patterns) if p == pat), + None, + ) + if existing_idx is not None: + affected_indices.add(existing_idx) + else: + loop_structures.patterns.append(pat) + loop_structures.sequences.append([]) + affected_indices.add(len(loop_structures.patterns) - 1) + + for pidx in sorted(affected_indices): + all_seqs = find_pattern_occurrences(loop_structures, pidx) + loop_structures.sequences[pidx] = all_seqs + + +def split_long_pattern( + loop_structures: LoopStructures, + pattern_idx: int, + min_pattern_length: int, + split_pattern_length_limit: int, +) -> bool: + sub_patterns = _analyze_pattern_for_splits( + loop_structures, + pattern_idx, + min_pattern_length, + split_pattern_length_limit, + ) + if not sub_patterns: + return False + _apply_pattern_split(loop_structures, pattern_idx, sub_patterns) + return True + + +# --------------------------------------------------------------------------- +# Pattern rotation +# --------------------------------------------------------------------------- + + +def rotate_pattern(loop_structures: LoopStructures, pattern_idx: int) -> bool: + pattern = loop_structures.patterns[pattern_idx] + sequences = loop_structures.sequences[pattern_idx] + pattern_len = len(pattern) + + shift_steps = 0 + while shift_steps < pattern_len: + all_match = all( + seq[-1] + shift_steps < len(loop_structures.operations) + and loop_structures.operation_name_to_idx[ + loop_structures.operations[seq[-1] + shift_steps].name + ] + == pattern[shift_steps] + for seq in sequences + ) + if all_match: + shift_steps += 1 + else: + break + + shift_steps -= 1 + if shift_steps < 0: + return False + + loop_structures.patterns[pattern_idx] = ( + pattern[shift_steps:] + pattern[:shift_steps] + ) + loop_structures.sequences[pattern_idx] = [ + list(range(seq[0] + shift_steps, seq[-1] + shift_steps + 1)) + for seq in sequences + ] + sequences = loop_structures.sequences[pattern_idx] + + covered = loop_structures.get_covered_indices( + [i for i in range(len(loop_structures.patterns)) if i != pattern_idx] + ) + + new_sequences = [] + for i, seq in enumerate(sequences): + if i > 0 and sequences[i - 1][-1] + 1 == seq[0]: + continue + potential_start = seq[0] - pattern_len + if potential_start >= 0 and all( + (potential_start + j) not in covered for j in range(pattern_len) + ): + if all( + loop_structures.operation_name_to_idx[ + loop_structures.operations[potential_start + j].name + ] + == loop_structures.patterns[pattern_idx][j] + for j in range(pattern_len) + ): + new_sequences.append( + list(range(potential_start, potential_start + pattern_len)) + ) + covered.update(range(potential_start, potential_start + pattern_len)) + + if new_sequences: + loop_structures.sequences[pattern_idx] = sorted( + sequences + new_sequences, + key=lambda seq: seq[0], + ) + return True + + +# --------------------------------------------------------------------------- +# Pattern occurrence finding +# --------------------------------------------------------------------------- + + +def find_pattern_occurrences( + loop_structures: LoopStructures, + pattern_idx: int, +) -> List[List[int]]: + pattern = loop_structures.patterns[pattern_idx] + pattern_len = len(pattern) + + covered = loop_structures.get_covered_indices( + [i for i in range(len(loop_structures.patterns)) if i != pattern_idx] + ) + + op_name_indices = [ + loop_structures.operation_name_to_idx[op.name] + for op in loop_structures.operations + ] + + first_element = pattern[0] + max_start = len(loop_structures.operations) - pattern_len + + candidates = [] + for i in range(max_start + 1): + if op_name_indices[i] != first_element or i in covered: + continue + if all((i + j) not in covered for j in range(1, pattern_len)) and all( + op_name_indices[i + j] == pattern[j] for j in range(1, pattern_len) + ): + candidates.append(i) + + sequences = [] + for i in candidates: + if sequences and i < sequences[-1][-1] + 1: + continue + sequences.append(list(range(i, i + pattern_len))) + + return sequences + + +# --------------------------------------------------------------------------- +# Seed pattern creation +# --------------------------------------------------------------------------- + + +def get_new_pattern_to_start( + loop_structures: LoopStructures, + operations: List[GPUOperation], + min_loop_count: int, + failed_ops: Optional[Set[int]] = None, +) -> Optional[Tuple[List[int], List[List[int]]]]: + if failed_ops is None: + failed_ops = set() + + covered = loop_structures.get_covered_indices() + unused_op_indices = [ + loop_structures.operation_name_to_idx[op.name] + for i, op in enumerate(operations) + if i not in covered + ] + op_frequencies = Counter(unused_op_indices) + if not op_frequencies: + return None + + for most_frequent_kernel, freq in op_frequencies.most_common(): + if most_frequent_kernel in failed_ops or freq < min_loop_count: + continue + kernel_name = loop_structures.pattern_operation_names[most_frequent_kernel] + new_pattern = [most_frequent_kernel] + new_sequences = [ + [i] + for i, op in enumerate(operations) + if op.name == kernel_name and i not in covered + ] + if len(new_sequences) >= min_loop_count: + return (new_pattern, new_sequences) + + return None + + +# =========================================================================== +# Loop detection main entry point +# =========================================================================== + + +def find_kernel_loops_single_stream( + operations: List[GPUOperation], + min_loop_count: int = 20, + min_pattern_length: int = 5, + split_pattern_length_limit: int = 35, + loop_structures: Optional[LoopStructures] = None, + stream_id: Optional[int] = None, + snippets: bool = False, +) -> LoopStructures: + """Find repeating kernel patterns in a single-stream operation list.""" + + if stream_id is None and operations: + streams_in_ops = {op.stream for op in operations} + if len(streams_in_ops) == 1: + stream_id = next(iter(streams_in_ops)) + + if loop_structures is not None: + if loop_structures.operations != operations: + raise ValueError("loop_structures must have the same operations list") + else: + loop_structures = LoopStructures(operations=operations, stream_id=stream_id) + + mpl_schedule = list(range(max(20, min_pattern_length), min_pattern_length - 1, -1)) + + if min_loop_count > 20: + pass_schedule: List[Tuple[int, int]] = [ + (mpl, min_loop_count) for mpl in mpl_schedule + ] + else: + mlc_start = max(20, min_loop_count) + pass_schedule = [(mpl, mlc_start) for mpl in mpl_schedule] + if min_loop_count < 20: + _mlc_relaxed = list(range(19, min_loop_count - 1, -1)) + pass_schedule += [(min_pattern_length, mlc) for mlc in _mlc_relaxed] + + max_iterations_per_pass = 50 + + extension_cache: Dict[int, Tuple[int, int, int]] = {} + split_analysis_cache: Dict[Tuple[int, ...], List[List[int]]] = {} + + for current_min_pattern_length, current_min_loop_count in pass_schedule: + if loop_structures._get_pattern_coverage() >= 1.0: + break + + failed_ops = { + op + for op, (alen, _, bagree) in extension_cache.items() + if alen < current_min_pattern_length and bagree < current_min_loop_count + } + + for pidx in range(len(loop_structures.patterns) - 1, -1, -1): + if len(loop_structures.patterns[pidx]) <= split_pattern_length_limit: + continue + pattern_key = tuple(loop_structures.patterns[pidx]) + sub_patterns = split_analysis_cache.get(pattern_key) + if not sub_patterns: + continue + min_sub_len = min(len(sp) for sp in sub_patterns) + if min_sub_len >= current_min_pattern_length: + _apply_pattern_split(loop_structures, pidx, sub_patterns) + + iteration = 0 + while iteration < max_iterations_per_pass: + pattern_idx = len(loop_structures.patterns) + + result = get_new_pattern_to_start( + loop_structures, + operations, + current_min_loop_count, + failed_ops, + ) + if result is None: + break + + new_pattern, new_sequences = result + pattern_init_op = new_pattern[0] + loop_structures.patterns.append(new_pattern) + loop_structures.sequences.append(new_sequences) + pattern_idx = len(loop_structures.patterns) - 1 + + if loop_structures._get_pattern_coverage() >= 1.0: + iteration += 1 + continue + + stop_best_agree = extend_pattern( + loop_structures, + pattern_idx, + current_min_loop_count, + current_min_pattern_length, + snippets=snippets, + ) + + pattern_len = len(loop_structures.patterns[pattern_idx]) + seq_count = len(loop_structures.sequences[pattern_idx]) + + if pattern_len < min_pattern_length: + if loop_structures._get_pattern_coverage() < 1.0: + loop_structures.patterns.pop() + loop_structures.sequences.pop() + extension_cache[pattern_init_op] = ( + pattern_len, + seq_count, + stop_best_agree, + ) + failed_ops.add(pattern_init_op) + + elif pattern_len > split_pattern_length_limit: + pattern_key = tuple(loop_structures.patterns[pattern_idx]) + if pattern_key not in split_analysis_cache: + sub_patterns = _analyze_pattern_for_splits( + loop_structures, + pattern_idx, + min_pattern_length, + split_pattern_length_limit, + ) + split_analysis_cache[pattern_key] = sub_patterns + else: + sub_patterns = split_analysis_cache[pattern_key] + + if sub_patterns: + min_sub_len = min(len(sp) for sp in sub_patterns) + if min_sub_len >= current_min_pattern_length: + _apply_pattern_split( + loop_structures, + pattern_idx, + sub_patterns, + ) + else: + all_seqs = find_pattern_occurrences(loop_structures, pattern_idx) + loop_structures.sequences[pattern_idx] = all_seqs + + iteration += 1 + + return loop_structures + + +# =========================================================================== +# Pattern finder – trace-level repeating-pattern discovery +# =========================================================================== + + +def _kernels_to_gpu_operations(kernels): + """Convert extracted kernel dicts to GPUOperation objects.""" + return [ + GPUOperation(name=k["name"], ts=k.get("ts", 0.0), dur=k.get("dur", 0.0)) + for k in kernels + ] + + +def _detect_primary_stream(kernels): + """Identify the primary stream when multiple streams are present. + + Returns (primary_stream_id, idx_map, secondary_indices) where + - primary_stream_id is the stream with the most kernels (None if single-stream) + - idx_map maps dense primary-only indices back to original indices + - secondary_indices is a sorted list of original indices on non-primary streams + If only one stream exists, returns (None, None, []). + """ + stream_counts = Counter(k.get("stream_id") for k in kernels) + non_null = {s: c for s, c in stream_counts.items() if s is not None} + + if len(non_null) <= 1: + return None, None, [] + + primary = max(non_null, key=non_null.get) + idx_map = [] + secondary = [] + for i, k in enumerate(kernels): + sid = k.get("stream_id") + if sid == primary or sid is None: + idx_map.append(i) + else: + secondary.append(i) + + return primary, idx_map, secondary + + +def find_repeating_pattern( + extracted_data, min_loop_count=10, min_pattern_length=3, max_pattern_length=100 +): + """Discover repeating kernel patterns in a trace. + + Automatically detects multi-stream traces and runs pattern discovery + on the primary stream only, remapping indices to the original list. + + Args: + extracted_data: dict from extract_trace_data.py with "kernels" list + min_loop_count: minimum number of pattern repetitions required + min_pattern_length: minimum ops in a pattern to keep + max_pattern_length: maximum ops before splitting into sub-patterns + + Returns: + dict with: + - patterns: list of discovered patterns, each a list of kernel names + - sequences: for each pattern, list of (start_idx, end_idx) tuples + (indices into the original kernel list) + - pattern_labels: letter labels for each pattern (A, B, C, ...) + - coverage: fraction of primary-stream kernels covered by patterns + - total_kernels: total number of kernels in the trace + - preamble_indices: kernel indices before the first pattern occurrence + that are not covered by any pattern (original indices) + - epilogue_indices: kernel indices after the last pattern occurrence + that are not covered by any pattern (original indices) + - primary_stream_id: stream used for pattern discovery (None if single-stream) + - secondary_stream_indices: kernel indices on non-primary streams + """ + kernels = extracted_data["kernels"] + n = len(kernels) + + primary_sid, idx_map, secondary_indices = _detect_primary_stream(kernels) + + if idx_map is not None: + primary_kernels = [kernels[i] for i in idx_map] + else: + primary_kernels = kernels + idx_map = list(range(n)) + + ops = _kernels_to_gpu_operations(primary_kernels) + + loop_structures = find_kernel_loops_single_stream( + ops, + min_loop_count=min_loop_count, + min_pattern_length=min_pattern_length, + split_pattern_length_limit=max_pattern_length, + ) + + patterns = [] + sequences = [] + pattern_labels = [] + for pidx, pat_indices in enumerate(loop_structures.patterns): + pat_names = [loop_structures.pattern_operation_names[i] for i in pat_indices] + patterns.append(pat_names) + pattern_labels.append(loop_structures.get_pattern_name(pidx)) + + seqs = [] + for seq in loop_structures.sequences[pidx]: + orig_start = idx_map[seq[0]] + orig_end = idx_map[seq[-1]] + seqs.append((orig_start, orig_end)) + sequences.append(seqs) + + coverage = loop_structures.get_pattern_coverage() + + covered_dense = loop_structures.get_covered_indices() + covered_orig = set(idx_map[i] for i in covered_dense) + + # Primary loop region spans the first to the last covered kernel across ALL + # detected patterns. Patterns are ordered by seed frequency, not position, + # so using sequences[0] alone would miss earlier/later patterns and mislabel + # their gaps. covered_orig already unions every pattern's covered indices. + if covered_orig: + primary_first = min(covered_orig) + primary_last = max(covered_orig) + else: + # No repeating pattern: there is no primary region. Put all uncovered + # kernels in preamble and leave epilogue empty so none are counted twice + # (primary_first=n -> preamble=range(n); primary_last=n-1 -> + # epilogue=range(n, n) is empty). + primary_first = n + primary_last = n - 1 + + secondary_set = set(secondary_indices) + preamble_indices = [ + i + for i in range(primary_first) + if i not in covered_orig and i not in secondary_set + ] + epilogue_indices = [ + i + for i in range(primary_last + 1, n) + if i not in covered_orig and i not in secondary_set + ] + + return { + "patterns": patterns, + "sequences": sequences, + "pattern_labels": pattern_labels, + "coverage": coverage, + "total_kernels": n, + "preamble_indices": preamble_indices, + "epilogue_indices": epilogue_indices, + "primary_stream_id": primary_sid, + "secondary_stream_indices": secondary_indices, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Discover repeating kernel patterns in a GPU trace" + ) + parser.add_argument("extracted_json", help="Path to extracted trace data JSON") + parser.add_argument("-o", "--output", help="Output JSON path (default: stdout)") + parser.add_argument( + "--min-loop-count", + type=int, + default=10, + help="Minimum pattern repetitions (default: 10)", + ) + parser.add_argument( + "--min-pattern-length", + type=int, + default=3, + help="Minimum ops per pattern (default: 3)", + ) + parser.add_argument( + "--max-pattern-length", + type=int, + default=100, + help="Maximum ops before splitting (default: 100)", + ) + args = parser.parse_args() + + with open(args.extracted_json) as f: + extracted = json.load(f) + + result = find_repeating_pattern( + extracted, + min_loop_count=args.min_loop_count, + min_pattern_length=args.min_pattern_length, + max_pattern_length=args.max_pattern_length, + ) + + if result["primary_stream_id"] is not None: + print( + f"Multi-stream: primary={result['primary_stream_id']}, " + f"{len(result['secondary_stream_indices'])} secondary kernels", + file=sys.stderr, + ) + + for i, (pat, label) in enumerate(zip(result["patterns"], result["pattern_labels"])): + n_seqs = len(result["sequences"][i]) + print( + f"Pattern {label}: {len(pat)} kernels, {n_seqs} occurrences", + file=sys.stderr, + ) + + print( + f"Coverage: {result['coverage']:.1%} of {result['total_kernels']} kernels", + file=sys.stderr, + ) + print( + f"Preamble: {len(result['preamble_indices'])} kernels, " + f"Epilogue: {len(result['epilogue_indices'])} kernels", + file=sys.stderr, + ) + + output = json.dumps(result, indent=2) + if args.output: + with open(args.output, "w") as f: + f.write(output) + print(f"Wrote {args.output}", file=sys.stderr) + else: + print(output) + + +if __name__ == "__main__": + main() diff --git a/TraceLens/Agent/Analysis/semantic_analyses/trace_split_adapter.py b/TraceLens/Agent/Analysis/semantic_analyses/trace_split_adapter.py new file mode 100644 index 000000000..d7d918114 --- /dev/null +++ b/TraceLens/Agent/Analysis/semantic_analyses/trace_split_adapter.py @@ -0,0 +1,207 @@ +############################################################################### +# Copyright (c) 2024 - 2025 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Adapter for TraceUtils split_inference_trace_annotation.py. + +Invokes TraceLens.TraceUtils.split_inference_trace_annotation as a subprocess +and converts its output to the format expected by extract_trace_data.py: +[(trace_dict, region_metadata), ...]. + +Usage: + from trace_split_adapter import split_vllm_trace, get_steady_state_key + result = split_vllm_trace("trace.json.gz") + if result: + for trace_dict, metadata in result: + key = get_steady_state_key(metadata) +""" + +import json +import os +import subprocess +import sys +import tempfile +from typing import List, Optional, Tuple + +from _helpers import load_json + + +def get_steady_state_key(metadata: dict) -> str: + """Return a key for grouping by steady-state region. + + For single-iteration data the key encodes (type, token_count) so that + corresponding iterations from two traces match by directory name. + """ + ctx = metadata.get("context_requests", 0) + gen = metadata.get("generation_requests", 0) + ctx_sum = metadata.get("context_sum", 0) + gen_sum = metadata.get("generation_sum", 0) + batch = metadata.get("batch_size", 0) + + if ctx > 0 and gen == 0: + return f"prefill_only_{ctx_sum}" + if ctx == 0 and gen > 0: + return f"decode_only_{gen_sum}" + if ctx > 0 and gen > 0: + return f"prefill_decode_{ctx_sum}_{gen_sum}" + return f"prefill_decode_{batch}_{batch}" + + +def _phase_to_region_meta(phase: dict) -> dict: + """Map TraceUtils phase dict to region_meta expected by extract_trace_data.""" + num_prefill = phase.get("num_prefill", 0) + num_prefilldecode = phase.get("num_prefilldecode", 0) + num_decode = phase.get("num_decode", 0) + avg_bs = phase.get("avg_bs", 0) + avg_conc = phase.get("avg_conc", 0) + + if num_prefill > 0 and num_prefilldecode == 0 and num_decode == 0: + # Prefill-only + return { + "context_requests": num_prefill, + "generation_requests": 0, + "context_sum": avg_bs, + "generation_sum": 0, + "batch_size": avg_bs, + "num_requests": avg_conc, + } + if num_prefill == 0 and num_prefilldecode == 0 and num_decode > 0: + # Decode-only + return { + "context_requests": 0, + "generation_requests": num_decode, + "context_sum": 0, + "generation_sum": avg_bs, + "batch_size": avg_bs, + "num_requests": avg_conc, + } + # Prefill-decode (or combined) + return { + "context_requests": num_prefill + num_prefilldecode, + "generation_requests": num_decode + num_prefilldecode, + "context_sum": avg_bs, + "generation_sum": avg_bs, + "batch_size": avg_bs, + "num_requests": avg_conc, + } + + +def _load_trace(path: str) -> dict: + """Load trace JSON from .json or .json.gz file.""" + return load_json(path) + + +def _is_single_iteration(phase: dict) -> bool: + """True if phase represents a single annotation iteration.""" + total = ( + phase.get("num_prefill", 0) + + phase.get("num_prefilldecode", 0) + + phase.get("num_decode", 0) + ) + return total <= 1 + + +def _iter_type_key(phase: dict) -> str: + """Group key for deduplicating single iterations by type + token count.""" + np = phase.get("num_prefill", 0) + npd = phase.get("num_prefilldecode", 0) + nd = phase.get("num_decode", 0) + bs = phase.get("avg_bs", 0) + if np > 0: + return f"prefill_{bs}" + if npd > 0: + return f"prefilldecode_{bs}" + if nd > 0: + return f"decode_{bs}" + return f"empty_{bs}" + + +def split_vllm_trace( + trace_path: str, +) -> Optional[List[Tuple[dict, dict]]]: # pragma: no cover + """ + Split a vLLM trace using TraceUtils split_inference_trace_annotation. + + Uses --store-single-iteration to get per-annotation-iteration data, + then selects one representative iteration per unique (type, token_count) + group. This matches the reference analysis style of showing per-step + averages rather than multi-step totals. + + Returns None if no annotation iterations are found (caller should fall back + to full trace). + """ + with tempfile.TemporaryDirectory(prefix="trace_split_") as tmpdir: + cmd = [ + sys.executable, + "-m", + "TraceLens.TraceUtils.split_inference_trace_annotation", + trace_path, + "-o", + tmpdir, + "--find-steady-state", + "--store-single-iteration", + "--iterations", + "all", + "--emit-gpu-op-uid", + ] + result = subprocess.run( + cmd, + capture_output=True, + text=True, + ) + if result.returncode != 0: + print( + f"TraceUtils split failed (exit {result.returncode}): {result.stderr}", + file=sys.stderr, + ) + return None + + details_path = os.path.join(tmpdir, "execution_details.json") + if not os.path.isfile(details_path): + return None + + with open(details_path, "r") as f: + execution_details = json.load(f) + + if not execution_details: + return None + + # Collect single-iteration entries, grouped by type+token_count + groups: dict = {} # iter_type_key -> list of (entry, phase) + for entry in execution_details: + out_path = entry.get("output_path") + if not out_path or not os.path.isfile(out_path): + continue + phase = entry.get("phase") or {} + if not _is_single_iteration(phase): + continue + total_steps = ( + phase.get("num_prefill", 0) + + phase.get("num_prefilldecode", 0) + + phase.get("num_decode", 0) + ) + if total_steps == 0: + continue + key = _iter_type_key(phase) + groups.setdefault(key, []).append((entry, phase)) + + # Pick the median-busy-time representative from each group + output = [] + for key, entries in groups.items(): + entries.sort(key=lambda ep: ep[0].get("gpu_busy_duration", 0)) + mid = len(entries) // 2 + entry, phase = entries[mid] + region_meta = _phase_to_region_meta(phase) + if "gpu_duration" in entry: + region_meta["traceutils_gpu_duration_us"] = entry["gpu_duration"] + if "gpu_busy_duration" in entry: + region_meta["traceutils_gpu_busy_duration_us"] = entry[ + "gpu_busy_duration" + ] + trace_dict = _load_trace(entry["output_path"]) + output.append((trace_dict, region_meta)) + + return output if output else None diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-coherence-agent.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-coherence-agent.md new file mode 100644 index 000000000..e1e75e97f --- /dev/null +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-coherence-agent.md @@ -0,0 +1,137 @@ + + +--- +name: kernel-coherence-agent +description: Second-pass cross-trace refinement. Consumes kernel_coherence_context.json (one-sided kernel buckets with their shared-neighbor context) and writes kernel_coherence_decisions.json, pairing one-sided buckets across traces by position and splitting names that occur in different contexts, so the comparison has no one-sided condensed symbols. +model: claude-opus-4-8 +--- + +# Kernel Coherence Agent + +Second pass over the name-first unification. The first pass leaves **one-sided** +buckets: kernels whose unified name appears in only one trace -- most importantly +vendor GEMM families that could not be paired by name ((platform 1 name) `(likely gemm name platform 1)_*` vs +(platform 2 name) `(likely gemm name platform 2)_*`). This pass uses the first-pass **shared** buckets as cross-trace +positional anchors and re-labels the one-sided buckets by their **neighbor +context**. + +Two capabilities: + +1. **Cross-trace pairing by position.** A one-sided bucket in trace A and a + one-sided bucket in trace B that sit in the *same* shared-neighbor context are + the same operation -- give them the **same new shared name**. (The GEMM + between `add_rmsnorm` and `rotary_embedding` is the QKV projection on both + traces, even though the raw kernels are `(likely gemm name platform 1)_...` vs `(likely gemm name platform 2)_...`.) +2. **Context-dependent splitting.** One name that appears in *different* contexts + should become *different* buckets (a GEMM between attention and norm vs a GEMM + between embedding and MoE). + +**Scripts directory:** `TraceLens/Agent/Analysis/semantic_analyses/` + +## Input + +The full `kernel_coherence_context.json` is provided inline. Key fields: + +- `condensed_sequence_a` / `_b` -- kernel-order semantic_block values with + consecutive duplicates collapsed (one symbol per run). +- `shared_blocks` -- symbols in both condensed sequences (your stable anchors). +- `one_sided_in_a` / `one_sided_in_b` -- the symbols you must resolve. +- `one_sided_details_a` / `_b` -- per one-sided symbol, a list of `contexts`, + each with `id`, `left_window` / `right_window` (nearest shared symbols), + `kernels_in_run`, `perf_categories`, `sample_input_dims`, and + `top_kernel_names_by_dur`. +- `context_catalog` -- flat list of every context with its `id` (use these ids + verbatim in your output). + +## How to decide + +For each one-sided context, work out the operation from: + +- **`left_window` / `right_window`** -- the primary signal. In a transformer + decoder layer the projections are pinned by their neighbors. Possible likely examples (not strict rules): + - `(*norm* | rotary_embedding)` -> QKV projection + - `(paged_attention | *norm*)` -> output projection + - `(*norm* | act_and_mul)` -> gate/up projection + - `(act_and_mul | *norm*)` -> down projection +- **kernel name** -- for example, GEMMs, attention, norms, and rotary embeddings often have characteristic names. +- **`sample_input_dims`** and **`kernels_in_run`** -- corroborate a pairing + (matching shapes / per-layer counts). + +**Pairing rule:** two contexts on opposite traces with the same +`(left_window, right_window)` and a compatible operation get the **same** final +name. Pick a short, vendor-neutral name (`qkv_projection`, `output_projection`, +`gate_up_projection`, `down_projection`, ...). + +**perf_category caveat:** the regex classifier may tag a vendor kernel as +`Others`. Do **not** require identical `perf_category` to pair +-- rely on the kernel name and the neighbor context. Still never pair operations +that are clearly different types. + +## Output file: `kernel_coherence_decisions.json` + +```json +{ + "context_renames": { + "(platform 1 name):6": "qkv_projection", + "(platform 2 name):6": "qkv_projection" + }, + "fallback_remap_a": { + "(likely gemm name platform 1)_..._(likely gemm signature)_...": "qkv_projection" + }, + "fallback_remap_b": { + "(likely gemm name platform 2)_..._(likely gemm signature)_...": "qkv_projection" + }, + "notes": "optional rationale" +} +``` + +Resolution order applied per kernel (see `apply`): + +1. **`context_renames[context_id]`** -- context-specific; wins. Use it to express + context-dependent behavior: the *same* first-pass symbol with *different* + `(left,right)` windows gets *different* finals via *different* context ids. +2. **`fallback_remap_a` / `_b[symbol]`** -- blanket per-symbol remap, applied when + no context id matched (empty windows at a sequence boundary) or when a symbol + means the same thing everywhere and you don't need to split it by context. +3. Otherwise the first-pass name is kept. + +### Guidance + +- For a GEMM shape that has one dominant role across the trace, a single + `fallback_remap` entry is cleaner than many identical `context_renames`; add + `context_renames` only where a symbol genuinely changes role by context. +- Prefer merging a one-sided bucket into an **existing shared** symbol when the + evidence supports it (e.g. an attention split-K/reduce kernel -> the shared + `paged_attention` bucket). +- When you introduce a new bucket for a paired concept, use the **same** name on + both traces so the condensed sets align. +- **Leave genuine one-offs alone.** Pre/post-layer, framework-specific setup + kernels (index build, dtype copies, prefix-scan) have no counterpart; do not + force them into a shared bucket. Residual one-sided singletons are acceptable. + +## Apply + verify + +The orchestrator runs: + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_coherence.py apply \ + --context /kernel_coherence_context.json \ + --decisions /kernel_coherence_decisions.json \ + --audit-csv-a /per_kernel_final_.csv \ + --audit-csv-b /per_kernel_final_.csv +``` + +Rewrites `semantic_block` on both label files, writes per-kernel audit CSVs, and +prints the residual one-sided condensed symbols. If meaningful (non-singleton) +symbols remain one-sided, revise the decisions and re-run; residual pre/post +singletons may be accepted. + +## Return Value + +Return: `status`, `pairs_created` (new shared names used on both sides), +`context_renames`, `fallback_a`, `fallback_b`, `shared_blocks_after`, and the +residual one-sided symbols with a one-line reason for leaving each. diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-stem-preprocessing-agent.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-stem-preprocessing-agent.md new file mode 100644 index 000000000..fdf3ec8c9 --- /dev/null +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-stem-preprocessing-agent.md @@ -0,0 +1,128 @@ + + +--- +name: kernel-stem-preprocessing-agent +description: Conditional pre-step for kernel unification when the combined unique kernel-name count is too large (> threshold, default 5000) to fit LLM context. Inspects a sample of names, authors custom regex rules that collapse high-cardinality families to stems, preserve families whose parameters matter for later analysis, and drop noise, then iterates until the stem count is manageable. +model: claude-opus-4-8 +--- + +# Kernel Stem Preprocessing Agent + +Reduce the number of distinct kernel names to a set small enough for the +`kernel-unification-agent` to reason over, **without losing information that +matters**. + +This step runs **only** when `kernel_unification.py prepare-context` reports +`needs_stem_preprocessing: true` (combined unique names exceed the threshold, +default 5000). Otherwise skip it entirely. + +**Scripts directory:** `TraceLens/Agent/Analysis/semantic_analyses/` + +## Why this is needed + +High-cardinality name families blow up the unique count. A single logical GEMM +family can appear as thousands of variants that differ only by an autotuner id +or embedded shape/tile suffix. These variants are the +same operation for the purpose of establishing cross-trace anchors, so +collapsing them to a **stem** (dropping the autotuner id and shape/tile suffix) drastically +shrinks the set. + +But **not every varying parameter is noise.** Some distinctions should be kept +for later analysis (e.g. GEMM tile / grid dimensions are useful downstream), and +some names are pure noise that should be dropped (profiler markers, one-off +setup kernels). Use judgment. + +## Input + +`prepare-context` output is provided inline. Key fields: + +- `summary.combined_unique`, `threshold` -- the size problem. +- `sample` -- a representative, impact-spanning sample of names, each with + `trace`, `name`, `perf_categories`, `kernel_count`, `total_dur_us`. Use it to + discover the naming families; do **not** assume it is exhaustive. + +## Your job + +1. **Identify families.** Group the sampled names into families that share a + structure (same prefix / vendor scheme, varying only in ids or shapes). +2. **Decide per family** one of three actions: + - **`collapse`** -- the variation is an id or a detail irrelevant to + matching. Author a regex that rewrites the name to a stable stem. + *Prefer collapsing the biggest families first* -- that is where the + cardinality lives. + - **`preserve`** -- the variation carries information a later stage needs + (e.g. GEMM dims). Keep the full name; still list the family so your intent + is explicit. + - **`drop`** -- the name is noise (profiler markers, debug kernels) and + should be excluded from unification. Dropped kernels keep their raw name + and simply will not unify. +3. **Author `stem_rules.json`** (see shape below). Rules are applied **in + order**; the first matching rule wins. Names matching no rule default to + `preserve`. +4. **Iterate.** Run `apply-stem-rules` and read the printed cardinality. If it + is still above threshold, broaden collapse patterns or drop more low-value + families and re-run until it reports `OK` (or is comfortably small). + +### Guidance for writing collapse regexes + +- Anchor to the family so a rule cannot over-match a different family. +- Collapse the varying token, not the whole name: replace digit runs / shape + tokens with a placeholder rather than erasing the identifying prefix. E.g. + `('_MT\d+x\d+x\d+', '_MT#')` keeps `(vendor gemm name)_..._MT#_...` distinguishable from a + different vendor GEMM variant, while `('_[0-9]+$', '')` strips a trailing id. +- Keep the stem **stable across both traces** where the operation is the same, + but you do not need to make two frameworks' stems identical here -- the + `kernel-unification-agent` maps stems across traces afterward. + +## Output file: `stem_rules.json` + +```json +{ + "rules": [ + {"pattern": "^gemm_[0-9]+$", "replacement": "gemm", + "action": "collapse", "note": "autotuner-id GEMM family"}, + {"pattern": "_MT[0-9]+x[0-9]+x[0-9]+", "replacement": "_MT#", + "action": "collapse", "note": "collapse vendor GEMM tile-size id, keep family"}, + {"pattern": "(?i)profiler|marker|nvtx", "replacement": "", + "action": "drop", "note": "profiler noise"}, + {"pattern": "^(vendor gemm prefix)", "replacement": "", + "action": "preserve", "note": "keep vendor GEMM dims for later analysis"} + ] +} +``` + +Fields per rule: +- `pattern` -- Python `re` regex, matched with `re.search`. +- `replacement` -- used by `re.sub` for `collapse` (may reference groups); + ignored for `preserve` / `drop`. +- `action` -- `collapse` | `preserve` | `drop`. +- `note` -- short rationale (required; this is the audit trail). + +## Apply loop + +The orchestrator runs: + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py apply-stem-rules \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a --name-b \ + --rules /stem_rules.json \ + --raw-to-stem /raw_to_stem.json \ + -o /kernel_unification_context.json +``` + +- Emits `raw_to_stem.json` (used later by `apply-map`) and a **stem-level** + `kernel_unification_context.json` for the `kernel-unification-agent`. +- Prints `raw -> stems` counts and per-action tallies. Re-run after editing + rules until the stem count is within budget. + +## Return Value + +Return: `status`, `raw_unique_before`, `stem_unique_after`, `action_counts` +(collapse/preserve/drop), a one-line rationale per family, and the path to +`stem_rules.json` and `raw_to_stem.json`. diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-unification-agent.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-unification-agent.md new file mode 100644 index 000000000..6764f0750 --- /dev/null +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-unification-agent.md @@ -0,0 +1,131 @@ + + +--- +name: kernel-unification-agent +description: Name-first cross-trace kernel-name unification. Reads kernel_unification_context.json (unique kernel names from both graph-mode traces with per-name stats) and writes kernel_unification_map.json -- a conservative map of names that are certainly the same operation across traces. Establishes matching anchors; does not resolve every ambiguity. +model: claude-opus-4-8 +--- + +# Kernel Unification Agent + +Unify raw GPU **kernel names** across two traces so equivalent kernels can be +matched cross-trace. + +Graph-mode traces collapse the CPU->GPU call stack under +`hip/cudaGraphLaunch`, so `nn_module` / `cpu_op` context is unavailable. The +only reliable cross-trace signal is the raw kernel name. Different frameworks +and vendors name the same operation differently (e.g. `moe_attn_vllm` vs +`sglang_moe_attention`), so this agent proposes a map that unifies such names +to a single shared label -- establishing clear **anchors** for the comparison. + +**Scripts directory:** `TraceLens/Agent/Analysis/semantic_analyses/` + +## Input + +The full `kernel_unification_context.json` is provided inline in your prompt. +Do NOT re-read it from disk. Key fields: + +- `name_a`, `name_b` -- short trace labels. +- `key_level` -- `raw_name` (map raw kernel names) or `stem` (stem + preprocessing was already applied; map the stems shown). +- `only_in_`, `only_in_` -- names present in only one trace. + **These are your unification candidates.** Each entry has `name`, + `kernel_count`, `total_dur_us`, `perf_categories`, and a `sample_input_dims` + (and `sample_raw_names` when `key_level` is `stem`). +- `in_both` -- names already identical in both traces. They are unified by + default; **do not** add map entries for them. +- `summary` -- counts for orientation. + +## Your job + +Produce a map that unifies names in `only_in_` with their counterparts +in `only_in_` when you are **certain** they are the same operation. + +### Signals for matching + +- **`perf_categories`** -- a GEMM only unifies with a GEMM, SDPA with SDPA, + etc. Never unify across different perf categories. +- **Name semantics** -- decode the mangled name. Vendor GEMM kernels often have + characteristic mangled names; `*paged_attention*`, + `*fmha*`, `*flash*` are attention; `*rmsnorm*`, `*layer_norm*` are + normalization; `*reduce*`, `*all_reduce*`, `*allgather*` are communication. +- **`sample_input_dims`** -- matching shapes across traces strengthen a pairing. +- **`kernel_count` / `total_dur_us`** -- an operation that runs N times per + layer on one trace usually runs a comparable number of times on the other. + +### Rules + +1. **Certainty only.** Map a pair only when the evidence is strong. When + unsure, leave both names unmapped -- they fall back to their raw name / stem + and simply remain unmatched. This pass builds anchors, not a full + resolution. +2. **Skip identical names.** Anything in `in_both` is already unified. Do NOT + add map entries for them. +3. **Preserve granularity.** Do not merge two functionally distinct kernels + because their names look similar. Do not collapse a family that a later + analysis stage may want to keep separate (e.g. distinct attention variants). +4. **Same value on both sides.** For a matched pair, use the **same** unified + string as the value in both `map_a` and `map_b`. Choose a short, neutral, + vendor-agnostic name (e.g. `moe_attn`, `qkv_projection`, `allreduce`). +5. **Exact keys.** Keys must be copied verbatim from the context lists + (`only_in_` entries' `name` field). Every key must exist in that + trace's list. +6. **perf_category is not yours to change.** You only unify names. + +## Output file: `kernel_unification_map.json` + +Write to the output directory. Exact shape: + +```json +{ + "name_a": "(platform 1 name)", + "name_b": "(platform 2 name)", + "map_a": { + "moe_attn_vllm": "moe_attn", + "(vendor1_gemm_name)_(shape1)": "expert_gemm" + }, + "map_b": { + "sglang_moe_attention": "moe_attn", + "(vendor2_gemm_name)_(shape1)": "expert_gemm" + } +} +``` + +Using the same `(shape1)` on both sides shows *why* this pairs: two vendors' +differently-mangled names for the same GEMM tile shape/role -- not license to +collapse every GEMM in a trace into one bucket; kernels with distinct shapes +or otherwise clearly serving distinct per-layer roles should stay separate +keys. + +- `map_a` keys are `name_a` names; `map_b` keys are `name_b` names. +- A pair is unified when a `map_a` value equals a `map_b` value. +- A one-sided rename (a name you want to relabel but that has no counterpart) + is allowed but usually unnecessary -- prefer leaving it unmapped. +- Either map may be empty if no confident unification exists on that side. + +## Apply + verify + +The orchestrator runs: + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py apply-map \ + --labels-a /semantic_labels.json \ + --labels-b /semantic_labels.json \ + --name-a --name-b \ + --map /kernel_unification_map.json \ + [--raw-to-stem /raw_to_stem.json] # only if stem preprocessing was used +``` + +This writes the unified name into each kernel's `semantic_block` field (default += raw name / stem when the map has no entry) and prints the resulting shared / +one-sided vocabulary counts. + +## Return Value + +Return: `status` (SUCCESS/ERROR), `pairs_unified` (count of matched values), +`names_mapped_a`, `names_mapped_b`, `shared_blocks_after_apply`, and any names +you deliberately left unmapped due to uncertainty (with a one-line reason). diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md new file mode 100644 index 000000000..8fe3e7c8b --- /dev/null +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md @@ -0,0 +1,298 @@ + + +--- +name: semantic-comparison-agent +description: End-to-end semantic comparison of two graph-mode GPU traces. Runs deterministic breakdown per trace (extraction + classification + pattern finding + label assembly), then a name-first LLM kernel-name unification pass that establishes cross-trace matching anchors in the semantic_block field, followed by a comparison pipeline. +model: claude-opus-4-8 +--- + +# Semantic Comparison + +Orchestrate end-to-end semantic comparison of two GPU traces. The user +provides two raw trace files and the orchestrator handles everything: +deterministic parallel breakdown (no LLM), a name-first LLM kernel-name +unification pass that establishes cross-trace matching anchors, and the +comparison pipeline. + +**Why name-first.** In graph mode the CPU->GPU call stack collapses under +`hip/cudaGraphLaunch`, so `nn_module` / `cpu_op` context is unavailable and +block-alignment harmonization cannot work. The only reliable cross-trace +signal is the raw GPU **kernel name**. This workflow unifies kernel names +across the two traces (e.g. `moe_attn_vllm` and `sglang_moe_attention` -> +`moe_attn`) and writes the unified name into each kernel's `semantic_block` +field, which the downstream comparison uses as its matching key. + +Use vendor-agnostic terminology (GPU kernels, vendor GEMM library, etc.) +except when quoting actual kernel names from traces. + +**Command prefix.** All commands below run through ``: read the command +prefix path from your execution context (`/cache/cmd_prefix.txt`) and +substitute the command for `{CMD}`. `` is blank for local runs and wraps +commands into the target environment otherwise. Shell control flow (`for`, `if`, +`wait`) and variable assignments run in the driving shell and are not prefixed. + +--- + +## Workflow Steps + +``` +0. Query User Inputs +1. Semantic Breakdown (PARALLEL shell commands, one per trace) +2. Kernel-Name Unification (name-first anchors + coherence refinement, LLM) +3. Generate TraceDiff Output (script) +4. Generate Comparison CSV (script) +``` + +--- + +## Step 0: Query User Inputs + +Ask the user for: + +**Required:** +- Trace A path (.json or .json.gz) +- Trace B path (.json or .json.gz) +- Short labels for each trace (e.g., MI355 / B200) + +**Optional:** +- Output directory (default: `comparison_output/`) + +**vLLM / annotated traces** are auto-detected by `extract_trace_data.py`. +No special flag is needed. + +--- + +## Step 1: Semantic Breakdown (Deterministic, PARALLEL) + +Breakdown is fully deterministic -- no LLM calls. Run both traces as +parallel shell commands. + +### 1.1 Per-trace Pipeline + +Run the full breakdown for **both traces in a single shell call** using +background jobs + `wait`. + +```bash +SCRIPTS=TraceLens/Agent/Analysis/semantic_analyses +CLASSIFY=TraceLens/Agent/Analysis/utils/classify_kernels.py +DIR_A=/work/ +DIR_B=/work/ + mkdir -p $DIR_A $DIR_B + +run_breakdown() { + local TRACE=$1 DIR=$2 + + # Extract (auto-splits vLLM traces into region subdirs) + python3 $SCRIPTS/extract_trace_data.py $TRACE -o $DIR/ + + # Check whether extraction produced region subdirs or a flat file. + # gpu_op_uid is stamped by extract_trace_data.py (raw-index UID aligned + # with the perf report); no trace-tree build is needed. + if ls $DIR/*/extracted.json >/dev/null 2>&1; then + for REGION in $DIR/*/; do + python3 $SCRIPTS/pattern_finder.py $REGION/extracted.json -o $REGION/pattern.json & + python3 $CLASSIFY $REGION/extracted.json -o $REGION/classified.json & + done + wait + for REGION in $DIR/*/; do + python3 $SCRIPTS/build_semantic_labels.py \ + $REGION/extracted.json $REGION/classified.json $REGION/pattern.json \ + -o $REGION/semantic_labels.json + done + else + python3 $SCRIPTS/pattern_finder.py $DIR/extracted.json -o $DIR/pattern.json & + python3 $CLASSIFY $DIR/extracted.json -o $DIR/classified.json & + wait + python3 $SCRIPTS/build_semantic_labels.py \ + $DIR/extracted.json $DIR/classified.json $DIR/pattern.json \ + -o $DIR/semantic_labels.json + fi +} + +run_breakdown $DIR_A & +run_breakdown $DIR_B & +wait +``` + +**Output directories:** +- Trace A: `/work//` +- Trace B: `/work//` + +For multi-region traces, each directory contains per-region subdirs +(e.g., `decode_only_3/`, `prefill_only_1024/`). + +### 1.2 Verify Breakdown Outputs + +**CRITICAL: DO NOT proceed to Step 2 until both breakdowns have +completed and outputs are verified.** + +After both breakdowns complete, verify that `semantic_labels.json` exists +(in each region subdir for multi-region traces, or directly in the trace +directory for single-trace). If either trace failed, report the error +and stop. + +--- + +## Step 2: Kernel-Name Unification (Name-First, LLM) + +The LLM unifies raw **kernel names** across the two traces. The unified name is written into each kernel's `semantic_block` field. Kernel names already identical in both traces unify by default (no map entry needed); the LLM only maps names that differ but denote the same operation + +Scripts: `TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py` + +### 2.1 Prepare Unification Context + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py prepare-context \ + --labels-a /work//semantic_labels.json \ + --labels-b /work//semantic_labels.json \ + --name-a --name-b \ + -o /work/kernel_unification_context.json +``` + +### 2.2 Stem Preprocessing (conditional, only if flagged) + +If Step 2.1 prints `STEM PREPROCESSING NEEDED` (combined unique names exceed +the threshold, default 5000), the raw name set is too large for the LLM. +Launch the subagent `TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-stem-preprocessing-agent.md`, then: + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py apply-stem-rules \ + --labels-a /work//semantic_labels.json \ + --labels-b /work//semantic_labels.json \ + --name-a --name-b \ + --rules /work/stem_rules.json \ + --raw-to-stem /work/raw_to_stem.json \ + -o /work/kernel_unification_context.json +``` + +Re-run until the printed stem count is within budget. + +### 2.3 Launch Kernel Unification Agent + +Read `TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-unification-agent.md` and +launch it with `kernel_unification_context.json` inline. For multi-region vLLM: run once per matching region pair. + +### 2.4 Apply the Map + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/kernel_unification.py apply-map \ + --labels-a /work//semantic_labels.json \ + --labels-b /work//semantic_labels.json \ + --name-a --name-b \ + --map /work/kernel_unification_map.json \ + --raw-to-stem /work/raw_to_stem.json # only if 2.2 ran +``` + +### 2.5 Verify Unification + +Check that `kernel_unification_context.json` and `kernel_unification_map.json` +exist in `/work/`, and that `apply-map` reported a non-empty +shared vocabulary. + +### 2.6 Coherence Pass (second pass, LLM) + +The first pass leaves **one-sided** buckets -- kernels whose unified name +appears in only one trace. This pass uses +the first-pass **shared** buckets as cross-trace positional anchors and +re-labels one-sided buckets by their shared-neighbor context, which (a) +pairs GEMMs across vendors by position and (b) splits a name that occurs in +different contexts. Skip it only if `apply-map` already reported no +meaningful one-sided buckets. + +### 2.6a Prepare coherence context + +```bash +KC=TraceLens/Agent/Analysis/semantic_analyses/kernel_coherence.py + python3 $KC prepare-context \ + --labels-a /work//semantic_labels.json \ + --labels-b /work//semantic_labels.json \ + --name-a --name-b \ + --neighbor-radius 1 \ + -o /work/kernel_coherence_context.json +``` + +### 2.6b Launch coherence agent + +Read +`TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/kernel-coherence-agent.md` and launch +it with `kernel_coherence_context.json` inline. It writes +`/work/kernel_coherence_decisions.json` (`context_renames` + +`fallback_remap_a` / `fallback_remap_b`), pairing same-context one-sided buckets +across traces into new shared names. + +### 2.6c Apply + +```bash + python3 $KC apply \ + --context /work/kernel_coherence_context.json \ + --decisions /work/kernel_coherence_decisions.json \ + --audit-csv-a /work/per_kernel_final_.csv \ + --audit-csv-b /work/per_kernel_final_.csv +``` + +`apply` rewrites `semantic_block` in place and prints residual one-sided symbols. +Revise decisions and re-run 2.6b--2.6c if meaningful (non-singleton) symbols remain (singleton setup/copy kernels may be accepted); raise `--neighbor-radius` in 2.6a for ambiguous contexts. + +--- + +## Step 3: Generate TraceDiff Output + +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/generate_semantic_diff.py \ + /work//semantic_labels.json \ + /work//semantic_labels.json \ + --name-a --name-b \ + -o /tracediff_output/ +``` + +Produces in `/tracediff_output/`: +- `diff_stats.csv` -- per-kernel rows matching TraceDiff schema +- `diff_stats_unique_args_summary.csv` -- aggregated by semantic block +- `cpu_op_map.json`, `cpu_op_map_trace1.json`, `cpu_op_map_trace2.json` +- `merged_tree_output.txt` + +This is a **final deliverable** directory for downstream TraceDiff consumers. + +**Perf-report enrichment compatibility:** `diff_stats.csv` carries per-kernel `gpu_op_uid` and per-LCA `busy_time`, consumable by `enrich_perf_report_dict_inplace` in `TraceLens/Reporting/tracediff_comparison_extension.py`. + +--- + +## Step 4: Generate Comparison CSV + +**Single-region mode:** +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/match_and_compare.py \ + /work//semantic_labels.json \ + /work//semantic_labels.json \ + --name-a --name-b \ + -o /work/comparison.csv +``` + +**Multi-region mode (vLLM):** +```bash + python3 TraceLens/Agent/Analysis/semantic_analyses/match_and_compare.py \ + --regions-dir-a /work/ \ + --regions-dir-b /work/ \ + --name-a --name-b \ + -o /work/comparison.csv +``` + +--- + +## Key Principles + +1. **Conservative anchors** -- map only certain equivalences; preserve + granularity and leave uncertain names unmapped +2. **No script creation** -- subagents use only existing scripts (the stem + preprocessing authors regex *rules*, not new scripts) + +--- + +## Final Deliverables + +- `/work/` -- per-trace `semantic_labels.json`, the unification/coherence JSON artifacts, `per_kernel_final_.csv`, and `comparison.csv` +- `/tracediff_output/` -- TraceDiff deliverables (see Step 3) diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/reference.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/reference.md index 18eccd385..a9cf8d1e1 100644 --- a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/reference.md +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/reference.md @@ -38,7 +38,7 @@ Use vendor-agnostic terminology throughout such as GPU kernels, collective commu optionally invoke agent_extension.py (when present), then embed the PNG into the report. ``` -**Subagent usage:** Only invoke Task subagents in steps that explicitly say "subagent" (Steps 6, 7, 9). All other steps (including Step 7.5) must be performed directly by the orchestrator using the command prefix. +**Subagent usage:** Only invoke Task subagents in steps that explicitly say "subagent" (Step 1.5 semantic diff, Steps 6, 7, 9). All other steps (including Step 7.5) must be performed directly by the orchestrator using the command prefix. --- @@ -80,9 +80,11 @@ Use vendor-agnostic terminology throughout such as GPU kernels, collective commu - If **Inference (vLLM/SGLang/ATOM)** is selected, ask **Execution Mode** → ``: 1. **Eager mode** (`` = `eager`) — only the trace file is needed 2. **Graph replay + capture** (`` = `graph_capture`) — also requires a capture folder path - - If **Graph replay + capture**, ask for **Capture Folder Path** → ``: - - Ask: "Please provide the full path to the graph capture traces folder" - - If **Graph replay + capture** and **comparative**, ask for **Trace2 Capture Folder Path** → `` + + - If **Graph replay + capture**, ask for the **Capture Folder Path(s)**: + - `standalone`: one folder → ``. Ask: "Please provide the full path to the graph capture traces folder" + - `comparative`: one folder per trace → `` (primary/trace1) and `` (comparison/trace2). Ask: "Please provide the graph capture traces folder for the primary trace and for the comparison trace." + - **Comparative + graph replay** (do not abort): collect capture folders for both traces when available. If capture is not available, the comparison uses the semantic path (see Step 0.5). 5. **Environment Setup** - Ask: "Are you running locally or on a cluster?" @@ -160,9 +162,20 @@ Do NOT proceed to Step 1 until validation passes. --- +## Step 0.5: Comparison Method (comparative only) + +For `standalone`, skip this step. + +For `comparative`, set `` directly from what was already collected in Step 0: + +- If a graph-replay trace is involved (`` = `graph_capture`) but capture folders were **not** collected for it in Step 0, set `` = `semantic`. +- Otherwise set `` = `tracediff`. + +--- + ## Step 1: Generate Performance Report -Use **``** to determine which CLI tool to run and then **``** to determine arguments. +Use **``** to determine which CLI tool to run and then **``** (and, for comparative, **``** from Step 0.5) to determine arguments. For all of these scripts below, look at the environment variable TL_EXTENSION to recursively search for a file called .json. Do not look for .json; it is not needed. If it is not found also look in TraceLens/Agent/Analysis/utils/arch/.json. @@ -255,6 +268,37 @@ All commands below append `` and ``, resolved by `` = `semantic` only) + +When `` = `comparative` and `` = `semantic`, run Step 1 in this order: + +1. **Trace2 report** — run the analysis-mode CLI above for trace2 using the `comparative` trace2 `` and empty `` (identical to the TraceDiff path). + +2. **Semantic diff (subagent)** — launch a Task subagent that reads and follows the FULL instructions in `TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md`. Prompt context: + +``` +Read and follow the FULL instructions in: + TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md + +**Execution Context:** +- Trace A (primary/trace1): (platform ) +- Trace B (comparison/trace2): (platform ) +- Labels: name-a trace1, name-b trace2 +- Output directory: /semantic/ +- Command prefix: read /cache/cmd_prefix.txt — substitute {CMD} + +Run the full semantic comparison through "Generate TraceDiff Output" so that +/semantic/tracediff_output/diff_stats.csv is produced. Return "DONE". +``` + + Verify `/semantic/tracediff_output/diff_stats.csv` exists before continuing. If it is missing, retry the subagent once; if it still fails, stop and report. + +3. **Trace1 report** — run the analysis-mode CLI for trace1 with the `comparative` trace1 `` and `` = `--precomputed_diff_stats_csv /semantic/tracediff_output/diff_stats.csv`. + +4. **Confirm** `/perf_report_trace1_csvs/diff_stats.csv` exists (written by the report script). If absent, copy `/semantic/tracediff_output/diff_stats.csv` to that path so the comparative fusion step (Steps 2-5) can read it. + +--- + ## Steps 2-5: Prepare Category Data Execute the TraceLens Agentic Mode orchestrator preparation script: @@ -573,7 +617,6 @@ The report at `/analysis.md` must use these exact `##` headers — d 5. `## Detailed Analysis` 6. `## Appendix` - ### 11.1 Validate Report Structure (Retry up to 2x) After writing `analysis.md`, validate that the report contains all required `##` section headers. If validation fails, modify the report with the missing sections. @@ -648,5 +691,6 @@ If the plot is skipped, the `{{PERF_PLOT}}` placeholder is removed so the report If Steps 1 or many of Steps 2-5 fail or produce unexpected results, check whether the trace uses the following features before retrying: - **GPU Graph Replay**: raw trace JSON contains `hipGraphLaunch` or `cudaGraphLaunch`. - - **Default mode** (analysis_mode = `default`): Inform the user with `[DIAG:trace_quality:GPU_GRAPH_REPLAY]` that GPU graph replay was detected and that the default analysis mode supports typical PyTorch traces. **Abort** -- do not retry or continue. - - **Inference mode** (analysis_mode = `inference`): Graph launches are expected and supported if graph capture folder is provided, do not abort. If inference_exec_mode is `eager` (no capture folder was provided), continue. + - **Comparative scope**: graph replay is **supported**; the comparison method (tracediff+capture vs semantic) was set in Step 0.5 from capture availability, with no trace classification. Ensure capture folders were collected for both traces when available. + - **Default mode, standalone** (analysis_mode = `default`): Inform the user with `[DIAG:trace_quality:GPU_GRAPH_REPLAY]` that GPU graph replay was detected and that the default analysis mode supports typical PyTorch traces. **Abort** -- do not retry or continue. + - **Inference mode, standalone**: graph launches are expected and supported; continue whether or not a capture folder was provided (eager mode has none). diff --git a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/templates/analysis_template.md b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/templates/analysis_template.md index 89e446265..d7d98644a 100644 --- a/TraceLens/Agent/Analysis/skills/analysis-orchestrator/templates/analysis_template.md +++ b/TraceLens/Agent/Analysis/skills/analysis-orchestrator/templates/analysis_template.md @@ -90,8 +90,10 @@ section that has STANDALONE / COMPARATIVE variants. Delete the unused variant. [1 paragraph comparative overview: summarize which trace is faster overall, by how much, and the dominant gap categories] - | Metric | Trace 1 - () | Trace 2 - () | Difference | |--------|----------------------------|-------------------------------|------------| @@ -144,6 +146,7 @@ One row per entry in `priority_data.json::priorities[]`, in array order (no mani `Trace 2 Time (ms)` = matching `manifest.trace2_ops_summary_by_category[]["total_direct_kernel_time_ms"]` where `"op category"` matches the row Category **case-insensitively**; use — if no match. `Difference (ms)` = Trace 2 Time − Trace 1 Time. +Note: for communication categories (e.g. collectives), `% of Compute Time` can exceed 100% because the kernel time is not compute-bound — report the value as computed. | Rank | Category | Trace 1 Time (ms) | Trace 2 Time (ms) | % of Compute Time | Ops | Difference (ms) | |------|----------|-------------------|-------------------|-------------------|-----|-----------------| diff --git a/TraceLens/Reporting/generate_perf_report_pytorch.py b/TraceLens/Reporting/generate_perf_report_pytorch.py index 16bb9b644..bf6e366b0 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch.py @@ -422,6 +422,10 @@ def generate_perf_report_pytorch( topk_ops: Optional[int] = None, topk_roofline_ops: Optional[int] = None, comparison_json_path: Optional[str] = None, + # precomputed diff_stats.csv (e.g. from the semantic comparison path); when + # set, the internal TraceDiff is skipped and this CSV is used as-is to + # enrich the report. Mutually exclusive with comparison_json_path. + precomputed_diff_stats_csv: Optional[str] = None, extension_file: Optional[str] = None, # for gemm simulator / Origami (Origami requires --enable_origami when arch is set) python_path: Optional[str] = None, @@ -799,7 +803,10 @@ def generate_perf_report_pytorch( # Add unified perf metrics table (ops with perf models + leaf ops with GPU kernels) df_unified_perf = perf_analyzer.build_df_unified_perf_table() - # Run TraceDiff when a comparison trace is provided. diff_stats_df is generated + # Obtain the comparison diff_stats. Normally generated by running + # TraceDiff against a comparison trace; alternatively a precomputed + # diff_stats.csv (e.g. from the semantic comparison path) can be + # supplied, in which case the internal TraceDiff is skipped. _tracediff_diff_stats: Optional[pd.DataFrame] = None if comparison_json_path and not df_unified_perf.empty: perf_analyzer2 = TreePerfAnalyzer.from_file( @@ -815,6 +822,8 @@ def generate_perf_report_pytorch( td = TraceDiff(perf_analyzer.tree, perf_analyzer2.tree) td.generate_tracediff_report() _tracediff_diff_stats = td.diff_stats_df + elif precomputed_diff_stats_csv and not df_unified_perf.empty: + _tracediff_diff_stats = pd.read_csv(precomputed_diff_stats_csv) if not df_unified_perf.empty: df_unified_perf_summary = perf_analyzer.summarize_df_unified_perf_table( @@ -1156,6 +1165,19 @@ def main(): ), ) + parser.add_argument( + "--precomputed_diff_stats_csv", + type=str, + default=None, + help=( + "Path to a precomputed TraceDiff-schema diff_stats.csv (e.g. from " + "the semantic comparison path). When set, the internal TraceDiff " + "is skipped and this CSV is used to enrich unified_perf_summary " + "and emit the diff_stats sheet. Mutually exclusive with " + "--comparison_json_path." + ), + ) + parser.add_argument( "--extension_file", type=str, @@ -1220,6 +1242,11 @@ def main(): ) args = parser.parse_args() + if args.comparison_json_path and args.precomputed_diff_stats_csv: + parser.error( + "--comparison_json_path and --precomputed_diff_stats_csv cannot be " + "used together; provide only one comparison diff_stats source." + ) generate_perf_report_pytorch( profile_json_path=args.profile_json_path, output_xlsx_path=args.output_xlsx_path, @@ -1238,6 +1265,7 @@ def main(): topk_ops=args.topk_ops, topk_roofline_ops=args.topk_roofline_ops, comparison_json_path=args.comparison_json_path, + precomputed_diff_stats_csv=args.precomputed_diff_stats_csv, extension_file=args.extension_file, python_path=args.python_path, gpu_arch_json_path=args.gpu_arch_json_path, diff --git a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py index 905193a34..7f712b6ca 100644 --- a/TraceLens/Reporting/generate_perf_report_pytorch_inference.py +++ b/TraceLens/Reporting/generate_perf_report_pytorch_inference.py @@ -533,6 +533,10 @@ def generate_perf_report_pytorch( topk_ops: Optional[int] = None, topk_roofline_ops: Optional[int] = None, comparison_json_path: Optional[str] = None, + # precomputed diff_stats.csv (e.g. from the semantic comparison path); when + # set, the internal TraceDiff is skipped and this CSV is used as-is to + # enrich the report. Mutually exclusive with comparison_json_path. + precomputed_diff_stats_csv: Optional[str] = None, comparison_augmented_tree: Optional[TraceToTree] = None, extension_file: Optional[str] = None, # for gemm simulator / Origami (Origami requires --enable_origami when arch is set) @@ -920,7 +924,10 @@ def generate_perf_report_pytorch( include_nccl=collective_analysis ) - # Run TraceDiff when a comparison trace is provided. diff_stats_df is generated + # Obtain the comparison diff_stats. Normally generated by running + # TraceDiff against a comparison trace; alternatively a precomputed + # diff_stats.csv (e.g. from the semantic comparison path) can be + # supplied, in which case the internal TraceDiff is skipped. _tracediff_diff_stats: Optional[pd.DataFrame] = None if comparison_json_path and not df_unified_perf.empty: if comparison_augmented_tree is not None: @@ -947,6 +954,8 @@ def generate_perf_report_pytorch( td = TraceDiff(perf_analyzer.tree, perf_analyzer2.tree) td.generate_tracediff_report() _tracediff_diff_stats = td.diff_stats_df + elif precomputed_diff_stats_csv and not df_unified_perf.empty: + _tracediff_diff_stats = pd.read_csv(precomputed_diff_stats_csv) if not df_unified_perf.empty: df_unified_perf_summary = perf_analyzer.summarize_df_unified_perf_table( @@ -1303,6 +1312,20 @@ def main(): ), ) + parser.add_argument( + "--precomputed_diff_stats_csv", + type=str, + default=None, + help=( + "Path to a precomputed TraceDiff-schema diff_stats.csv (e.g. from " + "the semantic comparison path). When set, the internal TraceDiff " + "is skipped and this CSV is used to enrich unified_perf_summary " + "and emit the diff_stats sheet. Mutually exclusive with " + "--comparison_json_path; unlike --comparison_json_path it may be " + "combined with --capture_folder." + ), + ) + parser.add_argument( "--extension_file", type=str, @@ -1358,6 +1381,14 @@ def main(): ) args = parser.parse_args() + if args.comparison_json_path and args.precomputed_diff_stats_csv: + parser.error( + "--comparison_json_path and --precomputed_diff_stats_csv cannot be " + "used together; provide only one comparison diff_stats source." + ) + # NOTE: --capture_folder + --precomputed_diff_stats_csv IS allowed: the + # precomputed (e.g. semantic) diff_stats does not run the internal + # TraceDiff, so graph capture traces can be compared via this path. if args.comparison_capture_folder and not args.comparison_json_path: parser.error("--comparison_capture_folder requires --comparison_json_path.") if args.capture_folder: @@ -1399,6 +1430,7 @@ def main(): topk_ops=args.topk_ops, topk_roofline_ops=args.topk_roofline_ops, comparison_json_path=args.comparison_json_path, + precomputed_diff_stats_csv=args.precomputed_diff_stats_csv, comparison_augmented_tree=comparison_graph_tree, extension_file=args.extension_file, python_path=args.python_path, diff --git a/TraceLens/TraceUtils/split_inference_trace_annotation.py b/TraceLens/TraceUtils/split_inference_trace_annotation.py index e8a928955..e483e1768 100644 --- a/TraceLens/TraceUtils/split_inference_trace_annotation.py +++ b/TraceLens/TraceUtils/split_inference_trace_annotation.py @@ -250,12 +250,28 @@ def main(): "output_dir/decode_only/. Each step is a separate trace file." ), ) + parser.add_argument( + "--emit-gpu-op-uid", + action="store_true", + default=False, + help=( + "Tag each event with a 'gpu_op_uid' field equal to its index in " + "the original (unfiltered) traceEvents array before splitting, " + "so downstream consumers can recover each extracted event's " + "position in the source trace without re-loading it. Off by " + "default to keep existing output byte-for-byte unchanged." + ), + ) args = parser.parse_args() execution_details = [] # Load trace trace_json = DataLoader.load_data(get_filename(args.trace_path)) events = trace_json.get("traceEvents", []) + if args.emit_gpu_op_uid: + for i, e in enumerate(events): + if isinstance(e, dict): + e["gpu_op_uid"] = i gpu_corr_map, flow_corr_map, meta_events = preprocess_trace(events) print(f"Loaded {len(events)} events") diff --git a/agent_evals/Analysis/eval_utils/compare_lca_partitions.py b/agent_evals/Analysis/eval_utils/compare_lca_partitions.py new file mode 100755 index 000000000..bacf179f1 --- /dev/null +++ b/agent_evals/Analysis/eval_utils/compare_lca_partitions.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Compare LCA (lowest-common-ancestor) partitions between the "with capture" +(gold standard) and "no capture" comparative analyses of the same traces, and +score the no-capture partition against semantics-free baselines. + +Both analyses emit a ``diff_stats.csv`` in which every row is a single GPU kernel +(identified by ``source`` = trace1/trace2 and ``gpu_op_uid``) tagged with a +``lowest_common_ancestor_id`` (LCA). LCA ids are opaque cluster labels: their +numeric value is only meaningful within a single analysis, so we never compare +ids across analyses directly. Instead we measure how well the two LCA +*partitions* of the shared kernel set agree, using bidirectional cluster purity. + +Forward purity (gold -> no-capture): + For every gold LCA group, take the majority no-capture LCA among its kernels + and count how many kernels in the group carry that majority label. Sum across + groups and divide by the number of matched kernels. + +Reverse purity (no-capture -> gold): the same metric with the roles swapped. + +Reporting both catches degenerate collapses: if the no-capture analysis dumped +every kernel into one LCA, forward purity stays high (each gold group trivially +agrees with the single label) but reverse purity collapses. + +Strict consistency +------------------ +An all-or-nothing variant of purity: a kernel is *consistent* iff its ENTIRE +group maps to a single group on the other side (no partial credit for the +majority). Strict forward = fraction of kernels whose whole gold group lands in +one no-capture LCA; strict reverse swaps the roles. Always <= the matching +purity, with equality only when every multi-kernel group is perfectly pure. + +Random baselines +---------------- +To judge whether the real no-capture purity reflects genuine semantic signal or +is merely an artifact of the bucket-size distribution and the metric's structure, +we compare against semantics-free baselines. All are evaluated on the SAME +matched kernels and PRESERVE the no-capture bucket sizes as observed on the +matched set -- only the kernel->bucket assignment changes. + +Baseline 1 (random): randomly shuffle which matched kernels fall into each + bucket, preserving bucket sizes. Reported as mean +/- std over many seeds. +Baseline 2 (sequential blocks, key string order): order buckets from most to + least popular, then walk the matched kernels ordered by the lexicographic key + string ("source:uid"), assigning the first k1 to bucket 1 (largest), the next + k2 to bucket 2, and so on. +Baseline 3 (sequential blocks, integer uid order): identical block assignment, + but kernels are ordered by (source, integer gpu_op_uid). + +Usage: + compare_lca_partitions.py +""" + +import argparse +from pathlib import Path + +import numpy as np +import pandas as pd + +KEY_COLS = ["source", "gpu_op_uid"] +LCA_COL = "lowest_common_ancestor_id" +LCA_NAME = "lowest_common_ancestor_name" + +# Number of random-baseline seeds (Baseline 1). Kept modest by default; the +# baseline loop is O(trials * kernels), so for large kernel sets it is capped to +# REDUCED_TRIALS (see the guard in main()). +DEFAULT_TRIALS = 200 +LARGE_KERNEL_THRESHOLD = 5000 +REDUCED_TRIALS = 2 +SEED = 0 + + +def load(path: Path) -> pd.DataFrame: + """Load a diff_stats.csv, validate identity columns, and add a unique key.""" + df = pd.read_csv(path) + missing = [c for c in KEY_COLS + [LCA_COL, LCA_NAME, "name"] if c not in df.columns] + if missing: + raise ValueError(f"{path} is missing required columns: {missing}") + df = df.copy() + df["key"] = df["source"].astype(str) + ":" + df["gpu_op_uid"].astype(str) + if df["key"].duplicated().any(): + n = int(df["key"].duplicated().sum()) + raise ValueError( + f"{path}: {n} duplicate (source, gpu_op_uid) keys; not a unique kernel id" + ) + return df + + +def purity(df: pd.DataFrame, group_col: str, label_col: str): + """Sum over each group of the majority-label count within that group. + + Returns (matched_count, total, fraction, per_group_records).""" + total = len(df) + matched = 0 + records = [] + for gid, grp in df.groupby(group_col, sort=True): + counts = grp[label_col].value_counts() # sorted desc by count + majority_label = counts.index[0] + majority_count = int(counts.iloc[0]) + matched += majority_count + records.append( + { + "group_id": gid, + "group_size": len(grp), + "majority_label": majority_label, + "majority_count": majority_count, + "group_purity": majority_count / len(grp), + } + ) + frac = matched / total if total else float("nan") + return matched, total, frac, records + + +def purity_frac(labels_a: np.ndarray, labels_b: np.ndarray) -> float: + """Fraction of items whose A-group carries the A-group's majority B-label. + + Sum over groups of A of the modal B-count, divided by the total item count. + """ + df = pd.DataFrame({"a": labels_a, "b": labels_b}) + matched = 0 + for _, grp in df.groupby("a", sort=False): + matched += int(grp["b"].value_counts().iloc[0]) + return matched / len(df) if len(df) else float("nan") + + +def both_purities(gold: np.ndarray, nc: np.ndarray): + """Return (forward, reverse). + + forward = gold group -> majority nc label + reverse = nc group -> majority gold label + """ + return purity_frac(gold, nc), purity_frac(nc, gold) + + +def strict_consistency_frac( + group_labels: np.ndarray, other_labels: np.ndarray +) -> float: + """Fraction of items whose entire group is homogeneous in ``other_labels``. + + Group items by ``group_labels``; an item is *consistent* iff every item that + shares its group carries the same ``other_labels`` value (so the whole group + maps to a single other-bin). Report the fraction of consistent items. + + This is an all-or-nothing variant of ``purity_frac``: a group contributes its + full size when perfectly pure, else 0 (no partial credit for the majority). + Hence ``strict_consistency_frac <= purity_frac`` always, with equality only + when every multi-item group is perfectly pure. + """ + df = pd.DataFrame({"g": group_labels, "o": other_labels}) + consistent = 0 + for _, grp in df.groupby("g", sort=False): + if grp["o"].nunique() == 1: + consistent += len(grp) + return consistent / len(df) + + +def both_strict(gold: np.ndarray, nc: np.ndarray): + """Return (strict_forward, strict_reverse). + + strict_forward = fraction of items whose gold group is fully pure in nc + strict_reverse = fraction of items whose nc group is fully pure in gold + """ + return strict_consistency_frac(gold, nc), strict_consistency_frac(nc, gold) + + +def agreement_report(gold: pd.DataFrame, ncap: pd.DataFrame) -> None: + """Print the LCA-partition agreement (forward/reverse purity) section.""" + print("=" * 78) + print("LCA assignment agreement: with-capture (GOLD) vs no-capture") + print("=" * 78) + print( + f"with-capture kernels : {len(gold)} | distinct gold LCAs: {gold[LCA_COL].nunique()}" + ) + print( + f"no-capture kernels : {len(ncap)} | distinct nc LCAs: {ncap[LCA_COL].nunique()}" + ) + + g = gold[["key", "name", LCA_COL]].rename( + columns={LCA_COL: "lca_gold", "name": "name_gold"} + ) + n = ncap[["key", "name", LCA_COL]].rename( + columns={LCA_COL: "lca_nc", "name": "name_nc"} + ) + merged = g.merge(n, on="key") + name_agree = ( + (merged["name_gold"] == merged["name_nc"]).mean() + if len(merged) + else float("nan") + ) + + gold_only = len(gold) - len(merged) + nc_only = len(ncap) - len(merged) + print("-" * 78) + print(f"matched kernels (in both) : {len(merged)}") + print(f"gold-only (unmatched) : {gold_only}") + print(f"no-capture-only (unmatched) : {nc_only}") + print( + f"kernel-name agreement on matched keys: {name_agree:.4f} (sanity: should be 1.0)" + ) + + df = merged[["key", "lca_gold", "lca_nc"]] + print("-" * 78) + print(f"matched-set distinct gold LCAs: {df['lca_gold'].nunique()}") + print(f"matched-set distinct nc LCAs: {df['lca_nc'].nunique()}") + + # Forward purity: gold groups -> majority no-capture LCA. + f_matched, f_total, f_frac, _ = purity(df, "lca_gold", "lca_nc") + # Reverse purity: no-capture groups -> majority gold LCA. + r_matched, r_total, r_frac, r_recs = purity(df, "lca_nc", "lca_gold") + + # Strict consistency (all-or-nothing per group; no partial credit). + gold_lbl = df["lca_gold"].to_numpy() + nc_lbl = df["lca_nc"].to_numpy() + sf = strict_consistency_frac(gold_lbl, nc_lbl) # gold groups fully pure in nc + sr = strict_consistency_frac(nc_lbl, gold_lbl) # nc groups fully pure in gold + n_tot = len(df) + + print("=" * 78) + print("RESULTS") + print("=" * 78) + print( + f"Forward purity (gold group -> majority no-capture LCA): " + f"{f_matched}/{f_total} = {f_frac:.4f}" + ) + print( + f"Reverse purity (no-capture group -> majority gold LCA): " + f"{r_matched}/{r_total} = {r_frac:.4f}" + ) + print( + f"Strict forward consistency (whole gold group -> one nc LCA): " + f"{int(round(sf * n_tot))}/{n_tot} = {sf:.4f}" + ) + print( + f"Strict reverse consistency (whole nc group -> one gold LCA): " + f"{int(round(sr * n_tot))}/{n_tot} = {sr:.4f}" + ) + print("=" * 78) + + # Extra context: the largest no-capture groups (where collapse would show up). + top_nc = pd.DataFrame(r_recs).sort_values("group_size", ascending=False).head(5) + print("Largest no-capture LCA groups (matched set):") + print(top_nc.to_string(index=False)) + + +def baseline_report(merged: pd.DataFrame, title: str, n_trials: int) -> None: + """Compute real + baseline purities for a given matched-kernel frame. + + ``merged`` must have columns: lca_gold, lca_nc, nc_name, source, gpu_op_uid, + key. Bucket sizes are recomputed from ``merged`` so the "total matched + preserved" property holds for whatever subset is passed in. + """ + merged = merged.sort_values("key").reset_index(drop=True) + n = len(merged) + + gold = merged["lca_gold"].to_numpy() + nc_real = merged["lca_nc"].to_numpy() + + # No-capture bucket sizes on this (sub)set (sum == n). + sizes = merged["lca_nc"].value_counts() # sorted desc by count + bucket_ids = sizes.index.to_numpy() + bucket_sizes = sizes.to_numpy() + names = merged.drop_duplicates("lca_nc").set_index("lca_nc")["nc_name"].to_dict() + assert bucket_sizes.sum() == n, (bucket_sizes.sum(), n) + + # Label pool: bucket id repeated by its size (length == n). + pool = np.repeat(bucket_ids, bucket_sizes) + assert len(pool) == n + + # ---- Real no-capture ---- + real_fwd, real_rev = both_purities(gold, nc_real) + real_sfwd, real_srev = both_strict(gold, nc_real) + + # ---- Baseline 1: random shuffle of the pool over matched kernels ---- + rng = np.random.default_rng(SEED) + fwds = np.empty(n_trials) + revs = np.empty(n_trials) + sfwds = np.empty(n_trials) + srevs = np.empty(n_trials) + for t in range(n_trials): + shuffled = pool.copy() + rng.shuffle(shuffled) + fwds[t], revs[t] = both_purities(gold, shuffled) + sfwds[t], srevs[t] = both_strict(gold, shuffled) + b1_fwd_m, b1_fwd_s = fwds.mean(), fwds.std() + b1_rev_m, b1_rev_s = revs.mean(), revs.std() + b1_sfwd_m, b1_sfwd_s = sfwds.mean(), sfwds.std() + b1_srev_m, b1_srev_s = srevs.mean(), srevs.std() + + # ---- Baseline 2: sequential blocks, buckets ordered most->least popular ---- + # Kernels ordered by key (lexicographic string: source then string uid). + seq = pool.copy() + b2_fwd, b2_rev = both_purities(gold, seq) + b2_sfwd, b2_srev = both_strict(gold, seq) + + # ---- Baseline 3: sequential blocks, kernels ordered by INT gpu_op_uid ---- + order3 = np.lexsort((merged["gpu_op_uid"].to_numpy(), merged["source"].to_numpy())) + gold3 = gold[order3] + b3_fwd, b3_rev = both_purities(gold3, pool) + b3_sfwd, b3_srev = both_strict(gold3, pool) + + # ---------------- report ---------------- + print("=" * 74) + print(f"{title}") + print("Matched kernels:", n, "| no-capture buckets:", len(bucket_ids)) + print("Distinct gold LCAs:", merged["lca_gold"].nunique()) + print("=" * 74) + print("No-capture bucket sizes (preserved by all baselines):") + for bid, sz in zip(bucket_ids, bucket_sizes): + print(f" LCA {bid:>3} {names.get(bid, '?'):<18} {sz:>4}") + print("-" * 74) + hdr = f"{'partition':<26}{'forward (gold->nc)':>24}{'reverse (nc->gold)':>24}" + print(hdr) + print("-" * 74) + print(f"{'Real no-capture':<26}{real_fwd:>24.4f}{real_rev:>24.4f}") + print( + f"{'Baseline 1 (random)':<26}" + f"{f'{b1_fwd_m:.4f} +/- {b1_fwd_s:.4f}':>24}" + f"{f'{b1_rev_m:.4f} +/- {b1_rev_s:.4f}':>24}" + ) + print(f"{'Baseline 2 (seq, key str)':<26}{b2_fwd:>24.4f}{b2_rev:>24.4f}") + print(f"{'Baseline 3 (seq, int uid)':<26}{b3_fwd:>24.4f}{b3_rev:>24.4f}") + print("-" * 74) + print("Strict consistency (fraction of items in a FULLY pure group):") + sh = f"{'partition':<26}{'strict fwd (gold pure)':>24}{'strict rev (nc pure)':>24}" + print(sh) + print("-" * 74) + print(f"{'Real no-capture':<26}{real_sfwd:>24.4f}{real_srev:>24.4f}") + print( + f"{'Baseline 1 (random)':<26}" + f"{f'{b1_sfwd_m:.4f} +/- {b1_sfwd_s:.4f}':>24}" + f"{f'{b1_srev_m:.4f} +/- {b1_srev_s:.4f}':>24}" + ) + print(f"{'Baseline 2 (seq, key str)':<26}{b2_sfwd:>24.4f}{b2_srev:>24.4f}") + print(f"{'Baseline 3 (seq, int uid)':<26}{b3_sfwd:>24.4f}{b3_srev:>24.4f}") + print("=" * 74) + + +def main() -> None: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "with_capture", type=Path, help="path to with-capture (GOLD) diff_stats.csv" + ) + ap.add_argument("no_capture", type=Path, help="path to no-capture diff_stats.csv") + ap.add_argument( + "--trials", + type=int, + default=DEFAULT_TRIALS, + help=f"random-baseline seeds (default: {DEFAULT_TRIALS}; auto-capped to " + f"{REDUCED_TRIALS} when matched kernels > {LARGE_KERNEL_THRESHOLD})", + ) + args = ap.parse_args() + + gold = load(args.with_capture) + ncap = load(args.no_capture) + + # Guard: with no shared identity keys the matched set is empty and every + # downstream report divides by zero / sorts an empty frame. Fail cleanly. + if not (set(gold["key"]) & set(ncap["key"])): + print( + f"No matched keys between the two files ({len(gold)} gold rows, " + f"{len(ncap)} no-capture rows, 0 shared keys). Nothing to compare." + ) + return + + # ---- Part 1: LCA-partition agreement ---- + agreement_report(gold, ncap) + print() + + # Matched set for the baselines: identity present in both. Deterministic by key. + merged = ( + gold[["key", LCA_COL]] + .rename(columns={LCA_COL: "lca_gold"}) + .merge( + ncap[["key", "source", "gpu_op_uid", LCA_COL, LCA_NAME]].rename( + columns={LCA_COL: "lca_nc", LCA_NAME: "nc_name"} + ), + on="key", + ) + ) + + # Guard: the random baseline is O(trials * kernels); cap trials on large sets. + n_trials = args.trials + if len(merged) > LARGE_KERNEL_THRESHOLD: + n_trials = REDUCED_TRIALS + print( + f"[guard] matched kernels = {len(merged)} > {LARGE_KERNEL_THRESHOLD}; " + f"reducing random-baseline trials {args.trials} -> {n_trials}" + ) + print() + + # ---- Part 2: random / sequential baselines ---- + baseline_report(merged, "ALL MATCHED KERNELS", n_trials) + + # Exclude singleton gold groups: gold LCAs with exactly one matched kernel. + gsz = merged.groupby("lca_gold")["key"].transform("size") + merged_ns = merged[gsz > 1].copy() + n_dropped = len(merged) - len(merged_ns) + print() + baseline_report( + merged_ns, + f"EXCLUDING SINGLETON GOLD GROUPS (dropped {n_dropped} kernels)", + n_trials, + ) + print(f"(Baseline 1 over {n_trials} random seeds; Baselines 2 & 3 deterministic.)") + + +if __name__ == "__main__": + main() diff --git a/agent_evals/Analysis/partial_tests/.gitignore b/agent_evals/Analysis/partial_tests/.gitignore new file mode 100644 index 000000000..314eaa1e0 --- /dev/null +++ b/agent_evals/Analysis/partial_tests/.gitignore @@ -0,0 +1,8 @@ +# Expanded fixtures -- regenerated from fixtures/*.tar.gz by run_partial_tests.sh. +# Only the tarballs under fixtures/ are committed (matches the e2e_tests_* convention). +/semantic_purity_deepseek_r1/ +/semantic_purity_qwen3_30b_a3b/ +# Run outputs and scratch +/partial_results/ +_gold_scratch_*/ +eval_utils/__pycache__/ diff --git a/agent_evals/Analysis/partial_tests/README.md b/agent_evals/Analysis/partial_tests/README.md new file mode 100644 index 000000000..b2f347288 --- /dev/null +++ b/agent_evals/Analysis/partial_tests/README.md @@ -0,0 +1,115 @@ + + +# Partial-workflow tests + +This folder holds tests that run only **part** of an analysis workflow — just +far enough to produce an artifact that a scripted eval scores against a +pre-baked reference. This is deliberately lighter than +`../eval_scripts/run_repeatability_parallel.sh`, which runs the full 11-step +analysis orchestrator per case. + +The first (and currently only) partial-workflow test is **`semantic_purity`**. +The folder is named generically so other partial-workflow tests can be added +beside it (give them a new value in the `workflow` column of +`partial_test_cases.csv` and a matching branch in `run_partial_tests.sh`). + +## Layout + +``` +partial_tests/ +├── run_partial_tests.sh # the runner +├── generate_updated_semantic_gold.sh # OPTIONAL, rare: regenerate gold +├── partial_test_cases.csv # manifest (id,workflow,trace_a,trace_b,reference_dir,platform_a,platform_b) +├── eval_utils/ +│ ├── semantic_partition_scripted_evals.py # per-run purity metrics (informational) +│ └── semantic_purity_aggregate.py # the regression gate +├── fixtures/*.tar.gz # DECODE trace pair + pre-baked gold (small) +└── / # expanded fixture (MI300/, B300/, analysis_output_ref/) +``` + +`compare_lca_partitions.py` — the shared purity/consistency math — lives in +`../eval_utils/` (it is also imported by `tests/test_compare_lca_partitions.py`). + +## What `semantic_purity` does + +For each model (DeepSeek-R1, Qwen3-30B-A3B) the fixture ships: +- the **DECODE** trace for MI300 and B300 (single batch-16 decode execution — the + only trace kind these tests use), and +- a pre-baked **gold** `analysis_output_ref/semantic_purity_gold_diff_stats.csv`. + +The runner invokes the **semantic-comparison workflow** (the no-capture, +name-first + coherence bucketing method) on the DECODE pair, running only +through its "Generate TraceDiff Output" step — equivalent to the main +orchestrator through Step 1.S. That produces +`analysis_output/tracediff_output/diff_stats.csv`, whose LCA partition is then +compared to gold. + +Gold itself is the *with-capture* TraceDiff partition on the same DECODE pair. +Because both sides derive from the same single-execution DECODE trace, their +`gpu_op_uid` numbering lines up, so the `(source, gpu_op_uid)` join used by +`compare_lca_partitions.py` is meaningful. + +## Running + +```bash +# from the repo root +bash agent_evals/Analysis/partial_tests/run_partial_tests.sh + +# one model, single run (quick check) +TEST_IDS=semantic_purity_deepseek_r1 NUM_REPEATS=1 \ + bash agent_evals/Analysis/partial_tests/run_partial_tests.sh + +# variance study +NUM_REPEATS=5 bash agent_evals/Analysis/partial_tests/run_partial_tests.sh +``` + +Env knobs: `NUM_REPEATS` (default 1), `TEST_IDS` (space-separated whitelist), +`MAX_PARALLEL` (default 3), `CONTAINER` (docker container to exec python in), +`AGENT_MODEL`. Results land under `partial_tests/partial_results//run_*/`. + +## Semantic-purity quality gate + +The semantic bucketing method uses an LLM for its unification/coherence steps, +so it has **real run-to-run variance**. Therefore: + +- `semantic_partition_scripted_evals.py` writes each run's metrics + (`forward_purity`, `strict_forward`, etc.) but is **informational only** — its + per-run `result` is PASS unless the pipeline itself failed (missing candidate, + no matched keys). +- `semantic_purity_aggregate.py` is the **actual gate**. It averages + `strict_forward` across the runs found and compares to a per-model floor. + The floors — and the observed-run data they're derived from — live in that + script's `MIN_STRICT_FORWARD_AVG` and its module docstring, which are the + single source of truth (not duplicated here). The floors are set below the + worst single observed run with margin, so a quick `NUM_REPEATS=1` check does + not false-alarm, while staying well above the random-shuffle baseline. Raise + them if the method is intentionally improved; do not lower them without a + documented reason. + +## Regenerating gold (optional, rare — maintainer only) + +Gold ships **pre-baked** inside the fixture tarballs, so normal test runs never +need this. Regenerate only after an intentional, reviewed change to the +with-capture TraceDiff path. + +**Precondition — not runnable from a clean clone.** Regeneration needs the +original **full-capture** traces (the DECODE traces *and* their +`capture_traces/` folders). Those are large and are **not committed to the +repo** — only the slimmed DECODE-only traces ship in the fixtures. You must +already have the full-capture traces available locally (on the machine where +they were captured). The script's source paths default to +`tests/traces/semantic/...` and are set in the `SEM_ROOT` / `SOURCES` entries at +the top of the script; edit them to point at wherever you keep the traces. The +script fails with a clear "missing source" error if they aren't present. + +```bash +bash agent_evals/Analysis/partial_tests/generate_updated_semantic_gold.sh +``` + +It runs the with-capture perf-report / TraceDiff path on the DECODE pair, writes +the refreshed CSV into each `analysis_output_ref/`, and repacks the fixture +tarball. diff --git a/agent_evals/Analysis/partial_tests/eval_utils/semantic_partition_scripted_evals.py b/agent_evals/Analysis/partial_tests/eval_utils/semantic_partition_scripted_evals.py new file mode 100644 index 000000000..e5c6790c8 --- /dev/null +++ b/agent_evals/Analysis/partial_tests/eval_utils/semantic_partition_scripted_evals.py @@ -0,0 +1,167 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Scripted eval: LCA-partition purity vs. a with-capture (gold) reference. + +Self-gating: only test cases whose reference_dir contains a +``semantic_purity_gold_diff_stats.csv`` (which ships pre-baked in each +partial-test fixture; regenerated only rarely by +``generate_updated_semantic_gold.sh``) are scored. Any other test case +produces zero rows here, so this eval is inert unless a gold reference is +present. + +For gated test cases, this reuses compare_lca_partitions.py's forward/reverse +purity and strict forward/reverse consistency metrics, computed between the +gold partition and the candidate's semantic-bucketing output +(``tracediff_output/diff_stats.csv``, produced by the semantic-comparison +workflow through its "Generate TraceDiff Output" step), and records them. + +IMPORTANT: this per-run result is informational only. Because the semantic +method has real run-to-run variance (observed on Qwen3-30B-A3B), a single +run's metrics are not a reliable regression signal by themselves -- the +actual pass/fail decision is made by semantic_purity_aggregate.py, which +averages strict_forward across NUM_REPEATS runs and compares against a +floor derived from currently-observed performance. See that script and the +partial_tests/README.md "Semantic-purity quality gate" section. + +Consequently the "result" written here is PASS whenever the metrics were +computed at all; it only turns FAIL for genuine pipeline errors (candidate +output missing, no matched keys). +""" + +import argparse +import csv +import os +import sys + +# compare_lca_partitions lives two dirs up in eval_utils/ (not a package). +_HERE = os.path.dirname(os.path.abspath(__file__)) +_ANALYSIS_EVAL_UTILS = os.path.abspath(os.path.join(_HERE, "..", "..", "eval_utils")) +sys.path.insert(0, _ANALYSIS_EVAL_UTILS) +from compare_lca_partitions import ( + LCA_COL, + both_purities, + both_strict, + load, +) # noqa: E402 + +CSV_COLUMNS = [ + "index", + "category", + "issue_summary", + "result", + "details", + "root_cause", + "recommended_fix", +] + +GOLD_FILENAME = "semantic_purity_gold_diff_stats.csv" +CANDIDATE_RELPATH = os.path.join("tracediff_output", "diff_stats.csv") + + +def _write(results_path: str, rows: list[dict]) -> None: + with open(results_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS) + writer.writeheader() + writer.writerows(rows) + + +def run(output_dir: str, reference_dir: str, results_path: str) -> list[dict]: + gold_path = os.path.join(reference_dir, GOLD_FILENAME) + if not os.path.isfile(gold_path): + # No gold reference here; self-gate to a no-op. + _write(results_path, []) + return [] + + candidate_path = os.path.join(output_dir, CANDIDATE_RELPATH) + if not os.path.isfile(candidate_path): + rows = [ + { + "index": "semantic_purity_1", + "category": "Quality", + "issue_summary": "LCA-partition purity vs with-capture gold", + "result": "FAIL", + "details": f"Candidate semantic diff-stats not found: {candidate_path}", + "root_cause": "pipeline", + "recommended_fix": "Semantic workflow did not produce tracediff_output/diff_stats.csv; check the run log", + } + ] + _write(results_path, rows) + return rows + + gold = load(gold_path) + cand = load(candidate_path) + + merged = ( + gold[["key", LCA_COL]] + .rename(columns={LCA_COL: "lca_gold"}) + .merge( + cand[["key", LCA_COL]].rename(columns={LCA_COL: "lca_cand"}), + on="key", + ) + ) + + if merged.empty: + rows = [ + { + "index": "semantic_purity_1", + "category": "Quality", + "issue_summary": "LCA-partition purity vs with-capture gold", + "result": "FAIL", + "details": "No matched (source, gpu_op_uid) keys between gold and candidate", + "root_cause": "data", + "recommended_fix": "Check that candidate and gold were generated from the same DECODE trace pair", + } + ] + _write(results_path, rows) + return rows + + gold_lbl = merged["lca_gold"].to_numpy() + cand_lbl = merged["lca_cand"].to_numpy() + fwd, rev = both_purities(gold_lbl, cand_lbl) + sfwd, srev = both_strict(gold_lbl, cand_lbl) + + # Informational only -- see module docstring. The regression gate is + # applied post-hoc, across repeats, by semantic_purity_aggregate.py. + details = ( + f"matched={len(merged)} forward_purity={fwd:.4f} reverse_purity={rev:.4f} " + f"strict_forward={sfwd:.4f} strict_reverse={srev:.4f} " + f"shared_buckets={merged['lca_cand'].nunique()} " + f"(informational only -- see semantic_purity_aggregate.py for the actual gate)" + ) + rows = [ + { + "index": "semantic_purity_1", + "category": "Quality", + "issue_summary": "LCA-partition purity vs with-capture gold", + "result": "PASS", + "details": details, + "root_cause": "", + "recommended_fix": "", + } + ] + _write(results_path, rows) + return rows + + +def main(): + parser = argparse.ArgumentParser( + description="LCA-partition purity vs with-capture gold" + ) + parser.add_argument("--output-dir", required=True) + parser.add_argument("--reference-dir", required=True) + parser.add_argument("--results", required=True) + args = parser.parse_args() + + rows = run(args.output_dir, args.reference_dir, args.results) + if not rows: + sys.exit(0) + passed = sum(1 for r in rows if r["result"] == "PASS") + sys.exit(0 if passed == len(rows) else 1) + + +if __name__ == "__main__": + main() diff --git a/agent_evals/Analysis/partial_tests/eval_utils/semantic_purity_aggregate.py b/agent_evals/Analysis/partial_tests/eval_utils/semantic_purity_aggregate.py new file mode 100644 index 000000000..b24084abc --- /dev/null +++ b/agent_evals/Analysis/partial_tests/eval_utils/semantic_purity_aggregate.py @@ -0,0 +1,241 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Aggregate the semantic-purity regression gate across repeated runs. + +semantic_partition_scripted_evals.py records per-run LCA-partition purity +metrics but does not gate on them (see that module's docstring) -- the +semantic-bucketing method has real run-to-run variance (observed on +Qwen3-30B-A3B), so a single run is not a reliable regression signal. + +This script is the actual gate. For each semantic-purity test case it: + 1. Scans results_root//run_*/semantic_purity_results.csv, + parsing out each run's strict_forward value (the "strict forward + consistency" metric from compare_lca_partitions.py: the fraction of + gold groups that map, in their entirety, to a single candidate LCA). + strict_forward is used as the sole gating metric -- it is more + discriminating than forward_purity because it does not give partial + credit to a majority-only match. + 2. Averages strict_forward across however many runs were found. + 3. Compares that average against MIN_STRICT_FORWARD_AVG[test_id] -- a + floor set BELOW the worst single independently-observed run (with + margin), so that even a quick NUM_REPEATS=1 check does not + false-alarm on ordinary variance, while still sitting far above the + random-shuffle baseline (Baseline 1 in compare_lca_partitions.py, + ~0.04-0.06 strict_forward for both models). The intent is regression + detection ("don't get materially worse than today"), not an absolute + quality bar. + +Writes one verdict row (7-column eval schema) per test id to +results_root//semantic_purity_aggregate_verdict.csv, and prints a +summary table. + +Usage: python3 semantic_purity_aggregate.py --results-root +""" + +import argparse +import csv +import glob +import os +import re +import statistics +import sys + +CSV_COLUMNS = [ + "index", + "category", + "issue_summary", + "result", + "details", + "root_cause", + "recommended_fix", +] + +# Floors are set below the worst single strict_forward run observed during +# decode-only method validation (with margin), NOT at the mean -- so a +# one-off NUM_REPEATS=1 quick check cannot false-alarm on ordinary +# variance. Averaging across repeats leaves even more headroom. +# +# Observed decode-only strict_forward (compare_lca_partitions.py vs the +# pre-baked gold): +# deepseek_r1: 0.7280, 0.9780, 0.9748, 0.9748, 0.9780 (worst 0.7280) +# qwen3_30b_a3b: 0.3710, 0.5671, 0.4677, 0.7304, 0.5691 (worst 0.3710) +# Both floors sit far above the random-shuffle baseline (~0.04-0.06). If the +# method is intentionally improved, raise these; do not lower them without a +# documented reason. +MIN_STRICT_FORWARD_AVG = { + "semantic_purity_deepseek_r1": 0.60, + "semantic_purity_qwen3_30b_a3b": 0.30, +} + +# Degenerate-collapse guard. strict_forward saturates to ~1.0 when the candidate +# partition collapses toward a single LCA bucket (every gold group is trivially +# homogeneous inside one bucket), so the floor above cannot catch total collapse. +# Distinct candidate LCA buckets on the matched set of observed healthy runs +# ranged 26-32 for both models; total collapse is 1. A floor of 10 sits ~2.6x +# below the worst healthy run and ~10x above collapse, so ordinary bucketing +# variance never false-alarms while any real collapse fails hard. +MIN_SHARED_BUCKETS = { + "semantic_purity_deepseek_r1": 10, + "semantic_purity_qwen3_30b_a3b": 10, +} + +DETAILS_RE = re.compile(r"strict_forward=([0-9.]+)") +FORWARD_RE = re.compile(r"forward_purity=([0-9.]+)") +BUCKETS_RE = re.compile(r"shared_buckets=([0-9]+)") + + +def find_run_csvs(results_root: str, test_id: str) -> list[str]: + pattern = os.path.join( + results_root, test_id, "run_*", "semantic_purity_results.csv" + ) + return sorted(glob.glob(pattern)) + + +def parse_run(csv_path: str): + with open(csv_path, newline="") as f: + rows = list(csv.DictReader(f)) + if not rows: + return None # self-gated no-op for this test id (shouldn't happen here) + row = rows[0] + details = row.get("details", "") + m_sfwd = DETAILS_RE.search(details) + m_fwd = FORWARD_RE.search(details) + m_buckets = BUCKETS_RE.search(details) + if not m_sfwd: + return None # pipeline failure row, no metrics to parse + return { + "strict_forward": float(m_sfwd.group(1)), + "forward_purity": float(m_fwd.group(1)) if m_fwd else None, + "shared_buckets": int(m_buckets.group(1)) if m_buckets else None, + "source": csv_path, + } + + +def aggregate_one(results_root: str, test_id: str) -> dict: + run_csvs = find_run_csvs(results_root, test_id) + parsed = [p for p in (parse_run(c) for c in run_csvs) if p is not None] + + if not parsed: + return { + "index": "semantic_purity_aggregate", + "category": "Quality", + "issue_summary": f"Semantic-purity regression gate ({test_id})", + "result": "FAIL", + "details": f"No usable runs found under {results_root}/{test_id}/run_*/semantic_purity_results.csv", + "root_cause": "pipeline", + "recommended_fix": "Ensure the run(s) completed for this test id before aggregating", + } + + values = [p["strict_forward"] for p in parsed] + avg = statistics.mean(values) + floor = MIN_STRICT_FORWARD_AVG.get(test_id) + if floor is None: + return { + "index": "semantic_purity_aggregate", + "category": "Quality", + "issue_summary": f"Semantic-purity regression gate ({test_id})", + "result": "FAIL", + "details": f"No floor configured for test id {test_id}; add one to MIN_STRICT_FORWARD_AVG", + "root_cause": "config", + "recommended_fix": "Add this test id to MIN_STRICT_FORWARD_AVG in semantic_purity_aggregate.py", + } + + passed = avg >= floor + + # Collapse guard: fail hard if the distinct candidate LCA count on the + # matched set drops below the floor, which would saturate strict_forward. + bucket_floor = MIN_SHARED_BUCKETS.get(test_id) + bucket_vals = [ + p["shared_buckets"] for p in parsed if p.get("shared_buckets") is not None + ] + min_buckets = min(bucket_vals) if bucket_vals else None + collapsed = ( + bucket_floor is not None + and min_buckets is not None + and min_buckets < bucket_floor + ) + passed = passed and not collapsed + + n_note = ( + "" + if len(parsed) > 1 + else " (single run; a multi-run average is a stronger estimate)" + ) + bucket_note = ( + f" min_shared_buckets={min_buckets} bucket_floor={bucket_floor}" + if bucket_floor is not None + else "" + ) + details = ( + f"n_runs={len(parsed)} strict_forward_values={[round(v, 4) for v in values]} " + f"avg_strict_forward={avg:.4f} floor={floor}{bucket_note}{n_note}" + ) + if collapsed: + details += ( + f" -- FAIL: candidate partition collapsed toward a single LCA bucket " + f"(min {min_buckets} < {bucket_floor}); strict_forward is unreliable" + ) + + if passed: + recommended_fix = "" + elif collapsed: + recommended_fix = ( + "Candidate LCA partition collapsed (near-single-bucket); strict_forward " + "is saturated and meaningless. Investigate the clustering/coherence step." + ) + else: + recommended_fix = ( + "Semantic-bucketing quality regressed below the observed-performance " + "floor; investigate the clustering change" + ) + + return { + "index": "semantic_purity_aggregate", + "category": "Quality", + "issue_summary": f"Semantic-purity regression gate ({test_id})", + "result": "PASS" if passed else "FAIL", + "details": details, + "root_cause": "" if passed else "quality", + "recommended_fix": recommended_fix, + } + + +def main(): + parser = argparse.ArgumentParser( + description="Aggregate semantic-purity gate across repeats" + ) + parser.add_argument("--results-root", required=True) + parser.add_argument( + "--test-ids", + default=",".join(MIN_STRICT_FORWARD_AVG.keys()), + help="Comma-separated test ids to aggregate (default: all known semantic-purity test ids)", + ) + args = parser.parse_args() + + test_ids = [t.strip() for t in args.test_ids.split(",") if t.strip()] + overall_pass = True + for test_id in test_ids: + case_dir = os.path.join(args.results_root, test_id) + if not os.path.isdir(case_dir): + print(f"[{test_id}] skipped (no results directory found)") + continue + result = aggregate_one(args.results_root, test_id) + out_path = os.path.join(case_dir, "semantic_purity_aggregate_verdict.csv") + with open(out_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=CSV_COLUMNS) + writer.writeheader() + writer.writerow(result) + print(f"[{test_id}] {result['result']}: {result['details']}") + print(f" -> {out_path}") + if result["result"] != "PASS": + overall_pass = False + + sys.exit(0 if overall_pass else 1) + + +if __name__ == "__main__": + main() diff --git a/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_deepseek_r1.tar.gz b/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_deepseek_r1.tar.gz new file mode 100644 index 000000000..69c88671f Binary files /dev/null and b/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_deepseek_r1.tar.gz differ diff --git a/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_qwen3_30b_a3b.tar.gz b/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_qwen3_30b_a3b.tar.gz new file mode 100644 index 000000000..22ab838e0 Binary files /dev/null and b/agent_evals/Analysis/partial_tests/fixtures/semantic_purity_qwen3_30b_a3b.tar.gz differ diff --git a/agent_evals/Analysis/partial_tests/generate_updated_semantic_gold.sh b/agent_evals/Analysis/partial_tests/generate_updated_semantic_gold.sh new file mode 100755 index 000000000..ec080b305 --- /dev/null +++ b/agent_evals/Analysis/partial_tests/generate_updated_semantic_gold.sh @@ -0,0 +1,105 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -uo pipefail + +# --------------------------------------------------------------------------- +# OPTIONAL, RARELY RUN. The gold reference ships pre-baked inside each +# partial-test fixture tarball -- normal test runs never call this script. +# +# Run it only to regenerate the gold reference (e.g. after an intentional, +# reviewed change to the with-capture TraceDiff path). It reproduces gold the +# same way it was originally produced: the with-capture perf-report / TraceDiff +# path on the single-execution DECODE trace pair (batch-16 decode), which is +# what makes gold's gpu_op_uid range line up with the no-capture candidate. +# +# It sources the DECODE traces AND their capture_traces folders from the +# original full-fidelity trace location (NOT from the slimmed fixture, which +# has no capture data), writes the refreshed CSV into each fixture's +# analysis_output_ref/, and rebuilds the fixture tarball so gold stays shipped. +# +# Usage: bash generate_updated_semantic_gold.sh [test_id ...] (default: all) +# --------------------------------------------------------------------------- + +REPO_ROOT="${REPO_ROOT:-$(pwd)}" +PARTIAL_DIR="$REPO_ROOT/agent_evals/Analysis/partial_tests" +GOLD_FILENAME="semantic_purity_gold_diff_stats.csv" +ARCH_JSON="${ARCH_JSON:-$REPO_ROOT/TraceLens/Agent/Analysis/utils/arch/MI300X.json}" +PERF_CLI="${PERF_CLI:-TraceLens_generate_perf_report_pytorch_inference}" + +# Original full-capture source locations (repo-relative). Override via env if +# the traces move. Format: " ". +SEM_ROOT="tests/traces/semantic" +declare -A SOURCES=( + [semantic_purity_deepseek_r1]="\ +$SEM_ROOT/deepseek_R1/Deepseek-R1-Distill-LLama-8B/MI300/torch_trace/1782841132.2905657-TP-0-DECODE.trace.json.gz \ +$SEM_ROOT/deepseek_R1/Deepseek-R1-Distill-LLama-8B/MI300/torch_trace/capture_traces \ +$SEM_ROOT/deepseek_R1/Deepseek-R1-Distill-LLama-8B/B300/torch_trace/1782843274.7730265-TP-0-DECODE.trace.json.gz \ +$SEM_ROOT/deepseek_R1/Deepseek-R1-Distill-LLama-8B/B300/torch_trace/capture_traces" + [semantic_purity_qwen3_30b_a3b]="\ +$SEM_ROOT/qwen3_30b_a3b/Qwen3-30B-A3B/MI300/torch_trace/1782859605.7198358-TP-0-DECODE.trace.json.gz \ +$SEM_ROOT/qwen3_30b_a3b/Qwen3-30B-A3B/MI300/torch_trace/capture_traces \ +$SEM_ROOT/qwen3_30b_a3b/Qwen3-30B-A3B/B300/torch_trace/1782799835.6885648-TP-0-DECODE.trace.json.gz \ +$SEM_ROOT/qwen3_30b_a3b/Qwen3-30B-A3B/B300/torch_trace/capture_traces" +) + +generate_one() { + local id="$1" + local spec="${SOURCES[$id]:-}" + if [[ -z "$spec" ]]; then + echo "ERROR: unknown test id '$id' (no source mapping)" >&2 + return 1 + fi + # shellcheck disable=SC2086 + set -- $spec + local decode_a="$REPO_ROOT/$1" cap_a="$REPO_ROOT/$2" decode_b="$REPO_ROOT/$3" cap_b="$REPO_ROOT/$4" + + for p in "$decode_a" "$cap_a" "$decode_b" "$cap_b" "$ARCH_JSON"; do + [[ -e "$p" ]] || { echo "ERROR [$id]: missing source $p" >&2; return 1; } + done + + local scratch="$PARTIAL_DIR/_gold_scratch_$id" + local out_csvs="$scratch/perf_report_trace1_csvs" + rm -rf "$scratch"; mkdir -p "$scratch" + + echo "[$id] running with-capture TraceDiff on the DECODE pair..." + # trace1 (MI300) with capture + comparison against trace2 (B300) with its + # capture triggers TraceDiff internally and writes diff_stats.csv into + # perf_report_trace1_csvs/. Only platform1's arch JSON is needed. + "$PERF_CLI" \ + --profile_json_path "$decode_a" --capture_folder "$cap_a" \ + --gpu_arch_json_path "$ARCH_JSON" \ + --group_by_parent_module --enable_pseudo_ops --group_by_num_kernels --include_call_stack \ + --comparison_json_path "$decode_b" --comparison_capture_folder "$cap_b" \ + --output_xlsx_path "$scratch/perf_report_trace1.xlsx" \ + --output_csvs_dir "$out_csvs" \ + || { echo "ERROR [$id]: perf-report CLI failed" >&2; return 1; } + + local gold_src="$out_csvs/diff_stats.csv" + [[ -f "$gold_src" ]] || { echo "ERROR [$id]: $gold_src not produced" >&2; return 1; } + + local ref_dir="$PARTIAL_DIR/$id/analysis_output_ref" + mkdir -p "$ref_dir" + cp "$gold_src" "$ref_dir/$GOLD_FILENAME" + echo "[$id] gold -> $ref_dir/$GOLD_FILENAME ($(wc -l < "$ref_dir/$GOLD_FILENAME") lines)" + + # Re-pack the fixture tarball so the refreshed gold ships with it. + tar czf "$PARTIAL_DIR/fixtures/${id}.tar.gz" -C "$REPO_ROOT" "agent_evals/Analysis/partial_tests/$id" + echo "[$id] fixture repacked -> fixtures/${id}.tar.gz" + rm -rf "$scratch" +} + +ids=("$@") +if [[ ${#ids[@]} -eq 0 ]]; then + ids=("${!SOURCES[@]}") +fi + +rc=0 +for id in "${ids[@]}"; do + generate_one "$id" || rc=1 +done +exit "$rc" diff --git a/agent_evals/Analysis/partial_tests/partial_test_cases.csv b/agent_evals/Analysis/partial_tests/partial_test_cases.csv new file mode 100644 index 000000000..a00b6bbee --- /dev/null +++ b/agent_evals/Analysis/partial_tests/partial_test_cases.csv @@ -0,0 +1,3 @@ +id,workflow,trace_a,trace_b,reference_dir,platform_a,platform_b +semantic_purity_deepseek_r1,semantic_comparison,agent_evals/Analysis/partial_tests/semantic_purity_deepseek_r1/MI300/1782841132.2905657-TP-0-DECODE.trace.json.gz,agent_evals/Analysis/partial_tests/semantic_purity_deepseek_r1/B300/1782843274.7730265-TP-0-DECODE.trace.json.gz,agent_evals/Analysis/partial_tests/semantic_purity_deepseek_r1/analysis_output_ref,MI300,B300 +semantic_purity_qwen3_30b_a3b,semantic_comparison,agent_evals/Analysis/partial_tests/semantic_purity_qwen3_30b_a3b/MI300/1782859605.7198358-TP-0-DECODE.trace.json.gz,agent_evals/Analysis/partial_tests/semantic_purity_qwen3_30b_a3b/B300/1782799835.6885648-TP-0-DECODE.trace.json.gz,agent_evals/Analysis/partial_tests/semantic_purity_qwen3_30b_a3b/analysis_output_ref,MI300,B300 diff --git a/agent_evals/Analysis/partial_tests/run_partial_tests.sh b/agent_evals/Analysis/partial_tests/run_partial_tests.sh new file mode 100755 index 000000000..fa7624564 --- /dev/null +++ b/agent_evals/Analysis/partial_tests/run_partial_tests.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +set -uo pipefail + +# --------------------------------------------------------------------------- +# Partial-workflow test runner. +# +# Unlike run_repeatability_parallel.sh (which runs the FULL 11-step analysis +# orchestrator per case), this harness runs only PART of a workflow -- just +# far enough to produce the artifact a scripted eval scores against a +# pre-baked reference. The first such test is semantic_purity, which runs the +# semantic-comparison workflow through its "Generate TraceDiff Output" step +# (== main orchestrator through Step 1.S) to produce +# /tracediff_output/diff_stats.csv +# and scores its LCA-partition purity against the committed gold. +# +# New partial-workflow tests can be added by giving them a distinct value in +# the `workflow` column of partial_test_cases.csv and a matching branch in +# run_single_job below. +# +# Usage: bash run_partial_tests.sh +# NUM_REPEATS repeats per case (default 1; raise for a variance study) +# TEST_IDS space-separated id whitelist (default: all) +# MAX_PARALLEL concurrent jobs (default 3) +# CONTAINER optional docker container to exec python in +# --------------------------------------------------------------------------- + +MAX_PARALLEL="${MAX_PARALLEL:-3}" +NUM_REPEATS="${NUM_REPEATS:-1}" +TEST_IDS="${TEST_IDS:-}" +CONTAINER="${CONTAINER:-}" +AGENT_MODEL="${AGENT_MODEL:-claude-opus-4-8-thinking-medium}" + +REPO_ROOT="${REPO_ROOT:-$(pwd)}" +ANALYSIS_DIR="$REPO_ROOT/TraceLens/Agent/Analysis" +EVALS_DIR="$REPO_ROOT/agent_evals/Analysis" +PARTIAL_DIR="$EVALS_DIR/partial_tests" +TEST_CASES_CSV="${TEST_CASES_CSV:-$PARTIAL_DIR/partial_test_cases.csv}" +RESULTS_ROOT="${RESULTS_ROOT:-$PARTIAL_DIR/partial_results}" + +if [[ -n "$CONTAINER" ]]; then + DEXEC=(docker exec -w "$REPO_ROOT" "$CONTAINER") +else + DEXEC=() +fi + +ts() { date "+%H:%M:%S"; } +log_status() { flock 1 echo "$@"; } + +# Expand a fixture tarball (paths inside are repo-relative) if not already present. +expand_fixture() { + local id="$1" + local archive="$PARTIAL_DIR/fixtures/${id}.tar.gz" + local target="$PARTIAL_DIR/$id" + if [[ -f "$archive" ]] && [[ ! -d "$target" ]]; then + echo "Expanding fixtures/${id}.tar.gz..." + tar xzf "$archive" -C "$REPO_ROOT" + fi +} + +# --------------------------------------------------------------------------- +# One (test_case, repeat) iteration. +# Args: id workflow repeat trace_a trace_b reference_dir platform_a platform_b +# --------------------------------------------------------------------------- +run_single_job() { + local id="$1" workflow="$2" repeat="$3" trace_a="$4" trace_b="$5" reference_dir="$6" platform_a="$7" platform_b="$8" + local tag="[$id|run_$repeat]" + + local CASE_RESULTS="$RESULTS_ROOT/$id/run_${repeat}" + local OUTPUT_DIR="$CASE_RESULTS/analysis_output" + rm -rf "$CASE_RESULTS" 2>/dev/null || true + mkdir -p "$OUTPUT_DIR" + + local abs_a="$REPO_ROOT/$trace_a" + local abs_b="$REPO_ROOT/$trace_b" + local abs_ref="$REPO_ROOT/$reference_dir" + + log_status " $tag [$(ts)] workflow=$workflow starting" + + case "$workflow" in + semantic_comparison) + local agent_success=false + local attempts=0 + while [ "$agent_success" = false ] && [ "$attempts" -lt 3 ]; do + attempts=$((attempts + 1)) + ( + cd "$ANALYSIS_DIR" || exit + timeout 1800 agent --model "$AGENT_MODEL" --print --force --trust --output-format stream-json \ + "Follow the semantic-comparison-agent workflow (TraceLens/Agent/Analysis/skills/analysis-orchestrator/agents/semantic-comparison-agent.md). Run it on trace A=$abs_a (platform $platform_a) and trace B=$abs_b (platform $platform_b), output to $OUTPUT_DIR. Run only through Step 3 'Generate TraceDiff Output' so that $OUTPUT_DIR/tracediff_output/diff_stats.csv is produced; do NOT run the full analysis orchestrator or any downstream report/category steps." + ) < /dev/null > "$CASE_RESULTS/semantic_stream.ndjson" 2>&1 + + if head -c 2048 "$CASE_RESULTS/semantic_stream.ndjson" | grep -qiE 'Error:.*unavailable|Service Unavailable'; then + log_status " $tag Attempt $attempts/3 failed (agent unavailable). Backing off 30s..." + sleep 30 + else + agent_success=true + fi + done + if [ "$agent_success" = false ]; then + log_status " $tag FAILED after 3 attempts (agent unavailable). Skipping eval." + return 1 + fi + + "${DEXEC[@]}" python3 "$PARTIAL_DIR/eval_utils/semantic_partition_scripted_evals.py" \ + --output-dir "$OUTPUT_DIR" --reference-dir "$abs_ref" \ + --results "$CASE_RESULTS/semantic_purity_results.csv" \ + > "$CASE_RESULTS/semantic_purity_eval.log" 2>&1 || true + ;; + *) + log_status " $tag ERROR: unknown workflow '$workflow'" + return 1 + ;; + esac + + log_status " $tag [$(ts)] done -> $CASE_RESULTS/semantic_purity_results.csv" +} + +# --------------------------------------------------------------------------- +# FIFO semaphore for concurrency control +# --------------------------------------------------------------------------- +FIFO="$RESULTS_ROOT/.job_fifo" +cleanup() { rm -f "$FIFO"; } +setup_semaphore() { + rm -f "$FIFO"; mkfifo "$FIFO"; exec 4<>"$FIFO" + for ((t = 0; t < MAX_PARALLEL; t++)); do echo >&4; done + trap cleanup EXIT +} + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- +mkdir -p "$RESULTS_ROOT" + +# Expand every fixture referenced by the manifest. +while IFS=, read -r id workflow trace_a trace_b reference_dir platform_a platform_b <&3; do + [[ -z "$id" ]] && continue + expand_fixture "$id" +done 3< <(tail -n +2 "$TEST_CASES_CSV"; echo) + +echo "=========================================" +echo " Partial-workflow tests" +echo " Repeats: $NUM_REPEATS" +echo " Max parallel: $MAX_PARALLEL" +echo " CSV: $TEST_CASES_CSV" +if [[ -n "$TEST_IDS" ]]; then echo " Test filter: $TEST_IDS"; fi +echo "=========================================" +echo "" + +_spawn_jobs() { + local id="$1" workflow="$2" trace_a="$3" trace_b="$4" reference_dir="$5" platform_a="$6" platform_b="$7" + if [[ -n "$TEST_IDS" ]]; then + case " $TEST_IDS " in + *" $id "*) ;; + *) return ;; + esac + fi + for ((i = 0; i < NUM_REPEATS; i++)); do + read -r -u4 + ( + run_single_job "$id" "$workflow" "$i" "$trace_a" "$trace_b" "$reference_dir" "$platform_a" "$platform_b" || true + echo >&4 + sleep 2 # stagger agent startup to avoid ~/.cursor/cli-config.json rename race + ) & + sleep 2 + done +} + +setup_semaphore + +# manifest: id,workflow,trace_a,trace_b,reference_dir,platform_a,platform_b +while IFS=, read -r id workflow trace_a trace_b reference_dir platform_a platform_b <&3; do + [[ -z "$id" ]] && continue + _spawn_jobs "$id" "$workflow" "$trace_a" "$trace_b" "$reference_dir" "$platform_a" "$platform_b" +done 3< <(tail -n +2 "$TEST_CASES_CSV"; echo) + +wait + +echo "" +echo "=========================================" +echo " Runs finished. Applying regression gate..." +echo "=========================================" +"${DEXEC[@]}" python3 "$PARTIAL_DIR/eval_utils/semantic_purity_aggregate.py" \ + --results-root "$RESULTS_ROOT" || true + +echo "" +echo " Results in: $RESULTS_ROOT" diff --git a/tests/test_compare_lca_partitions.py b/tests/test_compare_lca_partitions.py new file mode 100644 index 000000000..8d8e920c6 --- /dev/null +++ b/tests/test_compare_lca_partitions.py @@ -0,0 +1,151 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Fast, no-agent unit tests for agent_evals/Analysis/eval_utils/compare_lca_partitions.py. + +Uses small synthetic diff_stats.csv fixtures (no traces, no LLM) to pin down +the purity/strict-consistency metric semantics. +""" + +import os +import sys + +import pandas as pd +import pytest + +sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "agent_evals", + "Analysis", + "eval_utils", + ), +) +from compare_lca_partitions import both_purities, both_strict, load # noqa: E402 + +COLUMNS = [ + "source", + "gpu_op_uid", + "name", + "lowest_common_ancestor_id", + "lowest_common_ancestor_name", +] + + +def _write_csv(path, rows): + pd.DataFrame(rows, columns=COLUMNS).to_csv(path, index=False) + + +def _row(source, uid, name, lca_id): + return [source, uid, name, lca_id, f"lca_{lca_id}"] + + +def test_load_requires_key_columns(tmp_path): + path = tmp_path / "bad.csv" + pd.DataFrame({"foo": [1, 2]}).to_csv(path, index=False) + with pytest.raises(ValueError, match="missing required columns"): + load(path) + + +def test_load_rejects_duplicate_keys(tmp_path): + path = tmp_path / "dup.csv" + _write_csv( + path, + [ + _row("trace1", 1, "kernelA", 0), + _row("trace1", 1, "kernelA", 0), + ], + ) + with pytest.raises(ValueError, match="duplicate"): + load(path) + + +def test_perfect_agreement_gives_purity_one(): + gold = ["g0", "g0", "g1", "g1"] + cand = ["c0", "c0", "c1", "c1"] + fwd, rev = both_purities(gold, cand) + sfwd, srev = both_strict(gold, cand) + assert fwd == 1.0 + assert rev == 1.0 + assert sfwd == 1.0 + assert srev == 1.0 + + +def test_candidate_collapse_hurts_reverse_purity_not_forward(): + # Every gold group maps entirely into one giant candidate bucket: forward + # purity is trivially perfect (each gold group agrees with the sole label + # it sees), but reverse purity collapses because that one candidate bucket + # disagrees with most of the gold labels it swallowed. + gold = ["g0", "g0", "g1", "g1", "g2", "g2"] + cand = ["c0"] * 6 + fwd, rev = both_purities(gold, cand) + assert fwd == 1.0 + assert rev == pytest.approx(2 / 6) # majority gold label (any of the 3) covers 2/6 + + +def test_partial_split_purity_between_zero_and_one(): + # gold group g0 (4 items) splits 3-1 across two candidate buckets. + gold = ["g0", "g0", "g0", "g0", "g1", "g1"] + cand = ["c0", "c0", "c0", "c1", "c2", "c2"] + fwd, rev = both_purities(gold, cand) + sfwd, srev = both_strict(gold, cand) + # forward: g0 majority=3/4, g1 majority=2/2 -> (3+2)/6 + assert fwd == pytest.approx(5 / 6) + # strict forward: g0 not fully pure (0 credit), g1 fully pure (2 credit) -> 2/6 + assert sfwd == pytest.approx(2 / 6) + assert sfwd <= fwd + assert srev <= rev + + +def test_strict_never_exceeds_matching_purity_on_random_like_partition(): + gold = ["g0"] * 5 + ["g1"] * 3 + ["g2"] * 2 + cand = ["c0", "c0", "c1", "c1", "c1", "c2", "c0", "c1", "c2", "c2"] + fwd, rev = both_purities(gold, cand) + sfwd, srev = both_strict(gold, cand) + assert sfwd <= fwd + assert srev <= rev + + +def test_end_to_end_via_csv_files(tmp_path): + gold_path = tmp_path / "gold.csv" + cand_path = tmp_path / "cand.csv" + _write_csv( + gold_path, + [ + _row("trace1", 1, "kernelA", 0), + _row("trace1", 2, "kernelB", 0), + _row("trace2", 1, "kernelC", 1), + _row("trace2", 2, "kernelD", 1), + ], + ) + _write_csv( + cand_path, + [ + _row("trace1", 1, "kernelA", 10), + _row("trace1", 2, "kernelB", 10), + _row("trace2", 1, "kernelC", 11), + _row("trace2", 2, "kernelD", 11), + ], + ) + gold = load(gold_path) + cand = load(cand_path) + merged = ( + gold[["key", "lowest_common_ancestor_id"]] + .rename(columns={"lowest_common_ancestor_id": "lca_gold"}) + .merge( + cand[["key", "lowest_common_ancestor_id"]].rename( + columns={"lowest_common_ancestor_id": "lca_cand"} + ), + on="key", + ) + ) + assert len(merged) == 4 + fwd, rev = both_purities( + merged["lca_gold"].to_numpy(), merged["lca_cand"].to_numpy() + ) + assert fwd == 1.0 and rev == 1.0 diff --git a/tests/test_semantic_annotation_metadata.py b/tests/test_semantic_annotation_metadata.py new file mode 100644 index 000000000..4821a9926 --- /dev/null +++ b/tests/test_semantic_annotation_metadata.py @@ -0,0 +1,248 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/annotation_metadata.py.""" + +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import annotation_metadata + + +# --------------------------------------------------------------------------- +# parse_filename_metadata +# --------------------------------------------------------------------------- +def test_parse_filename_all_fields_with_path(): + path = "/data/traces/mi355_tp2_isl1024_osl8_conc4_opt.pt.trace.json.gz" + result = annotation_metadata.parse_filename_metadata(path) + assert result["isl"] == 1024 + assert result["osl"] == 8 + assert result["conc"] == 4 + assert result["tp"] == 2 + assert result["num_tokens_prefill"] == 1024 * 4 + assert result["num_tokens_decode"] == 4 + + +def test_parse_filename_all_fields_no_slash(): + result = annotation_metadata.parse_filename_metadata("tp1_isl16_osl2_conc3.json") + assert result["isl"] == 16 + assert result["osl"] == 2 + assert result["conc"] == 3 + assert result["tp"] == 1 + assert result["num_tokens_prefill"] == 48 + assert result["num_tokens_decode"] == 3 + + +def test_parse_filename_no_fields(): + assert annotation_metadata.parse_filename_metadata("randomfile.txt") == {} + + +def test_parse_filename_conc_only_no_isl(): + result = annotation_metadata.parse_filename_metadata("run_conc5.json") + assert result["conc"] == 5 + assert result["num_tokens_decode"] == 5 + assert "num_tokens_prefill" not in result + + +def test_parse_filename_isl_only_no_conc(): + result = annotation_metadata.parse_filename_metadata("run_isl99.json") + assert result["isl"] == 99 + assert "num_tokens_prefill" not in result + assert "num_tokens_decode" not in result + + +# --------------------------------------------------------------------------- +# parse_trace_input_dims +# --------------------------------------------------------------------------- +def test_parse_trace_input_dims_empty(): + assert annotation_metadata.parse_trace_input_dims([]) == {} + + +def test_parse_trace_input_dims_full(): + events = [ + # non-kernel event -> skipped + {"cat": "cpu_op", "name": "unified_attention"}, + # kernel but not attention -> skipped + {"cat": "kernel", "name": "gemm_kernel"}, + # attention kernel, q/k len == 3 + { + "cat": "kernel", + "name": "unified_attention", + "args": {"Input Dims": [[5, 1, 1], [7, 1, 1]]}, + }, + # attention kernel, q/k len > 3 (uses [-3]) + { + "cat": "kernel", + "name": "flash_attention_fwd", + "args": {"Input Dims": [[2, 3, 4, 5], [6, 7, 8, 9]]}, + }, + # attention kernel, dims too short -> skipped + { + "cat": "kernel", + "name": "unified_attention", + "args": {"Input Dims": [[1, 1]]}, + }, + # attention kernel, q/k not list-like -> skips nq/nkv appends + { + "cat": "kernel", + "name": "unified_attention", + "args": {"Input Dims": [5, 6]}, + }, + # attention kernel, q/k list but len < 3 -> skipped + { + "cat": "kernel", + "name": "unified_attention", + "args": {"Input Dims": [[1, 2], [3, 4]]}, + }, + # attention kernel, dims is a non-subscriptable set -> triggers except + { + "cat": "kernel", + "name": "unified_attention", + "args": {"Input Dims": {1, 2, 3}}, + }, + ] + result = annotation_metadata.parse_trace_input_dims(events) + # nq_values = [5, 3] -> avg = 4, count = 2 + assert result["trace_avg_nq"] == 4 + assert result["trace_attention_kernel_count"] == 2 + # nkv_values = [7, 7] -> avg = 7 + assert result["trace_avg_nkv"] == 7 + + +def test_parse_trace_input_dims_no_dims_key(): + events = [{"cat": "kernel", "name": "unified_attention", "args": {}}] + assert annotation_metadata.parse_trace_input_dims(events) == {} + + +# --------------------------------------------------------------------------- +# run_sanity_checks +# --------------------------------------------------------------------------- +def test_run_sanity_checks_all_warnings(): + warnings = annotation_metadata.run_sanity_checks( + annotation_meta={"batch_size": 100, "context_sum": 100}, + filename_meta={"isl": 10, "conc": 100}, + trace_meta={"trace_avg_nq": 200}, + user_meta={"num_tokens": 500}, + ) + assert len(warnings) == 3 + assert any("Batch size mismatch" in w for w in warnings) + assert any("num_tokens mismatch" in w for w in warnings) + assert any("Context sum mismatch" in w for w in warnings) + + +def test_run_sanity_checks_no_warnings_matching(): + warnings = annotation_metadata.run_sanity_checks( + annotation_meta={"batch_size": 1000, "context_sum": 100}, + filename_meta={"isl": 10, "conc": 100}, + trace_meta={"trace_avg_nq": 100}, + user_meta={"num_tokens": 1000}, + ) + assert warnings == [] + + +def test_run_sanity_checks_empty_all(): + assert annotation_metadata.run_sanity_checks({}, {}, {}, {}) == [] + + +def test_run_sanity_checks_filename_missing_conc(): + # isl present but conc missing -> batch_file stays None, no batch warning + warnings = annotation_metadata.run_sanity_checks( + annotation_meta={"batch_size": 100}, + filename_meta={"isl": 10}, + trace_meta={}, + user_meta={}, + ) + assert warnings == [] + + +# --------------------------------------------------------------------------- +# merge_metadata +# --------------------------------------------------------------------------- +def test_merge_metadata_all_none(): + merged = annotation_metadata.merge_metadata() + assert merged["annotation"] == {} + assert merged["filename"] == {} + assert merged["trace"] == {} + assert merged["user"] == {} + assert merged["num_tokens"] is None + assert merged["context_length"] is None + assert merged["batch_size"] is None + assert merged["context_sum"] is None + assert merged["generation_sum"] is None + assert merged["_warnings"] == [] + + +def test_merge_metadata_user_priority(): + merged = annotation_metadata.merge_metadata( + user_meta={"num_tokens": 42, "context_length": 77}, + ) + assert merged["num_tokens"] == 42 + assert merged["context_length"] == 77 + + +def test_merge_metadata_annotation_and_filename_fallbacks(): + merged = annotation_metadata.merge_metadata( + annotation_meta={"context_sum": 256, "generation_sum": 8}, + filename_meta={"num_tokens_prefill": 512, "isl": 128}, + ) + # num_tokens: no user/annotation batch_size -> prefill + assert merged["num_tokens"] == 512 + # context_length: annotation context_sum wins + assert merged["context_length"] == 256 + assert merged["batch_size"] == 512 + assert merged["context_sum"] == 256 + assert merged["generation_sum"] == 8 + + +def test_merge_metadata_context_from_num_tokens_fallback(): + # no context sources except num_tokens (from decode) + merged = annotation_metadata.merge_metadata( + filename_meta={"num_tokens_decode": 4}, + ) + assert merged["num_tokens"] == 4 + assert merged["context_length"] == 4 + + +def test_merge_metadata_emits_warning(): + merged = annotation_metadata.merge_metadata( + annotation_meta={"batch_size": 100}, + user_meta={"num_tokens": 500}, + ) + assert merged["_warnings"] + assert any("num_tokens mismatch" in w for w in merged["_warnings"]) + + +# --------------------------------------------------------------------------- +# gather_metadata +# --------------------------------------------------------------------------- +def test_gather_metadata_with_user_values(): + merged = annotation_metadata.gather_metadata( + trace_path="/x/tp1_isl1024_conc4.json", + events=None, + annotation_meta={"batch_size": 4096}, + num_tokens=4096, + context_length=1024, + ) + assert merged["num_tokens"] == 4096 + assert merged["context_length"] == 1024 + assert merged["filename"]["isl"] == 1024 + assert merged["trace"] == {} + + +def test_gather_metadata_no_user_values(): + merged = annotation_metadata.gather_metadata( + trace_path="plain.json", + events=None, + annotation_meta=None, + num_tokens=None, + context_length=None, + ) + assert merged["user"] == {} + assert merged["num_tokens"] is None diff --git a/tests/test_semantic_build_labels.py b/tests/test_semantic_build_labels.py new file mode 100644 index 000000000..0ef492a73 --- /dev/null +++ b/tests/test_semantic_build_labels.py @@ -0,0 +1,189 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/build_semantic_labels.py. + +Uses in-memory dict fixtures (no file I/O) to exercise the deterministic +labeling logic: positional block numbering, the layer-cycle detection and +the region (pre / body / post / secondary) helpers. +""" + +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import build_semantic_labels + + +# --------------------------------------------------------------------------- # +# _build_cycle_names +# --------------------------------------------------------------------------- # +def test_build_cycle_names_nonpositive_period(): + rle = [("GEMM", 1, [0], ["T"])] + assert build_semantic_labels._build_cycle_names(rle, 0, {}) == [] + + +def test_build_cycle_names_uses_global_counter(): + rle = [ + ("GEMM", 1, [0], ["T"]), + ("Normalization", 1, [1], ["T"]), + ("SDPA", 1, [2], ["T"]), + ] + counter = {"GEMM": 1} # pretend one GEMM was already numbered + names = build_semantic_labels._build_cycle_names(rle, 3, counter) + assert names == ["GEMM_1", "Normalization_0", "SDPA_0"] + assert counter == {"GEMM": 2, "Normalization": 1, "SDPA": 1} + + +# --------------------------------------------------------------------------- # +# _build_region_block_names +# --------------------------------------------------------------------------- # +def test_build_region_block_names_empty(): + assert build_semantic_labels._build_region_block_names(set(), {}, {}) == {} + + +def test_build_region_block_names_groups_consecutive(): + cls_by_idx = { + 0: {"perf_category": "GEMM"}, + 1: {"perf_category": "GEMM"}, + 2: {"perf_category": "Normalization"}, + # index 3 intentionally absent -> defaults to "Others" + } + counter = {} + result = build_semantic_labels._build_region_block_names( + {0, 1, 2, 3}, cls_by_idx, counter + ) + assert result == { + 0: "GEMM_0", + 1: "GEMM_0", + 2: "Normalization_0", + 3: "Others_0", + } + assert counter == {"GEMM": 1, "Normalization": 1, "Others": 1} + + +# --------------------------------------------------------------------------- # +# build_labels (integration over the pure labeling path) +# --------------------------------------------------------------------------- # +def _classified(cats): + return { + "classified_kernels": [ + {"index": i, "perf_category": cat, "kernel_type": "T"} + for i, cat in enumerate(cats) + ] + } + + +def _extracted(n, total_time=123.456, graph_mode=True): + return { + "source_file": "trace.json", + "metadata": { + "total_kernel_time_us": total_time, + "is_graph_mode": graph_mode, + }, + "kernels": [ + {"name": f"k{i}", "dur": float(i + 1), "gpu_op_uid": f"raw{i}"} + for i in range(n) + ], + } + + +def test_build_labels_positional_labels_and_layers(): + # 0,1 = preamble ; 2..13 = body (G,N,S repeated x4) ; 14,15 = epilogue ; 16 = secondary + cats = ( + ["GEMM", "GEMM"] + + ["GEMM", "Normalization", "SDPA"] * 4 + + ["Normalization", "Normalization"] + + ["GEMM"] + ) + extracted = _extracted(len(cats)) + classified = _classified(cats) + pattern = { + "preamble_indices": [0, 1], + "epilogue_indices": [14, 15], + "secondary_stream_indices": [16], + } + result = build_semantic_labels.build_labels(extracted, classified, pattern) + + info = result["model_info"] + assert info["period"] == 3 + assert info["num_layers"] == 4 + assert info["graph_mode"] is True + assert result["total_kernel_time_us"] == 123.46 + assert result["source_file"] == "trace.json" + + kernels = result["labeled_kernels"] + + # Preamble region: GEMM_0 (numbered before the body cycle). + assert kernels[0]["region"] == "pre" + assert kernels[0]["layer"] is None + assert kernels[0]["semantic_block"] == "GEMM_0" + + # Body cycle: GEMM in the body is GEMM_1 (global counter continues). + assert kernels[2]["region"] == "body" + assert kernels[2]["layer"] == 0 + assert kernels[2]["semantic_block"] == "GEMM_1" + assert kernels[3]["semantic_block"] == "Normalization_0" + assert kernels[4]["semantic_block"] == "SDPA_0" + + # Positional labels repeat across layers. + assert kernels[5]["semantic_block"] == "GEMM_1" + assert kernels[5]["layer"] == 1 + assert kernels[13]["semantic_block"] == "SDPA_0" + assert kernels[13]["layer"] == 3 + + # Epilogue + secondary reuse the global counter (unique numbering). + assert kernels[14]["region"] == "post" + assert kernels[14]["semantic_block"] == "Normalization_1" + assert kernels[16]["region"] == "secondary" + assert kernels[16]["semantic_block"] == "GEMM_2" + + # Enrichment fields are always empty (no trace-tree build); gpu_op_uid + # comes straight from the raw-index UID stamped by extract_trace_data. + assert kernels[2]["nn_module"] == "" + assert kernels[2]["cpu_op"] == "" + assert kernels[2]["input_dims"] == [] + assert kernels[2]["gpu_op_uid"] == "raw2" + assert kernels[4]["nn_module"] == "" + assert kernels[4]["gpu_op_uid"] == "raw4" + assert kernels[3]["nn_module"] == "" + assert kernels[3]["cpu_op"] == "" + assert kernels[3]["gpu_op_uid"] == "raw3" + + +def test_build_labels_no_body_period_zero(): + cats = ["GEMM", "GEMM", "Normalization"] + extracted = {"kernels": [{"name": f"k{i}", "dur": 1.0} for i in range(3)]} + classified = _classified(cats) + pattern = { + "preamble_indices": [0, 1, 2], + "epilogue_indices": [], + "secondary_stream_indices": [], + } + + result = build_semantic_labels.build_labels(extracted, classified, pattern) + + info = result["model_info"] + assert info["period"] == 0 + assert info["num_layers"] == 0 + assert info["graph_mode"] is False + assert result["total_kernel_time_us"] == 0 + assert result["source_file"] == "" + + kernels = result["labeled_kernels"] + assert [k["semantic_block"] for k in kernels] == [ + "GEMM_0", + "GEMM_0", + "Normalization_0", + ] + assert all(k["region"] == "pre" for k in kernels) + assert all(k["layer"] is None for k in kernels) + # No tree context and no raw uid -> gpu_op_uid is None. + assert kernels[0]["gpu_op_uid"] is None + assert kernels[0]["nn_module"] == "" diff --git a/tests/test_semantic_extract_trace_data.py b/tests/test_semantic_extract_trace_data.py new file mode 100644 index 000000000..5afc94269 --- /dev/null +++ b/tests/test_semantic_extract_trace_data.py @@ -0,0 +1,206 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/extract_trace_data.py extractors.""" + +import json +import os, sys +import tempfile + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) +import extract_trace_data + + +def _kernel(name, ts, dur, stream=None, cat="kernel", **extra): + ev = {"name": name, "ts": ts, "dur": dur, "cat": cat} + if stream is not None: + ev["args"] = {"stream": stream} + ev.update(extra) + return ev + + +def test_load_trace_from_dict_sorts_and_skips_nondict(): + data_in = { + "traceEvents": [ + {"name": "b", "cat": "kernel", "ts": 5}, + {"name": "a", "cat": "kernel", "ts": 1}, + "not-a-dict", + {"name": "c", "cat": "cpu_op", "ts": 3}, + ] + } + data, by_cat = extract_trace_data.load_trace(data_in) + assert data is data_in + # Non-dict element skipped; kernels sorted by ts. + assert [e["name"] for e in by_cat["kernel"]] == ["a", "b"] + assert [e["name"] for e in by_cat["cpu_op"]] == ["c"] + + +def test_load_trace_from_path(): + payload = {"traceEvents": [{"name": "k", "cat": "kernel", "ts": 1}]} + with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as tmp: + json.dump(payload, tmp) + path = tmp.name + try: + data, by_cat = extract_trace_data.load_trace(path) + assert data["traceEvents"][0]["name"] == "k" + assert [e["name"] for e in by_cat["kernel"]] == ["k"] + finally: + os.unlink(path) + + +def test_stamp_raw_uid(): + data = {"traceEvents": [{"name": "a"}, "skip", {"name": "b"}]} + extract_trace_data._stamp_raw_uid(data) + assert data["traceEvents"][0]["_gpu_op_uid"] == 0 + assert data["traceEvents"][2]["_gpu_op_uid"] == 2 + + +def test_get_stream_id_variants(): + # Valid string stream -> int. + assert extract_trace_data.get_stream_id({"args": {"stream": "3"}}) == 3 + # Unhashable stream value raises TypeError -> falls back to tid. + assert extract_trace_data.get_stream_id({"args": {"stream": [1]}, "tid": 7}) == 7 + # Non-numeric stream string raises ValueError -> falls back to tid. + assert extract_trace_data.get_stream_id({"args": {"stream": "x"}, "tid": 4}) == 4 + # No stream, tid present. + assert extract_trace_data.get_stream_id({"tid": 5}) == 5 + # Invalid tid -> None. + assert extract_trace_data.get_stream_id({"tid": "nope"}) is None + # Nothing usable -> None. + assert extract_trace_data.get_stream_id({}) is None + + +def test_filter_to_primary_stream_empty(): + by_cat = {} + extract_trace_data.filter_to_primary_stream(by_cat) + assert by_cat == {} + + +def test_filter_to_primary_stream_single_stream_noop(): + kernels = [_kernel("k", 1, 10, stream=0), _kernel("k", 2, 10, stream=0)] + by_cat = {"kernel": kernels} + extract_trace_data.filter_to_primary_stream(by_cat) + assert len(by_cat["kernel"]) == 2 + + +def test_filter_to_primary_stream_filters_minor_secondary(): + kernels = [_kernel("k", i, 10, stream=0) for i in range(9)] + kernels.append(_kernel("k", 100, 1, stream=1)) + by_cat = {"kernel": kernels} + extract_trace_data.filter_to_primary_stream(by_cat) + # Secondary stream is <5% of time -> dropped, keeping primary only. + assert all(k["args"]["stream"] == 0 for k in by_cat["kernel"]) + assert len(by_cat["kernel"]) == 9 + + +def test_filter_to_primary_stream_keeps_significant_secondary(): + kernels = [_kernel("k", i, 10, stream=0) for i in range(5)] + kernels.append(_kernel("k", 100, 10, stream=1)) + by_cat = {"kernel": kernels} + extract_trace_data.filter_to_primary_stream(by_cat) + # Secondary stream >5% of time -> keep all streams. + assert len(by_cat["kernel"]) == 6 + + +def test_extract_kernel_sequence(): + by_cat = { + "kernel": [ + _kernel("gemm", 5, 3, stream=0, _gpu_op_uid=11), + _kernel("relu", 1, 2, stream=0, gpu_op_uid=22), + ], + "gpu_memcpy": [_kernel("memcpy", 3, 1, cat="gpu_memcpy")], + } + seq = extract_trace_data.extract_kernel_sequence(by_cat) + # Sorted by ts across kernels + memcpy. + assert [k["name"] for k in seq] == ["relu", "memcpy", "gemm"] + # _gpu_op_uid preferred, else gpu_op_uid. + assert seq[2]["gpu_op_uid"] == 11 + assert seq[0]["gpu_op_uid"] == 22 + # memcpy has neither -> None. + assert seq[1]["gpu_op_uid"] is None + assert seq[1]["cat"] == "gpu_memcpy" + + +def test_detect_graph_mode(): + by_cat = {"cuda_runtime": [{"name": "hipGraphLaunch"}, {"name": "other"}]} + is_graph, launches = extract_trace_data.detect_graph_mode(by_cat) + assert is_graph is True + assert len(launches) == 1 + + is_graph2, launches2 = extract_trace_data.detect_graph_mode({}) + assert is_graph2 is False + assert launches2 == [] + + +def test_run_assertions_happy_path(): + data = {"traceEvents": [1]} + by_cat = {"kernel": [], "cpu_op": []} + kernels = [ + {"name": "a", "ts": 1, "dur": 2.0}, + {"name": "b", "ts": 3, "dur": 4.0}, + ] + errors = extract_trace_data.run_assertions(data, by_cat, kernels, False) + assert errors == [] + + +def test_run_assertions_missing_trace_events_and_categories(): + data = {} + by_cat = {"kernel": []} # cpu_op missing under strict. + kernels = [{"name": "a", "ts": 1, "dur": 2.0}] + errors = extract_trace_data.run_assertions(data, by_cat, kernels, False) + joined = " ".join(errors) + assert "A1.1 FAIL" in joined + assert "A1.2 FAIL" in joined + + +def test_run_assertions_no_kernels_and_zero_time(): + data = {"traceEvents": []} + by_cat = {"kernel": [], "cpu_op": []} + errors = extract_trace_data.run_assertions(data, by_cat, [], False) + joined = " ".join(errors) + assert "A1.3 FAIL" in joined + assert "A1.5 FAIL" in joined + + +def test_run_assertions_nonpositive_duration(): + data = {"traceEvents": [1]} + by_cat = {"kernel": [], "cpu_op": []} + kernels = [{"name": "bad", "ts": 1, "dur": 0}] + errors = extract_trace_data.run_assertions(data, by_cat, kernels, False) + assert any("A3.2 FAIL" in e for e in errors) + + +def test_run_assertions_nonmonotonic_timestamps(): + data = {"traceEvents": [1]} + by_cat = {"kernel": [], "cpu_op": []} + kernels = [ + {"name": "a", "ts": 5, "dur": 1.0}, + {"name": "b", "ts": 1, "dur": 1.0}, + ] + errors = extract_trace_data.run_assertions(data, by_cat, kernels, False) + assert any("A3.1 FAIL" in e for e in errors) + + +def test_extract_and_build_result_with_and_without_region_meta(): + by_cat = { + "kernel": [_kernel("gemm", 1, 3, stream=0, _gpu_op_uid=0)], + "cpu_op": [], + } + result, kernels = extract_trace_data.extract_and_build_result( + {"traceEvents": []}, by_cat, "trace.json" + ) + assert result["source_file"] == "trace.json" + assert result["metadata"]["total_kernels"] == 1 + assert result["metadata"]["total_kernel_time_us"] == 3.0 + assert "region_metadata" not in result + assert len(kernels) == 1 + + result2, _ = extract_trace_data.extract_and_build_result( + {"traceEvents": []}, by_cat, "trace.json", region_metadata={"region": "steady"} + ) + assert result2["region_metadata"] == {"region": "steady"} diff --git a/tests/test_semantic_generate_diff.py b/tests/test_semantic_generate_diff.py new file mode 100644 index 000000000..8b69fec95 --- /dev/null +++ b/tests/test_semantic_generate_diff.py @@ -0,0 +1,194 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/generate_semantic_diff.py builders.""" + +import os, sys + +import pandas as pd + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) +import generate_semantic_diff + + +def _labeled(): + labeled_a = [ + { + "name": "kA1", + "semantic_block": "QKV", + "nn_module": "attn", + "cpu_op": "aten::mm", + "dur": 10.0, + "input_dims": [[1, 2], (3, 4), 5], + "gpu_op_uid": 100, + }, + { + "name": "kA2", + "semantic_block": "MLP", + "perf_category": "Compute", + "dur": 5.0, + "input_dims": [], + "gpu_op_uid": 101, + }, + { + "name": "kA3", + "semantic_block": "OnlyA", + "dur": 2.0, + "gpu_op_uid": 102, + }, + ] + labeled_b = [ + { + "name": "kB1", + "semantic_block": "QKV", + "nn_module": "attn", + "cpu_op": "aten::mm", + "dur": 12.0, + "input_dims": [[1, 2]], + "gpu_op_uid": 200, + }, + { + "name": "kB2", + "semantic_block": "MLP", + "perf_category": "Compute", + "dur": 6.0, + "gpu_op_uid": 201, + }, + { + "name": "kB3", + "semantic_block": "OnlyB", + "dur": 3.0, + "gpu_op_uid": 202, + }, + ] + return labeled_a, labeled_b + + +def _diff_df(): + labeled_a, labeled_b = _labeled() + rows, _ = generate_semantic_diff.build_diff_stats(labeled_a, labeled_b) + df = pd.DataFrame(rows) + df["busy_time"] = ( + df.groupby(["source", "lowest_common_ancestor_id"])["kernel_time"] + .transform("sum") + .round(3) + ) + return df + + +def test_build_diff_stats(): + labeled_a, labeled_b = _labeled() + rows, block_id_map = generate_semantic_diff.build_diff_stats(labeled_a, labeled_b) + + assert len(rows) == 6 + assert block_id_map == {"QKV": 0, "MLP": 1, "OnlyA": 2, "OnlyB": 3} + + by_name = {r["name"]: r for r in rows} + # nn_module present. + assert by_name["kA1"]["nn_module_stack"] == "attn" + # falls back to perf_category. + assert by_name["kA2"]["nn_module_stack"] == "Compute" + # falls back to "Others". + assert by_name["kA3"]["nn_module_stack"] == "Others" + + # cpu_op present vs. falling back to the block name. + assert by_name["kA1"]["cpu_op_name"] == "aten::mm" + assert by_name["kA3"]["cpu_op_name"] == "OnlyA" + + # _format_dims: list, tuple, and scalar entries. + assert by_name["kA1"]["Input Dims"] == "(1, 2), (3, 4), 5" + # empty dims -> empty string. + assert by_name["kA2"]["Input Dims"] == "" + + assert by_name["kA1"]["source"] == "trace1" + assert by_name["kB1"]["source"] == "trace2" + assert by_name["kA1"]["kernel_time"] == 10.0 + + +def test_build_unique_args_summary(): + df = _diff_df() + summary = generate_semantic_diff.build_unique_args_summary(df) + + assert "kernel_time_sum" in summary.columns + assert "operation_count" in summary.columns + # 6 unique kernels -> 6 aggregated rows. + assert len(summary) == 6 + # sorted descending by kernel_time_sum. + vals = list(summary["kernel_time_sum"]) + assert vals == sorted(vals, reverse=True) + + +def test_build_unique_args_summary_unhashable_fallback(): + # A list-valued grouping column forces the TypeError -> str-repr branch. + df = pd.DataFrame( + [ + { + "name": "x", + "cpu_op_name": [1, 2], + "source": "trace1", + "kernel_time": 1.0, + }, + { + "name": "y", + "cpu_op_name": [3, 4], + "source": "trace1", + "kernel_time": 2.0, + }, + { + "name": "x", + "cpu_op_name": [1, 2], + "source": "trace1", + "kernel_time": 4.0, + }, + ] + ) + summary = generate_semantic_diff.build_unique_args_summary(df) + assert isinstance(summary, pd.DataFrame) + assert "kernel_time_sum" in summary.columns + assert "operation_count" in summary.columns + + +def test_build_cpu_op_maps(): + df = _diff_df() + cpu_op_map, t1, t2 = generate_semantic_diff.build_cpu_op_maps(df) + + assert "aten::mm" in cpu_op_map + assert set(cpu_op_map["aten::mm"].keys()) == {"trace1", "trace2"} + assert cpu_op_map["aten::mm"]["trace1"]["kernels"] == ["kA1"] + assert cpu_op_map["aten::mm"]["trace2"]["kernels"] == ["kB1"] + + # Per-source grouped frames. + t1_map = t1.to_dict()["name"] + t2_map = t2.to_dict()["name"] + assert t1_map["aten::mm"] == ["kA1"] + assert t2_map["OnlyB"] == ["kB3"] + assert "OnlyA" not in t2_map + + +def test_build_merged_tree_text(): + labeled_a, labeled_b = _labeled() + _, block_id_map = generate_semantic_diff.build_diff_stats(labeled_a, labeled_b) + text = generate_semantic_diff.build_merged_tree_text( + block_id_map, labeled_a, labeled_b, "MI355", "B200" + ) + + assert text.startswith("└── Root (MI355 vs B200)") + # Combined block (present in both) shows the plain label. + assert "QKV" in text + # trace1-only and trace2-only markers. + assert ">> trace1: OnlyA" in text + assert "<< trace2: OnlyB" in text + # Kernel of a combined block rendered plainly. + assert "kA1" in text + # Kernel of a trace1-only block carries the >> marker. + assert ">> trace1: kA3" in text + assert "<< trace2: kB3" in text + # Grouping by nn_module / perf_category. + assert "attn" in text + assert "Compute" in text + assert "Others" in text diff --git a/tests/test_semantic_helpers.py b/tests/test_semantic_helpers.py new file mode 100644 index 000000000..020412845 --- /dev/null +++ b/tests/test_semantic_helpers.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/_helpers.py.""" + +import gzip +import json +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +from _helpers import build_rle, detect_period, load_json, load_labels + + +# --------------------------------------------------------------------------- +# build_rle +# --------------------------------------------------------------------------- +def test_build_rle_empty(): + assert build_rle([], {}) == [] + + +def test_build_rle_single_group(): + cls = { + 0: {"perf_category": "Compute", "kernel_type": "gemm"}, + 1: {"perf_category": "Compute", "kernel_type": "gemm"}, + } + groups = build_rle([0, 1], cls) + assert groups == [("Compute", 2, [0, 1], ["gemm", "gemm"])] + + +def test_build_rle_multiple_groups_with_defaults(): + cls = { + 0: {"perf_category": "Compute", "kernel_type": "gemm"}, + 1: {"perf_category": "Compute", "kernel_type": "gemm"}, + 2: {"perf_category": "Memory", "kernel_type": "copy"}, + # idx 3 missing -> defaults ("Others", "Unknown") + } + groups = build_rle([0, 1, 2, 3], cls) + assert groups == [ + ("Compute", 2, [0, 1], ["gemm", "gemm"]), + ("Memory", 1, [2], ["copy"]), + ("Others", 1, [3], ["Unknown"]), + ] + + +# --------------------------------------------------------------------------- +# detect_period +# --------------------------------------------------------------------------- +def test_detect_period_short_returns_length(): + groups = [("A",), ("B",), ("C",), ("D",), ("E",)] # n = 5 < 6 + assert detect_period(groups) == 5 + + +def test_detect_period_repeating(): + # ABC repeated 4 times -> n = 12, period = 3 + cats = ["A", "B", "C"] * 4 + groups = [(c,) for c in cats] + assert detect_period(groups) == 3 + + +def test_detect_period_no_period(): + cats = ["A", "B", "C", "D", "E", "F"] # n = 6, no repetition + groups = [(c,) for c in cats] + assert detect_period(groups) == 6 + + +def test_detect_period_empty(): + assert detect_period([]) == 0 + + +# --------------------------------------------------------------------------- +# load_json / load_labels +# --------------------------------------------------------------------------- +def test_load_json_plain(tmp_path): + p = tmp_path / "data.json" + payload = {"a": 1, "b": [2, 3]} + p.write_text(json.dumps(payload)) + assert load_json(str(p)) == payload + + +def test_load_json_gzip(tmp_path): + p = tmp_path / "data.json.gz" + payload = {"x": "y"} + with gzip.open(str(p), "wt") as f: + json.dump(payload, f) + assert load_json(str(p)) == payload + + +def test_load_labels(tmp_path): + p = tmp_path / "semantic_labels.json" + payload = {"labeled_kernels": [{"semantic_block": "attn"}]} + p.write_text(json.dumps(payload)) + assert load_labels(str(p)) == payload + + +def test_detect_period_fuzzy_one_defect(): + # ABC repeated 4x with a single defect in the last group -> matches 8/9 > 0.85, + # exercising the non-exact fuzzy-match branch (still returns period 3). + cats = ["A", "B", "C"] * 4 + cats[-1] = "X" + groups = [(c,) for c in cats] + assert detect_period(groups) == 3 diff --git a/tests/test_semantic_kernel_coherence.py b/tests/test_semantic_kernel_coherence.py new file mode 100644 index 000000000..5c6c33060 --- /dev/null +++ b/tests/test_semantic_kernel_coherence.py @@ -0,0 +1,219 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/kernel_coherence.py deterministic core.""" + +import json +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import kernel_coherence # noqa: E402 +import kernel_runlength # noqa: E402 + +# --------------------------------------------------------------------------- +# _load / _dims_repr +# --------------------------------------------------------------------------- + + +def test_load_reads_json(tmp_path): + p = tmp_path / "doc.json" + p.write_text(json.dumps({"labeled_kernels": []})) + assert kernel_coherence._load(str(p)) == {"labeled_kernels": []} + + +def test_dims_repr_empty_and_compact_and_truncate(): + assert kernel_coherence._dims_repr(None) == "" + assert kernel_coherence._dims_repr([[1, 2]]) == "[[1,2]]" + long = kernel_coherence._dims_repr([[i] for i in range(200)], limit=10) + assert long.endswith("...") + assert len(long) == 13 + + +# --------------------------------------------------------------------------- +# _run_to_kernel_indices +# --------------------------------------------------------------------------- + + +def test_run_to_kernel_indices(): + mapping = kernel_coherence._run_to_kernel_indices([0, 0, 1, 2, 2]) + assert dict(mapping) == {0: [0, 1], 1: [2], 2: [3, 4]} + + +# --------------------------------------------------------------------------- +# _symbol_evidence +# --------------------------------------------------------------------------- + + +def _kernels_for_evidence(): + return [ + {"name": "gemm", "dur": 10.0, "perf_category": "GEMM", "input_dims": [[1, 2]]}, + {"name": "gemm", "dur": 5.0, "perf_category": "GEMM"}, + { + "name": "elem", + "dur": 3.0, + "perf_category": "Elementwise", + "input_dims": [[9]], + }, + ] + + +def test_symbol_evidence_ranks_and_aggregates(): + kernels = _kernels_for_evidence() + cats, dims, top = kernel_coherence._symbol_evidence(kernels, [0, 1, 2], 5) + assert cats == ["Elementwise", "GEMM"] + assert dims == "[[1,2]]" # first non-empty input_dims + assert top[0] == {"kernel_name": "gemm", "total_us": 15.0, "kernel_count": 2} + assert top[1] == {"kernel_name": "elem", "total_us": 3.0, "kernel_count": 1} + + +def test_symbol_evidence_top_kernels_limit(): + kernels = _kernels_for_evidence() + _, _, top = kernel_coherence._symbol_evidence(kernels, [0, 1, 2], 1) + assert [t["kernel_name"] for t in top] == ["gemm"] + + +# --------------------------------------------------------------------------- +# _collect_contexts +# --------------------------------------------------------------------------- + + +def test_collect_contexts_unique_neighbor_contexts(): + seq = ["A", "X", "B", "A", "X", "B", "A", "X", "C"] + kernels = [ + {"name": f"k{i}", "semantic_block": s, "dur": float(i + 1)} + for i, s in enumerate(seq) + ] + condensed = kernel_runlength.collapse_consecutive(seq) + run_per_kernel = kernel_runlength.run_index_per_kernel(seq) + shared = {"A", "B", "C"} + problematic = {"X"} + + detail = kernel_coherence._collect_contexts( + "A", kernels, condensed, run_per_kernel, shared, problematic, 1, 5 + ) + assert list(detail) == ["X"] + contexts = detail["X"]["contexts"] + # (A,B) appears twice but is deduped; (A,C) is distinct -> 2 contexts + assert detail["X"]["context_count"] == 2 + assert contexts[0]["id"] == "A:0" + assert contexts[0]["left_window"] == ["A"] + assert contexts[0]["right_window"] == ["B"] + assert contexts[1]["id"] == "A:1" + assert contexts[1]["right_window"] == ["C"] + assert contexts[0]["first_pass_block"] == "X" + + +# --------------------------------------------------------------------------- +# _context_lookup +# --------------------------------------------------------------------------- + + +def test_context_lookup_builds_key_map(): + catalog = [ + { + "id": "A:0", + "workload": "A", + "first_pass_block": "X", + "left_window": ["A"], + "right_window": ["B"], + }, + { + "id": "A:1", + "workload": "A", + "first_pass_block": "X", + "left_window": None, + "right_window": ["C"], + }, + ] + lookup = kernel_coherence._context_lookup(catalog) + assert lookup[("A", "X", ("A",), ("B",))] == "A:0" + # left_window None coerced to empty tuple + assert lookup[("A", "X", (), ("C",))] == "A:1" + + +# --------------------------------------------------------------------------- +# _final_blocks +# --------------------------------------------------------------------------- + + +def _final_kernels(): + return [ + {"name": "a", "semantic_block": "S1", "dur": 1.0, "index": 100}, + {"name": "b", "semantic_block": "P", "dur": 5.0, "index": 101}, + {"name": "c", "semantic_block": "S2", "dur": 2.0}, # no index -> uses i + ] + + +def test_final_blocks_context_rename(): + kernels = _final_kernels() + lookup = {("A", "P", ("S1",), ("S2",)): "A:0"} + finals, audit = kernel_coherence._final_blocks( + "A", kernels, {"S1", "S2"}, {"P"}, 1, lookup, {"A:0": "QKV"}, {} + ) + assert finals == ["S1", "QKV", "S2"] + assert audit[1]["context_id"] == "A:0" + assert audit[1]["first_pass_block"] == "P" + assert audit[1]["final_block"] == "QKV" + assert audit[1]["kernel_index"] == 101 + # kernel without "index" falls back to positional index + assert audit[2]["kernel_index"] == 2 + + +def test_final_blocks_fallback_remap(): + kernels = _final_kernels() + lookup = {("A", "P", ("S1",), ("S2",)): "A:0"} + finals, _ = kernel_coherence._final_blocks( + "A", kernels, {"S1", "S2"}, {"P"}, 1, lookup, {}, {"P": "FB"} + ) + # cid found but not in renames -> fallback used + assert finals == ["S1", "FB", "S2"] + + +def test_final_blocks_no_decision_keeps_first_pass(): + kernels = _final_kernels() + finals, audit = kernel_coherence._final_blocks( + "A", kernels, {"S1", "S2"}, {"P"}, 1, {}, {}, {} + ) + assert finals == ["S1", "P", "S2"] + assert audit[1]["context_id"] == "" + + +def test_final_blocks_missing_context_uses_fallback(): + kernels = _final_kernels() + # empty lookup -> cid "" -> falls through to fallback + finals, audit = kernel_coherence._final_blocks( + "A", kernels, {"S1", "S2"}, {"P"}, 1, {}, {}, {"P": "FB"} + ) + assert finals == ["S1", "FB", "S2"] + assert audit[1]["context_id"] == "" + + +# --------------------------------------------------------------------------- +# _residual_one_sided +# --------------------------------------------------------------------------- + + +def test_residual_one_sided(): + labels_a = { + "labeled_kernels": [ + {"semantic_block": "A"}, + {"semantic_block": "A"}, + {"semantic_block": "X"}, + ] + } + labels_b = { + "labeled_kernels": [ + {"semantic_block": "A"}, + {"semantic_block": "Y"}, + ] + } + res_a, res_b = kernel_coherence._residual_one_sided(labels_a, labels_b) + assert res_a == ["X"] + assert res_b == ["Y"] diff --git a/tests/test_semantic_kernel_runlength.py b/tests/test_semantic_kernel_runlength.py new file mode 100644 index 000000000..ee5629a21 --- /dev/null +++ b/tests/test_semantic_kernel_runlength.py @@ -0,0 +1,112 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/kernel_runlength.py.""" + +import json +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import kernel_runlength + + +# --------------------------------------------------------------------------- +# load_sequence +# --------------------------------------------------------------------------- +def test_load_sequence(tmp_path): + p = tmp_path / "semantic_labels.json" + payload = { + "labeled_kernels": [ + {"semantic_block": "A"}, + {"semantic_block": "B"}, + {}, # missing key -> "" + ] + } + p.write_text(json.dumps(payload)) + assert kernel_runlength.load_sequence(str(p)) == ["A", "B", ""] + + +def test_load_sequence_no_kernels(tmp_path): + p = tmp_path / "empty_labels.json" + p.write_text(json.dumps({})) + assert kernel_runlength.load_sequence(str(p)) == [] + + +# --------------------------------------------------------------------------- +# collapse_consecutive +# --------------------------------------------------------------------------- +def test_collapse_consecutive_empty(): + assert kernel_runlength.collapse_consecutive([]) == [] + + +def test_collapse_consecutive_basic(): + assert kernel_runlength.collapse_consecutive(["A", "A", "B", "D", "D"]) == [ + "A", + "B", + "D", + ] + + +def test_collapse_consecutive_no_dups(): + assert kernel_runlength.collapse_consecutive(["A", "B", "C"]) == ["A", "B", "C"] + + +# --------------------------------------------------------------------------- +# run_index_per_kernel +# --------------------------------------------------------------------------- +def test_run_index_per_kernel_empty(): + assert kernel_runlength.run_index_per_kernel([]) == [] + + +def test_run_index_per_kernel_basic(): + # A A B D D -> runs: A(0) A(0) B(1) D(2) D(2) + assert kernel_runlength.run_index_per_kernel(["A", "A", "B", "D", "D"]) == [ + 0, + 0, + 1, + 2, + 2, + ] + + +# --------------------------------------------------------------------------- +# shared_neighbor_windows_skip_non_shared +# --------------------------------------------------------------------------- +def test_shared_neighbor_windows_basic(): + # 0 1 2 3 4 5 6 + condensed = ["S1", "x", "S2", "C", "S3", "y", "S4"] + shared = {"S1", "S2", "S3", "S4"} + left, right = kernel_runlength.shared_neighbor_windows_skip_non_shared( + condensed, center_j=3, shared=shared, radius=2 + ) + # left: walk from idx 2 -> S2 (shared), idx 1 x (skip), idx 0 S1 (shared) + assert left == ["S2", "S1"] + # right: idx 4 S3 (shared), idx 5 y (skip), idx 6 S4 (shared) + assert right == ["S3", "S4"] + + +def test_shared_neighbor_windows_radius_limit(): + condensed = ["S1", "S2", "C", "S3", "S4"] + shared = {"S1", "S2", "S3", "S4"} + left, right = kernel_runlength.shared_neighbor_windows_skip_non_shared( + condensed, center_j=2, shared=shared, radius=1 + ) + assert left == ["S2"] + assert right == ["S3"] + + +def test_shared_neighbor_windows_boundaries(): + condensed = ["C", "S1"] + shared = {"S1"} + left, right = kernel_runlength.shared_neighbor_windows_skip_non_shared( + condensed, center_j=0, shared=shared, radius=3 + ) + assert left == [] + assert right == ["S1"] diff --git a/tests/test_semantic_kernel_unification.py b/tests/test_semantic_kernel_unification.py new file mode 100644 index 000000000..83de7a2a2 --- /dev/null +++ b/tests/test_semantic_kernel_unification.py @@ -0,0 +1,330 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/kernel_unification.py pure transforms.""" + +import json +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import kernel_unification # noqa: E402 + +# --------------------------------------------------------------------------- +# _load / _dims_repr +# --------------------------------------------------------------------------- + + +def test_load_reads_json(tmp_path): + p = tmp_path / "doc.json" + p.write_text(json.dumps({"a": 1, "b": [2, 3]})) + assert kernel_unification._load(str(p)) == {"a": 1, "b": [2, 3]} + + +def test_dims_repr_empty(): + assert kernel_unification._dims_repr(None) == "" + assert kernel_unification._dims_repr([]) == "" + + +def test_dims_repr_compact(): + assert kernel_unification._dims_repr([[1, 2], [3, 4]]) == "[[1,2],[3,4]]" + + +def test_dims_repr_truncates(): + dims = [[i, i + 1] for i in range(200)] + out = kernel_unification._dims_repr(dims, limit=20) + assert out.endswith("...") + assert len(out) == 23 # 20 + "..." + + +# --------------------------------------------------------------------------- +# aggregate_names +# --------------------------------------------------------------------------- + + +def _labels(kernels): + return {"labeled_kernels": kernels} + + +def test_aggregate_names_duplicates_and_ordering(): + labels = _labels( + [ + {"name": "small", "dur": 1.0, "perf_category": "Elementwise"}, + { + "name": "big", + "dur": 10.0, + "perf_category": "GEMM", + "input_dims": [[1, 2]], + }, + {"name": "big", "dur": 5.0, "perf_category": "GEMM"}, + {"name": "big", "dur": 2.0, "perf_category": "Attention"}, + ] + ) + agg = kernel_unification.aggregate_names(labels) + # ordered by descending total duration -> big (17) before small (1) + assert list(agg) == ["big", "small"] + big = agg["big"] + assert big["kernel_count"] == 3 + assert big["total_dur_us"] == 17.0 + # perf_categories deduped + sorted + assert big["perf_categories"] == ["Attention", "GEMM"] + # first non-empty input_dims sampled + assert big["sample_input_dims"] == "[[1,2]]" + # without key_fn there are no raw-name samples + assert "sample_raw_names" not in big + + +def test_aggregate_names_missing_and_none_dur(): + labels = _labels( + [ + {"dur": None}, # missing name -> "", None dur -> 0 + {"name": "", "dur": 4.0}, + ] + ) + agg = kernel_unification.aggregate_names(labels) + assert list(agg) == [""] + assert agg[""]["kernel_count"] == 2 + assert agg[""]["total_dur_us"] == 4.0 + assert agg[""]["sample_input_dims"] == "" + + +def test_aggregate_names_key_fn_collects_raw_names(): + labels = _labels( + [ + {"name": "moe_attn_vllm", "dur": 3.0}, + {"name": "sglang_moe_attention", "dur": 2.0}, + {"name": "moe_attn", "dur": 1.0}, # equals key -> not sampled + ] + ) + agg = kernel_unification.aggregate_names(labels, key_fn=lambda n: "moe_attn") + assert list(agg) == ["moe_attn"] + entry = agg["moe_attn"] + assert entry["kernel_count"] == 3 + # raw names differing from the key are recorded and sorted + assert entry["sample_raw_names"] == ["moe_attn_vllm", "sglang_moe_attention"] + + +def test_aggregate_names_key_fn_none_drops_kernel(): + labels = _labels( + [ + {"name": "keep", "dur": 1.0}, + {"name": "drop", "dur": 9.0}, + ] + ) + agg = kernel_unification.aggregate_names( + labels, key_fn=lambda n: None if n == "drop" else n + ) + assert list(agg) == ["keep"] + + +# --------------------------------------------------------------------------- +# _entry_list / _build_context +# --------------------------------------------------------------------------- + + +def test_entry_list_preserves_agg_order(): + labels = _labels( + [ + {"name": "b", "dur": 5.0}, + {"name": "a", "dur": 9.0}, + ] + ) + agg = kernel_unification.aggregate_names(labels) # order: a, b + picked = kernel_unification._entry_list(agg, {"a", "b"}) + assert [e["name"] for e in picked] == ["a", "b"] + + +def test_build_context_shape_and_extra(): + agg_a = kernel_unification.aggregate_names( + _labels([{"name": "shared", "dur": 2.0}, {"name": "onlyA", "dur": 1.0}]) + ) + agg_b = kernel_unification.aggregate_names( + _labels([{"name": "shared", "dur": 2.0}, {"name": "onlyB", "dur": 1.0}]) + ) + ctx = kernel_unification._build_context( + agg_a, agg_b, "MI300", "B300", "raw_name", extra={"flag": True} + ) + assert ctx["name_a"] == "MI300" + assert ctx["key_level"] == "raw_name" + assert ctx["summary"]["combined_unique"] == 3 + assert ctx["summary"]["in_both"] == 1 + assert ctx["in_both"] == ["shared"] + assert [e["name"] for e in ctx["only_in_MI300"]] == ["onlyA"] + assert [e["name"] for e in ctx["only_in_B300"]] == ["onlyB"] + assert ctx["flag"] is True + + +def test_build_context_no_extra(): + agg = kernel_unification.aggregate_names(_labels([{"name": "x", "dur": 1.0}])) + ctx = kernel_unification._build_context(agg, agg, "a", "b", "stem") + assert ctx["summary"]["combined_unique"] == 1 + assert "flag" not in ctx + + +# --------------------------------------------------------------------------- +# _sample_names +# --------------------------------------------------------------------------- + + +def test_sample_names_small_returns_all_tagged(): + agg_a = kernel_unification.aggregate_names(_labels([{"name": "a", "dur": 1.0}])) + agg_b = kernel_unification.aggregate_names(_labels([{"name": "b", "dur": 1.0}])) + sample = kernel_unification._sample_names(agg_a, agg_b, "A", "B", 10) + traces = {row["trace"] for row in sample} + assert traces == {"A", "B"} + assert len(sample) == 2 + + +def test_sample_names_spaced_subset(): + agg_a = kernel_unification.aggregate_names( + _labels([{"name": n, "dur": float(i)} for i, n in enumerate("abcd")]) + ) + agg_b = kernel_unification.aggregate_names( + _labels([{"name": n, "dur": float(i)} for i, n in enumerate("wxyz")]) + ) + sample = kernel_unification._sample_names(agg_a, agg_b, "A", "B", 2) + # sample_size 2 -> half=1 from A, 1 from B + assert len(sample) == 2 + assert [r["trace"] for r in sample] == ["A", "B"] + + +# --------------------------------------------------------------------------- +# _compile_rules +# --------------------------------------------------------------------------- + + +def test_compile_rules_valid(): + compiled = kernel_unification._compile_rules( + [ + {"pattern": r"foo_\d+", "action": "collapse", "replacement": "foo"}, + {"pattern": "bar", "action": "preserve"}, + {"pattern": "baz"}, # default action collapse + ] + ) + assert [c["action"] for c in compiled] == ["collapse", "preserve", "collapse"] + assert compiled[0]["regex"].search("foo_123") + + +def test_compile_rules_bad_action(): + try: + kernel_unification._compile_rules([{"pattern": "x", "action": "nope"}]) + assert False, "expected SystemExit" + except SystemExit as e: + assert "invalid action" in str(e) + + +def test_compile_rules_bad_regex(): + try: + kernel_unification._compile_rules([{"pattern": "(", "action": "preserve"}]) + assert False, "expected SystemExit" + except SystemExit as e: + assert "bad regex" in str(e) + + +# --------------------------------------------------------------------------- +# stem_for +# --------------------------------------------------------------------------- + + +def test_stem_for_collapse(): + compiled = kernel_unification._compile_rules( + [{"pattern": r"gemm_[a-z0-9]+", "action": "collapse", "replacement": "gemm"}] + ) + assert kernel_unification.stem_for("gemm_fp16x8", compiled) == ("gemm", "collapse") + + +def test_stem_for_preserve_and_drop(): + compiled = kernel_unification._compile_rules( + [ + {"pattern": "keep", "action": "preserve"}, + {"pattern": "trash", "action": "drop"}, + ] + ) + assert kernel_unification.stem_for("keep_me", compiled) == ("keep_me", "preserve") + assert kernel_unification.stem_for("trash_me", compiled) == ("trash_me", "drop") + + +def test_stem_for_no_match_defaults_preserve(): + compiled = kernel_unification._compile_rules( + [{"pattern": "zzz", "action": "collapse", "replacement": "z"}] + ) + assert kernel_unification.stem_for("other", compiled) == ("other", "preserve") + + +def test_stem_for_bad_replacement_falls_back(capsys): + # replacement references a non-existent capture group -> re.error at sub time + compiled = kernel_unification._compile_rules( + [{"pattern": r"(a)b", "action": "collapse", "replacement": r"\2"}] + ) + # first call warns and preserves + assert kernel_unification.stem_for("ab", compiled) == ("ab", "preserve") + err = capsys.readouterr().err + assert "ignoring stem rule" in err + assert compiled[0]["_warned"] is True + # second call is already warned -> no new warning, still preserves + assert kernel_unification.stem_for("ab", compiled) == ("ab", "preserve") + assert capsys.readouterr().err == "" + + +# --------------------------------------------------------------------------- +# _load_map_side +# --------------------------------------------------------------------------- + + +def test_load_map_side_by_side_key(): + doc = {"map_a": {"foo": "bar"}} + assert kernel_unification._load_map_side(doc, "a", "MI300") == {"foo": "bar"} + + +def test_load_map_side_by_name_key(): + doc = {"map_MI300": {"foo": "bar"}} + assert kernel_unification._load_map_side(doc, "a", "MI300") == {"foo": "bar"} + + +def test_load_map_side_nested_map(): + doc = {"a": {"map": {"foo": "bar"}, "notes": "x"}} + assert kernel_unification._load_map_side(doc, "a", "MI300") == {"foo": "bar"} + + +def test_load_map_side_missing_returns_empty(): + assert kernel_unification._load_map_side({}, "a", "MI300") == {} + + +# --------------------------------------------------------------------------- +# _apply_side +# --------------------------------------------------------------------------- + + +def test_apply_side_no_stem_map(): + labels = _labels( + [ + {"name": "foo", "dur": 1.0}, + {"name": "baz", "dur": 1.0}, + ] + ) + stats = kernel_unification._apply_side(labels, {"foo": "FOO"}, None) + blocks = [k["semantic_block"] for k in labels["labeled_kernels"]] + assert blocks == ["FOO", "baz"] # unmapped falls back to raw + assert stats == {"kernels": 2, "mapped": 1, "stemmed": 0} + + +def test_apply_side_with_stem_map(): + labels = _labels( + [ + {"name": "gemm_v1", "dur": 1.0}, + {"name": "plain", "dur": 1.0}, + ] + ) + raw_to_stem = {"gemm_v1": "gemm", "plain": "plain"} + unified = {"gemm": "GEMM_UNIFIED"} + stats = kernel_unification._apply_side(labels, unified, raw_to_stem) + blocks = [k["semantic_block"] for k in labels["labeled_kernels"]] + assert blocks == ["GEMM_UNIFIED", "plain"] + # gemm_v1 was stemmed (base != raw); plain was not + assert stats == {"kernels": 2, "mapped": 1, "stemmed": 1} diff --git a/tests/test_semantic_match_and_compare.py b/tests/test_semantic_match_and_compare.py new file mode 100644 index 000000000..44696d586 --- /dev/null +++ b/tests/test_semantic_match_and_compare.py @@ -0,0 +1,225 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/match_and_compare.py. + +Uses in-memory dict fixtures (no file I/O) to exercise the pure +aggregation / comparison / assertion helpers. +""" + +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +from match_and_compare import aggregate, build_comparison, run_assertions + + +def _labeled(spec): + """Build a list of labeled-kernel dicts from compact tuples. + + Each tuple is (semantic_block, name, dur, perf_category, nn_module). + """ + return [ + { + "name": name, + "dur": dur, + "semantic_block": block, + "perf_category": pc, + "nn_module": nm, + } + for (block, name, dur, pc, nm) in spec + ] + + +# Standard, self-consistent A/B fixture used by several tests. +LABELED_A = _labeled( + [ + ("GEMM_0", "aa", 10.0, "GEMM", "Blk"), + ("GEMM_0", "bb", 30.0, "GEMM", "Blk"), + ("GEMM_0", "aa", 10.0, "GEMM", "Blk"), # duplicate name -> deduped in set + ("Norm_0", "cc", 20.0, "Normalization", "Blk"), + ("Extra_0", "dd", 10.0, "Others", ""), # block only present in A + ] +) +TOTAL_A = 80.0 + +LABELED_B = _labeled( + [ + ("GEMM_0", "dd", 20.0, "GEMM", "Blk"), + ("Norm_0", "ee", 10.0, "Normalization", "Blk"), + ("Norm_0", "ff", 10.0, "Normalization", "Blk"), + ("SDPA_0", "gg", 5.0, "SDPA", "Blk"), # block only present in B + ] +) +TOTAL_B = 45.0 + + +def _agg_pair(): + return aggregate(LABELED_A), aggregate(LABELED_B) + + +def _rows(**kwargs): + agg_a, agg_b = _agg_pair() + return build_comparison(agg_a, agg_b, TOTAL_A, TOTAL_B, "MI355", "B200", **kwargs) + + +# --------------------------------------------------------------------------- # +# aggregate +# --------------------------------------------------------------------------- # +def test_aggregate_groups_dedups_names_and_counts(): + agg = aggregate(LABELED_A) + assert list(agg.keys()) == ["GEMM_0", "Norm_0", "Extra_0"] + gemm = agg["GEMM_0"] + assert gemm["count"] == 3 + assert gemm["durs"] == [10.0, 30.0, 10.0] + assert gemm["names"] == {"aa", "bb"} # duplicate "aa" collapsed + + +def test_aggregate_carries_perf_category_and_nn_module(): + agg = aggregate(LABELED_A) + assert agg["GEMM_0"]["perf_category"] == "GEMM" + assert agg["GEMM_0"]["nn_module"] == "Blk" + + +def test_aggregate_empty_input(): + assert aggregate([]) == {} + + +# --------------------------------------------------------------------------- # +# build_comparison +# --------------------------------------------------------------------------- # +def test_build_comparison_rows_ratios_and_order(): + rows = _rows() + blocks = [r["semantic_block"] for r in rows] + assert blocks == ["GEMM_0", "Norm_0", "Extra_0", "SDPA_0"] + + by_block = {r["semantic_block"]: r for r in rows} + + gemm = by_block["GEMM_0"] + assert gemm["MI355_kernel_count"] == 3 + assert gemm["MI355_total_us"] == 50.0 + assert gemm["MI355_avg_us"] == round(50.0 / 3, 2) + assert gemm["B200_total_us"] == 20.0 + assert gemm["MI355_vs_B200_ratio"] == 2.5 + assert gemm["MI355_kernel_names"] == "aa | bb" + assert gemm["algorithm_order"] == 1 + + # Block only in A -> B side empty, ratio "inf", zero avg on B. + extra = by_block["Extra_0"] + assert extra["B200_kernel_count"] == 0 + assert extra["B200_avg_us"] == 0 + assert extra["MI355_vs_B200_ratio"] == "inf" + + # Block only in B -> A side empty. + sdpa = by_block["SDPA_0"] + assert sdpa["MI355_kernel_count"] == 0 + assert sdpa["MI355_avg_us"] == 0 + assert sdpa["MI355_vs_B200_ratio"] == 0.0 + + +def test_build_comparison_region_field_present(): + rows = _rows(region="prefill_only_3072") + assert rows[0]["region"] == "prefill_only_3072" + assert list(rows[0].keys())[0] == "region" + + +def test_build_comparison_gpu_timeline_fields(): + rows = _rows( + gpu_timeline_a={"busy_time_us": 2000.0, "idle_pct": 10.0}, + gpu_timeline_b={"busy_time_us": 4000.0, "idle_pct": 25.0}, + ) + r = rows[0] + assert r["MI355_busy_ms"] == 2.0 + assert r["MI355_idle_pct"] == 10.0 + assert r["B200_busy_ms"] == 4.0 + assert r["B200_idle_pct"] == 25.0 + + +def test_build_comparison_no_gpu_timeline_when_one_missing(): + rows = _rows(gpu_timeline_a={"busy_time_us": 2000.0, "idle_pct": 10.0}) + assert "MI355_busy_ms" not in rows[0] + + +def test_build_comparison_empty_inputs(): + assert build_comparison({}, {}, 0, 0, "MI355", "B200") == [] + + +def test_build_comparison_perf_category_and_module_fallback(): + labeled = _labeled([("X_0", "x", 5.0, None, None)]) + agg = aggregate(labeled) + rows = build_comparison(agg, agg, 5.0, 5.0, "A", "B") + assert rows[0]["perf_category"] == "Others" + assert rows[0]["nn_module"] == "" + + +def test_build_comparison_zero_totals_give_zero_pct(): + agg_a, agg_b = _agg_pair() + rows = build_comparison(agg_a, agg_b, 0, 0, "MI355", "B200") + assert all(r["MI355_pct"] == 0 for r in rows) + assert all(r["B200_pct"] == 0 for r in rows) + + +# --------------------------------------------------------------------------- # +# run_assertions +# --------------------------------------------------------------------------- # +def test_run_assertions_pass_on_consistent_data(): + rows = _rows() + errors = run_assertions( + rows, LABELED_A, LABELED_B, TOTAL_A, TOTAL_B, "MI355", "B200" + ) + assert errors == [] + + +def test_run_assertions_count_mismatch_a(): + rows = _rows() + padded_a = LABELED_A + [{"name": "z", "dur": 0.0, "semantic_block": "GEMM_0"}] + errors = run_assertions( + rows, padded_a, LABELED_B, TOTAL_A, TOTAL_B, "MI355", "B200" + ) + assert any("A6.1" in e and "MI355" in e for e in errors) + assert not any("A6.2" in e for e in errors) + + +def test_run_assertions_count_mismatch_b(): + rows = _rows() + padded_b = LABELED_B + [{"name": "z", "dur": 0.0, "semantic_block": "SDPA_0"}] + errors = run_assertions( + rows, LABELED_A, padded_b, TOTAL_A, TOTAL_B, "MI355", "B200" + ) + assert any("A6.2" in e and "B200" in e for e in errors) + + +def test_run_assertions_time_mismatch_both_sides(): + rows = _rows() + errors = run_assertions( + rows, LABELED_A, LABELED_B, TOTAL_A + 5.0, TOTAL_B + 5.0, "MI355", "B200" + ) + assert any("A6.3" in e and "MI355" in e for e in errors) + assert any("A6.3" in e and "B200" in e for e in errors) + + +def test_run_assertions_pct_mismatch_both_sides(): + # Feed doubled totals into build_comparison so percentages sum to ~50%. + agg_a, agg_b = _agg_pair() + rows = build_comparison(agg_a, agg_b, TOTAL_A * 2, TOTAL_B * 2, "MI355", "B200") + errors = run_assertions( + rows, LABELED_A, LABELED_B, TOTAL_A * 2, TOTAL_B * 2, "MI355", "B200" + ) + assert any("A7.2" in e and "MI355" in e for e in errors) + assert any("A7.2" in e and "B200" in e for e in errors) + + +def test_run_assertions_ratio_mismatch(): + rows = _rows() + # Corrupt the stored ratio of the GEMM_0 row (expected 2.5). + rows[0]["MI355_vs_B200_ratio"] = 9.999 + errors = run_assertions( + rows, LABELED_A, LABELED_B, TOTAL_A, TOTAL_B, "MI355", "B200" + ) + assert any("A7.5" in e for e in errors) diff --git a/tests/test_semantic_precomputed_diff_stats.py b/tests/test_semantic_precomputed_diff_stats.py new file mode 100644 index 000000000..bfaa959fe --- /dev/null +++ b/tests/test_semantic_precomputed_diff_stats.py @@ -0,0 +1,75 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Integration test for the semantic->report handoff via precomputed diff_stats. + +The semantic comparison path emits a TraceDiff-schema ``diff_stats.csv`` and +feeds it to the perf-report generators through ``precomputed_diff_stats_csv``. +When set (and no ``comparison_json_path``), the generator skips the internal +TraceDiff and loads the CSV as-is, then enriches the report and emits a +``diff_stats`` sheet. These tests exercise that branch for both the training +and inference report variants. +""" + +import json +import os + +import pytest + +from tests.fixtures.reporting import _build_synthetic_trace +from TraceLens.Reporting.generate_perf_report_pytorch import ( + generate_perf_report_pytorch, +) +from TraceLens.Reporting.generate_perf_report_pytorch_inference import ( + generate_perf_report_pytorch as generate_inference_report, +) + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +_DIFF_STATS_CSV = os.path.join( + REPO_ROOT, "tests", "traces", "tracediff_test", "diff_stats.csv" +) + +pytestmark = pytest.mark.filterwarnings( + "ignore:Source column .* not found.*:UserWarning", + "ignore:There are hipgraph launches.*:UserWarning", + "ignore:Found .* events with failed performance metric.*:UserWarning", + "ignore:Input list of events is empty.*:UserWarning", +) + + +def _write_trace(tmp_path, specs): + path = tmp_path / "trace.json" + path.write_text(json.dumps(_build_synthetic_trace(specs))) + return str(path) + + +def test_precomputed_diff_stats_csv_is_used_pytorch(tmp_path): + trace = _write_trace( + tmp_path, + [("aten::mm", "gemm_kernel", 100), ("aten::relu", "relu_kernel", 20)], + ) + result = generate_perf_report_pytorch( + profile_json_path=trace, + output_csvs_dir=str(tmp_path / "csvs"), + collective_analysis=False, + precomputed_diff_stats_csv=_DIFF_STATS_CSV, + ) + # The elif branch loaded the CSV and emitted the diff_stats sheet without + # running the internal TraceDiff (no comparison_json_path was supplied). + assert "diff_stats" in result + assert not result["diff_stats"].empty + + +def test_precomputed_diff_stats_csv_is_used_inference(tmp_path): + trace = _write_trace(tmp_path, [("aten::mm", "gemm_kernel", 100)]) + result = generate_inference_report( + profile_json_path=trace, + output_csvs_dir=str(tmp_path / "csvs"), + collective_analysis=False, + precomputed_diff_stats_csv=_DIFF_STATS_CSV, + ) + assert "diff_stats" in result + assert not result["diff_stats"].empty diff --git a/tests/test_semantic_trace_split_adapter.py b/tests/test_semantic_trace_split_adapter.py new file mode 100644 index 000000000..59313859e --- /dev/null +++ b/tests/test_semantic_trace_split_adapter.py @@ -0,0 +1,166 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for semantic_analyses/trace_split_adapter pure helpers. + +Covers the deterministic phase/metadata mapping used to convert TraceUtils +annotation-split output into the region_metadata shape expected by +extract_trace_data. The subprocess driver ``split_vllm_trace`` is excluded +from coverage (it shells out to TraceLens.TraceUtils.split_inference_trace_ +annotation and is exercised only in the end-to-end path). +""" + +import json +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +SEM = os.path.join(REPO_ROOT, "TraceLens", "Agent", "Analysis", "semantic_analyses") +sys.path.insert(0, SEM) + +import trace_split_adapter as tsa # noqa: E402 + + +# --------------------------------------------------------------------------- # +# get_steady_state_key +# --------------------------------------------------------------------------- # +def test_steady_state_key_prefill_only(): + meta = {"context_requests": 4, "generation_requests": 0, "context_sum": 128} + assert tsa.get_steady_state_key(meta) == "prefill_only_128" + + +def test_steady_state_key_decode_only(): + meta = {"context_requests": 0, "generation_requests": 8, "generation_sum": 64} + assert tsa.get_steady_state_key(meta) == "decode_only_64" + + +def test_steady_state_key_prefill_decode(): + meta = { + "context_requests": 2, + "generation_requests": 3, + "context_sum": 100, + "generation_sum": 50, + } + assert tsa.get_steady_state_key(meta) == "prefill_decode_100_50" + + +def test_steady_state_key_fallback_uses_batch(): + # ctx == 0 and gen == 0 -> batch/batch fallback branch + meta = {"context_requests": 0, "generation_requests": 0, "batch_size": 16} + assert tsa.get_steady_state_key(meta) == "prefill_decode_16_16" + + +def test_steady_state_key_defaults_when_empty(): + assert tsa.get_steady_state_key({}) == "prefill_decode_0_0" + + +# --------------------------------------------------------------------------- # +# _phase_to_region_meta +# --------------------------------------------------------------------------- # +def test_phase_to_region_meta_prefill_only(): + phase = { + "num_prefill": 5, + "num_prefilldecode": 0, + "num_decode": 0, + "avg_bs": 32, + "avg_conc": 5, + } + out = tsa._phase_to_region_meta(phase) + assert out == { + "context_requests": 5, + "generation_requests": 0, + "context_sum": 32, + "generation_sum": 0, + "batch_size": 32, + "num_requests": 5, + } + + +def test_phase_to_region_meta_decode_only(): + phase = { + "num_prefill": 0, + "num_prefilldecode": 0, + "num_decode": 7, + "avg_bs": 8, + "avg_conc": 7, + } + out = tsa._phase_to_region_meta(phase) + assert out == { + "context_requests": 0, + "generation_requests": 7, + "context_sum": 0, + "generation_sum": 8, + "batch_size": 8, + "num_requests": 7, + } + + +def test_phase_to_region_meta_combined(): + phase = { + "num_prefill": 2, + "num_prefilldecode": 3, + "num_decode": 4, + "avg_bs": 10, + "avg_conc": 9, + } + out = tsa._phase_to_region_meta(phase) + assert out["context_requests"] == 5 # num_prefill + num_prefilldecode + assert out["generation_requests"] == 7 # num_decode + num_prefilldecode + assert out["context_sum"] == 10 + assert out["generation_sum"] == 10 + assert out["batch_size"] == 10 + assert out["num_requests"] == 9 + + +def test_phase_to_region_meta_defaults_empty(): + # empty phase -> falls through to combined branch with all zeros + out = tsa._phase_to_region_meta({}) + assert out["context_requests"] == 0 + assert out["generation_requests"] == 0 + assert out["batch_size"] == 0 + + +# --------------------------------------------------------------------------- # +# _is_single_iteration +# --------------------------------------------------------------------------- # +def test_is_single_iteration_true(): + assert tsa._is_single_iteration({"num_prefill": 1}) is True + assert tsa._is_single_iteration({}) is True # total 0 + + +def test_is_single_iteration_false(): + assert tsa._is_single_iteration({"num_prefill": 1, "num_decode": 2}) is False + + +# --------------------------------------------------------------------------- # +# _iter_type_key +# --------------------------------------------------------------------------- # +def test_iter_type_key_prefill(): + assert tsa._iter_type_key({"num_prefill": 1, "avg_bs": 4}) == "prefill_4" + + +def test_iter_type_key_prefilldecode(): + assert ( + tsa._iter_type_key({"num_prefilldecode": 1, "avg_bs": 2}) == "prefilldecode_2" + ) + + +def test_iter_type_key_decode(): + assert tsa._iter_type_key({"num_decode": 1, "avg_bs": 8}) == "decode_8" + + +def test_iter_type_key_empty(): + assert tsa._iter_type_key({"avg_bs": 3}) == "empty_3" + + +# --------------------------------------------------------------------------- # +# _load_trace (thin wrapper over _helpers.load_json) +# --------------------------------------------------------------------------- # +def test_load_trace_reads_plain_json(tmp_path): + p = tmp_path / "trace.json" + p.write_text(json.dumps({"traceEvents": [{"name": "a"}]})) + out = tsa._load_trace(str(p)) + assert out == {"traceEvents": [{"name": "a"}]}