diff --git a/metainfer/tasks/evolve_kernel/form.yaml b/metainfer/tasks/evolve_kernel/form.yaml index f117584b..73daae93 100644 --- a/metainfer/tasks/evolve_kernel/form.yaml +++ b/metainfer/tasks/evolve_kernel/form.yaml @@ -1,10 +1,10 @@ # Question bank for the `evolve-kernel` task type . # LLM-guided iterative GPU kernel optimization. -# The user provides a Triton kernel file; the orchestrator generates -# test harnesses and iteratively optimizes the kernel. +# Supports Triton kernel optimization and full HIP C++ rewrites. +# The orchestrator generates test harnesses and iteratively optimizes the kernel. - key: kernel_file_path - question: "Path to the Triton kernel .py file (must contain a @triton.jit decorated function):" + question: "Path to the kernel .py file (Triton @triton.jit or plain Python wrapper for HIP C++):" header: "Kernel file" required: true @@ -28,8 +28,50 @@ - label: "50" description: "Deep optimization, long-running" +- key: optimizer_mode + question: "Optimization strategy?" + header: "Mode" + required: false + multi: false + default: "Triton (standard)" + options: + - label: "Triton (standard)" + description: "Pure Triton optimization — autotuning, tile size tuning, memory coalescing" + - label: "HIP C++ (from scratch)" + description: "Full HIP C++ rewrite w/ inline GCN assembly — max perf, replaces Triton entirely" + +- key: enable_profiling + question: "Enable hipprof GPU profiling (kernel duration + bandwidth estimation)?" + header: "Profiling" + required: false + multi: false + default: "No" + options: + - label: "No" + description: "Standard perf harness only (end-to-end exec time)" + - label: "Yes" + description: "Run hipprof for per-kernel timing + data-movement analysis" + +- key: multi_gpu + question: "Distribute optimization across multiple GPUs?" + header: "Multi-GPU" + required: false + multi: false + default: "no" + options: + - label: "Single GPU" + description: "Run on one GPU only (default)" + - label: "All GPUs (auto)" + description: "Split target shapes across all available GPUs" + +- key: gpu_count + question: "Number of GPUs to use (only if Multi-GPU is on, leave blank for auto-detect):" + header: "GPU Count" + required: false + default: "" + - key: extra_notes - question: "Extra constraints or notes for the optimizer?" + question: "Extra constraints, target shapes (for multi-GPU splitting), or notes for the optimizer?" header: "Notes" required: false diff --git a/metainfer/tasks/evolve_kernel/orchestrator/_parallel.py b/metainfer/tasks/evolve_kernel/orchestrator/_parallel.py new file mode 100644 index 00000000..2e402fcb --- /dev/null +++ b/metainfer/tasks/evolve_kernel/orchestrator/_parallel.py @@ -0,0 +1,645 @@ +"""Unified multi-GPU orchestrator — ONE task, N GPU workers. + +The MultiGpuOrchestrator spawns one standard evolve-kernel orchestrator +subprocess per GPU, each pinned to ``CUDA_VISIBLE_DEVICES=N`` and +assigned a subset of target shapes. All state lives under a single +``state_dir`` / ``workspace_dir`` pair so the WebUI sees ONE task with +a unified multi-GPU dashboard. + +Lifecycle: + 1. Parse shapes → split across GPUs + 2. Write per-GPU requirements under ``state_dir/gpu_N/`` + 3. Spawn subprocesses (``python -m ...cli run ... --gpu-device N``) + 4. Poll per-GPU ``run.json`` → update parent ``run.json`` + 5. When all workers finish: merge kernel_library.json + shape_bench.json +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from metainfer.tasks.evolve_kernel.server._multi_gpu import ( + detect_gpus, + split_shapes_for_gpus, + shapes_to_notes, +) + + +# =========================================================================== # +# Config +# =========================================================================== # + + +@dataclass +class GpuWorkerConfig: + gpu_idx: int + label: str # "GPU 0" + shapes: list # List[ShapeSpec] + state_dir: Path + workspace_dir: Path + requirements_path: Path + + +# =========================================================================== # +# MultiGpuOrchestrator +# =========================================================================== # + + +class MultiGpuOrchestrator: + """Spawns and monitors N parallel GPU kernel-optimization workers.""" + + def __init__( + self, + req: Dict[str, Any], + state_dir: Path, + workspace_dir: Path, + claude_bin: str = "ccb", + model: Optional[str] = None, + permission_mode: str = "bypassPermissions", + effort: str = "max", + ) -> None: + self.req = req + self.state_dir = state_dir + self.workspace_dir = workspace_dir + self.claude_bin = claude_bin + self.model = model + self.permission_mode = permission_mode + self.effort = effort + + num_gpus_str = req.get("gpu_count", "") + try: + self.num_gpus = int(num_gpus_str) if num_gpus_str else detect_gpus() + except (ValueError, TypeError): + self.num_gpus = detect_gpus() + self.num_gpus = max(1, min(self.num_gpus, 8)) + + self.worker_procs: list[Tuple[GpuWorkerConfig, subprocess.Popen]] = [] + self._start_ts = time.time() + + # ------------------------------------------------------------------ # + # Public entry + # ------------------------------------------------------------------ # + + def run(self) -> None: + """Main orchestration: setup → spawn → monitor → aggregate → done.""" + # 1. Setup + workers = self._prepare_workers() + + # Write parent run.json with multi-gpu metadata + self._write_parent_run({"gpu_count": len(workers), "phase": "starting"}) + + # 2. Spawn + self._spawn_all(workers) + + # Refresh PID immediately after spawn — closes the gap between + # the initial write and the first monitor loop iteration. + self._refresh_pid_file() + + # 3. Monitor + self._monitor_loop() + + # 4. Check results + self._check_failures() + + # 5. Aggregate + self._merge_results() + + # 6. Summary report + summary = self._generate_summary() + + # 7. Done + self._write_parent_run({ + "phase": "finished", + "finished": True, + "final_status": "success", + "summary": summary, + }) + + # Print summary to orchestrator log for quick inspection + print("\n" + "=" * 60) + print("TASK COMPLETE — Summary Report") + print("=" * 60) + print(summary) + print("=" * 60) + + # ------------------------------------------------------------------ # + # Worker setup + # ------------------------------------------------------------------ # + + def _prepare_workers(self) -> List[GpuWorkerConfig]: + """Split shapes and create per-GPU worker configs.""" + extra_notes = self.req.get("extra_notes", "") + groups = split_shapes_for_gpus(extra_notes, self.num_gpus) + + workers: List[GpuWorkerConfig] = [] + for gpu_label, shapes in groups: + gpu_idx = int(gpu_label.replace("GPU ", "")) + gpu_state = self.state_dir / f"gpu_{gpu_idx}" + gpu_workspace = self.workspace_dir / f"gpu_{gpu_idx}" + + gpu_state.mkdir(parents=True, exist_ok=True) + gpu_workspace.mkdir(parents=True, exist_ok=True) + + # Build per-GPU requirements (copy parent, override) + gpu_req = dict(self.req) + gpu_req.pop("multi_gpu", None) + gpu_req.pop("gpu_count", None) + gpu_req["gpu_device"] = str(gpu_idx) + gpu_req["extra_notes"] = ( + "{}\nGPU={} shapes:\n{}".format( + extra_notes.split("\n")[0] if extra_notes else "", + gpu_label, + shapes_to_notes(shapes), + ) + ) + + req_path = gpu_state / "requirements.json" + req_path.write_text(json.dumps(gpu_req, indent=2), encoding="utf-8") + + # Copy reference kernel to per-GPU workspace + ref_src = self.workspace_dir / "reference" / "original_kernel.py" + if ref_src.is_file(): + ref_dst = gpu_workspace / "reference" + ref_dst.mkdir(parents=True, exist_ok=True) + (ref_dst / "original_kernel.py").write_text( + ref_src.read_text(encoding="utf-8"), + encoding="utf-8", + ) + + workers.append(GpuWorkerConfig( + gpu_idx=gpu_idx, + label=gpu_label, + shapes=shapes, + state_dir=gpu_state, + workspace_dir=gpu_workspace, + requirements_path=req_path, + )) + + return workers + + # ------------------------------------------------------------------ # + # Spawn + # ------------------------------------------------------------------ # + + def _spawn_all(self, workers: List[GpuWorkerConfig]) -> None: + """Launch one orchestrator subprocess per GPU.""" + for w in workers: + log_path = w.state_dir / "orchestrator.log" + log_fp = open(str(log_path), "ab", buffering=0) + + # Resolve python executable — avoid the pip entry-point wrapper + # which points to the wrong module path. + python_bin = os.environ.get("METAINFER_PYTHON", sys.executable) + # If sys.executable is the metainfer-orchestrator entry point, + # fall back to the standard python3 binary. + if "orchestrator" in python_bin.lower(): + import shutil + fallback = shutil.which("python3") or "/usr/bin/python3" + python_bin = fallback + + cmd = [ + python_bin, + "-m", "metainfer.tasks.evolve_kernel.orchestrator.cli", + "run", str(w.requirements_path), + "--state-dir", str(w.state_dir), + "--workspace-dir", str(w.workspace_dir), + "--gpu-device", str(w.gpu_idx), + "--claude-bin", self.claude_bin, + "--permission-mode", self.permission_mode, + "--effort", self.effort, + ] + if self.model: + cmd += ["--model", self.model] + + # Set PYTHONPATH so the worker can import metainfer + worker_env = dict(os.environ) + python_path = os.pathsep.join( + p for p in sys.path + if p and p not in worker_env.get("PYTHONPATH", "").split(os.pathsep) + ) + if python_path: + existing = worker_env.get("PYTHONPATH", "") + worker_env["PYTHONPATH"] = ( + f"{python_path}{os.pathsep}{existing}".rstrip(os.pathsep) + if existing else python_path + ) + + proc = subprocess.Popen( + cmd, + stdout=log_fp, + stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, + cwd=str(w.state_dir), + env=worker_env, + start_new_session=True, + ) + log_fp.close() + + # Write PID file for status monitoring + pid_data = { + "pid": proc.pid, + "task_id": self.req.get("task_id", ""), + "started_at": time.time(), + } + (w.state_dir / "orchestrator.pid").write_text( + json.dumps(pid_data, indent=2), encoding="utf-8", + ) + + self.worker_procs.append((w, proc)) + print(f"[multi-gpu] spawned GPU {w.gpu_idx} worker (PID={proc.pid})") + + # ------------------------------------------------------------------ # + # Monitor loop + # ------------------------------------------------------------------ # + + def _monitor_loop(self) -> None: + """Poll per-GPU run.json and update parent state until all finish.""" + poll_s = 3.0 + _crashed_handled: set = set() + while True: + all_done = True + gpu_statuses: List[Dict[str, Any]] = [] + + for w, proc in self.worker_procs: + # Detect crashed/zombie workers: if process exited but + # run.json was never updated to "finished", mark as crashed. + if w.gpu_idx not in _crashed_handled and proc.poll() is not None: + _crashed_handled.add(w.gpu_idx) + run_path = w.state_dir / "run.json" + if run_path.is_file(): + try: + run = json.loads(run_path.read_text(encoding="utf-8")) + if run.get("current_phase") not in ("finished", "crashed"): + exit_code = proc.poll() + print(f"[multi-gpu] GPU {w.gpu_idx} worker exited with code {exit_code} " + f"at phase {run.get('current_phase')} — marking as crashed") + run["finished"] = True + run["final_status"] = "crashed" + run["crash_exit_code"] = exit_code + run_path.write_text(json.dumps(run, indent=2), encoding="utf-8") + except Exception: + pass + + status = self._read_gpu_status(w) + gpu_statuses.append(status) + + if status["phase"] not in ("finished", "crashed", ""): + all_done = False + + # Update parent run.json with live status + self._write_parent_run({ + "phase": "running", + "gpu_status": gpu_statuses, + "gpu_count": len(self.worker_procs), + }) + + # Refresh parent orchestrator.pid every cycle so the liveness + # scanner never sees a stale PID file. The liveness checker + # runs every 10s in the WebUI and will stamp finished_at on any + # PID file it misdiagnoses as dead. + self._refresh_pid_file() + + if all_done: + print("[multi-gpu] all workers finished") + break + + time.sleep(poll_s) + + # Wait for processes to actually exit + for w, proc in self.worker_procs: + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + def _refresh_pid_file(self) -> None: + """Rewrite parent orchestrator.pid to keep liveness scanner happy. + + Preserves the original started_at from the launcher's placeholder + (which matches /proc//stat) — overwriting it causes the + liveness scanner to falsely detect PID reuse. + """ + import json as _json + pid_path = self.state_dir / "orchestrator.pid" + + # Read existing started_at (written by launcher at spawn time) + started_at = None + my_pid = os.getpid() + if pid_path.is_file(): + try: + prev = _json.loads(pid_path.read_text(encoding="utf-8")) + if prev.get("pid") == my_pid and prev.get("started_at"): + started_at = prev["started_at"] + except Exception: + pass + + payload = { + "pid": my_pid, + "task_id": self.req.get("task_id", ""), + "started_at": started_at or self._start_ts, + } + try: + tmp = pid_path.with_suffix(".tmp") + tmp.write_text(_json.dumps(payload, indent=2), encoding="utf-8") + tmp.replace(pid_path) + except OSError: + pass + + def _read_gpu_status(self, w: GpuWorkerConfig) -> Dict[str, Any]: + """Read one GPU worker's current state from disk.""" + run_path = w.state_dir / "run.json" + if not run_path.is_file(): + return { + "gpu_idx": w.gpu_idx, + "label": w.label, + "phase": "starting", + "iteration": 0, + "exec_time_ms": 0, + "speedup": 0, + "running": True, + "pid": 0, + } + + try: + run = json.loads(run_path.read_text(encoding="utf-8")) + except Exception: + return {"gpu_idx": w.gpu_idx, "label": w.label, "phase": "error"} + + # Check if process is alive + pid_file = w.state_dir / "orchestrator.pid" + running = False + pid = 0 + if pid_file.is_file(): + try: + pid_data = json.loads(pid_file.read_text(encoding="utf-8")) + pid = pid_data.get("pid", 0) + if pid and pid_data.get("finished_at") is None: + # Verify process is alive + try: + os.kill(pid, 0) + running = True + except OSError: + running = False + except Exception: + pass + + # Read phase/iteration from run.json + phase = run.get("current_phase", "idle") + iteration = run.get("current_iteration", 0) + + # Read best kernel from per-GPU library + exec_time_ms = 0.0 + speedup = 0.0 + lib_path = w.workspace_dir / "kernel_library.json" + if lib_path.is_file(): + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + if lib: + lib.sort(key=lambda k: k.get("exec_time_ms", float("inf"))) + best = lib[0] + exec_time_ms = best.get("exec_time_ms", 0) + # Speedup vs seed (exec_time of iteration_added=0) + for k in lib: + if k.get("iteration_added") == 0: + seed_time = k.get("exec_time_ms", 0) + if seed_time > 0 and exec_time_ms > 0: + speedup = seed_time / exec_time_ms + break + except Exception: + pass + + return { + "gpu_idx": w.gpu_idx, + "label": w.label, + "phase": phase, + "iteration": iteration, + "exec_time_ms": round(exec_time_ms, 4) if exec_time_ms else 0, + "speedup": round(speedup, 2) if speedup else 0, + "running": running, + "pid": pid, + } + + # ------------------------------------------------------------------ # + # Result aggregation + # ------------------------------------------------------------------ # + + def _check_failures(self) -> None: + """Log which workers crashed.""" + for w, _proc in self.worker_procs: + run_path = w.state_dir / "run.json" + if run_path.is_file(): + try: + run = json.loads(run_path.read_text(encoding="utf-8")) + if run.get("final_status") == "crashed": + print(f"[multi-gpu] GPU {w.gpu_idx} ended with crash") + except Exception: + pass + + def _merge_results(self) -> None: + """Merge all per-GPU kernel libraries and shape benchmarks into parent.""" + # Merge kernel libraries + all_kernels: List[Dict[str, Any]] = [] + for w, _proc in self.worker_procs: + lib_path = w.workspace_dir / "kernel_library.json" + if lib_path.is_file(): + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + for k in lib: + k["gpu_id"] = w.gpu_idx + k["gpu_label"] = w.label + all_kernels.extend(lib) + except Exception: + pass + + all_kernels.sort(key=lambda k: k.get("exec_time_ms", float("inf"))) + parent_lib = self.workspace_dir / "kernel_library.json" + parent_lib.write_text(json.dumps(all_kernels, indent=2, ensure_ascii=False), + encoding="utf-8") + + # Merge shape benchmarks + all_bench: List[Dict[str, Any]] = [] + for w, _proc in self.worker_procs: + bench_path = w.workspace_dir / "shape_bench.json" + if bench_path.is_file(): + try: + data = json.loads(bench_path.read_text(encoding="utf-8")) + gpu_source = f"GPU{w.gpu_idx}" + for r in data.get("results", []): + r["gpu_source"] = gpu_source + all_bench.extend(data.get("results", [])) + except Exception: + pass + + if all_bench: + parent_bench = self.workspace_dir / "shape_bench.json" + parent_bench.write_text(json.dumps({ + "results": all_bench, + "best_kernel_id": "merged", + "cached": False, + }, indent=2, ensure_ascii=False), encoding="utf-8") + + print(f"[multi-gpu] merged {len(all_kernels)} kernels, {len(all_bench)} benchmarks") + + # ------------------------------------------------------------------ # + # Parent run.json + # ------------------------------------------------------------------ # + + def _generate_summary(self) -> str: + """Generate a human-readable summary report after all workers finish. + + Writes the summary to ``summary.txt`` in the state dir and returns + the markdown-formatted report string. + """ + import time as _time + + elapsed_s = _time.time() - self._start_ts + elapsed_str = f"{elapsed_s:.0f}s" + if elapsed_s > 3600: + elapsed_str = f"{elapsed_s / 3600:.1f}h" + elif elapsed_s > 60: + elapsed_str = f"{elapsed_s / 60:.1f}min" + + lines = [] + lines.append("") + lines.append(f"## Task: {self.req.get('label', self.req.get('task_id', '?'))}") + lines.append(f"**Elapsed:** {elapsed_str} | **Mode:** {self.req.get('optimizer_mode', 'Triton')} | **GPUs:** {self.num_gpus}") + lines.append(f"**Max iterations:** {self.req.get('max_iterations', '?')} | **Profiling:** {self.req.get('enable_profiling', 'No')}") + lines.append("") + lines.append("| GPU | Shape Task | Seed (ms) | Best (ms) | Speedup | Kernels | Stop Reason |") + lines.append("|-----|-----------|----------|----------|---------|---------|-------------|") + + all_ok = True + for w, _proc in self.worker_procs: + run_path = w.state_dir / "run.json" + lib_path = w.workspace_dir / "kernel_library.json" + + stop_reason = "unknown" + final_status = "?" + if run_path.is_file(): + try: + run = json.loads(run_path.read_text(encoding="utf-8")) + final_status = run.get("final_status", "?") + if run.get("crash_exit_code"): + stop_reason = f"crashed (exit {run['crash_exit_code']})" + all_ok = False + elif final_status == "success": + stop_reason = "converged / max iters" + elif final_status == "crashed": + stop_reason = f"crashed: {run.get('crash_reason', '?')[:60]}" + all_ok = False + except Exception: + pass + + seed_ms = 0.0 + best_ms = 0.0 + speedup = 0.0 + kernel_count = 0 + if lib_path.is_file(): + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + kernel_count = len(lib) + if lib: + best = min(lib, key=lambda k: k.get("exec_time_ms", float("inf"))) + best_ms = best.get("exec_time_ms", 0) + for k in lib: + if k.get("iteration_added") == 0: + seed_ms = k.get("exec_time_ms", 0) + break + if seed_ms > 0 and best_ms > 0: + speedup = seed_ms / best_ms + except Exception: + pass + + # Extract shape task summary from requirements + shapes_notes = "" + req_path = w.state_dir / "requirements.json" + if req_path.is_file(): + try: + gpu_req = json.loads(req_path.read_text(encoding="utf-8")) + extra = gpu_req.get("extra_notes", "") + # Extract first line that looks like a shape description + for line in extra.split("\n"): + if "(" in line and "@" in line and ("M=" in line or "TP=" in line): + shapes_notes = line.strip()[:60] + break + except Exception: + pass + + su_str = f"{speedup:.2f}×" if speedup > 0 else "—" + seed_str = f"{seed_ms:.4f}" if seed_ms > 0 else "—" + best_str = f"{best_ms:.4f}" if best_ms > 0 else "—" + + lines.append( + f"| GPU {w.gpu_idx} | {shapes_notes} | {seed_str}ms | {best_str}ms | " + f"{su_str} | {kernel_count} | {stop_reason} |" + ) + + lines.append("") + + # Warnings + if not all_ok: + lines.append("### ⚠ Warnings") + lines.append("Some workers did not finish cleanly. Check per-GPU orchestrator logs for details.") + lines.append("") + + # Export paths + exp_dir = self.workspace_dir / "optimized_kernels" + if exp_dir.is_dir(): + cpp_count = len(list(exp_dir.glob("*.cpp"))) + py_count = len(list(exp_dir.glob("*.py"))) + lines.append(f"### Exported Kernels") + lines.append(f"`{exp_dir}` — {py_count} `.py` wrappers + {cpp_count} `.cpp` sources") + lines.append("") + + report = "\n".join(lines) + + # Write to summary.txt + summary_path = self.state_dir / "summary.txt" + try: + summary_path.write_text(report, encoding="utf-8") + except Exception: + pass + + return report + + def _write_parent_run(self, updates: Dict[str, Any]) -> None: + """Update parent run.json with current multi-GPU state. + + Always ensures ``finished``, ``final_status`` reflect the real state. + Stale ``finished: true`` from a prior crash/restart is reset. + """ + run_path = self.state_dir / "run.json" + current: Dict[str, Any] = { + "task_id": self.req.get("task_id", ""), + "current_iteration": 0, + "current_phase": "running", + "finished": False, + "final_status": None, + "multi_gpu": True, + "last_update": time.time(), + } + if run_path.is_file(): + try: + loaded = json.loads(run_path.read_text(encoding="utf-8")) + # Merge but NEVER carry over stale finished/stopped + for k, v in loaded.items(): + if k in ("finished", "final_status"): + continue + if k not in current: + current[k] = v + except Exception: + pass + # Apply latest updates (including override of phase, gpu_status, etc.) + current.update(updates) + current["last_update"] = time.time() + # Only set finished if explicitly passed True + if updates.get("finished") is not True: + current["finished"] = False + run_path.write_text(json.dumps(current, indent=2), encoding="utf-8") diff --git a/metainfer/tasks/evolve_kernel/orchestrator/cli.py b/metainfer/tasks/evolve_kernel/orchestrator/cli.py index 98ed4c50..458c6ede 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/cli.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/cli.py @@ -60,6 +60,8 @@ def main(argv: list[str] | None = None) -> int: run_p.add_argument("--model", default=None, help="Override model for sub-agents") run_p.add_argument("--effort", default=None, choices=_VALID_EFFORTS, help=f"Claude Code effort level (default: {DEFAULT_EFFORT!r})") + run_p.add_argument("--gpu-device", type=str, default=None, + help="CUDA visible device(s), e.g. '0' or '0,1'. Sets CUDA_VISIBLE_DEVICES env var before importing torch.") run_p.add_argument("--max-iterations", type=int, default=None, help="Override max iterations") run_p.add_argument("--extra-claude-arg", action="append", default=[], @@ -68,6 +70,10 @@ def main(argv: list[str] | None = None) -> int: args = parser.parse_args(argv) if args.cmd == "run": + # Set GPU device BEFORE importing torch/triton inside the orchestrator + if args.gpu_device is not None: + import os + os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_device from .orchestrator import run_with_requirements return run_with_requirements( requirements_path=args.requirements, diff --git a/metainfer/tasks/evolve_kernel/orchestrator/headroom.py b/metainfer/tasks/evolve_kernel/orchestrator/headroom.py new file mode 100644 index 00000000..40ab896c --- /dev/null +++ b/metainfer/tasks/evolve_kernel/orchestrator/headroom.py @@ -0,0 +1,725 @@ +"""Roofline-model headroom analysis for GPU kernel optimization. + +Computes the kernel's position on the roofline using the standard model: + + P_achieved = FLOPs / T + AI = FLOPs / Bytes_HBM + P_bandwidth_roof = BW_peak × AI + P_max = min(P_compute_peak, P_bandwidth_roof) + roofline_efficiency = P_achieved / P_max + +where FLOPs and Bytes are the *theoretical minimum* for the algorithm +(algorithmic FLOPs and compulsory HBM traffic). When hipprof profiling +is enabled, measured HBM bytes can replace the theoretical estimate for +a measured-AI comparison. + +Reference: Williams, Waterman, Patterson. "Roofline: an insightful visual +performance model for multicore architectures." CACM 52(4), 2009. + +Integrated into Phase H (measure perf) — runs after the perf harness but +before the kernel enters the library. Results are stored in KernelEntry +and displayed in the WebUI to guide future optimization iterations. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field, asdict +from typing import Any, Dict, List, Optional, Tuple + + +# =========================================================================== # +# GPU Spec Database +# =========================================================================== # + + +@dataclass +class GpuSpec: + name: str + peak_hbm_bw_gbps: float # peak HBM bandwidth (GB/s) + peak_tflops_fp32: float = 0.0 + peak_tflops_fp16: float = 0.0 + peak_tflops_bf16: float = 0.0 + peak_tflops_int8: float = 0.0 + warp_size: int = 64 + smem_per_cu_kb: int = 64 + cu_count: int = 120 + # Approximate effective L2 bandwidth (GB/s) — used for hierarchical roofline + l2_bw_gbps: float = 2000.0 + # Approximate shared memory bandwidth per CU (TB/s) + smem_bw_tbps: float = 10.0 + + +GPU_SPECS: Dict[str, GpuSpec] = { + "gfx928": GpuSpec( + name="AMD DCU K500SM_AI (gfx928)", + peak_hbm_bw_gbps=700.0, + peak_tflops_fp32=55.0, + peak_tflops_fp16=110.0, + peak_tflops_bf16=110.0, + peak_tflops_int8=220.0, + warp_size=64, + smem_per_cu_kb=64, + cu_count=120, + l2_bw_gbps=2000.0, + ), + "gfx942": GpuSpec( + name="AMD Instinct MI300X (gfx942)", + peak_hbm_bw_gbps=5300.0, + peak_tflops_fp32=81.7, + peak_tflops_fp16=163.4, + peak_tflops_bf16=163.4, + peak_tflops_int8=326.8, + warp_size=64, + smem_per_cu_kb=64, + cu_count=304, + ), + "gfx90a": GpuSpec( + name="AMD Instinct MI250X (gfx90a)", + peak_hbm_bw_gbps=1600.0, + peak_tflops_fp32=47.9, + peak_tflops_fp16=191.6, + peak_tflops_bf16=191.6, + peak_tflops_int8=383.2, + warp_size=64, + smem_per_cu_kb=64, + cu_count=220, + ), + "sm90": GpuSpec( + name="NVIDIA H100 (sm90)", + peak_hbm_bw_gbps=3350.0, + peak_tflops_fp32=67.0, + peak_tflops_fp16=989.0, + peak_tflops_bf16=989.0, + peak_tflops_int8=1979.0, + warp_size=32, + smem_per_cu_kb=228, + cu_count=132, + ), + "sm89": GpuSpec( + name="NVIDIA RTX 4090 (sm89)", + peak_hbm_bw_gbps=1008.0, + peak_tflops_fp32=82.6, + peak_tflops_fp16=165.2, + peak_tflops_bf16=165.2, + peak_tflops_int8=330.3, + warp_size=32, + smem_per_cu_kb=128, + cu_count=128, + ), +} + +_FALLBACK_SPEC = GpuSpec( + name="Unknown GPU", + peak_hbm_bw_gbps=500.0, + peak_tflops_fp32=10.0, + peak_tflops_fp16=20.0, + peak_tflops_bf16=20.0, + peak_tflops_int8=40.0, +) + + +def detect_gpu(req: Dict[str, Any], kernel_code: str = "") -> GpuSpec: + """Detect the GPU model from requirements or kernel code.""" + search_text = "" + for key in ("extra_notes", "notes", "target_hardware", "gpu_model"): + val = req.get(key, "") + if val: + search_text += str(val) + "\n" + search_text += kernel_code[:2000] + + for gpu_id, spec in GPU_SPECS.items(): + if gpu_id.lower() in search_text.lower(): + return spec + + name_map = { + "mi300x": "gfx942", "mi250x": "gfx90a", "mi250": "gfx90a", + "h100": "sm90", "h200": "sm90", "h800": "sm90", + "rtx 4090": "sm89", "rtx4090": "sm89", + "dcu": "gfx928", "k500sm": "gfx928", + } + for name_key, gpu_id in name_map.items(): + if name_key in search_text.lower(): + return GPU_SPECS[gpu_id] + + return _FALLBACK_SPEC + + +# =========================================================================== # +# Shape extraction +# =========================================================================== # + +_PLAIN_SHAPE_PATTERN = re.compile(r'\b(\d+)\s*[×x]\s*(\d+)\s*[×x]\s*(\d+)\b') +_PERGPU_SHAPE_PATTERN = re.compile( + r'\(\s*(?:M|m)\s*,\s*(\d+)\s*\)\s*@\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)' +) + + +@dataclass +class ProblemShape: + M: Optional[int] # None means variable (benchmarked across multiple M) + N: int + K: int + label: str = "" + + +def extract_shapes(req: Dict[str, Any]) -> List[ProblemShape]: + """Extract target problem shapes from requirements text.""" + shapes: List[ProblemShape] = [] + seen: set = set() + + search_text = "" + for key in ("extra_notes", "notes", "problem_shapes", "target_shapes"): + val = req.get(key, "") + if val: + search_text += str(val) + "\n" + + if not search_text: + return shapes + + # Pattern "M×K × K×N" (variable M) + for match in re.finditer( + r'(?:M|m)\s*[×x]\s*(\d+)\s*[×x]\s*(\d+)', + search_text + ): + k_val = int(match.group(1)) + n_val = int(match.group(2)) + key = (None, k_val, n_val) + if key not in seen: + seen.add(key) + shapes.append(ProblemShape(M=None, N=n_val, K=k_val, + label=f"M×{k_val} (K)×{n_val}")) + + # Pattern "(M, K) @ (K, N)" (per-GPU splitter format) + for match in _PERGPU_SHAPE_PATTERN.finditer(search_text): + k_val = int(match.group(1)) + n_val = int(match.group(3)) + key = (None, k_val, n_val) + if key not in seen: + seen.add(key) + shapes.append(ProblemShape(M=None, N=n_val, K=k_val, + label=f"M×{k_val} (K)×{n_val}")) + + # Pattern "m×k×n" + for match in re.finditer( + r'(?:m)\s*[×x]\s*(\d+)\s*[×x]\s*(\d+)', + search_text + ): + k_val = int(match.group(1)) + n_val = int(match.group(2)) + key = (None, k_val, n_val) + if key not in seen: + seen.add(key) + shapes.append(ProblemShape(M=None, N=n_val, K=k_val, + label=f"m×{k_val} (K)×{n_val}")) + + # Plain M×N×K + for match in _PLAIN_SHAPE_PATTERN.finditer(search_text): + m_val = int(match.group(1)) + n_val = int(match.group(2)) + k_val = int(match.group(3)) + key = (m_val, n_val, k_val) + if key not in seen: + seen.add(key) + shapes.append(ProblemShape(M=m_val, N=n_val, K=k_val, + label=f"{m_val}×{n_val}×{k_val}")) + + return shapes + + +def _resolve_shape_dim(val, default: int) -> int: + if val is None: + return default + if isinstance(val, int): + return val + try: + return int(val) + except (ValueError, TypeError): + return default + + +# =========================================================================== # +# FLOPs & Memory Traffic +# =========================================================================== # + + +def estimate_gemm_flops(M: int, N: int, K: int) -> int: + """GEMM FLOPs: 2*M*N*K (one multiply + one add per output element). + + For int8 w8a8 GEMM, the core dot-product is int8→int32 with fp32 scaling. + We count the matmul FLOPs; scaling adds ~M*N FLOPs (negligible for large K). + """ + return 2 * M * N * K + + +def estimate_hbm_traffic(M: int, N: int, K: int, + a_dtype_bytes: int = 1, + b_dtype_bytes: int = 1, + out_dtype_bytes: int = 2, + has_scales: bool = True) -> Tuple[int, int, int]: + """Theoretical *minimum* HBM traffic for a GEMM kernel. + + Assumes each matrix element is read from HBM exactly once and the output + is written once. Real kernels will exceed this due to cache misses, + non-coalesced access, write-allocate, and instruction overhead. + + Returns (bytes_read, bytes_write, total_bytes). + """ + bytes_read = M * K * a_dtype_bytes + K * N * b_dtype_bytes + if has_scales: + bytes_read += M * 4 + N * 4 + bytes_write = M * N * out_dtype_bytes + return bytes_read, bytes_write, bytes_read + bytes_write + + +# =========================================================================== # +# Roofline Analysis +# =========================================================================== # +# Notation: +# P_achieved = FLOPs / T (achieved TFLOPS) +# AI = FLOPs / Bytes_HBM (arithmetic intensity, FLOP/byte) +# P_bw_roof = BW_peak × AI (bandwidth ceiling, in TFLOPS) +# AI_ridge = P_compute_peak / BW_peak (ridge point, FLOP/byte) +# P_max = min(P_compute_peak, P_bw_roof) (roofline ceiling) +# η = P_achieved / P_max (roofline efficiency) + + +@dataclass +class HeadroomResult: + """Result of roofline headroom analysis for one representative shape.""" + + # -- Input / problem -- + shape_label: str = "" + M: int = 0 + N: int = 0 + K: int = 0 + exec_time_ms: float = 0.0 + + # -- Algorithmic quantities (theoretical minimum) -- + total_flops: int = 0 + total_bytes_hbm: int = 0 # theoretical compulsory HBM bytes + bytes_read: int = 0 + bytes_write: int = 0 + arithmetic_intensity: float = 0.0 # FLOP / byte (HBM-level, theoretical) + + # -- Achieved (derived from exec_time + theoretical bytes/FLOPs) -- + achieved_tflops: float = 0.0 # P_achieved + achieved_bw_gbps: float = 0.0 # total_bytes / T (not from profiler) + + # -- When hipprof data is available, these hold measured HBM bytes -- + measured_hbm_bytes: float = 0.0 # profiler-measured HBM traffic (0 = no data) + measured_ai: float = 0.0 # AI from measured bytes + + # -- Hardware peaks -- + peak_bw_gbps: float = 0.0 # HBM bandwidth + peak_tflops: float = 0.0 # compute peak (matched to dtype) + gpu_name: str = "" + + # -- Roofline model -- + ai_ridge: float = 0.0 # P_compute / BW_hbm (FLOP/byte) + p_bandwidth_roof_tflops: float = 0.0 # BW_peak × AI + p_max_tflops: float = 0.0 # min(P_compute, P_bw_roof) + roofline_efficiency_pct: float = 0.0 # P_achieved / P_max × 100 + + # -- Bottleneck classification -- + bottleneck: str = "unknown" + headroom_pct: float = 0.0 # 100 - roofline_efficiency (simplified) + shape_is_approximate: bool = False + + # -- Recommendations -- + suggestions: List[str] = field(default_factory=list) + optimization_advice: str = "" + + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- + + +def analyze_headroom( + kernel_code: str, + kernel_fn_name: str, + exec_time_ms: float, + req: Dict[str, Any], + gpu_spec: Optional[GpuSpec] = None, + a_dtype_bytes: int = 1, + b_dtype_bytes: int = 1, + out_dtype_bytes: int = 2, + profiled_hbm_bytes: float = 0.0, # from hipprof (0 = no data) +) -> HeadroomResult: + """Run roofline headroom analysis for a kernel. + + Uses the original Roofline model: + - Algorithmic FLOPs and theoretical minimum HBM bytes define AI. + - Exec time gives P_achieved. + - P_max = min(P_compute, BW_hbm × AI). + - Efficiency = P_achieved / P_max. + + Args: + kernel_code: Optimized kernel source (for shape fallback). + kernel_fn_name: Kernel wrapper function name. + exec_time_ms: Measured execution time from the perf harness. + req: Task requirements dict (shapes, GPU model). + gpu_spec: Pre-detected GPU spec. + a_dtype_bytes, b_dtype_bytes: Element sizes for A, B matrices. + out_dtype_bytes: Element size for output. + profiled_hbm_bytes: Actual HBM bytes from hipprof (0 = use theoretical). + + Returns: + HeadroomResult with all computed metrics. + """ + # 1. Detect GPU + if gpu_spec is None: + gpu_spec = detect_gpu(req, kernel_code) + + # 2. Extract shapes — pick the largest-FLOPs representative shape + shapes = extract_shapes(req) + if not shapes: + shapes = [_guess_shape_from_code(kernel_code)] + + explicit_shapes = [s for s in shapes if s.M is not None] + candidates = explicit_shapes if explicit_shapes else shapes + + best_shape = candidates[0] + best_flops = 0 + for s in candidates: + m = _resolve_shape_dim(s.M, 4096) + n = s.N + k = s.K + flops = estimate_gemm_flops(m, n, k) + if flops > best_flops: + best_flops = flops + best_shape = s + + M = _resolve_shape_dim(best_shape.M, 4096) + N = best_shape.N + K = best_shape.K + + # 3. Pick compute peak for the kernel's dtype + if a_dtype_bytes == 1 and b_dtype_bytes == 1: + peak_tflops = gpu_spec.peak_tflops_int8 + elif out_dtype_bytes == 2: + peak_tflops = gpu_spec.peak_tflops_fp16 + else: + peak_tflops = gpu_spec.peak_tflops_fp32 + if peak_tflops <= 0: + peak_tflops = gpu_spec.peak_tflops_fp32 + + peak_bw = gpu_spec.peak_hbm_bw_gbps + + # 4. Algorithmic quantities + total_flops = estimate_gemm_flops(M, N, K) + bytes_read, bytes_write, total_bytes = estimate_hbm_traffic( + M, N, K, + a_dtype_bytes=a_dtype_bytes, + b_dtype_bytes=b_dtype_bytes, + out_dtype_bytes=out_dtype_bytes, + ) + + # Arithmetic intensity: FLOP / byte (HBM level, theoretical minimum) + ai = total_flops / max(total_bytes, 1) + + # ---------------------------------------------------------------- + # 5. Roofline model + # ---------------------------------------------------------------- + time_s = max(exec_time_ms / 1000.0, 1e-9) + + # Achieved throughput (from theoretical FLOPs / measured time) + p_achieved = total_flops / time_s / 1e12 # TFLOPS + bw_achieved = total_bytes / time_s / 1e9 # GB/s + + # Bandwidth ceiling expressed as TFLOPS + p_bw_roof = peak_bw * ai / 1e3 # TFLOPS + # Overall roofline ceiling + p_max = min(peak_tflops, p_bw_roof) + + # Roofline efficiency: how close is P_achieved to P_max? + roofline_efficiency = (p_achieved / p_max * 100.0) if p_max > 0 else 0.0 + roofline_efficiency = min(100.0, roofline_efficiency) + + # Ridge point: where the bandwidth ceiling meets the compute ceiling + ai_ridge = peak_tflops * 1e3 / peak_bw if peak_bw > 0 else 0.0 + + # ---------------------------------------------------------------- + # 6. Measured AI (if profiler data available) + # ---------------------------------------------------------------- + measured_ai = 0.0 + if profiled_hbm_bytes > 0: + measured_ai = total_flops / profiled_hbm_bytes + + # ---------------------------------------------------------------- + # 7. Bottleneck classification + # ---------------------------------------------------------------- + bottleneck, headroom_pct = _classify_bottleneck( + ai=ai, + ai_ridge=ai_ridge, + roofline_efficiency=roofline_efficiency, + p_achieved=p_achieved, + p_max=p_max, + peak_bw=peak_bw, + peak_tflops=peak_tflops, + bw_achieved=bw_achieved, + ) + + # ---------------------------------------------------------------- + # 8. Suggestions + # ---------------------------------------------------------------- + suggestions, optimization_advice = _generate_suggestions( + bottleneck=bottleneck, + roofline_efficiency=roofline_efficiency, + ai=ai, ai_ridge=ai_ridge, + p_achieved=p_achieved, p_max=p_max, + p_bw_roof=p_bw_roof, + bw_achieved=bw_achieved, peak_bw=peak_bw, + peak_tflops=peak_tflops, + measured_ai=measured_ai, + M=M, N=N, K=K, + gpu_spec=gpu_spec, + ) + + shape_is_approximate = best_shape.M is None + + return HeadroomResult( + shape_label=best_shape.label, + M=M, N=N, K=K, + exec_time_ms=exec_time_ms, + total_flops=total_flops, + total_bytes_hbm=total_bytes, + bytes_read=bytes_read, + bytes_write=bytes_write, + arithmetic_intensity=round(ai, 2), + achieved_tflops=round(p_achieved, 4), + achieved_bw_gbps=round(bw_achieved, 2), + measured_hbm_bytes=profiled_hbm_bytes, + measured_ai=round(measured_ai, 2) if measured_ai > 0 else 0.0, + peak_bw_gbps=peak_bw, + peak_tflops=peak_tflops, + gpu_name=gpu_spec.name, + ai_ridge=round(ai_ridge, 2), + p_bandwidth_roof_tflops=round(p_bw_roof, 4), + p_max_tflops=round(p_max, 4), + roofline_efficiency_pct=round(roofline_efficiency, 1), + bottleneck=bottleneck, + headroom_pct=round(headroom_pct, 1), + shape_is_approximate=shape_is_approximate, + suggestions=suggestions, + optimization_advice=optimization_advice, + ) + + +def _guess_shape_from_code(kernel_code: str, default_M: int = 4096, + default_N: int = 4096, default_K: int = 4096) -> ProblemShape: + """Fallback: guess shape from kernel comments or docstring.""" + for match in _PLAIN_SHAPE_PATTERN.finditer(kernel_code[:5000]): + return ProblemShape( + M=int(match.group(1)), + N=int(match.group(2)), + K=int(match.group(3)), + label=f"{match.group(1)}×{match.group(2)}×{match.group(3)}", + ) + return ProblemShape(M=default_M, N=default_N, K=default_K, + label=f"{default_M}×{default_N}×{default_K} (guessed)") + + +# =========================================================================== # +# Bottleneck Classification +# =========================================================================== # + + +def _classify_bottleneck( + ai: float, + ai_ridge: float, + roofline_efficiency: float, + p_achieved: float, + p_max: float, + peak_bw: float, + peak_tflops: float, + bw_achieved: float, +) -> Tuple[str, float]: + """Classify the performance bottleneck using the roofline model. + + The regime is determined by AI relative to the ridge point: + - AI < AI_ridge : memory-bound regime → ceiling is BW_peak × AI + - AI ≥ AI_ridge : compute-bound regime → ceiling is P_compute_peak + + Within each regime, the achieved throughput relative to that ceiling + determines whether the kernel is efficient or not. + """ + NEAR_OPTIMAL = 90.0 # roofline efficiency ≥ 90% → near_optimal + EFFICIENT = 60.0 # roofline efficiency ≥ 60% → well-bound + INEFFICIENT = 30.0 # roofline efficiency < 30% → clearly inefficient + + if roofline_efficiency >= NEAR_OPTIMAL: + bottleneck = "near_optimal" + elif ai < ai_ridge: + # Memory-bound regime — bottleneck is HBM bandwidth + if roofline_efficiency >= EFFICIENT: + bottleneck = "memory_bound" + else: + bottleneck = "inefficient" + else: + # Compute-bound regime — bottleneck is compute throughput + if roofline_efficiency >= EFFICIENT: + bottleneck = "compute_bound" + else: + bottleneck = "inefficient" + + headroom = max(0.0, 100.0 - roofline_efficiency) + return bottleneck, headroom + + +# =========================================================================== # +# Suggestions +# =========================================================================== # + + +def _generate_suggestions( + bottleneck: str, + roofline_efficiency: float, + ai: float, + ai_ridge: float, + p_achieved: float, + p_max: float, + p_bw_roof: float, + bw_achieved: float, + peak_bw: float, + peak_tflops: float, + measured_ai: float, + M: int, N: int, K: int, + gpu_spec: GpuSpec, +) -> Tuple[List[str], str]: + """Generate optimization suggestions based on roofline analysis.""" + suggestions: List[str] = [] + parts: List[str] = [] + + # Always show the roofline equation breakdown + suggestions.append( + f"Roofline: P_achieved = {p_achieved:.2f} TFLOPS, " + f"AI = {ai:.1f} FLOP/byte, " + f"P_bw_roof = BW×AI = {p_bw_roof:.2f} TFLOPS, " + f"P_max = min(P_compute={peak_tflops:.0f}, P_bw_roof={p_bw_roof:.2f}) = {p_max:.2f} TFLOPS" + ) + if measured_ai > 0: + suggestions.append( + f"Measured AI (from hipprof HBM counters): {measured_ai:.1f} FLOP/byte " + f"(theoretical minimum: {ai:.1f}). " + f"{'Memory access is efficient' if measured_ai >= ai * 0.7 else 'Memory access is suboptimal — consider coalescing and cache blocking'}" + ) + + if bottleneck == "near_optimal": + suggestions.append( + f"Kernel is at {roofline_efficiency:.0f}% of the hardware roofline " + f"(P_max = {p_max:.2f} TFLOPS). " + f"Further optimization on this problem size is unlikely to yield significant gains." + ) + suggestions.append( + "Consider: (a) radically different algorithm (e.g., split-K, persistent kernel), " + "(b) larger problem size where this kernel's optimizations scale better, " + "(c) rewriting in HIP C++ with hand-tuned assembly, or " + "(d) declaring convergence." + ) + parts.append(f"This kernel achieves {roofline_efficiency:.0f}% of the hardware roofline (P_max={p_max:.2f} TFLOPS).") + parts.append("Near the limit. Consider a different algorithm, HIP C++ rewrite, or declaring convergence.") + + elif bottleneck == "memory_bound": + bw_util_pct = bw_achieved / peak_bw * 100 if peak_bw > 0 else 0 + suggestions.append( + f"HBM bandwidth is the bottleneck (AI={ai:.1f} < ridge={ai_ridge:.1f}). " + f"Achieved BW: {bw_achieved:.0f} GB/s ({bw_util_pct:.0f}% of {peak_bw:.0f} GB/s peak). " + f"Roofline efficiency: {roofline_efficiency:.0f}%." + ) + suggestions.append( + f"To improve: increase AI by raising data reuse. " + f"Larger tile sizes (BLOCK_M, BLOCK_N) reduce HBM reads per FLOP." + ) + suggestions.append( + f"Current AI={ai:.1f}. Raising it to {ai_ridge:.1f} would move to the " + f"compute-bound regime (needs {ai_ridge/ai:.1f}× more reuse)." + ) + if gpu_spec.smem_per_cu_kb > 0: + suggestions.append( + f"Use shared memory ({gpu_spec.smem_per_cu_kb} KB/CU) to cache " + f"B-tile across M iterations — avoids re-reading B from HBM." + ) + parts.append("This kernel is memory-bound.") + parts.append(f"AI={ai:.1f}, ridge={ai_ridge:.1f}. Increase data reuse to improve.") + + elif bottleneck == "compute_bound": + suggestions.append( + f"Compute throughput is the bottleneck (AI={ai:.1f} > ridge={ai_ridge:.1f}). " + f"Roofline efficiency: {roofline_efficiency:.0f}%." + ) + suggestions.append( + f"Check tensor core utilization: ensure tile sizes (M={M}, N={N}, K={K}) " + f"are compatible with gfx928 MMA instructions (16×16×32 int8)." + ) + suggestions.append( + f"Consider warp-level tuning: {gpu_spec.warp_size}-wide " + f"operations for coalesced execution on {gpu_spec.cu_count} CUs." + ) + parts.append("This kernel is compute-bound.") + parts.append("Focus on instruction-level optimizations and tensor core utilization.") + + else: # inefficient + suggestions.append( + f"Kernel is far from the roofline (efficiency={roofline_efficiency:.0f}%). " + f"Achieved: {p_achieved:.2f} TFLOPS vs ceiling of {p_max:.2f} TFLOPS." + ) + suggestions.append( + "Check occupancy: may be limited by register pressure or shared memory allocation. " + f"Target: {gpu_spec.cu_count} CUs × {gpu_spec.warp_size} threads." + ) + suggestions.append( + "Verify launch configuration: grid size, block size, wavefront occupancy. " + "Profile with hardware counters to identify the specific stall reason." + ) + parts.append("This kernel is inefficient — far from the roofline ceiling.") + parts.append("Investigate occupancy, register pressure, tile alignment, and launch configuration.") + + optimization_advice = ( + f"**Bottleneck:** {bottleneck.replace('_', ' ').title()}\n\n" + + " ".join(parts) + ) + + return suggestions, optimization_advice + + +# =========================================================================== # +# Serialization helpers for pipeline.py +# =========================================================================== # + + +def headroom_result_to_dict(result: HeadroomResult) -> Dict[str, Any]: + """Serialize HeadroomResult to a JSON-serializable dict.""" + return asdict(result) + + +def headroom_dict_to_summary(d: Dict[str, Any]) -> Dict[str, Any]: + """Extract a compact summary for storage in kernel_library.json. + + Key roofline fields: + roofline_efficiency_pct : P_achieved / P_max × 100 + p_max_tflops : the kernel's roofline ceiling + p_bandwidth_roof_tflops : BW_peak × AI + """ + return { + "bottleneck": d.get("bottleneck", "unknown"), + "roofline_efficiency_pct": d.get("roofline_efficiency_pct", 0.0), + "headroom_pct": d.get("headroom_pct", 0), + "achieved_tflops": d.get("achieved_tflops", 0), + "achieved_bw_gbps": d.get("achieved_bw_gbps", 0), + "peak_bw_gbps": d.get("peak_bw_gbps", 0), + "peak_tflops": d.get("peak_tflops", 0), + "ai_ridge": d.get("ai_ridge", 0), + "p_bandwidth_roof_tflops": d.get("p_bandwidth_roof_tflops", 0), + "p_max_tflops": d.get("p_max_tflops", 0), + "arithmetic_intensity": d.get("arithmetic_intensity", 0), + "measured_ai": d.get("measured_ai", 0), + "bw_util_pct": 0.0, # deprecated — use roofline_efficiency_pct + "compute_util_pct": 0.0, # deprecated — use roofline_efficiency_pct + "shape_is_approximate": d.get("shape_is_approximate", False), + "shape_label": d.get("shape_label", ""), + "gpu_name": d.get("gpu_name", ""), + "suggestions": d.get("suggestions", []), + "optimization_advice": d.get("optimization_advice", ""), + } diff --git a/metainfer/tasks/evolve_kernel/orchestrator/kernel_library.py b/metainfer/tasks/evolve_kernel/orchestrator/kernel_library.py index 86f9ded4..ddecba1a 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/kernel_library.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/kernel_library.py @@ -33,6 +33,40 @@ class KernelEntry: iteration_added: int = 0 parent_id: Optional[str] = None + # Headroom analysis (populated during Phase H after perf measurement) + headroom_bottleneck: Optional[str] = None # memory_bound | compute_bound | inefficient | near_optimal + headroom_roofline_efficiency_pct: float = 0.0 # P_achieved / P_max × 100 — primary metric + headroom_pct: float = 0.0 # 100 - roofline_efficiency (for backward compat) + headroom_p_max_tflops: float = 0.0 # min(P_compute, BW_peak × AI) + headroom_p_bw_roof_tflops: float = 0.0 # BW_peak × AI + headroom_ai_ridge: float = 0.0 # P_compute / BW_hbm ridge point + headroom_suggestions_json: Optional[str] = None # JSON-encoded list of suggestion strings + headroom_advice: Optional[str] = None # human-readable optimization advice paragraph + headroom_achieved_bw_gbps: float = 0.0 # achieved HBM bandwidth (GB/s) — from theoretical bytes + headroom_achieved_tflops: float = 0.0 # achieved compute throughput (TFLOPS) — from theoretical FLOPs + headroom_peak_bw_gbps: float = 0.0 # peak HBM bandwidth of the GPU + headroom_peak_tflops: float = 0.0 # peak TFLOPS for the kernel's compute dtype + headroom_arithmetic_intensity: float = 0.0 # FLOP / byte (HBM-level, theoretical) + headroom_measured_ai: float = 0.0 # FLOP / byte from profiler-measured HBM bytes (0=none) + headroom_shape_label: str = "" # shape used for roofline analysis (e.g. "M×2048 (K)×4096") + headroom_M: int = 0 + headroom_N: int = 0 + headroom_K: int = 0 + # Deprecated fields kept for backward compat with old kernel_library.json + headroom_bw_util_pct: float = 0.0 + headroom_compute_util_pct: float = 0.0 + + # HIP C++ kernel source (from scratch mode — the real .cpp file content) + cpp_code: Optional[str] = None + + # hipprof profiling (populated during Phase H when enable_profiling=True) + profiled: bool = False + profiling_kernel_duration_us: float = 0.0 + profiling_achieved_bw_gbps: float = 0.0 + profiling_occupancy_pct: float = 0.0 + profiling_l2_cache_hit_pct: float = 0.0 + profiling_advice: Optional[str] = None # profiling-based optimization advice + def recompute_combined(self) -> float: """Recompute combined_score from exec_time and complexity. @@ -146,6 +180,36 @@ def from_list(cls, data: List[Dict[str, Any]]) -> "KernelLibrary": combined_score=d.get("combined_score", 0.0), iteration_added=d.get("iteration_added", 0), parent_id=d.get("parent_id"), + headroom_bottleneck=d.get("headroom_bottleneck"), + headroom_roofline_efficiency_pct=( + d.get("headroom_roofline_efficiency_pct", 0.0) + or max(d.get("headroom_bw_util_pct", 0), d.get("headroom_compute_util_pct", 0)) + ), + headroom_pct=d.get("headroom_pct", 0.0), + headroom_p_max_tflops=d.get("headroom_p_max_tflops", 0.0), + headroom_p_bw_roof_tflops=d.get("headroom_p_bw_roof_tflops", 0.0), + headroom_ai_ridge=d.get("headroom_ai_ridge", 0.0), + headroom_suggestions_json=d.get("headroom_suggestions_json"), + headroom_advice=d.get("headroom_advice"), + headroom_achieved_bw_gbps=d.get("headroom_achieved_bw_gbps", 0.0), + headroom_achieved_tflops=d.get("headroom_achieved_tflops", 0.0), + headroom_peak_bw_gbps=d.get("headroom_peak_bw_gbps", 0.0), + headroom_peak_tflops=d.get("headroom_peak_tflops", 0.0), + headroom_arithmetic_intensity=d.get("headroom_arithmetic_intensity", 0.0), + headroom_measured_ai=d.get("headroom_measured_ai", 0.0), + headroom_shape_label=d.get("headroom_shape_label", ""), + headroom_M=d.get("headroom_M", 0), + headroom_N=d.get("headroom_N", 0), + headroom_K=d.get("headroom_K", 0), + headroom_bw_util_pct=d.get("headroom_bw_util_pct", 0.0), + headroom_compute_util_pct=d.get("headroom_compute_util_pct", 0.0), + cpp_code=d.get("cpp_code"), + profiled=d.get("profiled", False), + profiling_kernel_duration_us=d.get("profiling_kernel_duration_us", 0.0), + profiling_achieved_bw_gbps=d.get("profiling_achieved_bw_gbps", 0.0), + profiling_occupancy_pct=d.get("profiling_occupancy_pct", 0.0), + profiling_l2_cache_hit_pct=d.get("profiling_l2_cache_hit_pct", 0.0), + profiling_advice=d.get("profiling_advice"), ) for d in data ] diff --git a/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py b/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py index e2f0bcbb..c8ecd49b 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/orchestrator.py @@ -63,18 +63,57 @@ def run_with_requirements( req: Dict[str, Any] = json.loads(requirements_path.read_text(encoding="utf-8")) task_id = req.get("task_id", "task") - set_process_name("metainfer-orch-evolve-kernel") - + # Resolve dirs early — needed by both single-GPU and multi-GPU paths if state_dir is None: state_dir = Path.cwd() / "nodes" / "localhost" / ".metainfer" / "tasks" / task_id if workspace_dir is None: workspace_dir = Path.cwd() / "nodes" / "localhost" / "workspaces" / task_id paths = _task_subdirs(state_dir, workspace_dir) + # Copy requirements if needed target_req = paths["requirements"] if requirements_path.resolve() != target_req.resolve(): target_req.write_text(requirements_path.read_text(encoding="utf-8"), encoding="utf-8") + # Also copy reference kernel from kernel_file_path to workspace + kernel_path_str = req.get("kernel_file_path", "") + if kernel_path_str: + kernel_path = Path(kernel_path_str) + if kernel_path.is_file(): + ref_dir = workspace_dir / "reference" + ref_dir.mkdir(parents=True, exist_ok=True) + ref_path = ref_dir / "original_kernel.py" + if not ref_path.is_file(): + ref_path.write_text(kernel_path.read_text(encoding="utf-8"), encoding="utf-8") + + # Multi-GPU dispatch: one orchestrator spawns N GPU workers internally + multi_gpu = req.get("multi_gpu", "no") + if multi_gpu in ("All GPUs (auto)", "yes", "true", "1"): + set_process_name("metainfer-orch-evolve-kernel-multi") + print(f"[metainfer-evolve-kernel] MULTI-GPU mode") + print(f"[metainfer-evolve-kernel] task_id = {task_id}") + print(f"[metainfer-evolve-kernel] state dir = {state_dir}") + print(f"[metainfer-evolve-kernel] workspace dir = {workspace_dir}") + + from ._parallel import MultiGpuOrchestrator + write_pid_file(paths["pid_file"], task_id) + orch_multi = MultiGpuOrchestrator( + req=req, + state_dir=state_dir, + workspace_dir=workspace_dir, + claude_bin=claude_bin, + model=model, + permission_mode=permission_mode, + effort=effort, + ) + try: + orch_multi.run() + finally: + clear_pid_file(paths["pid_file"]) + return 0 + + set_process_name("metainfer-orch-evolve-kernel") + write_pid_file(paths["pid_file"], task_id) repo_root = _repo_root() diff --git a/metainfer/tasks/evolve_kernel/orchestrator/pipeline.py b/metainfer/tasks/evolve_kernel/orchestrator/pipeline.py index 874a3490..64b34f75 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/pipeline.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/pipeline.py @@ -31,7 +31,14 @@ from typing import Any, Dict, List, Optional, Tuple from . import phases as P +from .headroom import analyze_headroom, headroom_result_to_dict, headroom_dict_to_summary from .kernel_library import KernelEntry, KernelLibrary, MAX_LIBRARY_SIZE +from .profiling import ( + run_hipprof_profile, + profile_result_to_dict, + profile_to_advice, + profile_summary_for_storage, +) from .harness import ( build_correctness_harness_template, build_perf_harness_template, @@ -311,6 +318,11 @@ def _loop(self, is_resume: bool = False) -> None: if perf: phase_rec["perf"] = perf + # Persist phases to disk IMMEDIATELY so the failure log in the + # WebUI sees results even for intra-iteration retries (F→G→F + # loops where consume_iteration=False keeps the iteration open). + self.store.write_iteration(iter_num, iter_rec.to_dict()) + # Get next transition t = P.next_transition(phase, outcome) if t is None: @@ -601,6 +613,10 @@ def _do_optimize( failure_feedback=ctx.failure if ctx.failure and "correctness" in (ctx.failure or "").lower() else None, iteration=n, logs_dir=self._logs_dir_for(n), + headroom_bottleneck=ctx.selected_kernel.headroom_bottleneck or "", + headroom_bw_util_pct=ctx.selected_kernel.headroom_roofline_efficiency_pct, + headroom_compute_util_pct=ctx.selected_kernel.headroom_roofline_efficiency_pct, + headroom_pct=ctx.selected_kernel.headroom_pct, ), timeout=self.cfg.agent_timeout_s, resume_session_id=ctx.session_id, @@ -617,6 +633,28 @@ def _do_optimize( if not opt_path.exists(): return P.LOGIC_FAIL, None, "Agent did not produce optimized_kernel.py" + # HIP C++ mode: validate that the .cpp file exists and compiles + if self._get_optimizer_mode() == "hip_cpp": + cpp_path = iter_dir / "optimized_kernel.cpp" + if not cpp_path.exists(): + # Agent might have inlined C++ code in the Python wrapper, + # which is acceptable but not ideal. Log a warning. + self.store.append_timeline("hip_cpp_missing_source", { + "iteration": n, + "warning": "optimized_kernel.cpp not found; agent may have inlined C++ in .py", + }) + else: + # Validate HIP compilation + cpp_ok, cpp_error = self._validate_hip_kernel(cpp_path, n) + if not cpp_ok: + self.store.append_timeline("hip_compile_error", { + "iteration": n, + "error": cpp_error[:500], + }) + # Don't fail the phase — let G phase (correctness) catch + # runtime compilation errors. The harness will report + # import failures with details. + ctx.session_id = None return P.OK, None, None @@ -724,19 +762,104 @@ def _do_measure_perf( ctx.no_improvement_count += 1 return P.LOGIC_FAIL, None, f"Perf measurement failed: {perf_result.get('error', 'unknown')}" - # 2. Complexity evaluation + # Read optimized kernel code for headroom analysis and complexity evaluation opt_code = evolved_path.read_text(encoding="utf-8") + + # 1.5. Headroom analysis (roofline model) + headroom = analyze_headroom( + kernel_code=opt_code, + kernel_fn_name=ctx.kernel_fn_name, + exec_time_ms=exec_time_ms, + req=self.req, + ) + headroom_dict = headroom_result_to_dict(headroom) + headroom_summary = headroom_dict_to_summary(headroom_dict) + + self.store.append_timeline("headroom_analysis", { + "iteration": n, + "bottleneck": headroom.bottleneck, + "bw_util_pct": headroom_summary.get("bw_util_pct", 0.0), + "compute_util_pct": headroom_summary.get("compute_util_pct", 0.0), + "headroom_pct": headroom.headroom_pct, + }) + + # 1.6. Fine-grained profiling via hipprof (optional, controlled by requirements) + profile_result = None + profile_advice = "" + enable_profiling = self.req.get("enable_profiling", False) + if isinstance(enable_profiling, str): + enable_profiling = enable_profiling.lower() in ("true", "yes", "1") + if enable_profiling: + try: + profile_result = run_hipprof_profile( + kernel_script=evolved_path, + shape_args={"M": headroom.M, "N": headroom.N, "K": headroom.K}, + kernel_fn_name=ctx.kernel_fn_name, + output_dir=iter_dir / "profiling", + timeout_s=120, + ) + profile_advice = profile_to_advice( + profile_result, headroom.M, headroom.N, headroom.K, + ) + self.store.append_timeline("profiling_complete", { + "iteration": n, + "success": profile_result.success, + "kernel_duration_us": profile_result.kernel_duration_us, + "achieved_bw_gbps": profile_result.achieved_bw_total_gbps, + "occupancy_pct": profile_result.achieved_occupancy_pct, + }) + except Exception as e: + self.store.append_timeline("profiling_error", { + "iteration": n, + "error": str(e)[:500], + }) + + # 2. Complexity evaluation complexity = self._evaluate_complexity(opt_code, n, iter_dir) + # 2.5. In HIP mode, also capture the real .cpp kernel source + cpp_code: Optional[str] = None + optimizer_mode = self._get_optimizer_mode() + if optimizer_mode == "hip_cpp": + cpp_path = iter_dir / "optimized_kernel.cpp" + if cpp_path.exists(): + try: + cpp_code = cpp_path.read_text(encoding="utf-8") + except Exception: + pass + # 3. Create kernel entry and try to add to library entry = KernelEntry( id=str(uuid.uuid4()), code=opt_code, + cpp_code=cpp_code, exec_time_ms=exec_time_ms, complexity_score=complexity, combined_score=0.0, iteration_added=n, parent_id=ctx.selected_kernel.id if ctx.selected_kernel else None, + headroom_bottleneck=headroom.bottleneck, + headroom_roofline_efficiency_pct=headroom.roofline_efficiency_pct, + headroom_pct=headroom.headroom_pct, + headroom_p_max_tflops=headroom.p_max_tflops, + headroom_p_bw_roof_tflops=headroom.p_bandwidth_roof_tflops, + headroom_ai_ridge=headroom.ai_ridge, + headroom_suggestions_json=json.dumps(headroom.suggestions, ensure_ascii=False), + headroom_advice=headroom.optimization_advice, + headroom_achieved_bw_gbps=headroom.achieved_bw_gbps, + headroom_achieved_tflops=headroom.achieved_tflops, + headroom_peak_bw_gbps=headroom.peak_bw_gbps, + headroom_peak_tflops=headroom.peak_tflops, + headroom_arithmetic_intensity=headroom.arithmetic_intensity, + headroom_measured_ai=headroom.measured_ai, + headroom_shape_label=headroom.shape_label, + headroom_M=headroom.M, headroom_N=headroom.N, headroom_K=headroom.K, + profiled=(profile_result is not None and profile_result.success), + profiling_kernel_duration_us=profile_result.kernel_duration_us if profile_result else 0.0, + profiling_achieved_bw_gbps=profile_result.achieved_bw_total_gbps if profile_result else 0.0, + profiling_occupancy_pct=profile_result.achieved_occupancy_pct if profile_result else 0.0, + profiling_l2_cache_hit_pct=profile_result.l2_cache_hit_pct if profile_result else 0.0, + profiling_advice=profile_advice if profile_advice else None, ) entry.recompute_combined() @@ -747,10 +870,13 @@ def _do_measure_perf( ctx.current_complexity = complexity # Track optimization history - ctx.optimization_history.append( + history_line = ( f"Iter {n}: exec={exec_time_ms:.4f}ms, complexity={complexity:.2f}, " f"combined={entry.combined_score:.4f}, added={'YES' if added else 'NO'}" ) + if profile_advice: + history_line += f"\n Profile: {profile_advice[:300]}" + ctx.optimization_history.append(history_line) # Update best if exec_time_ms < ctx.best_exec_time_ms: @@ -770,6 +896,18 @@ def _do_measure_perf( shared_kernel.parent.mkdir(parents=True, exist_ok=True) shared_kernel.write_text(opt_code, encoding="utf-8") + # HIP C++ mode: also persist the .cpp source alongside the .py wrapper + if self._get_optimizer_mode() == "hip_cpp": + cpp_path = iter_dir / "optimized_kernel.cpp" + if cpp_path.exists(): + shared_cpp = self.cfg.workspace_dir / "optimized_kernels" / f"{entry.id}.cpp" + shared_cpp.write_text(cpp_path.read_text(encoding="utf-8")) + else: + self.store.append_timeline("hip_cpp_source_missing_on_save", { + "iteration": n, + "kernel_id": entry.id, + }) + perf_dict = { "exec_time_ms": exec_time_ms, "speedup": speedup, @@ -777,6 +915,10 @@ def _do_measure_perf( "combined_score": entry.combined_score, "added_to_library": added, "library_size": ctx.library.size, + "headroom_bottleneck": headroom_summary.get("bottleneck"), + "headroom_pct": headroom_summary.get("headroom_pct"), + "headroom_bw_util_pct": headroom_summary.get("bw_util_pct"), + "headroom_compute_util_pct": headroom_summary.get("compute_util_pct"), } return P.OK, perf_dict, None @@ -1075,6 +1217,42 @@ def _load_kernel_from_req(self, ctx: IterationContext) -> None: "Create a new task using the current form to provide a kernel file." ) + def _get_optimizer_mode(self) -> str: + """Read and normalize optimizer_mode from task requirements. + + Returns one of: 'triton', 'hip_cpp'. + """ + mode = self.req.get("optimizer_mode", "Triton (standard)") + if isinstance(mode, list): + mode = mode[0] if mode else "Triton (standard)" + if mode in ("HIP C++ (from scratch)", "hip_cpp", "hip"): + return "hip_cpp" + else: + return "triton" # default + + def _validate_hip_kernel( + self, cpp_path: Path, iteration: int, + ) -> Tuple[bool, str]: + """Run hipcc syntax check on a HIP C++ kernel file. + + Returns (success: bool, error_output: str). + """ + try: + proc = subprocess.run( + ["hipcc", "-fsyntax-only", "-x", "hip", str(cpp_path)], + capture_output=True, text=True, timeout=60, + ) + if proc.returncode == 0: + return True, "" + error = (proc.stderr or proc.stdout or "unknown compilation error")[:2000] + return False, error + except subprocess.TimeoutExpired: + return False, "hipcc syntax check timed out after 60s" + except FileNotFoundError: + return False, "hipcc not found in PATH — cannot validate HIP kernel" + except Exception as e: + return False, f"hipcc validation error: {e!r}" + def _resolve_max_iterations(self) -> int: from metainfer.orchestrator.requirements import req_field_int return req_field_int(self.req, "max_iterations", self.cfg.max_iterations) diff --git a/metainfer/tasks/evolve_kernel/orchestrator/profiling.py b/metainfer/tasks/evolve_kernel/orchestrator/profiling.py new file mode 100644 index 00000000..c72e55f1 --- /dev/null +++ b/metainfer/tasks/evolve_kernel/orchestrator/profiling.py @@ -0,0 +1,385 @@ +"""Fine-grained GPU profiling via hipprof/rocprof for kernel optimization. + +Provides hardware counter-level profiling that complements the roofline +model in headroom.py. hipprof gives actual achieved bandwidth, occupancy, +cache hit rates, and stall reasons — not just theoretical estimates. + +Integrated into Phase H (measure perf) when ``enable_profiling`` is set +in the task requirements. +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import tempfile +from dataclasses import dataclass, field, asdict +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + + +# =========================================================================== # +# Profiling result +# =========================================================================== # + + +@dataclass +class ProfileResult: + """Fine-grained profiling data from hipprof for one kernel execution.""" + + # Kernel identification + kernel_name: str = "" + kernel_duration_us: float = 0.0 + + # Bandwidth + achieved_bw_read_gbps: float = 0.0 + achieved_bw_write_gbps: float = 0.0 + achieved_bw_total_gbps: float = 0.0 + + # Occupancy + achieved_occupancy_pct: float = 0.0 + theoretical_occupancy_pct: float = 0.0 + + # Cache + l2_cache_hit_pct: float = 0.0 + + # Stalls (percentage of cycles) + stall_memory_pct: float = 0.0 + stall_dependency_pct: float = 0.0 + stall_sync_pct: float = 0.0 + stall_other_pct: float = 0.0 + + # Raw data + raw_stats: Dict[str, Any] = field(default_factory=dict) + trace_file: str = "" + success: bool = False + error: str = "" + + +# =========================================================================== # +# hipprof runner +# =========================================================================== # + + +_HIPPROF_BIN = "/opt/dtk/bin/hipprof" + + +def run_hipprof_profile( + kernel_script: Path, + shape_args: Dict[str, int], + kernel_fn_name: str = "matmul_int8", + output_dir: Optional[Path] = None, + timeout_s: int = 120, +) -> ProfileResult: + """Run hipprof on a kernel script and parse the output. + + Args: + kernel_script: Path to the .py file containing the kernel. + shape_args: Dict with M, N, K dimensions. + kernel_fn_name: Name of the kernel wrapper function. + output_dir: Where to write hipprof output files. + timeout_s: Timeout in seconds. + + Returns: + ProfileResult with extracted metrics. + """ + if not os.path.isfile(_HIPPROF_BIN): + return ProfileResult( + success=False, + error=f"hipprof not found at {_HIPPROF_BIN}", + ) + + if output_dir is None: + output_dir = Path(tempfile.mkdtemp(prefix="hipprof_")) + + output_dir.mkdir(parents=True, exist_ok=True) + output_prefix = str(output_dir / "hipprof") + + # Write a minimal runner script + M = shape_args.get("M", 4096) + N = shape_args.get("N", 4096) + K = shape_args.get("K", 4096) + + runner_script = output_dir / "_prof_runner.py" + runner_script.write_text(f'''"""Auto-generated hipprof profiling runner.""" +import sys +sys.path.insert(0, "{kernel_script.parent}") +import importlib.util + +spec = importlib.util.spec_from_file_location("_prof_kernel", "{kernel_script}") +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) +fn = getattr(mod, "{kernel_fn_name}") + +import torch +device = torch.device("cuda") + +# Generate inputs matching the target shape +a = torch.randint(-128, 127, ({M}, {K}), device=device, dtype=torch.int8) +a_scale = torch.randn({M}, 1, device=device, dtype=torch.float32).abs() / 127.0 +b = torch.randint(-128, 127, ({K}, {N}), device=device, dtype=torch.int8) +b_scale = torch.randn({N}, 1, device=device, dtype=torch.float32).abs() / 127.0 + +# Warmup +for _ in range(5): + fn(a, a_scale, b, b_scale, torch.bfloat16) + +torch.cuda.synchronize() + +# Timed run +out = fn(a, a_scale, b, b_scale, torch.bfloat16) +torch.cuda.synchronize() + +# Verify output is valid +assert out.shape == ({M}, {N}), f"Bad shape: {{out.shape}}" +assert not torch.isnan(out).any(), "NaN in output" +print("PROFILE_OK") +''') + + try: + proc = subprocess.run( + [ + _HIPPROF_BIN, + "--stats", + "--hip-trace", + "-o", output_prefix, + "python3", str(runner_script), + ], + capture_output=True, text=True, + timeout=timeout_s, + env={**os.environ, "PYTHONUNBUFFERED": "1"}, + ) + except subprocess.TimeoutExpired: + return ProfileResult( + success=False, + error=f"hipprof timed out after {timeout_s}s", + ) + except Exception as e: + return ProfileResult( + success=False, + error=f"Failed to run hipprof: {e!r}", + ) + + stdout = proc.stdout or "" + stderr = proc.stderr or "" + + # Check if hipprof produced kernel CSV (reliable indicator of success) + kernel_csv = output_dir / "hipprof.hipkernel.csv" + if not kernel_csv.is_file(): + return ProfileResult( + success=False, + error=f"hipprof did not produce output CSV. stderr: {stderr[-500:]!r}", + ) + + # Parse hipprof CSV output files + result = _parse_hipprof_stats(output_dir) + + # Find trace file + trace_file = "" + for f in output_dir.iterdir(): + if f.suffix in (".csv", ".db", ".json") and f.name.startswith("hipprof"): + trace_file = str(f) + break + + result.trace_file = trace_file + result.success = True + result.raw_stats["runner_script"] = str(runner_script) + result.raw_stats["stdout_tail"] = stdout[-2000:] + result.raw_stats["stderr_tail"] = stderr[-2000:] + + return result + + +# =========================================================================== # +# hipprof output parsing +# =========================================================================== # + + +def _parse_hipprof_stats(output_dir: Path) -> ProfileResult: + """Parse hipprof CSV output files into structured ProfileResult. + + hipprof produces two key CSV files: + - hipprof.hipkernel.csv: kernel Name, Calls, TotalDurationNs, AverageNs, Percentage + - hipprof.hiptrace.csv: HIP API calls with durations (used for BW estimation) + + We extract the dominant compute kernel's duration and derive bandwidth from + memory copy sizes divided by transfer time. + """ + import csv as _csv + + result = ProfileResult() + + kernel_csv = output_dir / "hipprof.hipkernel.csv" + trace_csv = output_dir / "hipprof.hiptrace.csv" + + # ---- Parse kernel CSV ---- + if kernel_csv.is_file(): + try: + with open(kernel_csv, newline="", encoding="utf-8") as fh: + reader = _csv.DictReader(fh) + kernel_rows = list(reader) + + # Filter to real compute kernels (exclude torch/pytorch internals, + # elementwise, reduce, random, memset, etc.) + _IGNORE_PATTERNS = ( + "at::native::vectorized_elementwise", + "at::native::reduce_kernel", + "at::native::distribution_", + "at::native::Bitwise", + "at::native::Fill", + "at::native::AUnary", + "at::native::BUnary", + "at::native::CUDA_tensor_apply", + "at::native::tensor_kernel", + "at::native::unrolled_elementwise", + "hipMemset", + "memset", + "Total", + ) + + candidates: List[Tuple[str, float, float]] = [] # (name, avg_ns, pct) + for row in kernel_rows: + name = row.get("Name", "") + if not name or any(p in name for p in _IGNORE_PATTERNS): + continue + try: + avg_ns = float(row.get("AverageNs", "0").replace(",", "")) + pct = float(row.get("Percentage", "0").replace(",", "")) + except (ValueError, KeyError): + continue + if avg_ns > 0: + candidates.append((name, avg_ns, pct)) + + if candidates: + # Pick the kernel with highest percentage (dominant compute) + candidates.sort(key=lambda x: x[2], reverse=True) + result.kernel_name = candidates[0][0] + result.kernel_duration_us = candidates[0][1] / 1000.0 + + result.raw_stats["kernel_csv"] = [ + {"name": n, "avg_ns": a, "pct": p} for n, a, p in candidates[:10] + ] + except Exception as e: + result.raw_stats["kernel_csv_error"] = str(e)[:500] + + # ---- Parse trace CSV for bandwidth estimation ---- + if trace_csv.is_file(): + try: + with open(trace_csv, newline="", encoding="utf-8") as fh: + reader = _csv.DictReader(fh) + trace_rows = list(reader) + + # Sum memory copy durations (hipMemcpyWithStream, hipMemcpy) + copy_total_ns = 0.0 + for row in trace_rows: + name = row.get("Name", "") + if name in ("hipMemcpyWithStream", "hipMemcpy", "hipMemcpyAsync"): + try: + copy_total_ns += float(row.get("TotalDurationNs", "0").replace(",", "")) + except (ValueError, KeyError): + pass + + # Bandwidth estimation: if we know the data size from the trace, + # compute achieved BW. Use the kernel's known data footprint. + # For now, store copy time so callers can compute BW with shape info. + result.raw_stats["copy_total_ns"] = copy_total_ns + + # Also extract hipLaunchKernel total + for row in trace_rows: + if row.get("Name") == "hipLaunchKernel": + try: + launch_ns = float(row.get("TotalDurationNs", "0").replace(",", "")) + result.raw_stats["launch_total_ns"] = launch_ns + except (ValueError, KeyError): + pass + break + except Exception as e: + result.raw_stats["trace_csv_error"] = str(e)[:500] + + # Mark as successful if we found kernel data + if result.kernel_duration_us > 0: + result.success = True + + return result + + +# =========================================================================== # +# Profile summary for optimizer feedback +# =========================================================================== # + + +def profile_result_to_dict(result: ProfileResult) -> Dict[str, Any]: + """Serialize ProfileResult to a JSON-serializable dict.""" + d = asdict(result) + # Truncate raw output for storage + if "raw_stats" in d and isinstance(d.get("raw_stats"), dict): + raw = d["raw_stats"] + if "output" in raw and isinstance(raw["output"], str): + raw["output"] = raw["output"][-2000:] # keep last 2K + return d + + +def profile_to_advice(result: ProfileResult, M: int, N: int, K: int) -> str: + """Generate optimization advice from profiling results. + + Returns a paragraph suitable for the optimizer agent's context. + """ + if not result.success and result.kernel_duration_us <= 0: + return f"(Profiling failed: {result.error})" + + parts = [f"**hipprof profile for ({M}×{N}×{K}):**"] + + if result.kernel_duration_us > 0: + parts.append(f"Kernel duration: {result.kernel_duration_us:.1f} µs") + # Estimate achieved BW from data footprint + # int8 w8a8: A=(M,K) int8, a_scale=(M,1) fp32, B=(K,N) int8, b_scale=(N,1) fp32, out=(M,N) bf16 + bytes_read = (M * K * 1) + (M * 1 * 4) + (K * N * 1) + (N * 1 * 4) + bytes_write = M * N * 2 # bf16 output + total_bytes = bytes_read + bytes_write + if result.kernel_duration_us > 0: + achieved_bw = total_bytes / (result.kernel_duration_us / 1e6) / 1e9 # GB/s + parts.append(f"Estimated BW: {achieved_bw:.1f} GB/s ({total_bytes/1024:.0f} KiB moved)") + + if result.achieved_bw_total_gbps > 0: + parts.append(f"Measured BW: {result.achieved_bw_total_gbps:.1f} GB/s") + peak_bw = 700.0 # gfx928 + util = result.achieved_bw_total_gbps / peak_bw * 100 + parts.append(f"Achieved HBM BW: {result.achieved_bw_total_gbps:.1f} GB/s ({util:.0f}% of {peak_bw:.0f} GB/s peak)") + + if result.achieved_occupancy_pct > 0: + parts.append(f"Occupancy: {result.achieved_occupancy_pct:.0f}%") + + if result.l2_cache_hit_pct > 0: + parts.append(f"L2 cache hit: {result.l2_cache_hit_pct:.0f}%") + + # Bottleneck-specific advice + if result.achieved_bw_total_gbps > 0 and result.achieved_bw_total_gbps < 200: + parts.append("→ CRITICAL: Bandwidth is very low. Focus on memory coalescing and vectorized loads.") + elif result.achieved_occupancy_pct > 0 and result.achieved_occupancy_pct < 40: + parts.append("→ Occupancy is low. Reduce register pressure or increase tile size.") + elif result.l2_cache_hit_pct > 0 and result.l2_cache_hit_pct < 30: + parts.append("→ Poor cache hit rate. Improve data reuse via larger tiles or shared memory caching.") + + return "\n".join(parts) + + +# =========================================================================== # +# Integration helper: generate profiling info for KernelEntry +# =========================================================================== # + + +def profile_summary_for_storage(result: ProfileResult) -> Dict[str, Any]: + """Extract compact profiling summary for storage in kernel_library.json.""" + return { + "profiled": result.success, + "kernel_duration_us": round(result.kernel_duration_us, 2), + "achieved_bw_gbps": round(result.achieved_bw_total_gbps, 2), + "occupancy_pct": round(result.achieved_occupancy_pct, 1), + "l2_cache_hit_pct": round(result.l2_cache_hit_pct, 1), + "stall_memory_pct": round(result.stall_memory_pct, 1), + "stall_dependency_pct": round(result.stall_dependency_pct, 1), + "kernel_name": result.kernel_name, + "error": result.error if not result.success else "", + } diff --git a/metainfer/tasks/evolve_kernel/orchestrator/prompts.py b/metainfer/tasks/evolve_kernel/orchestrator/prompts.py index 4bfd2cab..107ef922 100644 --- a/metainfer/tasks/evolve_kernel/orchestrator/prompts.py +++ b/metainfer/tasks/evolve_kernel/orchestrator/prompts.py @@ -272,6 +272,10 @@ def optimize_kernel_prompt( failure_feedback: Optional[str] = None, iteration: int = 0, logs_dir: Optional[Path] = None, + headroom_bottleneck: str = "", + headroom_bw_util_pct: float = 0.0, + headroom_compute_util_pct: float = 0.0, + headroom_pct: float = 0.0, ) -> str: fb = "" if failure_feedback: @@ -291,6 +295,125 @@ def optimize_kernel_prompt( ``` """ + # Build headroom context for the optimizer to guide escalation + headroom_section = "" + if headroom_bottleneck: + headroom_blk = ( + f"## Headroom Analysis\n" + f"- **Bottleneck:** {headroom_bottleneck}\n" + f"- **BW usage:** {headroom_bw_util_pct:.0f}% of peak HBM\n" + f"- **Compute usage:** {headroom_compute_util_pct:.0f}% of peak TFLOPS\n" + f"- **Estimated remaining headroom:** {headroom_pct:.0f}%\n" + ) + if headroom_bottleneck == "near_optimal" or headroom_pct < 10: + headroom_blk += ( + "\n**⚠ CRITICAL: This kernel is at the hardware roofline.** " + "Further Triton-level tuning is unlikely to yield meaningful gains. " + "You MUST attempt a radically different approach:\n" + "- Rewrite the entire kernel as a hand-tuned HIP C++ kernel with inline GCN assembly\n" + "- Drop Triton entirely — direct HIP gives control over register allocation, " + "shared memory banks, and instruction scheduling that Triton cannot match\n" + "- Output BOTH `optimized_kernel.cpp` (HIP kernel) AND `optimized_kernel.py` " + "(Python wrapper with same function signature for harness compatibility)\n" + ) + elif headroom_bw_util_pct < 30 and headroom_compute_util_pct < 30: + headroom_blk += ( + "\n**⚠ Kernel is inefficient.** Both bandwidth and compute utilization are low. " + "Consider rewriting in HIP C++ with direct register control and SMEM management " + "to overcome the Triton codegen overhead.\n" + ) + headroom_section = headroom_blk + + mode = req.get("optimizer_mode", "Triton (standard)") + if isinstance(mode, list): + mode = mode[0] if mode else "Triton (standard)" + # Normalize: form sends display labels, not internal keys + if mode in ("HIP C++ (from scratch)", "hip_cpp", "hip"): + mode = "hip_cpp" + else: + mode = "triton" # default: Triton + + return _build_optimizer_prompt( + mode=mode, + iteration=iteration, + original_kernel_code=original_kernel_code, + current_kernel_code=current_kernel_code, + current_exec_time_ms=current_exec_time_ms, + current_complexity=current_complexity, + best_exec_time_ms=best_exec_time_ms, + best_kernel_code=best_kernel_code, + best_section=best_section, + failure_block=fb, + optimization_history=optimization_history, + iter_dir=iter_dir, + kernel_fn_name=kernel_fn_name, + headroom_section=headroom_section, + ) + + +def _build_optimizer_prompt( + mode: str, + iteration: int, + original_kernel_code: str, + current_kernel_code: str, + current_exec_time_ms: float, + current_complexity: float, + best_exec_time_ms: float, + best_kernel_code: str, + best_section: str, + failure_block: str, + optimization_history: str, + iter_dir: Path, + kernel_fn_name: str, + headroom_section: str = "", +) -> str: + if mode == "hip_cpp": + return _hip_cpp_optimizer_prompt( + iteration=iteration, + original_kernel_code=original_kernel_code, + current_kernel_code=current_kernel_code, + current_exec_time_ms=current_exec_time_ms, + best_section=best_section, + failure_block=failure_block, + optimization_history=optimization_history, + iter_dir=iter_dir, + kernel_fn_name=kernel_fn_name, + headroom_section=headroom_section, + ) + else: # triton (default) + return _triton_optimizer_prompt( + iteration=iteration, + original_kernel_code=original_kernel_code, + current_kernel_code=current_kernel_code, + current_exec_time_ms=current_exec_time_ms, + current_complexity=current_complexity, + best_section=best_section, + failure_block=failure_block, + optimization_history=optimization_history, + iter_dir=iter_dir, + kernel_fn_name=kernel_fn_name, + headroom_section=headroom_section, + ) + + +# ========================================================================== # +# Mode: Triton (standard — default) +# ========================================================================== # + + +def _triton_optimizer_prompt( + iteration: int, + original_kernel_code: str, + current_kernel_code: str, + current_exec_time_ms: float, + current_complexity: float, + best_section: str, + failure_block: str, + optimization_history: str, + iter_dir: Path, + kernel_fn_name: str, + headroom_section: str = "", +) -> str: return f"""You are the **KERNEL OPTIMIZER** for GPU kernel optimization, iteration #{iteration}. # Goal @@ -306,34 +429,44 @@ def optimize_kernel_prompt( {current_kernel_code} ``` {best_section} -{fb} +{failure_block} +{headroom_section} # Optimization History {optimization_history} -# Optimization Guide (Triton-specific, for DCU/AMD GPU) -1. **Tile size tuning**: Adjust BLOCK_SIZE_M/N/K. Larger tiles = more parallelism but more register pressure. -2. **Thread coarsening**: Have each thread compute multiple output elements. -3. **Memory coalescing**: Ensure adjacent threads access adjacent memory. -4. **Software pipelining**: Use `num_stages` > 2 to overlap compute and memory. -5. **Warp count**: Adjust `num_warps` to balance occupancy vs register usage. -6. **Precision**: Use `tl.float16` or `tl.bfloat16` compute types where safe. -7. **Loop ordering**: Reorder loops for better memory access patterns. -8. **Reduce shared memory**: Use less SMEM to increase occupancy. +# Optimization Strategy: MEMORY-FIRST, THEN COMPUTE + +Most GPU kernels are memory-bound. Fix memory access patterns FIRST, +then fill idle compute slots SECOND. For AMD DCU gfx928 (warp_size=64, 120 CUs). + +## Phase 1 — Memory Access Optimization (do this first!) +1. **Coalesce global memory**: Adjacent threads in a warp → adjacent addresses. +2. **Eliminate shared memory bank conflicts**: Pad SMEM to avoid same-bank collisions. +3. **Maximize data reuse**: Larger BLOCK_SIZE_M/N increases FLOPs/byte. +4. **Vectorized loads**: 128-bit loads (4×32-bit) reduce load count 4×. +5. **Double-buffering**: num_stages >= 2 to overlap load with compute. +6. **Prefetching**: Load first iteration before the loop. + +## Phase 2 — Compute Optimization (after memory is fixed) +7. **Dual-issue**: gfx928 dual-issues scalar + vector in same cycle. +8. **ILP**: Process 2-4 output elements per thread. +9. **Reduce conversions**: Fuse dot→int32→fp32→scale path. +10. **Loop unrolling**: tl.constexpr with EVEN_K for compile-time unrolling. +11. **Warp count**: Adjust `num_warps` (4/8) to balance occupancy vs registers. +12. **Reduced BLOCK_SIZE_K**: Smaller K tiles (64/128) → more wavefronts in flight. # Constraints - MUST maintain the same function signature as the original - MUST produce numerically equivalent results (within 1e-3 tolerance) -- MUST remain a Triton kernel (no CUDA/HIP inline assembly) -- Test on the ACTUAL hardware — detect GPU type before optimizing -- Write the COMPLETE optimized kernel file (not a diff) +- Pure Triton — NO inline assembly, NO CUDA/HIP code +- Target hardware: AMD DCU gfx928, warp_size=64, 120 CUs +- Write the COMPLETE optimized kernel file # Deliverable Write the COMPLETE optimized kernel to `{iter_dir}/optimized_kernel.py`. -Include the FULL kernel with all imports, the `@triton.jit` decorated function, and any helper functions. ```python -# optimized_kernel.py — complete, runnable Triton kernel file import torch import triton import triton.language as tl @@ -343,8 +476,156 @@ def {kernel_fn_name}(...): # Your optimized implementation ... ``` +Write ONLY the optimized kernel file. +""" + + +# ========================================================================== # +# Mode: HIP C++ (full rewrite) +# ========================================================================== # + + +def _hip_cpp_optimizer_prompt( + iteration: int, + original_kernel_code: str, + current_kernel_code: str, + current_exec_time_ms: float, + best_section: str, + failure_block: str, + optimization_history: str, + iter_dir: Path, + kernel_fn_name: str, + headroom_section: str = "", +) -> str: + return f"""You are the **KERNEL OPTIMIZER** for GPU kernel optimization, iteration #{iteration}. + +# Goal +Rewrite the kernel as a HIP C++ kernel with a Python wrapper for MAXIMUM performance +on AMD DCU gfx928. HIP C++ gives direct control over register allocation, shared memory, +and instruction scheduling — no Triton codegen overhead. + +# Original (Reference) Triton Kernel — analyze its algorithm, inputs, outputs +```python +{original_kernel_code} +``` + +# Current Implementation (exec_time={current_exec_time_ms:.4f}ms) +```python +{current_kernel_code} +``` +{best_section} +{failure_block} +{headroom_section} + +# Optimization History +{optimization_history} + +# Step 1: Analyze the Triton Kernel +Before writing any HIP code, carefully analyze the original kernel to identify: +1. **Algorithm**: What computation does this kernel do? (GEMM, attention, reduction, convolution, element-wise, etc.) +2. **Input tensors**: Names, shapes, dtypes — trace through the function body to understand each argument +3. **Output tensor**: What dtype and shape does the kernel produce? +4. **Key computation pattern**: The inner loop structure, accumulation pattern, and any non-trivial logic +5. **Function signature**: The exact Python function signature — you MUST replicate this exactly in the wrapper + +# Step 2: Map to HIP C++ for gfx928 (warp_size=64, 120 CUs, peak HBM ~1.2 TB/s) + +## General HIP Kernel Structure +```cpp +#include +#include +#include + +// Launch grid configuration +// gridDim.x = ceil(N / BLOCK_N), gridDim.y = ceil(M / BLOCK_M) +// blockDim.x = warp_size * num_warps_per_block +// Typical: blockDim = 256 (4 warps), BLOCK_M=128, BLOCK_N=128 + +__global__ void kernel_name( + // Map each Triton tensor pointer to a HIP pointer + // Use appropriate types: float, __half, __hip_bfloat16, int8_t, etc. +) {{ + // Thread/work-item indexing + int tid = threadIdx.x; + int wid = tid / warpSize; // warp index within block + int lane = tid % warpSize; // lane within warp (0-63) + + // Block-level output tile: blockIdx.y → M tile, blockIdx.x → N tile + // Each thread computes a fragment of the output tile + + // Shared memory declarations (use extern __shared__ or fixed size) + // Accumulator in registers (fp32 or int32 depending on algorithm) + + // Main loop over reduction dimension + for (int k = 0; k < K; k += BLOCK_K) {{ + // 1. Cooperative global→shared load (vectorized, coalesced) + // 2. __syncthreads() + // 3. Compute tile operation + // 4. __syncthreads() + }} + + // Epilogue: activation, scaling, type conversion, etc. + // Cooperative store to global memory (coalesced) +}} +``` + +## Key Optimizations for gfx928 +1. **Vectorized loads**: Use `__builtin_amdgcn_global_load_dwordx4` or `float4`/`uint4` for 16-byte loads — critical for HBM BW utilization +2. **Shared memory banking**: Pad SMEM arrays by 1 element to avoid bank conflicts on gfx928's 32-bank LDS +3. **Register blocking**: Each thread computes an 4×4 or 8×8 output tile to amortize load instructions — use #pragma unroll +4. **Double buffering**: Two SMEM buffers with async copy (`__builtin_amdgcn_sched_barrier`) to overlap loads and compute +5. **Assembly intrinsics**: `__builtin_amdgcn_mfma_*` for matrix ops, `__builtin_amdgcn_ds_*` for LDS ops, `__builtin_amdgcn_s_barrier` for fine-grained sync +6. **Occupancy management**: Target 4+ waves/CU (waves_per_cu = (120 * 4) / num_workgroups). Keep VGPR ≤ 128, LDS ≤ 32KB per block +7. **Coalesced memory access**: Ensure adjacent threads access adjacent addresses for all global load/store operations +8. **ILP (Instruction-Level Parallelism)**: Interleave independent operations to hide latency without consuming VGPRs + +# Step 3: Write the TWO deliverable files + +## File 1: `{iter_dir}/optimized_kernel.cpp` +Complete, self-contained HIP C++ kernel: +- `#include ` and any needed fp16/bf16 headers +- `__global__` kernel function with the actual algorithm from Step 1 +- Host-side launch function (`void launch_kernel(...)`) that computes grid/block dims and calls the kernel +- HIP error checking wrapper (`HIP_CHECK(call)`) +- Use `extern "C"` for the host function so it can be called from Python via ctypes + +## File 2: `{iter_dir}/optimized_kernel.py` +Python wrapper using `torch.utils.cpp_extension.load_inline()`: +```python +import os +import torch +import torch.utils.cpp_extension + +# JIT-compile the HIP kernel at import time +_src_dir = os.path.dirname(os.path.abspath(__file__)) +_cpp_path = os.path.join(_src_dir, "optimized_kernel.cpp") +_cpp_source = open(_cpp_path).read() + +_hip_module = torch.utils.cpp_extension.load_inline( + name="evolved_hip_kernel", + cpp_sources=[_cpp_source], + functions=["launch_kernel"], # the host-side C function name + extra_cflags=["-O3", "-ffast-math"], + with_cuda=True, # PyTorch uses CUDA/HIP interchangeably + verbose=False, +) + +# IMPORTANT: Replicate the ORIGINAL kernel's function name and signature EXACTLY. +# The correctness/perf harness calls this function by name. +# Convert PyTorch tensors to raw pointers, call launch_kernel, return output. +def {kernel_fn_name}(...): + # ... same signature as original ... + raise NotImplementedError("Replace with actual implementation") +``` + +## Critical Constraints +- **SAME function name**: The Python wrapper MUST export a function named `{kernel_fn_name}` — the test harness imports it by this name +- **SAME function signature**: Arguments, return type, and semantics MUST match the original kernel — the harness calls it with the same inputs +- **Numerical equivalence**: Results must match the reference within 1e-3 relative tolerance +- **Self-contained**: The `.py` file reads the `.cpp` file at import time and JIT-compiles it — no manual build step +- **Error handling**: Check all HIP calls with `hipGetErrorString` and raise Python RuntimeError on failure -Write ONLY the optimized kernel file. Do NOT modify any other files. +Write BOTH files now. Do NOT omit either file. """ diff --git a/metainfer/tasks/evolve_kernel/server/_multi_gpu.py b/metainfer/tasks/evolve_kernel/server/_multi_gpu.py new file mode 100644 index 00000000..7d75a427 --- /dev/null +++ b/metainfer/tasks/evolve_kernel/server/_multi_gpu.py @@ -0,0 +1,362 @@ +"""Multi-GPU parallel task launcher for evolve-kernel. + +When the user selects "All GPUs (auto)" in the form, this module: + 1. Detects available GPUs + 2. Splits target shapes from extra_notes across GPUs + 3. Creates N child tasks (one per GPU), each optimizing a subset of shapes + 4. Creates a virtual "parent" task that aggregates results + +Child tasks are normal evolve-kernel tasks, each with: + - gpu_device: "0" / "1" / "2" / "3" + - multi_gpu_parent: parent task ID + - extra_notes: subset of shapes for this GPU + +The parent task has a requirements.json with: + - multi_gpu_children: ["child-id-1", "child-id-2", ...] + - multi_gpu_mode: true +""" + +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from metainfer.server import launcher as _launcher +from metainfer.server import paths as _paths +from metainfer.server.registry import get as _get_web_plugin +from ._shape_bench import parse_shapes_from_extra_notes, ShapeSpec + + +# =========================================================================== # +# GPU detection +# =========================================================================== # + + +def detect_gpus() -> int: + """Detect number of available GPUs. Tries torch, then rocm-smi, then env.""" + try: + import torch + return torch.cuda.device_count() + except Exception: + pass + try: + import subprocess + result = subprocess.run( + ["rocm-smi", "--showid", "--csv"], + capture_output=True, text=True, timeout=5, + ) + lines = [l for l in result.stdout.strip().split("\n") if l.strip() and "GPU" not in l] + return len(lines) if lines else 1 + except Exception: + pass + cu = _os_environ("CUDA_VISIBLE_DEVICES") + if cu: + return max(1, cu.count(",") + 1) + return 1 + + +def _os_environ(key: str) -> str: + import os + return os.environ.get(key, "") + + +# =========================================================================== # +# Shape splitting +# =========================================================================== # + + +def split_shapes_for_gpus( + extra_notes: str, + num_gpus: int, +) -> List[Tuple[str, List[ShapeSpec]]]: + """Split parsed shapes across GPUs. + + Strategy: + 1. Parse shapes from extra_notes + 2. If multiple distinct shapes (different label), assign roughly evenly + 3. If only one shape, each GPU gets the same shape (competing optimization) + 4. Returns list of (gpu_label, shapes) tuples + """ + shapes = parse_shapes_from_extra_notes(extra_notes) + if not shapes: + return [] + + # If only one shape, replicate across all GPUs + if len(shapes) == 1: + return [ + ("GPU {}".format(i), list(shapes)) + for i in range(num_gpus) + ] + + # Multiple shapes: distribute evenly + groups: List[List[ShapeSpec]] = [[] for _ in range(num_gpus)] + for i, s in enumerate(shapes): + groups[i % num_gpus].append(s) + + return [ + ("GPU {}".format(i), grp) + for i, grp in enumerate(groups) + if grp # skip empty groups + ] + + +def shapes_to_notes(shapes: List[ShapeSpec]) -> str: + """Convert ShapeSpec list back to extra_notes text.""" + lines = [] + for s in shapes: + m_str = ",".join(str(m) for m in s.M_values) + lines.append( + "{}: M={} (M, {}) @ ({}, {})".format( + s.label, m_str, s.K, s.K, s.N, + ) + ) + return "\n".join(lines) + + +# =========================================================================== # +# Child task creation +# =========================================================================== # + + +def create_child_task( + parent_id: str, + gpu_index: int, + child_shapes: List[ShapeSpec], + base_requirements: Dict[str, Any], + label_suffix: str, +) -> str: + """Create one child task pinned to a specific GPU. + + Args: + parent_id: The parent task ID for linkage. + gpu_index: Which GPU index to bind to (0, 1, 2, 3). + child_shapes: Subset of shapes for this child. + base_requirements: The user's original form answers. + label_suffix: Display label suffix, e.g. "GPU 0". + + Returns: + child task ID. + """ + import metainfer.server.tasks as _tasks + + child_id = _tasks.gen_task_id("evolve-kernel", "{}-gpu{}".format(parent_id, gpu_index)) + + sd = _paths.task_dir(child_id) + sd.mkdir(parents=True, exist_ok=True) + wd = _paths.workspace_dir(child_id) + wd.mkdir(parents=True, exist_ok=True) + + # Build child requirements: copy base but override shapes and GPU. + # Strip multi_gpu so child tasks don't recursively trigger more launches. + child_req = dict(base_requirements) + child_req.pop("multi_gpu", None) + child_req.pop("gpu_count", None) + child_req["task_id"] = child_id + child_req["gpu_device"] = str(gpu_index) + child_req["multi_gpu_parent"] = parent_id + child_req["extra_notes"] = shapes_to_notes(child_shapes) + + plugin = _get_web_plugin("evolve-kernel") + default_label = (plugin.label if plugin else "") or "evolve-kernel" + label = "{} — {}".format(default_label, label_suffix) + + # Register child task + entry = _tasks.TaskEntry( + id=child_id, type="evolve-kernel", label=label, + state_dir=str(sd), workspace_dir=str(wd), + created_at=time.time(), launcher="local", + ) + _tasks.add_task(entry) + + # Spawn child orchestrator + child_launcher = _launcher.get_default_launcher() + child_launcher.start(child_id, child_req, sd, wd) + + return child_id + + +def launch_multi_gpu( + base_requirements: Dict[str, Any], + num_gpus: Optional[int] = None, + label: str = "", +) -> Dict[str, Any]: + """Create a multi-GPU optimization batch. + + Args: + base_requirements: The user's form answers (as if creating one task). + num_gpus: Number of GPUs to use. Auto-detected if None. + label: User-provided label for the parent task. + + Returns: + {"parent_id": "...", "children": ["child1", "child2", ...]} + """ + import metainfer.server.tasks as _tasks + + if num_gpus is None: + num_gpus = detect_gpus() + num_gpus = max(1, min(num_gpus, 8)) # clamp 1-8 + + extra_notes = base_requirements.get("extra_notes", "") + + # Split shapes + groups = split_shapes_for_gpus(extra_notes, num_gpus) + if not groups: + raise ValueError( + "No shapes found in extra_notes. Cannot split across GPUs. " + "Add shapes using the format: shape_name: (M, K) @ (K, N)" + ) + + # Create parent task ID + parent_id = _tasks.gen_task_id("evolve-kernel", "{}-multi".format(label or "kernel")) + + # Create parent task entries (a virtual task — no orchestrator) + plugin = _get_web_plugin("evolve-kernel") + default_label = (plugin.label if plugin else "") or "evolve-kernel" + parent_label = "{} [{} GPUs]".format(label or default_label, num_gpus) + + sd = _paths.task_dir(parent_id) + sd.mkdir(parents=True, exist_ok=True) + wd = _paths.workspace_dir(parent_id) + wd.mkdir(parents=True, exist_ok=True) + + # Write parent requirements + parent_req = dict(base_requirements) + parent_req["task_id"] = parent_id + parent_req["multi_gpu_mode"] = True + parent_req["multi_gpu_children"] = [] + req_path = sd / "requirements.json" + req_path.write_text(json.dumps(parent_req, indent=2), encoding="utf-8") + + # Write parent run.json (virtual task, marked finished immediately) + run_data = { + "task_id": parent_id, + "current_phase": "idle", + "finished": False, + "final_status": None, + } + (sd / "run.json").write_text(json.dumps(run_data, indent=2), encoding="utf-8") + + # Register parent + entry = _tasks.TaskEntry( + id=parent_id, type="evolve-kernel", label=parent_label, + state_dir=str(sd), workspace_dir=str(wd), + created_at=time.time(), launcher="local", + ) + _tasks.add_task(entry) + + # Create children + children: List[str] = [] + for gpu_label, child_shapes in groups: + gpu_idx = int(gpu_label.replace("GPU ", "")) + try: + child_id = create_child_task( + parent_id=parent_id, + gpu_index=gpu_idx, + child_shapes=child_shapes, + base_requirements=base_requirements, + label_suffix=gpu_label, + ) + children.append(child_id) + except Exception: + # Continue creating other children + import traceback + traceback.print_exc() + + # Update parent with child IDs + parent_req["multi_gpu_children"] = children + req_path.write_text(json.dumps(parent_req, indent=2), encoding="utf-8") + + return {"parent_id": parent_id, "children": children, "num_gpus": num_gpus} + + +# =========================================================================== # +# Aggregated results +# =========================================================================== # + + +def get_child_task_ids(state_dir: Path) -> List[str]: + """Read the child task IDs from a parent task's requirements.""" + req_path = state_dir / "requirements.json" + if not req_path.is_file(): + return [] + try: + req = json.loads(req_path.read_text(encoding="utf-8")) + if req.get("multi_gpu_mode"): + return req.get("multi_gpu_children", []) + except Exception: + pass + return [] + + +def aggregate_status(task_ids: List[str]) -> List[Dict[str, Any]]: + """Collect status for each child task.""" + launcher = _launcher.get_default_launcher() + results = [] + for tid in task_ids: + entry_data = None + from metainfer.server import tasks as _tasks + entry = _tasks.get_task(tid) + if entry: + entry_data = { + "id": entry.id, + "label": entry.label, + "type": entry.type, + "state_dir": entry.state_dir, + "workspace_dir": entry.workspace_dir, + } + status = launcher.status(tid) + results.append({ + "task_id": tid, + "entry": entry_data, + "running": status.running, + "pid": status.pid, + }) + return results + + +def aggregate_shape_benchmarks(workspace_dirs: List[str]) -> List[Dict[str, Any]]: + """Aggregate shape benchmark results from all child tasks. + + Reads shape_bench.json from each child's workspace, merges into one table + with a 'source' column indicating which task produced each row. + """ + all_results: List[Dict[str, Any]] = [] + for wd_str in workspace_dirs: + bench_path = Path(wd_str) / "shape_bench.json" + if not bench_path.is_file(): + continue + try: + data = json.loads(bench_path.read_text(encoding="utf-8")) + except Exception: + continue + task_id = data.get("best_kernel_id", "?")[:8] + for r in data.get("results", []): + r["source"] = task_id + all_results.append(r) + + # Sort: by shape_label, then by M + all_results.sort(key=lambda r: (r.get("shape_label", ""), r.get("M", 0))) + return all_results + + +def aggregate_best_kernels(workspace_dirs: List[str]) -> List[Dict[str, Any]]: + """Collect the best kernel from each child's library.""" + kernels = [] + for wd_str in workspace_dirs: + lib_path = Path(wd_str) / "kernel_library.json" + if not lib_path.is_file(): + continue + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + except Exception: + continue + if not lib: + continue + # Sort by exec_time_ms + lib.sort(key=lambda k: k.get("exec_time_ms", float("inf"))) + best = lib[0] + best["source_workspace"] = str(Path(wd_str).parent.name) + kernels.append(best) + return kernels diff --git a/metainfer/tasks/evolve_kernel/server/_shape_bench.py b/metainfer/tasks/evolve_kernel/server/_shape_bench.py new file mode 100644 index 00000000..9521b191 --- /dev/null +++ b/metainfer/tasks/evolve_kernel/server/_shape_bench.py @@ -0,0 +1,333 @@ +"""On-demand shape benchmark for the evolve-kernel WebUI. + +Runs the best kernel from the library against the reference kernel on a set +of target shapes extracted from the task's extra_notes. Results are cached +in workspace/shape_bench.json and refreshed when the best kernel changes. + +Parsed shape formats: + - Table rows: wq_b: (M, 1024) @ (1024, 4096) # TP=8 + - Multi-line: gate_up_proj: (M, 4096) @ (4096, 1024) → M=1,4,16,4096 + - Inline: shapes=(M,1024,4096) (M,4096,1024) +""" + +from __future__ import annotations + +import importlib.util +import json +import re +import statistics +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional + +try: # torch optional so CPU-only CI can import/collect; runtime requires it + import torch +except ImportError: # pragma: no cover - CPU-only environment + torch = None # type: ignore[assignment] + + +# =========================================================================== # +# Shape spec +# =========================================================================== # + + +@dataclass +class ShapeSpec: + label: str # e.g. "gate_up_proj (TP=4)" + M_values: List[int] # e.g. [1, 2, 4, 8, 16, 4096] + K: int + N: int + + +# Default M values to benchmark per shape +_DEFAULT_M_VALUES = [1, 2, 4, 8, 16, 4096] + +# Commonly-used shapes from DeepSeek V4 Flash TP=4/TP=8 +_PRESET_SHAPES: List[ShapeSpec] = [ + ShapeSpec("wq_b (TP=4)", [1, 2, 4, 8, 16, 4096], K=1024, N=8192), + ShapeSpec("wq_b (TP=8)", [1, 2, 4, 8, 16, 4096], K=1024, N=4096), + ShapeSpec("wo_b (TP=4)", [1, 2, 4, 8, 16, 4096], K=2048, N=4096), + ShapeSpec("wo_b (TP=8)", [1, 2, 4, 8, 16, 4096], K=1024, N=4096), + ShapeSpec("gate_up_proj (TP=4)", [1, 2, 4, 8, 16, 4096], K=4096, N=1024), + ShapeSpec("gate_up_proj (TP=8)", [1, 2, 4, 8, 16, 4096], K=4096, N=512), + ShapeSpec("down_proj (TP=4)", [1, 2, 4, 8, 16, 4096], K=512, N=4096), + ShapeSpec("down_proj (TP=8)", [1, 2, 4, 8, 16, 4096], K=256, N=4096), +] + + +# =========================================================================== # +# Shape parsing +# =========================================================================== # + + +def parse_shapes_from_extra_notes(extra_notes: str) -> List[ShapeSpec]: + """Parse target shapes from extra_notes or other requirements text. + + Recognized formats: + + 1. Named rows (table format): + wq_b: (M, 1024) @ (1024, 8192) + gate_up_proj: (M, 4096) @ (4096, 1024) + + 2. Explicit M values on same line: + gate_up_proj: M=1,4,8,16,4096 (M, 4096) @ (4096, 1024) + + 3. Compact format: + (M,1024,4096) → K=1024, N=4096 + + If no shapes are found, returns the PRESET_SHAPES. + """ + if not extra_notes: + return list(_PRESET_SHAPES) + + shapes: List[ShapeSpec] = [] + + # Pattern 1: named shapes with (M, K) @ (K, N) notation + named_pattern = re.compile( + r'([\w_]+(?:\s*\([^)]*\))?)\s*[::]\s*' + r'\(?\s*M\s*,\s*(\d+)\s*\)\s*@\s*\(?\s*(\d+)\s*,\s*(\d+)\s*\)?', + re.IGNORECASE, + ) + for match in named_pattern.finditer(extra_notes): + label = match.group(1).strip() + k_val = int(match.group(2)) + n_val = int(match.group(4)) # group 3 is K from @(K, N), group 4 is N + + # Make label descriptive if it's generic + if label.upper() in ("SHAPE", "TARGET SHAPE", "TARGET", "SHAPES"): + label = f"({k_val}×{n_val})" + + # Collect M values from nearby context (up to 500 chars) + nearby = extra_notes[match.start():match.start() + 500] + m_vals_set = set() + # Match patterns: M=1,2,4,8 M=4096 M <=16 M=1,4,8,16,4096 + for m_match in re.finditer( + r'(?:M|m)\s*[=::<=]+\s*([\d,\s]+)', + nearby, + ): + for num_str in m_match.group(1).split(","): + num_str = num_str.strip() + if num_str.isdigit(): + m_vals_set.add(int(num_str)) + + if m_vals_set: + m_vals = sorted(m_vals_set) + else: + # Fallback: look for "M <=N" patterns in full text + m_range = re.search( + r'(?:M|m)\s*<=\s*(\d+)', + extra_notes, + ) + if m_range: + max_m = int(m_range.group(1)) + m_vals = [1, 2, 4, 8, 16, max_m] + else: + m_vals = list(_DEFAULT_M_VALUES) + + shapes.append(ShapeSpec(label=label, M_values=m_vals, K=k_val, N=n_val)) + + # Pattern 2: compact (M, K, N) format — pick the first one that looks significant + if not shapes: + compact = re.compile(r'\(\s*(?:M|m)\s*,\s*(\d+)\s*,\s*(\d+)\s*\)') + for match in compact.finditer(extra_notes): + k_val = int(match.group(1)) + n_val = int(match.group(2)) + shapes.append(ShapeSpec( + label=f"({k_val}×{n_val})", + M_values=[1, 2, 4, 8, 16, 4096], + K=k_val, N=n_val, + )) + + if not shapes: + return list(_PRESET_SHAPES) + + return shapes + + +# =========================================================================== # +# Benchmark runner +# =========================================================================== # + + +@dataclass +class ShapeResult: + shape_label: str + M: int + N: int + K: int + ref_ms: float = 0.0 + best_ms: float = 0.0 + speedup: float = 1.0 + error: str = "" + + +def _load_kernel_fn(filepath: str, fn_name: str): + spec = importlib.util.spec_from_file_location("_shape_bench_mod", filepath) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + if not hasattr(mod, fn_name): + raise ValueError(f"Function {fn_name!r} not found in {filepath}") + return getattr(mod, fn_name) + + +def _measure_one(kernel_fn, M: int, N: int, K: int, + warmup: int = 5, repeat: int = 20) -> float: + """Measure median execution time in ms for one (M,N,K) shape.""" + device = torch.device("cuda") + + a = torch.randint(-128, 127, (M, K), device=device, dtype=torch.int8) + a_scale = torch.randn(M, 1, device=device, dtype=torch.float32).abs() / 127.0 + b = torch.randint(-128, 127, (K, N), device=device, dtype=torch.int8) + b_scale = torch.randn(N, 1, device=device, dtype=torch.float32).abs() / 127.0 + + # Warmup + for _ in range(warmup): + kernel_fn(a, a_scale, b, b_scale, torch.bfloat16) + torch.cuda.synchronize() + + # Timed runs + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_ms = [] + for _ in range(repeat): + start.record() + kernel_fn(a, a_scale, b, b_scale, torch.bfloat16) + end.record() + torch.cuda.synchronize() + times_ms.append(start.elapsed_time(end)) + + return statistics.median(times_ms) + + +def _measure_one_safe(kernel_fn, M: int, N: int, K: int, + warmup: int = 2, repeat: int = 5) -> float: + """Measure with adaptive repetitions; catches OOM and crashes. + + Uses more repetitions for small M (fast) and fewer for large M (slow). + """ + try: + # Adaptive: more repeats for fast kernels (M <= 16), fewer for slow ones + actual_repeat = repeat + if M <= 16: + actual_repeat = 10 # small M is sub-ms, so more repeats for accuracy + actual_warmup = max(warmup, 3) + else: + actual_repeat = 3 # large M is slower, fewer repeats + actual_warmup = max(warmup, 1) + + device = torch.device("cuda") + a = torch.randint(-128, 127, (M, K), device=device, dtype=torch.int8) + a_scale = torch.randn(M, 1, device=device, dtype=torch.float32).abs() / 127.0 + b = torch.randint(-128, 127, (K, N), device=device, dtype=torch.int8) + b_scale = torch.randn(N, 1, device=device, dtype=torch.float32).abs() / 127.0 + + for _ in range(actual_warmup): + kernel_fn(a, a_scale, b, b_scale, torch.bfloat16) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + times_ms = [] + for _ in range(actual_repeat): + start.record() + kernel_fn(a, a_scale, b, b_scale, torch.bfloat16) + end.record() + torch.cuda.synchronize() + times_ms.append(start.elapsed_time(end)) + + return statistics.median(times_ms) + except Exception: + return 0.0 + + +def run_shape_benchmark( + ref_kernel_path: str, + best_kernel_path: str, + kernel_fn_name: str, + shapes: List[ShapeSpec], +) -> List[ShapeResult]: + """Run benchmarks across all shapes for both reference and best kernel. + + Returns one ShapeResult per (shape, M_value) combination. + """ + try: + ref_fn = _load_kernel_fn(ref_kernel_path, kernel_fn_name) + except Exception as e: + return [ShapeResult(shape_label="error", M=0, N=0, K=0, error=f"Failed to load ref kernel: {e}")] + + try: + best_fn = _load_kernel_fn(best_kernel_path, kernel_fn_name) + except Exception as e: + return [ShapeResult(shape_label="error", M=0, N=0, K=0, error=f"Failed to load best kernel: {e}")] + + results: List[ShapeResult] = [] + + for spec in shapes: + for m in spec.M_values: + # Skip shapes that are too large for this M + if m * spec.N > 2 ** 25: # > 32M elements output + continue + + ref_ms = _measure_one_safe(ref_fn, m, spec.N, spec.K) + best_ms = _measure_one_safe(best_fn, m, spec.N, spec.K) + + if ref_ms <= 0 or best_ms <= 0: + results.append(ShapeResult( + shape_label=spec.label, M=m, N=spec.N, K=spec.K, + ref_ms=ref_ms, best_ms=best_ms, + error="Measurement failed (OOM or kernel crash)", + )) + continue + + speedup = ref_ms / max(best_ms, 1e-6) + results.append(ShapeResult( + shape_label=spec.label, M=m, N=spec.N, K=spec.K, + ref_ms=ref_ms, best_ms=best_ms, speedup=speedup, + )) + + return results + + +def shape_results_to_dict(results: List[ShapeResult]) -> List[Dict[str, Any]]: + return [ + { + "shape_label": r.shape_label, + "M": r.M, + "N": r.N, + "K": r.K, + "ref_ms": round(r.ref_ms, 4), + "best_ms": round(r.best_ms, 4), + "speedup": round(r.speedup, 3), + "error": r.error, + } + for r in results + ] + + +# =========================================================================== # +# Cache management +# =========================================================================== # + + +def load_cached_benchmark(cache_path: Path, best_kernel_id: str) -> Optional[List[Dict[str, Any]]]: + """Load cached benchmark results if they match the current best kernel.""" + if not cache_path.is_file(): + return None + try: + data = json.loads(cache_path.read_text(encoding="utf-8")) + if data.get("best_kernel_id") == best_kernel_id: + return data.get("results", []) + except (json.JSONDecodeError, KeyError): + pass + return None + + +def save_cached_benchmark(cache_path: Path, best_kernel_id: str, + results: List[Dict[str, Any]]) -> None: + cache_path.parent.mkdir(parents=True, exist_ok=True) + cache_path.write_text(json.dumps({ + "best_kernel_id": best_kernel_id, + "results": results, + "timestamp": time.time(), + }, indent=2), encoding="utf-8") diff --git a/metainfer/tasks/evolve_kernel/server/_state_readers.py b/metainfer/tasks/evolve_kernel/server/_state_readers.py index 27effcb6..74364713 100644 --- a/metainfer/tasks/evolve_kernel/server/_state_readers.py +++ b/metainfer/tasks/evolve_kernel/server/_state_readers.py @@ -6,8 +6,11 @@ from __future__ import annotations import json +import re from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set + +import os from ..orchestrator import phases as _phases @@ -87,6 +90,19 @@ def read_charts(state_dir: Path) -> Dict[str, Any]: # --------------------------------------------------------------------------- # +def _find_retrospective_md(state_dir: Path, n: int) -> Optional[Path]: + """Find retrospective.md for iteration n, trying multiple locations.""" + # 1) logs//retrospective.md (where the agent writes it) + logs_path = state_dir / "logs" / f"{n:03d}" / "retrospective.md" + if logs_path.is_file(): + return logs_path + # 2) state_dir / "logs" / str(n) / "retrospective.md" (unpadded variant) + logs_path2 = state_dir / "logs" / str(n) / "retrospective.md" + if logs_path2.is_file(): + return logs_path2 + return None + + def read_retrospective(state_dir: Path, n: int) -> Dict[str, Any]: rec = read_iteration(state_dir, n) if rec is None: @@ -98,12 +114,25 @@ def read_retrospective(state_dir: Path, n: int) -> Dict[str, Any]: path_str = rec.get("retrospective_path") markdown = "" has = False + actual_path: Optional[str] = None if path_str: p = Path(path_str) if p.is_file(): try: markdown = p.read_text(encoding="utf-8", errors="replace") has = True + actual_path = path_str + except OSError: + markdown = "" + # Fallback: look for retrospective.md in logs dir (pipeline writes it there + # but doesn't always record retrospective_path in the iteration record) + if not has: + found = _find_retrospective_md(state_dir, n) + if found is not None: + try: + markdown = found.read_text(encoding="utf-8", errors="replace") + has = True + actual_path = str(found) except OSError: markdown = "" if not has: @@ -113,7 +142,7 @@ def read_retrospective(state_dir: Path, n: int) -> Dict[str, Any]: ) return { "has_retrospective": has, - "path": path_str, + "path": actual_path, "markdown": markdown, "this_perf": this_perf, "prev_perf": prev_perf, @@ -141,6 +170,15 @@ def read_state_graph(state_dir: Path) -> Dict[str, Any]: # --------------------------------------------------------------------------- # +def read_optimizer_mode(state_dir: Path) -> str: + """Read optimizer_mode from requirements.json. Returns display label or 'triton'.""" + req = _load_json(state_dir / "requirements.json", {}) + mode = req.get("optimizer_mode", "Triton (standard)") + if isinstance(mode, list): + mode = mode[0] if mode else "Triton (standard)" + return mode + + def read_kernel_library(workspace_dir: Path) -> Dict[str, Any]: path = workspace_dir / "kernel_library.json" if not path.is_file(): @@ -152,16 +190,58 @@ def read_kernel_library(workspace_dir: Path) -> Dict[str, Any]: if not isinstance(data, list): return {"kernels": [], "size": 0} - # Enrich each kernel entry with a preview (first 20 lines) + # Build a lookup from kernel id → exec_time_ms for speedup calculation + time_by_id: Dict[str, float] = {} + for k in data: + et = k.get("exec_time_ms", 0) + if et > 0: + time_by_id[k["id"]] = et + + # Enrich each kernel entry with a preview (first 20 lines) and speedup vs parent enriched = [] for k in data: code = k.get("code", "") preview = "\n".join(code.splitlines()[:20]) if code else "" - enriched.append({ + entry = { **k, "code_preview": preview, "code_lines": len(code.splitlines()) if code else 0, - }) + "speedup_vs_parent": None, + # Headroom: extract from kernel entry or from nested dict + "headroom": { + "bottleneck": k.get("headroom_bottleneck"), + "roofline_efficiency_pct": k.get("headroom_roofline_efficiency_pct", + # Fallback: compute from old bw_util_pct / compute_util_pct + max(k.get("headroom_bw_util_pct", 0), k.get("headroom_compute_util_pct", 0))), + "headroom_pct": k.get("headroom_pct", 100 - max(k.get("headroom_bw_util_pct", 0), k.get("headroom_compute_util_pct", 0))), + "p_max_tflops": k.get("headroom_p_max_tflops", 0), + "p_bw_roof_tflops": k.get("headroom_p_bw_roof_tflops", 0), + "ai_ridge": k.get("headroom_ai_ridge", 0), + "bw_util_pct": k.get("headroom_bw_util_pct", 0), + "compute_util_pct": k.get("headroom_compute_util_pct", 0), + "suggestions": _parse_headroom_suggestions(k.get("headroom_suggestions_json")), + "advice": k.get("headroom_advice"), + "achieved_bw_gbps": k.get("headroom_achieved_bw_gbps", 0), + "achieved_tflops": k.get("headroom_achieved_tflops", 0), + "peak_bw_gbps": k.get("headroom_peak_bw_gbps", 0), + "peak_tflops": k.get("headroom_peak_tflops", 0), + "arithmetic_intensity": k.get("headroom_arithmetic_intensity", 0), + "measured_ai": k.get("headroom_measured_ai", 0), + "shape_label": k.get("headroom_shape_label", ""), + "M": k.get("headroom_M", 0), + "N": k.get("headroom_N", 0), + "K": k.get("headroom_K", 0), + "has_data": bool(k.get("headroom_bottleneck")), + }, + } + # Compute speedup vs parent + pid = k.get("parent_id") + child_time = k.get("exec_time_ms", 0) + if pid and child_time > 0: + parent_time = time_by_id.get(pid, 0) + if parent_time > 0: + entry["speedup_vs_parent"] = round(parent_time / child_time, 2) + enriched.append(entry) return {"kernels": enriched, "size": len(enriched)} @@ -213,3 +293,867 @@ def read_reference_kernel(workspace_dir: Path) -> Dict[str, Any]: "path": str(path), "lines": len(code.splitlines()), } + + +# --------------------------------------------------------------------------- # +# Kernel Lineage +# --------------------------------------------------------------------------- # + + +def _parse_headroom_suggestions(raw: Optional[str]) -> List[str]: + """Parse JSON-encoded headroom suggestions list.""" + if not raw: + return [] + try: + parsed = json.loads(raw) + if isinstance(parsed, list): + return [str(s) for s in parsed] + except (json.JSONDecodeError, TypeError): + pass + return [] + + +def _auto_summary_items(kernel: Dict[str, Any], parent: Optional[Dict[str, Any]], + retro_md: Optional[str]) -> List[str]: + """Generate improvement summary as a list of items from retrospective sections. + + Extracts bullet points and paragraph breaks from ``## What Changed`` and + ``## Why Perf Moved`` sections. Falls back to metric-based summary if no + retrospective is available. + """ + items: List[str] = [] + + if retro_md: + lines = retro_md.splitlines() + sections: Dict[str, List[str]] = {} + current_section: Optional[str] = None + for line in lines: + if line.startswith("## "): + current_section = line[3:].strip().lower() + sections.setdefault(current_section, []) + elif current_section: + sections[current_section].append(line) + + # Process "What Changed" + what_lines = sections.get("what changed", []) + if what_lines: + # Group consecutive non-empty lines into items, splitting on + # bullet markers or double-newline gaps + text = "\n".join(what_lines).strip() + # Split by double-newline (paragraph breaks) or explicit bullet markers + parts = re.split(r'\n\n+|\n(?=(?:- |\* |\d+\. ))', text) + for part in parts: + part = part.strip() + if not part: + continue + # If part contains `- ` bullets, split further + if re.match(r'- ', part, re.MULTILINE): + sub = [s.strip()[2:].strip() for s in re.split(r'\n- ', part) if s.strip()] + items.extend(sub) + elif re.match(r'\* ', part, re.MULTILINE): + sub = [s.strip()[2:].strip() for s in re.split(r'\n\* ', part) if s.strip()] + items.extend(sub) + else: + items.append(part) + + # Process "Why Perf Moved" + why_lines = sections.get("why perf moved", []) + if why_lines: + text = "\n".join(why_lines).strip() + parts = re.split(r'\n\n+|\n(?=(?:- |\* |\d+\. ))', text) + for part in parts: + part = part.strip() + if not part: + continue + if re.match(r'- ', part, re.MULTILINE): + sub = [s.strip()[2:].strip() for s in re.split(r'\n- ', part) if s.strip()] + items.extend(sub) + elif re.match(r'\* ', part, re.MULTILINE): + sub = [s.strip()[2:].strip() for s in re.split(r'\n\* ', part) if s.strip()] + items.extend(sub) + else: + items.append(part) + + if items: + return items + + # Fallback: metric-based summary + if parent and parent.get("exec_time_ms", 0) > 0: + parent_time = parent["exec_time_ms"] + child_time = kernel.get("exec_time_ms", 0) + if child_time > 0: + delta = parent_time - child_time + speedup = parent_time / child_time if child_time > 0 else 1.0 + faster = "faster" if delta > 0 else "slower" + items.append( + f"Exec time: {parent_time:.4f}ms → {child_time:.4f}ms " + f"({abs(delta):.4f}ms {faster}, {speedup:.2f}× vs parent)" + ) + return items + + if parent is None: + items.append("Initial seed kernel — baseline for all subsequent optimizations.") + return items + + items.append("No performance data available for comparison.") + return items + + +def _auto_summary(kernel: Dict[str, Any], parent: Optional[Dict[str, Any]], + retro_md: Optional[str]) -> str: + """Generate a human-readable improvement summary (legacy string format). + + Prefer ``_auto_summary_items`` for structured display; this is kept for + backward compatibility with the string-based ``improvement.summary`` field. + """ + items = _auto_summary_items(kernel, parent, retro_md) + return "\n\n".join(items) + + +def read_kernel_lineage(workspace_dir: Path, state_dir: Path, + kernel_id: str) -> Dict[str, Any]: + """Build lineage info for one kernel: parent, improvement metrics, retrospective.""" + lib_data = read_kernel_library(workspace_dir) + kernels = lib_data.get("kernels", []) + + # Find the target kernel + target: Optional[Dict[str, Any]] = None + for k in kernels: + if k.get("id") == kernel_id: + target = k + break + + if target is None: + return {"error": f"kernel {kernel_id[:8]}… not found in library"} + + # Find parent + parent: Optional[Dict[str, Any]] = None + parent_id = target.get("parent_id") + if parent_id: + for k in kernels: + if k.get("id") == parent_id: + parent = k + break + + # Build ancestor chain + ancestors: List[str] = [] + current_pid = target.get("parent_id") + seen: set = set() + while current_pid and current_pid not in seen: + seen.add(current_pid) + ancestors.append(current_pid) + # Find parent's parent + found_parent = False + for k in kernels: + if k.get("id") == current_pid: + current_pid = k.get("parent_id") + found_parent = True + break + if not found_parent: + break + + # Improvement metrics + parent_time = parent.get("exec_time_ms", 0) if parent else 0 + child_time = target.get("exec_time_ms", 0) + improvement: Dict[str, Any] = { + "exec_time_delta_ms": round(child_time - parent_time, 6) if parent_time > 0 else None, + "speedup_vs_parent": round(parent_time / child_time, 4) if parent_time > 0 and child_time > 0 else None, + "complexity_delta": round(target.get("complexity_score", 0) - (parent.get("complexity_score", 0) if parent else 0), 2), + "summary": "", + } + + # Try retrospective for the iteration this kernel was added + iter_num = target.get("iteration_added", 0) + retro_md: Optional[str] = None + if iter_num > 0: + retro = read_retrospective(state_dir, iter_num) + if retro.get("has_retrospective"): + retro_md = retro.get("markdown", "") + + improvement["summary"] = _auto_summary(target, parent, retro_md) + improvement["summary_items"] = _auto_summary_items(target, parent, retro_md) + + # Headroom data + headroom = { + "bottleneck": target.get("headroom_bottleneck"), + "headroom_pct": target.get("headroom_pct", 0), + "bw_util_pct": target.get("headroom_bw_util_pct", 0), + "compute_util_pct": target.get("headroom_compute_util_pct", 0), + "suggestions": _parse_headroom_suggestions(target.get("headroom_suggestions_json")), + "advice": target.get("headroom_advice"), + "achieved_bw_gbps": target.get("headroom_achieved_bw_gbps", 0), + "achieved_tflops": target.get("headroom_achieved_tflops", 0), + "peak_bw_gbps": target.get("headroom_peak_bw_gbps", 0), + "peak_tflops": target.get("headroom_peak_tflops", 0), + "arithmetic_intensity": target.get("headroom_arithmetic_intensity", 0), + "has_data": bool(target.get("headroom_bottleneck")), + } + + return { + "kernel": target, + "parent": parent, + "improvement": improvement, + "retrospective": { + "has_retrospective": retro_md is not None, + "markdown": retro_md or "", + "iteration": iter_num, + }, + "ancestor_chain": ancestors, + "headroom": headroom, + } + + +# --------------------------------------------------------------------------- # +# Failure Log +# --------------------------------------------------------------------------- # + +PHASE_LABEL_MAP = { + "A_gen_correctness_harness": "A: Gen Correctness Harness", + "B_review_correctness_harness": "B: Review Correctness Harness", + "C_gen_perf_harness": "C: Gen Perf Harness", + "D_review_perf_harness": "D: Review Perf Harness", + "E_select_kernel": "E: Select Kernel", + "F_optimize": "F: Optimize", + "G_verify_correctness": "G: Verify Correctness", + "H_measure_perf": "H: Measure Perf", +} + + +def _read_correctness_failure_detail(state_dir: Path, n: int) -> Optional[str]: + """Read detailed correctness failure from logs dir.""" + log_path = state_dir / "logs" / f"{n:03d}" / "correctness_failure.json" + if not log_path.is_file(): + return None + try: + data = json.loads(log_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return None + if not isinstance(data, dict): + return None + + lines: List[str] = [] + error_msg = data.get("error", "") + if error_msg: + lines.append(f"Error: {error_msg}") + + results = data.get("results", []) + if results: + lines.append(f"\nFailed test cases ({len(results)}):") + for r in results: + if isinstance(r, dict): + name = r.get("name", r.get("test", "unknown")) + passed = r.get("passed", False) + err = r.get("error", "") + status = "✓" if passed else "✗" + lines.append(f" {status} {name}") + if err and not passed: + lines.append(f" {err}") + + return "\n".join(lines) if lines else None + + +def read_failures(state_dir: Path) -> Dict[str, Any]: + """Extract all failures from iteration records. + + Returns a list of failure entries, each with: + - iteration, phase, phase_label, outcome, severity, summary, detail, attempts, timestamp + """ + all_iters = read_iterations(state_dir) + failures: List[Dict[str, Any]] = [] + + for rec in all_iters: + iter_num = rec.get("iteration", 0) + phases = rec.get("phases", {}) + + # 1) Iteration-level failures + if rec.get("status") == "failed" and rec.get("failure_reason"): + failures.append({ + "iteration": iter_num, + "phase": None, + "phase_label": "Iteration", + "outcome": rec.get("outcome", "failed"), + "severity": "error", + "summary": _truncate(rec["failure_reason"], 120), + "detail": rec["failure_reason"], + "attempts": 1, + "timestamp": rec.get("ended_at", rec.get("started_at", 0)), + }) + + # 2) Phase-level failures (including recovered ones where failure text persists) + for phase_id, pdata in phases.items(): + if not isinstance(pdata, dict): + continue + failure_text = pdata.get("failure") + if not failure_text: + continue + outcome = pdata.get("outcome", "unknown") + attempts = pdata.get("attempts", 1) + + # Determine severity: error if final outcome is not ok, warning if recovered + severity = "error" if outcome != "ok" else "warning" + phase_label = PHASE_LABEL_MAP.get(phase_id, phase_id) + + # For correctness failures, try to read detailed test results + detail = failure_text + if phase_id == "G_verify_correctness": + detail_extra = _read_correctness_failure_detail(state_dir, iter_num) + if detail_extra: + detail = failure_text + "\n\n" + detail_extra + + failures.append({ + "iteration": iter_num, + "phase": phase_id, + "phase_label": phase_label, + "outcome": outcome, + "severity": severity, + "summary": _truncate(failure_text, 120), + "detail": detail, + "attempts": attempts, + "timestamp": pdata.get("ended_at", rec.get("started_at", 0)), + }) + + # Sort newest first + failures.sort(key=lambda f: f["timestamp"], reverse=True) + + # Summary counts + error_count = sum(1 for f in failures if f["severity"] == "error") + warning_count = sum(1 for f in failures if f["severity"] == "warning") + + return { + "failures": failures, + "total": len(failures), + "errors": error_count, + "warnings": warning_count, + } + + +def _truncate(text: str, max_len: int) -> str: + """Truncate text to max_len characters, adding ellipsis if needed.""" + text = text.strip() + if len(text) <= max_len: + return text + return text[:max_len - 3].rstrip() + "…" + + +# --------------------------------------------------------------------------- # +# Shape Benchmark +# --------------------------------------------------------------------------- # + +from ._shape_bench import ( + parse_shapes_from_extra_notes, + run_shape_benchmark, + shape_results_to_dict, + load_cached_benchmark, + save_cached_benchmark, +) +from ._multi_gpu import ( + aggregate_shape_benchmarks, + aggregate_best_kernels, +) + + +def read_shape_benchmark(state_dir: Path, workspace_dir: Path) -> Dict[str, Any]: + """Read or run shape benchmarks comparing best kernel vs reference. + + Returns cached results if available and best kernel hasn't changed. + Otherwise runs benchmarks on all target shapes — this is slow (~30s-2min). + """ + # 1. Find best kernel + lib_path = workspace_dir / "kernel_library.json" + if not lib_path.is_file(): + return {"results": [], "error": "No kernel library yet", "running": False} + + try: + lib_data = json.loads(lib_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return {"results": [], "error": "Failed to read kernel library", "running": False} + + if not lib_data: + return {"results": [], "error": "Kernel library is empty", "running": False} + + # Sort by exec_time_ms ascending + lib_data.sort(key=lambda k: k.get("exec_time_ms", float("inf"))) + best_kernel = lib_data[0] + best_id = best_kernel.get("id", "") + best_iter = best_kernel.get("iteration_added", 0) + + # 2. Find reference kernel + ref_path = workspace_dir / "reference" / "original_kernel.py" + if not ref_path.is_file(): + return {"results": [], "error": "No reference kernel found", "running": False} + + # 3. Find best kernel code + best_kernel_path = None + # Try iteration dir first + iter_path = workspace_dir / f"{best_iter:03d}" / "optimized_kernel.py" + if iter_path.is_file(): + best_kernel_path = iter_path + else: + # Try shared kernels dir + shared = workspace_dir / "optimized_kernels" / f"{best_id}.py" + if shared.is_file(): + best_kernel_path = shared + if best_kernel_path is None: + return {"results": [], "error": f"Best kernel file not found (id={best_id[:8]})", "running": False} + + # 4. Check cache + cache_path = workspace_dir / "shape_bench.json" + cached = load_cached_benchmark(cache_path, best_id) + if cached is not None: + return { + "results": cached, + "best_kernel_id": best_id, + "best_exec_time_ms": best_kernel.get("exec_time_ms", 0), + "cached": True, + "running": False, + } + + # 5. Parse shapes from requirements + req_path = state_dir / "requirements.json" + extra_notes = "" + kernel_fn_name = "matmul_int8" + if req_path.is_file(): + try: + req = json.loads(req_path.read_text(encoding="utf-8")) + extra_notes = req.get("extra_notes", "") + kernel_fn_name = req.get("kernel_function_name", "") or kernel_fn_name + except (json.JSONDecodeError, OSError): + pass + + shapes = parse_shapes_from_extra_notes(extra_notes) + + # 6. Run benchmarks + try: + results = run_shape_benchmark( + ref_kernel_path=str(ref_path), + best_kernel_path=str(best_kernel_path), + kernel_fn_name=kernel_fn_name, + shapes=shapes, + ) + except Exception as e: + return {"results": [], "error": f"Benchmark failed: {e}", "running": False, + "best_kernel_id": best_id, "best_exec_time_ms": best_kernel.get("exec_time_ms", 0)} + + result_dicts = shape_results_to_dict(results) + + # 7. Cache + save_cached_benchmark(cache_path, best_id, result_dicts) + + return { + "results": result_dicts, + "best_kernel_id": best_id, + "best_exec_time_ms": best_kernel.get("exec_time_ms", 0), + "cached": False, + "running": False, + } + + +def refresh_shape_benchmark(state_dir: Path, workspace_dir: Path) -> Dict[str, Any]: + """Force re-run shape benchmarks, ignoring cache.""" + cache_path = workspace_dir / "shape_bench.json" + if cache_path.is_file(): + cache_path.unlink() + return read_shape_benchmark(state_dir, workspace_dir) + + +# --------------------------------------------------------------------------- # +# Multi-GPU status +# --------------------------------------------------------------------------- # + + +def read_gpu_status(state_dir: Path, workspace_dir: Path) -> Dict[str, Any]: + """Read live status of all GPU workers from gpu_N/ subdirectories. + + Scans state_dir for gpu_0/, gpu_1/, etc. and returns each worker's + current phase, iteration, exec_time, and running status. + """ + workers: List[Dict[str, Any]] = [] + + # Check parent run.json for multi-GPU metadata + parent_run = _load_json(state_dir / "run.json", {}) + is_multi = parent_run.get("multi_gpu") or parent_run.get("multi_gpu_mode") + + if not is_multi: + # Also check requirements + req = _load_json(state_dir / "requirements.json", {}) + is_multi = req.get("multi_gpu") in ("All GPUs (auto)", "yes", "true", "1") + + if not is_multi: + return {"workers": [], "is_multi_gpu": False} + + # Scan gpu_N directories + for entry in sorted(state_dir.iterdir()): + if not entry.is_dir(): + continue + if not entry.name.startswith("gpu_"): + continue + try: + gpu_idx = int(entry.name.split("_")[1]) + except (IndexError, ValueError): + continue + + # Read run.json + run = _load_json(entry / "run.json", {}) + phase = run.get("current_phase", "idle") + iteration = run.get("current_iteration", 0) + + # Check process liveness via PID file + pid_data = _load_json(entry / "orchestrator.pid", {}) + pid = pid_data.get("pid", 0) + running = False + if pid and pid_data.get("finished_at") is None: + try: + os.kill(pid, 0) + running = True + except OSError: + running = False + + # Read best time from per-GPU library + exec_time_ms = 0.0 + speedup = 0.0 + kernel_count = 0 + best_kernel_id_short = "" + gpu_workspace = workspace_dir / entry.name + lib_path = gpu_workspace / "kernel_library.json" + if lib_path.is_file(): + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + kernel_count = len(lib) + if lib: + lib.sort(key=lambda k: k.get("exec_time_ms", float("inf"))) + best = lib[0] + exec_time_ms = best.get("exec_time_ms", 0) + best_kernel_id_short = best.get("id", "")[:8] + # Find seed (iteration_added=0) for speedup calc + for k in lib: + if k.get("iteration_added") == 0: + seed_time = k.get("exec_time_ms", 0) + if seed_time > 0 and exec_time_ms > 0: + speedup = seed_time / exec_time_ms + break + except Exception: + pass + + # Read agents from per-GPU agents.json + agents_data = _load_json(entry / "agents.json", {"ts": 0, "agents": []}) + agents_summary: List[Dict[str, Any]] = [] + for a in agents_data.get("agents", []): + if not isinstance(a, dict): + continue + agents_summary.append({ + "name": a.get("name", ""), + "role": a.get("role", ""), + "phase": a.get("phase", "unknown"), + "elapsed_s": round(a.get("elapsed_s", 0), 1), + "attempt": a.get("attempt", 1), + "success": a.get("success"), + "error": a.get("error"), + }) + + # Read shapes from requirements + gpu_req = _load_json(entry / "requirements.json", {}) + shapes_notes = gpu_req.get("extra_notes", "")[:200] + + workers.append({ + "gpu_idx": gpu_idx, + "label": f"GPU {gpu_idx}", + "phase": phase, + "iteration": iteration, + "exec_time_ms": round(exec_time_ms, 4) if exec_time_ms else 0, + "speedup": round(speedup, 2) if speedup else 0, + "running": running, + "pid": pid, + "shapes": shapes_notes, + "agents": agents_summary, + "kernel_count": kernel_count, + "best_kernel_id": best_kernel_id_short, + }) + + return { + "workers": workers, + "is_multi_gpu": True, + "num_workers": len(workers), + } + + +def read_aggregate_bench(state_dir: Path, workspace_dir: Path) -> Dict[str, Any]: + """Merge shape benchmarks from all per-GPU workspaces.""" + wds: List[str] = [] + for entry in sorted(workspace_dir.iterdir()): + if entry.is_dir() and entry.name.startswith("gpu_"): + wds.append(str(entry)) + + if not wds: + # Try parent workspace directly + bench_path = workspace_dir / "shape_bench.json" + if bench_path.is_file(): + try: + data = json.loads(bench_path.read_text(encoding="utf-8")) + return { + "results": data.get("results", []), + "kernels": [], + "is_multi_gpu": True, + "num_workers": 0, + } + except Exception: + pass + return {"results": [], "kernels": [], "is_multi_gpu": False} + + results = aggregate_shape_benchmarks(wds) + kernels = aggregate_best_kernels(wds) + + return { + "results": results, + "kernels": kernels, + "is_multi_gpu": True, + "num_workers": len(wds), + } + + +# --------------------------------------------------------------------------- # +# Combined timeline (across all GPU workers + parent) +# --------------------------------------------------------------------------- # + +from metainfer.server.state_reader import read_timeline as _read_timeline_file + + +# --------------------------------------------------------------------------- # +# Per-GPU detail (state graph + kernel library for one GPU worker) +# --------------------------------------------------------------------------- # + + +def _gpu_state_dir(state_dir: Path, gpu_idx: int) -> Path: + return state_dir / f"gpu_{gpu_idx}" + + +def _gpu_workspace_dir(workspace_dir: Path, gpu_idx: int) -> Path: + return workspace_dir / f"gpu_{gpu_idx}" + + +def read_gpu_state_graph(state_dir: Path, gpu_idx: int) -> Dict[str, Any]: + """Read state graph for a single GPU worker.""" + gpu_sd = _gpu_state_dir(state_dir, gpu_idx) + run = _load_json(gpu_sd / "run.json", {}) + current = run.get("current_phase", "idle") + last_outcome = run.get("last_outcome") + last_label = run.get("last_transition_label") + if hasattr(_phases, "graph_payload"): + return _phases.graph_payload(current, last_outcome, last_label) + return {"error": "phases module does not export graph_payload()"} + + +def read_gpu_kernel_library(workspace_dir: Path, gpu_idx: int) -> Dict[str, Any]: + """Read kernel library for a single GPU worker.""" + return read_kernel_library(_gpu_workspace_dir(workspace_dir, gpu_idx)) + + +def read_gpu_detail(state_dir: Path, workspace_dir: Path, gpu_idx: int) -> Dict[str, Any]: + """Aggregated per-GPU detail: state graph + kernel library + agents + phase info. + + Returns everything needed to render the single-GPU view for one worker. + """ + gpu_sd = _gpu_state_dir(state_dir, gpu_idx) + run = _load_json(gpu_sd / "run.json", {}) + + # State graph + graph = read_gpu_state_graph(state_dir, gpu_idx) + + # Kernel library + library = read_gpu_kernel_library(workspace_dir, gpu_idx) + + # Agents + agents_data = _load_json(gpu_sd / "agents.json", {"ts": 0, "agents": []}) + agents_summary: List[Dict[str, Any]] = [] + for a in agents_data.get("agents", []): + if not isinstance(a, dict): + continue + agents_summary.append({ + "name": a.get("name", ""), + "role": a.get("role", ""), + "phase": a.get("phase", "unknown"), + "elapsed_s": round(a.get("elapsed_s", 0), 1), + "attempt": a.get("attempt", 1), + "success": a.get("success"), + "error": a.get("error"), + }) + + # Iteration data + iteration = run.get("current_iteration", 0) + phase = run.get("current_phase", "idle") + + # Check liveness + pid_data = _load_json(gpu_sd / "orchestrator.pid", {}) + pid = pid_data.get("pid", 0) + running = False + if pid and pid_data.get("finished_at") is None: + try: + import os as _os + _os.kill(pid, 0) + running = True + except OSError: + running = False + + # Best perf from library + exec_time_ms = 0.0 + speedup = 0.0 + kernel_count = 0 + if library.get("kernels"): + kernels = library["kernels"] + kernel_count = len(kernels) + sorted_kernels = sorted(kernels, key=lambda k: k.get("exec_time_ms", float("inf"))) + best = sorted_kernels[0] + exec_time_ms = best.get("exec_time_ms", 0) + for k in kernels: + if k.get("iteration_added") == 0: + seed_time = k.get("exec_time_ms", 0) + if seed_time > 0 and exec_time_ms > 0: + speedup = seed_time / exec_time_ms + break + + # Read shapes from per-GPU requirements + gpu_req = _load_json(gpu_sd / "requirements.json", {}) + shapes_notes = gpu_req.get("extra_notes", "") + + # Read optimizer mode from per-GPU requirements + optimizer_mode = gpu_req.get("optimizer_mode", "Triton (standard)") + if isinstance(optimizer_mode, list): + optimizer_mode = optimizer_mode[0] if optimizer_mode else "Triton (standard)" + + return { + "gpu_idx": gpu_idx, + "label": f"GPU {gpu_idx}", + "phase": phase, + "iteration": iteration, + "running": running, + "pid": pid, + "exec_time_ms": round(exec_time_ms, 4) if exec_time_ms else 0, + "speedup": round(speedup, 2) if speedup else 0, + "kernel_count": kernel_count, + "shapes": shapes_notes, + "graph": graph, + "library": library, + "agents": agents_summary, + "optimizer_mode": optimizer_mode, + } + + +def read_gpu_harness(workspace_dir: Path, gpu_idx: int, harness_type: str) -> Dict[str, Any]: + """Read correctness or perf harness for a single GPU worker.""" + return read_harness(_gpu_workspace_dir(workspace_dir, gpu_idx), harness_type) + + +def read_gpu_kernel_lineage(workspace_dir: Path, state_dir: Path, + gpu_idx: int, kernel_id: str) -> Dict[str, Any]: + """Read kernel lineage for a single GPU worker.""" + return read_kernel_lineage( + _gpu_workspace_dir(workspace_dir, gpu_idx), + _gpu_state_dir(state_dir, gpu_idx), + kernel_id, + ) + + +def read_combined_timeline(state_dir: Path, since: float = 0.0) -> List[Dict[str, Any]]: + """Merge timeline events from parent + all gpu_N/ subdirectories. + + Each event gets a ``source`` field (``"parent"`` or ``"gpu_N"``) so + the frontend can color-code or filter by source. + """ + all_events: List[Dict[str, Any]] = [] + + # Parent timeline + parent_events = _read_timeline_file(state_dir, since) + for ev in parent_events: + ev["source"] = "parent" + all_events.append(ev) + + # Per-GPU timelines + for entry in sorted(state_dir.iterdir()): + if not entry.is_dir() or not entry.name.startswith("gpu_"): + continue + gpu_events = _read_timeline_file(entry, since) + for ev in gpu_events: + ev["source"] = entry.name + all_events.append(ev) + + # Sort by timestamp + all_events.sort(key=lambda e: e.get("ts", 0)) + return all_events + + +def read_summary(state_dir: Path, workspace_dir: Path) -> Dict[str, Any]: + """Read the end-of-task summary report. + + Looks for: + 1. ``summary.txt`` in state_dir (multi-GPU orchestrator writes this) + 2. ``summary`` field in ``run.json`` + 3. Falls back to auto-generating from kernel libraries + """ + # Method 1: summary.txt file + summary_path = state_dir / "summary.txt" + if summary_path.is_file(): + return { + "report": summary_path.read_text(encoding="utf-8"), + "source": "summary.txt", + } + + # Method 2: summary field in run.json + run = _load_json(state_dir / "run.json", {}) + if run.get("summary"): + return { + "report": run["summary"], + "source": "run.json", + } + + # Method 3: auto-generate from available data + gpu_summaries: List[Dict[str, Any]] = [] + for gpu_id in range(8): # max 8 GPUs + gpu_sd = state_dir / f"gpu_{gpu_id}" + run_path = gpu_sd / "run.json" + if not run_path.is_file(): + break + gpu_run = _load_json(run_path, {}) + gpu_ws = workspace_dir / f"gpu_{gpu_id}" + lib_path = gpu_ws / "kernel_library.json" + + seed_ms = 0.0 + best_ms = 0.0 + kernel_count = 0 + if lib_path.is_file(): + try: + lib = json.loads(lib_path.read_text(encoding="utf-8")) + kernel_count = len(lib) + if lib: + best = min(lib, key=lambda k: k.get("exec_time_ms", float("inf"))) + best_ms = best.get("exec_time_ms", 0) + for k in lib: + if k.get("iteration_added") == 0: + seed_ms = k.get("exec_time_ms", 0) + break + except Exception: + pass + + gpu_summaries.append({ + "gpu": gpu_id, + "phase": gpu_run.get("current_phase", "?"), + "status": gpu_run.get("final_status", "running"), + "seed_ms": seed_ms, + "best_ms": best_ms, + "speedup": round(seed_ms / best_ms, 2) if seed_ms > 0 and best_ms > 0 else 0, + "kernels": kernel_count, + }) + + # Build a simple text report + lines = ["## Auto-generated summary (task still running or summary not yet written)", ""] + lines.append("| GPU | Seed (ms) | Best (ms) | Speedup | Kernels | Status |") + lines.append("|-----|----------|----------|---------|---------|--------|") + for s in gpu_summaries: + lines.append( + f"| GPU {s['gpu']} | {s['seed_ms']:.4f} | {s['best_ms']:.4f} | " + f"{s['speedup']:.2f}× | {s['kernels']} | {s['status']} |" + ) + + return { + "report": "\n".join(lines), + "source": "auto-generated", + "gpu_summaries": gpu_summaries, + } diff --git a/metainfer/tasks/evolve_kernel/server/plugin.py b/metainfer/tasks/evolve_kernel/server/plugin.py index efc83cdb..dc5a0713 100644 --- a/metainfer/tasks/evolve_kernel/server/plugin.py +++ b/metainfer/tasks/evolve_kernel/server/plugin.py @@ -19,11 +19,17 @@ "Feed in a Triton GEMM kernel — the LLM generates test harnesses " "and iteratively optimizes it for GPU performance." ), - detail_view_module="app/ok-detail", + detail_view_module="app/ok-evolve-detail", qa_config=_QA_CONFIG, build_router=build_router, frontend_dir=Path(__file__).resolve().parent.parent / "static", - importmap_entries={}, + importmap_entries={ + "app/ok-evolve-detail": "/static/plugins/evolve-kernel/ok-evolve-detail.js?v=CACHE_BUST", + "app/ok-evolve-multi-gpu": "/static/plugins/evolve-kernel/ok-evolve-multi-gpu.js?v=CACHE_BUST", + "app/ok-evolve-state-graph": "/static/plugins/evolve-kernel/ok-evolve-state-graph.js?v=CACHE_BUST", + "app/ok-evolve-kernel-library": "/static/plugins/evolve-kernel/ok-evolve-kernel-library.js?v=CACHE_BUST", + "app/ok-evolve-runtime-api": "/static/plugins/evolve-kernel/ok-evolve-runtime-api.js?v=CACHE_BUST", + }, extra_stylesheets=["ok.css"], ) diff --git a/metainfer/tasks/evolve_kernel/server/routes.py b/metainfer/tasks/evolve_kernel/server/routes.py index 82b91d58..b5bdea64 100644 --- a/metainfer/tasks/evolve_kernel/server/routes.py +++ b/metainfer/tasks/evolve_kernel/server/routes.py @@ -71,7 +71,10 @@ def ok_state_graph(task_id: str) -> Dict[str, Any]: def ok_kernel_library(task_id: str) -> Dict[str, Any]: entry = task_or_404(task_id) require_task_type(entry, PLUGIN_TYPE) - return _state_readers.read_kernel_library(workspace_dir_for(entry)) + result = _state_readers.read_kernel_library(workspace_dir_for(entry)) + # Attach optimizer_mode from requirements for smart code display + result["optimizer_mode"] = _state_readers.read_optimizer_mode(state_dir_for(entry)) + return result # ---- Harnesses ---- @@ -95,8 +98,116 @@ def ok_reference_kernel(task_id: str) -> Dict[str, Any]: require_task_type(entry, PLUGIN_TYPE) return _state_readers.read_reference_kernel(workspace_dir_for(entry)) - # ---- QA ---- + # ---- Kernel Lineage ---- + + @router.get("/kernel-library/{kernel_id}/lineage") + def ok_kernel_lineage(task_id: str, kernel_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + result = _state_readers.read_kernel_lineage( + workspace_dir_for(entry), state_dir_for(entry), kernel_id, + ) + if "error" in result: + raise HTTPException(404, result["error"]) + return result + + # ---- Failures ---- + + @router.get("/failures") + def ok_failures(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_failures(state_dir_for(entry)) + + # ---- Shape Benchmark ---- + + @router.get("/shape-benchmark") + def ok_shape_benchmark(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_shape_benchmark( + state_dir_for(entry), workspace_dir_for(entry), + ) + + @router.post("/shape-benchmark/refresh") + def ok_shape_benchmark_refresh(task_id: str) -> Dict[str, Any]: + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.refresh_shape_benchmark( + state_dir_for(entry), workspace_dir_for(entry), + ) + + # ---- Multi-GPU ---- + + @router.get("/gpu-status") + def ok_gpu_status(task_id: str) -> Dict[str, Any]: + """Return live status of all GPU workers for a multi-GPU task.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_gpu_status( + state_dir_for(entry), workspace_dir_for(entry), + ) - register_qa_routes(router, plugin, prefix="/qa") + @router.get("/aggregate-bench") + def ok_aggregate_bench(task_id: str) -> Dict[str, Any]: + """Aggregated shape benchmarks across all GPUs.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_aggregate_bench( + state_dir_for(entry), workspace_dir_for(entry), + ) + + @router.get("/combined-timeline") + def ok_combined_timeline(task_id: str, since: float = 0.0) -> Dict[str, Any]: + """Timeline events from parent + all GPU workers.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + events = _state_readers.read_combined_timeline( + state_dir_for(entry), since=since, + ) + return {"events": events} + + # ---- Per-GPU Detail ---- + + @router.get("/gpu/{gpu_idx}/detail") + def ok_gpu_detail(task_id: str, gpu_idx: int) -> Dict[str, Any]: + """Aggregated detail for one GPU worker: state graph + kernel library + agents.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_gpu_detail( + state_dir_for(entry), workspace_dir_for(entry), gpu_idx, + ) + + @router.get("/gpu/{gpu_idx}/harnesses/{harness_type}") + def ok_gpu_harness(task_id: str, gpu_idx: int, harness_type: str) -> Dict[str, Any]: + """Read correctness or perf harness for one GPU worker.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_gpu_harness( + workspace_dir_for(entry), gpu_idx, harness_type, + ) + + @router.get("/gpu/{gpu_idx}/kernel-library/{kernel_id}/lineage") + def ok_gpu_kernel_lineage(task_id: str, gpu_idx: int, kernel_id: str) -> Dict[str, Any]: + """Kernel lineage for one GPU worker.""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + result = _state_readers.read_gpu_kernel_lineage( + workspace_dir_for(entry), state_dir_for(entry), gpu_idx, kernel_id, + ) + if "error" in result: + raise HTTPException(404, result["error"]) + return result + + # ---- Summary Report ---- + + @router.get("/summary") + def ok_summary(task_id: str) -> Dict[str, Any]: + """Return the end-of-task summary report (Markdown text).""" + entry = task_or_404(task_id) + require_task_type(entry, PLUGIN_TYPE) + return _state_readers.read_summary(state_dir_for(entry), workspace_dir_for(entry)) + + # ---- QA ---- return router diff --git a/metainfer/tasks/evolve_kernel/static/ok-detail.js b/metainfer/tasks/evolve_kernel/static/ok-detail.js deleted file mode 100644 index 9e316db0..00000000 --- a/metainfer/tasks/evolve_kernel/static/ok-detail.js +++ /dev/null @@ -1,233 +0,0 @@ -// evolve-kernel task detail body. -// -// Renders the 8-phase kernel optimization flow: -// Bootstrap: A(BGen Correctness Harness)→B(Review)→C(Gen Perf Harness)→D(Review) -// Loop: E(Select)→F(Optimize)→G(Verify)→H(Measure)⟲ -// -// Composed of: state graph, kernel library, harness status, iteration metrics. - -import { html } from "htm/preact"; -import { useCallback, useEffect, useState } from "preact/hooks"; -import { StateGraph } from "app/ok-state-graph"; -import { KernelLibrary } from "app/ok-kernel-library"; -import { AgentsPanel } from "app/agents-panel"; -import { Timeline } from "app/timeline"; -import { - getIterations, getCharts, getStateGraph, - getKernelLibrary, getCorrectnessHarness, getPerfHarness, - getReferenceKernel, -} from "app/ok-runtime-api"; - -const withTimeout = (p, ms = 8000) => - Promise.race([ - p, - new Promise((_, rej) => setTimeout(() => rej(new Error("timeout")), ms)), - ]); - -function useRuntimeData(taskId) { - const [data, setData] = useState({ - iterations: [], charts: null, graph: null, - library: null, correctnessHarness: null, perfHarness: null, - refKernel: null, - }); - const refresh = useCallback(async () => { - if (!taskId) return; - try { - const [it, ch, g, lib, chs, phs, ref] = await Promise.all([ - withTimeout(getIterations(taskId)).catch(() => []), - withTimeout(getCharts(taskId)).catch(() => null), - withTimeout(getStateGraph(taskId)).catch(() => null), - withTimeout(getKernelLibrary(taskId)).catch(() => null), - withTimeout(getCorrectnessHarness(taskId)).catch(() => null), - withTimeout(getPerfHarness(taskId)).catch(() => null), - withTimeout(getReferenceKernel(taskId)).catch(() => null), - ]); - setData({ - iterations: it || [], charts: ch, graph: g, - library: lib, correctnessHarness: chs, perfHarness: phs, - refKernel: ref, - }); - } catch (e) { - console.warn("evolve-kernel runtime fetch error:", e); - } - }, [taskId]); - useEffect(() => { refresh(); }, [refresh]); - useEffect(() => { - if (!taskId) return; - const id = setInterval(refresh, 5000); - return () => clearInterval(id); - }, [taskId, refresh]); - return { ...data, refresh }; -} - -// ---- Helper: flow indicator showing the active phase ---- - -function FlowIndicator({ graph }) { - if (!graph) return null; - const cur = graph.current; - if (cur === "idle" || cur === "finished") { - return html`
- Status: - ${cur === "finished" ? "Finished" : "Idle (waiting)"} -
`; - } - - const phases = [ - "A_gen_correctness_harness", "B_review_correctness_harness", - "C_gen_perf_harness", "D_review_perf_harness", - "E_select_kernel", "F_optimize", "G_verify_correctness", "H_measure_perf", - ]; - const labels = { - "A_gen_correctness_harness": "A: Gen Correctness Harness", - "B_review_correctness_harness": "B: Review Correctness Harness", - "C_gen_perf_harness": "C: Gen Perf Harness", - "D_review_perf_harness": "D: Review Perf Harness", - "E_select_kernel": "E: Select Kernel", - "F_optimize": "F: Optimize", - "G_verify_correctness": "G: Verify Correctness", - "H_measure_perf": "H: Measure Perf", - }; - const curIdx = phases.indexOf(cur); - - return html`
- Active: - ${phases.map((p, i) => html` - curIdx ? 0.3 : 1.0 }}> - ${labels[p]} - - ${i < phases.length - 1 ? html`` : null} - `)} -
`; -} - -// ---- Metric cards ---- - -function MetricCards({ library, iterations }) { - const last = iterations && iterations.length > 0 ? iterations[iterations.length - 1] : null; - const perf = last && last.perf ? last.perf : {}; - const bestKernel = library && library.kernels && library.kernels.length > 0 ? library.kernels[0] : null; - - return html`
-
-
${library ? library.size : 0}
-
Kernels in Library
-
-
-
${bestKernel && bestKernel.exec_time_ms ? bestKernel.exec_time_ms.toFixed(3) + " ms" : "—"}
-
Best Exec Time
-
-
-
${perf.speedup ? perf.speedup.toFixed(2) + "×" : "—"}
-
Last Speedup
-
-
-
${iterations ? iterations.length : 0}
-
Optimization Iterations
-
-
`; -} - -// ---- Harness status badges ---- - -function HarnessStatus({ harness, label }) { - if (!harness) { - return html`${label}: pending`; - } - if (harness.exists) { - return html`${label}: ready (${harness.lines} lines)`; - } - return html`${label}: not generated`; -} - -// ---- Main view ---- - -export default function OptKernelDetailView({ - taskId, - run, - status, - data, -}) { - const [selectedKernel, setSelectedKernel] = useState(null); - const [showHarness, setShowHarness] = useState(null); // "correctness" | "perf" | null - const { timeline, agents, loadState, lastErr } = data; - const rt = useRuntimeData(taskId); - - if (loadState === "error" && lastErr) { - return html` -
- Refresh failed: ${lastErr} - (auto-retry) -
- `; - } - - return html` - <${FlowIndicator} graph=${rt.graph} /> - - <${MetricCards} library=${rt.library} iterations=${rt.iterations} /> - -
-
-

State Machine

- <${StateGraph} graph=${rt.graph} /> -
- -
-

Kernel Library - (click for code) -

- <${KernelLibrary} - library=${rt.library} - selectedKernelId=${selectedKernel && selectedKernel.id} - onSelectKernel=${(k) => setSelectedKernel(k)} /> -
-
- - ${selectedKernel ? html` -
-

Kernel: ${selectedKernel.id.slice(0, 8)}… - -

-
${selectedKernel.code || selectedKernel.code_preview || "Code not available"}
-
- ` : null} - -
-
-

Harnesses

-
- <${HarnessStatus} harness=${rt.correctnessHarness} label="Correctness" /> - <${HarnessStatus} harness=${rt.perfHarness} label="Performance" /> -
-
- - -
- - ${showHarness === "correctness" && rt.correctnessHarness && rt.correctnessHarness.exists ? html` -
${rt.correctnessHarness.code}
- ` : null} - ${showHarness === "perf" && rt.perfHarness && rt.perfHarness.exists ? html` -
${rt.perfHarness.code}
- ` : null} -
- -
-

Live Sub-agents

- <${AgentsPanel} agents=${agents} /> -
-
- -
-

Event Timeline

- <${Timeline} events=${timeline.events} /> -
- `; -} diff --git a/metainfer/tasks/evolve_kernel/static/ok-evolve-detail.js b/metainfer/tasks/evolve_kernel/static/ok-evolve-detail.js new file mode 100644 index 00000000..faf85202 --- /dev/null +++ b/metainfer/tasks/evolve_kernel/static/ok-evolve-detail.js @@ -0,0 +1,619 @@ +// evolve-kernel task detail body. +// +// Renders the 8-phase kernel optimization flow: +// Bootstrap: A(BGen Correctness Harness)→B(Review)→C(Gen Perf Harness)→D(Review) +// Loop: E(Select)→F(Optimize)→G(Verify)→H(Measure)⟲ +// +// Composed of: state graph, kernel library, harness status, iteration metrics. + +import { html } from "htm/preact"; +import { useCallback, useEffect, useState } from "preact/hooks"; +import { StateGraph } from "app/ok-evolve-state-graph"; +import { KernelLibrary, KernelCodeView } from "app/ok-evolve-kernel-library"; +import { AgentsPanel } from "app/agents-panel"; +import { MultiGpuDashboard } from "app/ok-evolve-multi-gpu"; +import { + getIterations, getCharts, getStateGraph, + getKernelLibrary, getCorrectnessHarness, getPerfHarness, + getReferenceKernel, getKernelLineage, getFailures, + getShapeBenchmark, refreshShapeBenchmark, +} from "app/ok-evolve-runtime-api"; + +const withTimeout = (p, ms = 8000) => + Promise.race([ + p, + new Promise((_, rej) => setTimeout(() => rej(new Error("timeout")), ms)), + ]); + +function useRuntimeData(taskId) { + const [data, setData] = useState({ + iterations: [], charts: null, graph: null, + library: null, correctnessHarness: null, perfHarness: null, + refKernel: null, failures: null, + }); + const refresh = useCallback(async () => { + if (!taskId) return; + if (run?.multi_gpu === false) { setIsMultiGpu(false); return; } + try { + const [it, ch, g, lib, chs, phs, ref, fails] = await Promise.all([ + withTimeout(getIterations(taskId)).catch(() => []), + withTimeout(getCharts(taskId)).catch(() => null), + withTimeout(getStateGraph(taskId)).catch(() => null), + withTimeout(getKernelLibrary(taskId)).catch(() => null), + withTimeout(getCorrectnessHarness(taskId)).catch(() => null), + withTimeout(getPerfHarness(taskId)).catch(() => null), + withTimeout(getReferenceKernel(taskId)).catch(() => null), + withTimeout(getFailures(taskId)).catch(() => null), + ]); + setData({ + iterations: it || [], charts: ch, graph: g, + library: lib, correctnessHarness: chs, perfHarness: phs, + refKernel: ref, failures: fails, + }); + } catch (e) { + console.warn("evolve-kernel runtime fetch error:", e); + } + }, [taskId]); + useEffect(() => { refresh(); }, [refresh]); + useEffect(() => { + if (!taskId) return; + if (run?.multi_gpu === false) { setIsMultiGpu(false); return; } + const id = setInterval(refresh, 5000); + return () => clearInterval(id); + }, [taskId, refresh]); + return { ...data, refresh }; +} + +// ---- Helper: flow indicator showing the active phase ---- + +function FlowIndicator({ graph }) { + if (!graph) return null; + const cur = graph.current; + if (cur === "idle" || cur === "finished") { + return html`
+ Status: + ${cur === "finished" ? "Finished" : "Idle (waiting)"} +
`; + } + + const phases = [ + "A_gen_correctness_harness", "B_review_correctness_harness", + "C_gen_perf_harness", "D_review_perf_harness", + "E_select_kernel", "F_optimize", "G_verify_correctness", "H_measure_perf", + ]; + const labels = { + "A_gen_correctness_harness": "A: Gen Correctness Harness", + "B_review_correctness_harness": "B: Review Correctness Harness", + "C_gen_perf_harness": "C: Gen Perf Harness", + "D_review_perf_harness": "D: Review Perf Harness", + "E_select_kernel": "E: Select Kernel", + "F_optimize": "F: Optimize", + "G_verify_correctness": "G: Verify Correctness", + "H_measure_perf": "H: Measure Perf", + }; + const curIdx = phases.indexOf(cur); + + return html`
+ Active: + ${phases.map((p, i) => html` + curIdx ? 0.3 : 1.0 }}> + ${labels[p]} + + ${i < phases.length - 1 ? html`` : null} + `)} +
`; +} + +// ---- Metric cards ---- + +function MetricCards({ library, iterations }) { + const last = iterations && iterations.length > 0 ? iterations[iterations.length - 1] : null; + const perf = last && last.perf ? last.perf : {}; + const bestKernel = library && library.kernels && library.kernels.length > 0 ? library.kernels[0] : null; + + return html`
+
+
${library ? library.size : 0}
+
Kernels in Library
+
+
+
${bestKernel && bestKernel.exec_time_ms ? bestKernel.exec_time_ms.toFixed(3) + " ms" : "—"}
+
Best Exec Time
+
+
+
${perf.speedup ? perf.speedup.toFixed(2) + "×" : "—"}
+
Last Speedup
+
+
+
${iterations ? iterations.length : 0}
+
Optimization Iterations
+
+
`; +} + +// ---- Harness status badges ---- + +function HarnessStatus({ harness, label }) { + if (!harness) { + return html`${label}: pending`; + } + if (harness.exists) { + return html`${label}: ready (${harness.lines} lines)`; + } + return html`${label}: not generated`; +} + +// ---- Headroom Analysis Card ---- + +function bottleneckClass(bottleneck) { + if (!bottleneck) return ""; + if (bottleneck === "near_optimal") return "ok-score-good"; + if (bottleneck === "memory_bound" || bottleneck === "compute_bound") return "ok-score-ok"; + return "ok-score-bad"; +} + +function bottleneckLabel(bottleneck) { + if (!bottleneck) return "Unknown"; + const labels = { + memory_bound: "Memory-Bound", + compute_bound: "Compute-Bound", + near_optimal: "Near-Optimal", + inefficient: "Inefficient", + }; + return labels[bottleneck] || bottleneck; +} + +function HeadroomCard({ headroom }) { + if (!headroom || !headroom.has_data) return null; + + const bwW = Math.min(100, Math.max(0, headroom.bw_util_pct || 0)); + const compW = Math.min(100, Math.max(0, headroom.compute_util_pct || 0)); + + return html` +
+
+ Roofline Analysis + + ${bottleneckLabel(headroom.bottleneck)} + +
+ +
+
+
+ HBM Bandwidth + ${headroom.achieved_bw_gbps != null ? headroom.achieved_bw_gbps.toFixed(0) : "—"} / ${headroom.peak_bw_gbps != null ? headroom.peak_bw_gbps.toFixed(0) : "—"} GB/s +
+
+
+
+ ${headroom.bw_util_pct != null ? headroom.bw_util_pct.toFixed(0) : "—"}% +
+ +
+
+ Compute (TFLOPS) + ${headroom.achieved_tflops != null ? headroom.achieved_tflops.toFixed(2) : "—"} / ${headroom.peak_tflops != null ? headroom.peak_tflops.toFixed(0) : "—"} TFLOPS +
+
+
+
+ ${headroom.compute_util_pct != null ? headroom.compute_util_pct.toFixed(0) : "—"}% +
+
+ +
+ Estimated Headroom: ${headroom.headroom_pct != null ? headroom.headroom_pct.toFixed(0) + "%" : "—"} + Arithmetic Intensity: ${headroom.arithmetic_intensity != null ? headroom.arithmetic_intensity.toFixed(1) : "—"} FLOP/byte +
+ + ${headroom.advice ? html` +
${headroom.advice}
+ ` : null} + + ${headroom.suggestions && headroom.suggestions.length > 0 ? html` +
+ Optimization Suggestions (${headroom.suggestions.length}) +
    + ${headroom.suggestions.map(s => html`
  • ${s}
  • `)} +
+
+ ` : null} +
+ `; +} + +// ---- Kernel Lineage Panel (shown below kernel code when a kernel is selected) ---- + +function KernelLineagePanel({ kernelId, library, onSelectKernel, taskId }) { + const [lineage, setLineage] = useState(null); + const [loading, setLoading] = useState(false); + + useEffect(() => { + if (!taskId || !kernelId) return; + let cancelled = false; + setLoading(true); + getKernelLineage(taskId, kernelId) + .then((d) => { if (!cancelled) setLineage(d); }) + .catch((e) => { if (!cancelled) console.warn("lineage fetch error:", e); }) + .finally(() => { if (!cancelled) setLoading(false); }); + return () => { cancelled = true; }; + }, [taskId, kernelId]); + + if (loading && !lineage) { + return html`

Loading lineage…

`; + } + if (!lineage || lineage.error) { + return null; + } + + const { parent, improvement, ancestor_chain, retrospective, headroom } = lineage; + + return html` +
+
+ Optimization Lineage + ${ancestor_chain && ancestor_chain.length > 0 ? html` + + Ancestors: ${ancestor_chain.map(id => html`${id.slice(0, 8)}`)} + + ` : html`Seed kernel (no ancestors)`} +
+ +
+
+ +
+ Parent Time + ${parent && parent.exec_time_ms ? parent.exec_time_ms.toFixed(4) + " ms" : "—"} +
+
+ Speedup vs Parent + 1 ? "ok-score-good" : improvement.speedup_vs_parent && improvement.speedup_vs_parent < 1 ? "ok-score-bad" : "")}> + ${improvement.speedup_vs_parent ? improvement.speedup_vs_parent.toFixed(2) + "×" : "—"} + +
+
+ Time Delta + + ${improvement.exec_time_delta_ms != null ? (improvement.exec_time_delta_ms > 0 ? "+" : "") + improvement.exec_time_delta_ms.toFixed(4) + " ms" : "—"} + +
+
+ + ${improvement.summary_items && improvement.summary_items.length > 0 ? html` +
+
Improvement Summary
+
    + ${improvement.summary_items.map(item => html`
  • ${item}
  • `)} +
+
+ ` : improvement.summary ? html` +
+
Improvement Summary
+
${improvement.summary}
+
+ ` : null} + + ${headroom && headroom.has_data ? html` + <${HeadroomCard} headroom=${headroom} /> + ` : null} + + ${retrospective && retrospective.has_retrospective ? html` +
+ Full Retrospective (iteration ${retrospective.iteration}) +
${retrospective.markdown}
+
+ ` : null} +
+
+ `; +} + +// ---- Shape Benchmark Panel ---- + +function speedupClass(val) { + if (val == null || val <= 0) return ""; + if (val >= 1.2) return "ok-score-good"; + if (val >= 1.0) return "ok-score-ok"; + return "ok-score-bad"; +} + +function formatMs2(ms) { + if (ms == null || ms === 0) return "—"; + if (ms < 1) return (ms * 1000).toFixed(1) + " μs"; + if (ms < 1000) return ms.toFixed(3) + " ms"; + return (ms / 1000).toFixed(2) + " s"; +} + +function ShapeBenchmarkPanel({ taskId }) { + const [data, setData] = useState(null); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [refreshing, setRefreshing] = useState(false); + + const load = useCallback(async (forceRefresh) => { + if (!taskId) return; + if (run?.multi_gpu === false) { setIsMultiGpu(false); return; } + setLoading(true); + setError(null); + try { + const d = forceRefresh + ? await refreshShapeBenchmark(taskId) + : await getShapeBenchmark(taskId); + setData(d); + if (d.error) setError(d.error); + } catch (e) { + setError(e.message); + setData(null); + } finally { + setLoading(false); + setRefreshing(false); + } + }, [taskId]); + + useEffect(() => { load(false); }, [load]); + + if (loading && !data) { + return html`
+

Shape Benchmark

+

Loading benchmarks…

+
`; + } + + if (error && (!data || !data.results || data.results.length === 0)) { + return html`
+

Shape Benchmark

+

${error}

+
`; + } + + if (!data || !data.results || data.results.length === 0) { + return null; + } + + const results = data.results; + const bestId = data.best_kernel_id ? data.best_kernel_id.slice(0, 8) : "?"; + + // Group results by shape + const shapeMap = {}; + for (const r of results) { + const key = r.shape_label; + if (!shapeMap[key]) shapeMap[key] = []; + shapeMap[key].push(r); + } + + return html` +
+

Shape Benchmark — Best Kernel vs Baseline + + (kernel ${bestId}, ${data.cached ? "cached" : "fresh"}) + + +

+ +
+ ${Object.entries(shapeMap).map(([label, rows]) => html` +
+
${label}
+ + + + + + + + + + + + ${rows.map(r => html` + + + + + + + + `)} + +
M(M×K×N)BaselineBest KernelSpeedup
${r.M}${r.M}×${r.K}×${r.N}${formatMs2(r.ref_ms)}${formatMs2(r.best_ms)} + ${r.error ? html`err` + : r.speedup.toFixed(2) + "×"} +
+
+ `)} +
+
+ `; +} + + +// ---- Failure Log Panel ---- + +function FailureLog({ failures }) { + const [expandedId, setExpandedId] = useState(null); + + if (!failures || !failures.failures || failures.failures.length === 0) { + return null; + } + + const { failures: items, total, errors, warnings } = failures; + + return html` +
+

Failure Log + + ${errors} errors + ${warnings > 0 ? html`${warnings} warnings` : null} + +

+ +
+ ${items.map((f, i) => { + const isExpanded = expandedId === i; + const entryId = `fail-${f.iteration}-${f.phase || "iter"}-${i}`; + return html` +
+
setExpandedId(isExpanded ? null : i)} + > + ${f.severity === "error" ? "✗" : "⚠"} + Iter ${f.iteration} + ${f.phase_label} + ${f.summary} + ${f.attempts > 1 ? html`${f.attempts} attempts` : null} + ${isExpanded ? "▲" : "▼"} +
+ ${isExpanded ? html` +
+
+ Outcome: ${f.outcome} + ${f.attempts > 1 ? html`Attempts: ${f.attempts}` : null} +
+
${f.detail}
+
+ ` : null} +
+ `; + })} +
+
+ `; +} + +// ---- Main view ---- + +export default function OptKernelDetailView({ + taskId, + run, + status, + data, +}) { + const [selectedKernel, setSelectedKernel] = useState(null); + const [showHarness, setShowHarness] = useState(null); // "correctness" | "perf" | null + const { agents, loadState, lastErr } = data; + const rt = useRuntimeData(taskId); + const [isMultiGpu, setIsMultiGpu] = useState(run?.multi_gpu ? true : null); + + // Check if this task is a multi-GPU task (initialized from run prop) + useEffect(() => { + if (!taskId) return; + if (run?.multi_gpu === false) { setIsMultiGpu(false); return; } + fetch(`/api/evolve-kernel/${taskId}/gpu-status`) + .then(r => r.json()) + .then(d => setIsMultiGpu(d.is_multi_gpu ? true : false)) + .catch(() => setIsMultiGpu(false)); + }, [taskId]); + + if (loadState === "error" && lastErr) { + return html` +
+ Refresh failed: ${lastErr} + (auto-retry) +
+ `; + } + + // If this is a multi-GPU task, show the unified dashboard + if (isMultiGpu) { + return html` + <${MultiGpuDashboard} taskId=${taskId} /> + `; + } + + return html` + <${FlowIndicator} graph=${rt.graph} /> + + <${MetricCards} library=${rt.library} iterations=${rt.iterations} /> + +
+
+

State Machine

+ <${StateGraph} graph=${rt.graph} /> +
+ +
+

Kernel Library + (click for details) +

+ <${KernelLibrary} + library=${rt.library} + selectedKernelId=${selectedKernel && selectedKernel.id} + onSelectKernel=${(k) => setSelectedKernel(k)} /> +
+
+ + ${selectedKernel ? html` + <${KernelLineagePanel} + kernelId=${selectedKernel.id} + library=${rt.library} + onSelectKernel=${(k) => setSelectedKernel(k)} + taskId=${taskId} + /> +
+

Kernel Source: ${selectedKernel.id.slice(0, 8)}… + +

+ <${KernelCodeView} kernel=${selectedKernel} optimizerMode=${(rt.library && rt.library.optimizer_mode) || "Triton (standard)"} /> +
+ ` : null} + +
+
+

Harnesses

+
+ <${HarnessStatus} harness=${rt.correctnessHarness} label="Correctness" /> + <${HarnessStatus} harness=${rt.perfHarness} label="Performance" /> +
+
+ + +
+ + ${showHarness === "correctness" && rt.correctnessHarness && rt.correctnessHarness.exists ? html` +
+ ${rt.correctnessHarness.code.split("\n").map(line => html`
${line}
`)} +
+ ` : null} + ${showHarness === "perf" && rt.perfHarness && rt.perfHarness.exists ? html` +
+ ${rt.perfHarness.code.split("\n").map(line => html`
${line}
`)} +
+ ` : null} +
+ +
+

Live Sub-agents

+ <${AgentsPanel} agents=${agents} /> +
+
+ + <${ShapeBenchmarkPanel} taskId=${taskId} /> + + <${FailureLog} failures=${rt.failures} /> + `; +} diff --git a/metainfer/tasks/evolve_kernel/static/ok-evolve-kernel-library.js b/metainfer/tasks/evolve_kernel/static/ok-evolve-kernel-library.js new file mode 100644 index 00000000..5913041c --- /dev/null +++ b/metainfer/tasks/evolve_kernel/static/ok-evolve-kernel-library.js @@ -0,0 +1,158 @@ +// Kernel library table component for evolve-kernel. +// Shows ranked kernels with exec_time, complexity, combined_score. +// Color-coded rows by speedup, headroom badge, combined score bar. + +import { html } from "htm/preact"; +import { useState } from "preact/hooks"; + +function scoreClass(val) { + if (val == null || val === 0) return ""; + if (val > 0.8) return "ok-score-good"; + if (val > 0.4) return "ok-score-ok"; + return "ok-score-bad"; +} + +function formatMs(ms) { + if (ms == null || ms === 0) return "—"; + if (ms < 1) return (ms * 1000).toFixed(1) + " μs"; + if (ms < 1000) return ms.toFixed(3) + " ms"; + return (ms / 1000).toFixed(2) + " s"; +} + +function speedupRowClass(k) { + const su = k.speedup_vs_parent; + if (su == null || su <= 0) return ""; + if (su >= 1.2) return "speedup-good"; + if (su < 1.0) return "speedup-bad"; + return ""; +} + +function headroomBadge(headroom) { + if (!headroom || !headroom.bottleneck) return null; + const b = headroom.bottleneck; + const labels = { + near_optimal: "Near-Opt", + memory_bound: "Mem-Bound", + compute_bound: "Comp-Bound", + inefficient: "Ineff", + }; + return html`${labels[b] || b}`; +} + +function combinedBar(score, maxScore) { + const pct = maxScore > 0 ? Math.min(100, (score / maxScore) * 100) : 0; + return html` + + ${score.toFixed(2)} + `; +} + +export function KernelLibrary({ library, selectedKernelId, onSelectKernel }) { + if (!library || !library.kernels || library.kernels.length === 0) { + return html`

Kernel library is empty — waiting for the first optimization result.

`; + } + + const kernels = library.kernels; + const maxScore = Math.max(...kernels.map(k => k.combined_score || 0), 1); + + return html` +
+ + + + + + + + + + + + + + + + ${kernels.map((k, i) => { + const headroom = k.headroom; + const hasProfile = k.profiled; + return html` + onSelectKernel && onSelectKernel(k)} + style=${{ cursor: onSelectKernel ? "pointer" : "default" }} + > + + + + + + + + + + + `; + })} + +
#IDExec TimeΔ ParentScoreCmplxIterLinesHeadroom
${i + 1}${k.id.slice(0, 8)}${formatMs(k.exec_time_ms)} + ${k.speedup_vs_parent + ? html`= 1.2 ? "ok-score-good" : k.speedup_vs_parent >= 1.0 ? "ok-score-ok" : "ok-score-bad"}>${k.speedup_vs_parent.toFixed(2)}×` + : html``} + ${combinedBar(k.combined_score, maxScore)}${((k.complexity_score || 0) * 100).toFixed(0)}%${k.iteration_added || 0}${k.code_lines || 0}${headroomBadge(headroom)}${hasProfile ? html`📊` : null}
+
+

${library.size} kernel(s) ranked by combined score — + ≥1.2× speedup, + regression +

+ `; +} + +// Kernel code viewer: +// - HIP mode: defaults to .cpp source, toggle to .py wrapper +// - Triton mode: shows .py source directly +export function KernelCodeView({ kernel, optimizerMode }) { + const isHip = optimizerMode && ( + optimizerMode === "HIP C++ (from scratch)" || + optimizerMode === "hip_cpp" || + optimizerMode === "hip" + ); + const cppCode = kernel.cpp_code || null; + const pyCode = kernel.code || kernel.code_preview || null; + const hasBoth = cppCode && pyCode; + + // Default: .cpp for HIP mode, .py for Triton mode + const [showCpp, setShowCpp] = useState(isHip); + + const activeCode = showCpp && cppCode ? cppCode : pyCode; + const cppLines = cppCode ? cppCode.split("\n").length : 0; + const pyLines = pyCode ? pyCode.split("\n").length : 0; + + if (!activeCode) return html`

Code not available

`; + + const lines = activeCode.split("\n"); + + return html` +
+ ${hasBoth ? html` +
+ + +
+ ` : cppCode ? html` +

+ HIP C++ kernel source (${cppLines} lines) +

+ ` : null} +
+ ${lines.map(line => html`
${line}
`)} +
+
+ `; +} diff --git a/metainfer/tasks/evolve_kernel/static/ok-evolve-multi-gpu.js b/metainfer/tasks/evolve_kernel/static/ok-evolve-multi-gpu.js new file mode 100644 index 00000000..ffcc6321 --- /dev/null +++ b/metainfer/tasks/evolve_kernel/static/ok-evolve-multi-gpu.js @@ -0,0 +1,719 @@ +// Multi-GPU unified dashboard for evolve-kernel. + +import { html } from "htm/preact"; +import { useCallback, useEffect, useRef, useState } from "preact/hooks"; +import { StateGraph } from "app/ok-evolve-state-graph"; +import { KernelLibrary, KernelCodeView } from "app/ok-evolve-kernel-library"; + +const BASE = (taskId) => `/api/evolve-kernel/${taskId}`; + +async function fetchJSON(url) { + const res = await fetch(url); + if (!res.ok) throw new Error(`${res.status}`); + return res.json(); +} + +// ---- Helpers ---- + +function formatMs(ms) { + if (ms == null || ms <= 0) return "—"; + if (ms < 1) return (ms * 1000).toFixed(1) + " μs"; + if (ms < 1000) return ms.toFixed(3) + " ms"; + return (ms / 1000).toFixed(2) + " s"; +} + +function formatAge(s) { + if (s == null || s <= 0) return ""; + if (s < 60) return Math.round(s) + "s"; + if (s < 3600) return Math.round(s / 60) + "m"; + return (s / 3600).toFixed(1) + "h"; +} + +function speedupClass(val) { + if (val == null || val <= 0) return ""; + if (val >= 1.2) return "ok-score-good"; + if (val >= 1.0) return "ok-score-ok"; + return "ok-score-bad"; +} + +function phaseLabel(p) { + const labels = { + "A_gen_correctness_harness": "A:Gen Harness", "B_review_correctness_harness": "B:Review Harness", + "C_gen_perf_harness": "C:Gen Perf", "D_review_perf_harness": "D:Review Perf", + "E_select_kernel": "E:Select", "F_optimize": "F:Optimize", + "G_verify_correctness": "G:Verify", "H_measure_perf": "H:Measure", + "finished": "Done", "starting": "Starting", "idle": "Idle", "running": "Running", + }; + return labels[p] || (p || "?").slice(0, 20); +} + +function phaseDotClass(p) { + if (!p) return "idle"; + if (p.startsWith("A_") || p.startsWith("B_") || p.startsWith("C_") || p.startsWith("D_")) return "bootstrap"; + if (p.startsWith("E_") || p.startsWith("F_") || p.startsWith("G_") || p.startsWith("H_")) return "optimize"; + if (p === "finished") return "finished"; + return "idle"; +} + +function phaseCardClass(p) { + if (!p) return ""; + if (p.startsWith("A") || p.startsWith("B") || p.startsWith("C") || p.startsWith("D")) return "bootstrap"; + if (p.startsWith("E") || p.startsWith("F") || p.startsWith("G") || p.startsWith("H")) return "optimize"; + if (p === "finished") return "done"; + return ""; +} + +function roleLabel(r) { + const labels = { + "correctness_harness_generator": "Gen Harness", "correctness_harness_reviewer": "Review Harness", + "perf_harness_generator": "Gen Perf", "perf_harness_reviewer": "Review Perf", + "kernel_optimizer": "Optimizer", "correctness_verifier": "Verify", + "perf_measurer": "Measure", "headroom_analyzer": "Headroom", "seed_generator": "Seed", + }; + return labels[r] || (r || "?").replace(/_/g, " ").slice(0, 16); +} + +// ---- Parse shapes from extra_notes ---- + +function parseShapes(notes) { + if (!notes) return []; + const shapes = []; + // Match: name (TP=N): M=a,b,c,... (M, K) @ (K, N) + // Lines are indented within the multi-line string, so don't use ^. + const re = /([\w_]+)\s*\(TP=(\d+)\):\s*M=([\d,]+)\s*\(\s*(?:M|m)\s*,\s*(\d+)\s*\)\s*@\s*\(\s*(\d+)\s*,\s*(\d+)\s*\)/g; + let m; + while ((m = re.exec(notes)) !== null) { + const name = m[1]; + const tp = parseInt(m[2]); + const mValues = m[3].split(',').map(s => parseInt(s.trim())).filter(Boolean); + const K = parseInt(m[4]); + const K2 = parseInt(m[5]); + const N = parseInt(m[6]); + // Skip if K values don't match (malformed line or different format) + if (K !== K2) continue; + shapes.push({ name, tp, mValues, K, N, label: `${name} (TP=${tp}) ${K}×${N}` }); + } + return shapes; +} + +// ---- GPU Status Card ---- + +function GpuCard({ worker, onClick, isSelected }) { + const phase = worker.phase || "starting"; + const running = worker.running; + const gpuLabel = worker.label || `GPU ${worker.gpu_idx}`; + const agents = worker.agents || []; + const activeAgent = agents.find(a => a.phase === "running"); + + return html` +
+
+ ${gpuLabel} + + ${running ? "● Running" : "○ Idle"} + +
+
+
+
+ Phase + + ${phaseLabel(phase)} + +
+
+ Iter + ${worker.iteration || 0} +
+
+
+
+ Best + ${formatMs(worker.exec_time_ms)} +
+
+ Speedup + + ${worker.speedup > 0 ? worker.speedup.toFixed(2) + "×" : "—"} + +
+
+
+
+ Kernels + ${worker.kernel_count || 0} +
+
+
+ ${worker.iteration > 0 ? html` +
+ ` : null} +
+ `; +} + +function GpuCards({ workers, selectedGpu, onSelectGpu }) { + if (!workers || workers.length === 0) { + return html`

No GPU workers found.

`; + } + return html` +
+ ${workers.map(w => html`<${GpuCard} key=${w.gpu_idx} worker=${w} isSelected=${selectedGpu === w.gpu_idx} onClick=${() => onSelectGpu(selectedGpu === w.gpu_idx ? null : w.gpu_idx)} />`)} +
+ `; +} + +// ---- Collapsible Agent List ---- + +function PerGpuAgents({ workers }) { + const hasAnyAgents = workers.some(w => (w.agents || []).length > 0); + if (!hasAnyAgents) return null; + + const [collapsedSet, setCollapsedSet] = useState(() => { + const s = new Set(); + workers.forEach((w, i) => { if (i > 0) s.add(w.gpu_idx); }); + return s; + }); + + return html` +
+

Live Sub-agents + +

+
+ ${workers.map(w => { + const agents = w.agents || []; + if (agents.length === 0) return null; + const isCollapsed = collapsedSet.has(w.gpu_idx); + return html` +
+
{ + setCollapsedSet(prev => { const n = new Set(prev); if (n.has(w.gpu_idx)) n.delete(w.gpu_idx); else n.add(w.gpu_idx); return n; }); + }}> + ${w.label} + +
+ ${agents.map(a => html` +
+ ${roleLabel(a.role)} + + ${a.phase === "running" ? html`● ${formatAge(a.elapsed_s)}` + : a.success === false ? html`` + : a.success === true ? html`` + : html`${a.phase}`} + + ${a.error ? html`${a.error.slice(0, 30)}` : null} +
+ `)} +
+ `; + })} +
+
+ `; +} + +// ---- Stats Bar ---- + +function StatsBar({ detail }) { + if (!detail) return null; + return html` +
+
Phase ${phaseLabel(detail.phase)}
+ | +
Iter ${detail.iteration}
+ | +
Best ${formatMs(detail.exec_time_ms)}
+ | +
Speedup ${detail.speedup > 0 ? detail.speedup.toFixed(2) + "×" : "—"}
+ | +
Library ${detail.kernel_count}/10
+
+ `; +} + +// ---- Roofline callout (clean, no plot) ---- + +const BOTTLENECK_META = { + memory_bound: { color: "#58a6ff", icon: "▦", label: "Memory-Bound", desc: "HBM bandwidth limits performance. The data movement per FLOP (AI) is below the ridge point." }, + compute_bound: { color: "#d29922", icon: "◉", label: "Compute-Bound", desc: "Instruction throughput limits performance. The kernel has enough data reuse to saturate compute." }, + near_optimal: { color: "#3fb950", icon: "✓", label: "Near-Optimal", desc: "The kernel is close to the hardware roofline. Further tuning yields diminishing returns." }, + inefficient: { color: "#f85149", icon: "⚠", label: "Inefficient", desc: "Neither bandwidth nor compute are near peak. Check occupancy, register pressure, or tile alignment." }, +}; + +function RooflineCallout({ headroom }) { + if (!headroom || !headroom.has_data) return null; + + const b = headroom.bottleneck || "inefficient"; + const meta = BOTTLENECK_META[b] || BOTTLENECK_META.inefficient; + const eff = headroom.roofline_efficiency_pct || 0; + const ai = headroom.arithmetic_intensity || 0; + const ridge = headroom.ai_ridge || 0; + const pAchieved = headroom.achieved_tflops || 0; + const pMax = headroom.p_max_tflops || 0; + const peak = headroom.peak_tflops || 220; + const bwPeak = headroom.peak_bw_gbps || 700; + const pBwRoof = headroom.p_bw_roof_tflops; + + return html` +
+
+ ${meta.icon} + ${meta.label} + ${ridge > 0 ? html` + + AI = ${ai.toFixed(1)} ${ai < ridge ? "<" : "≥"} Ridge = ${ridge.toFixed(1)} FLOP/byte + + ` : html` + AI = ${ai.toFixed(1)} FLOP/byte + `} +
+

${meta.desc}

+ +
+
+ P_achieved + ${pAchieved.toFixed(2)} TFLOPS + = FLOPs / exec_time +
+ ${pBwRoof ? html` +
+ P_bw_roof + ${pBwRoof.toFixed(2)} TFLOPS + = BW_peak × AI = ${bwPeak.toFixed(0)} × ${ai.toFixed(0)} / 1000 +
+ ` : null} +
+ P_max + ${pMax > 0 ? pMax.toFixed(2) : `min(${peak.toFixed(0)}T compute, BW×AI)`} TFLOPS + = roofline ceiling +
+
+ +
+ Roofline Eff. +
+
+
+ ${eff.toFixed(0)}% +
+ + ${headroom.advice ? html`

${headroom.advice}

` : null} +
+ `; +} + +// ---- Shape Summary (what this GPU optimizes) ---- + +function ShapeSummary({ detail }) { + // Parse shapes from the agents data (the detail doesn't have extra_notes itself) + // We try to guess from the shapes field or leave it for the caller + return null; // Will be rendered inline with detail.shapes data +} + +// ---- Code with line numbers ---- + +function CodeBlock({ code }) { + if (!code) return html`

Code not available

`; + const lines = code.split("\n"); + return html` +
+ ${lines.map(line => html`
${line}
`)} +
+ `; +} + +// ---- Per-GPU Detail Panel ---- + +function GpuDetailPanel({ taskId, gpuIdx, onClose }) { + const [detail, setDetail] = useState(null); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + const [selectedKernel, setSelectedKernel] = useState(null); + const [kernelLineage, setKernelLineage] = useState(null); + const [lineageLoading, setLineageLoading] = useState(false); + const [showHarness, setShowHarness] = useState(null); + const [harnessData, setHarnessData] = useState({}); + const mountedRef = useRef(true); + + useEffect(() => { + mountedRef.current = true; + setLoading(true); + async function fetchDetail() { + if (!taskId || gpuIdx == null) return; + try { + const res = await fetch(`${BASE(taskId)}/gpu/${gpuIdx}/detail`); + if (!res.ok) throw new Error(`${res.status}`); + const d = await res.json(); + if (!mountedRef.current) return; + setDetail(d); + setError(null); + } catch (e) { + if (!mountedRef.current) return; + setError(e.message); + } finally { + if (mountedRef.current) setLoading(false); + } + } + fetchDetail(); + const pollId = setInterval(fetchDetail, 3000); + return () => { mountedRef.current = false; clearInterval(pollId); }; + }, [taskId, gpuIdx]); + + useEffect(() => { + if (!showHarness || !taskId || gpuIdx == null) return; + const htype = showHarness; + if (harnessData[htype]) return; + let cancelled = false; + (async () => { + try { + const res = await fetch(`${BASE(taskId)}/gpu/${gpuIdx}/harnesses/${htype}`); + if (!res.ok) throw new Error(`${res.status}`); + const d = await res.json(); + if (!cancelled) setHarnessData(prev => ({ ...prev, [htype]: d })); + } catch (e) { /* ignore */ } + })(); + return () => { cancelled = true; }; + }, [showHarness, taskId, gpuIdx]); + + useEffect(() => { + if (!selectedKernel || !taskId || gpuIdx == null) { + setKernelLineage(null); + return; + } + let cancelled = false; + setLineageLoading(true); + (async () => { + try { + const res = await fetch(`${BASE(taskId)}/gpu/${gpuIdx}/kernel-library/${selectedKernel.id}/lineage`); + if (!res.ok) throw new Error(`${res.status}`); + const d = await res.json(); + if (!cancelled) setKernelLineage(d); + } catch (e) { /* ignore */ } + finally { if (!cancelled) setLineageLoading(false); } + })(); + return () => { cancelled = true; }; + }, [selectedKernel, taskId, gpuIdx]); + + if (loading && !detail && !error) { + return html`

Loading GPU ${gpuIdx} detail…

`; + } + + if (error && !detail) { + return html`

Error: ${error}

`; + } + + if (!detail) return null; + + const headroom = kernelLineage && kernelLineage.headroom && kernelLineage.headroom.has_data + ? kernelLineage.headroom : null; + const improvement = kernelLineage ? kernelLineage.improvement : null; + const parent = kernelLineage ? kernelLineage.parent : null; + const profiling = kernelLineage ? kernelLineage.kernel : null; + const correctnessReady = harnessData["correctness"] ? harnessData["correctness"].exists : null; + const perfReady = harnessData["perf"] ? harnessData["perf"].exists : null; + + // Parse target shapes from detail + const shapes = parseShapes(detail.shapes || ""); + + // Categorize M values + const decodeMs = (() => { + const s = new Set(); + shapes.forEach(sh => sh.mValues.forEach(v => { if (v <= 16) s.add(v); })); + return [...s].sort((a, b) => a - b); + })(); + const largeMs = (() => { + const s = new Set(); + shapes.forEach(sh => sh.mValues.forEach(v => { if (v > 16) s.add(v); })); + return [...s].sort((a, b) => a - b); + })(); + + return html` +
+

GPU ${gpuIdx} Detail + ${shapes.map(s => `TP=${s.tp}`).filter((v,i,a) => a.indexOf(v)===i).join(', ')} + +

+ + <${StatsBar} detail=${detail} /> + + + ${shapes.length > 0 ? html` +
+ ${shapes.map(s => { + const hasDecode = s.mValues.some(v => v <= 16); + const hasLarge = s.mValues.some(v => v > 16); + return html` +
+
+ ${s.name} + TP=${s.tp} + ${s.K}×${s.N} +
+
+ ${hasDecode ? html`⚡ M≤16: ${s.mValues.filter(v => v <= 16).join(', ')}` : null} + ${hasLarge ? html`📦 M=${s.mValues.filter(v => v > 16).join(', ')}` : null} +
+
+ `; + })} +
+ ` : null} + + +
+
+

State Machine

+ <${StateGraph} graph=${detail.graph} /> +
+
+

Kernel Library + (${detail.kernel_count} kernels) +

+ <${KernelLibrary} + library=${detail.library} + selectedKernelId=${selectedKernel && selectedKernel.id} + onSelectKernel=${(k) => setSelectedKernel(selectedKernel && selectedKernel.id === k.id ? null : k)} + /> +
+
+ + + ${selectedKernel ? html` +
+
+

Kernel: ${selectedKernel.id.slice(0, 8)}… + + ${formatMs(selectedKernel.exec_time_ms)} | Cmplx ${((selectedKernel.complexity_score || 0) * 100).toFixed(0)}% | Iter ${selectedKernel.iteration_added || 0} + + +

+ <${KernelCodeView} kernel=${selectedKernel} optimizerMode=${detail && detail.optimizer_mode} /> +
+
+ + ${lineageLoading ? html`

Loading analysis…

` + : kernelLineage && !kernelLineage.error ? html` +
+ +
+

Roofline + ${headroom && headroom.shape_label ? html`(on ${headroom.shape_label})` : null} +

+ <${RooflineCallout} headroom=${headroom} /> + + ${profiling && profiling.profiled ? html` +
+
hipprof ${profiling.profiling_kernel_duration_us ? profiling.profiling_kernel_duration_us.toFixed(0) + " µs" : "—"}
+
BW ${profiling.profiling_achieved_bw_gbps ? profiling.profiling_achieved_bw_gbps.toFixed(1) + " GB/s" : "—"}
+
Occ ${profiling.profiling_occupancy_pct ? profiling.profiling_occupancy_pct.toFixed(0) + "%" : "—"}
+
L2$ ${profiling.profiling_l2_cache_hit_pct ? profiling.profiling_l2_cache_hit_pct.toFixed(0) + "%" : "—"}
+
+ ` : null} +
+ + + ${improvement ? html` +
+

Optimization + vs ${parent ? parent.id.slice(0,8) + '…' : 'seed'} +

+
+
${parent && parent.exec_time_ms ? parent.exec_time_ms.toFixed(4) + " ms" : "—"}Parent Time
+
1 ? "ok-score-good" : improvement.speedup_vs_parent && improvement.speedup_vs_parent < 1 ? "ok-score-bad" : "")}>${improvement.speedup_vs_parent ? improvement.speedup_vs_parent.toFixed(2) + "×" : "—"}Speedup
+
${improvement.exec_time_delta_ms != null ? (improvement.exec_time_delta_ms > 0 ? "+" : "") + improvement.exec_time_delta_ms.toFixed(4) + " ms" : "—"}Delta
+
+ + ${improvement.summary_items && improvement.summary_items.length > 0 ? html` +
+
    + ${improvement.summary_items.map(item => html`
  • ${item}
  • `)} +
+
+ ` : improvement.summary ? html` +
+
${improvement.summary}
+
+ ` : null} +
+ ` : html`

Optimization

Seed kernel — no parent to compare.

`} +
+ ` : kernelLineage && kernelLineage.error ? html` +

Analysis not available for this kernel.

+ ` : null} + ` : null} + + +
+
+

Harnesses

+
+ ${correctnessReady ? "✓ Correctness" : correctnessReady === false ? "✗ Correctness" : "… Correctness"} + ${perfReady ? "✓ Perf" : perfReady === false ? "✗ Perf" : "… Perf"} +
+
+ + +
+ ${showHarness && harnessData[showHarness] && harnessData[showHarness].exists ? html` +
+ ${(harnessData[showHarness].code || "").split("\n").map(line => html`
${line}
`)} +
+ ` : showHarness && harnessData[showHarness] && !harnessData[showHarness].exists ? html` +

Harness not yet generated.

+ ` : showHarness ? html`

Loading…

` : null} +
+ +
+

Agents

+
+ ${(detail.agents || []).length === 0 ? html`

No agents yet.

` : null} + ${(detail.agents || []).map(a => html` +
+ ${roleLabel(a.role)} + + ${a.phase === "running" ? html`● ${formatAge(a.elapsed_s)}` + : a.success === false ? html`` + : a.success === true ? html`` + : html`${a.phase}`} + +
+ `)} +
+
+
+
+ `; +} + +// ---- Aggregated Shape Benchmark ---- + +function AggregatedBenchTable({ benchData }) { + if (!benchData || !benchData.results || benchData.results.length === 0) { + return html`

No benchmark data yet. Run shape benchmarks from the task page.

`; + } + const results = benchData.results; + const groups = {}; + for (const r of results) { + const key = r.shape_label || r.gpu_source || "?"; + if (!groups[key]) groups[key] = []; + groups[key].push(r); + } + return html` +
+ ${Object.entries(groups).map(([label, rows]) => html` +
+
${label}
+ + + + ${rows.map(r => html` + = 4096 ? "row-large" : ""}> + + + + + + `)} + +
MBaselineBestSpeedup
${r.M}${formatMs(r.ref_ms)}${formatMs(r.best_ms)}${r.error ? html`err` : ((r.speedup || 0) > 0 ? r.speedup.toFixed(2) + "×" : "—")}
+
+ `)} +
+ `; +} + +// ---- Best Kernels Summary ---- + +function BestKernelsSummary({ benchData }) { + if (!benchData || !benchData.kernels || benchData.kernels.length === 0) return null; + return html` +
+

Best Kernels by GPU

+
+ + + + ${benchData.kernels.map(k => html` + + + + + + + + `)} + +
SourceIDExec TimeCmplxIter
${k.source_workspace || k.gpu_label || "?"}${(k.id || "").slice(0, 8)}${formatMs(k.exec_time_ms)}${((k.complexity_score || 0) * 100).toFixed(0)}%${k.iteration_added || 0}
+
+
+ `; +} + +// ---- Main Multi-GPU Dashboard ---- + +export function MultiGpuDashboard({ taskId }) { + const [gpuData, setGpuData] = useState(null); + const [benchData, setBenchData] = useState(null); + const [loading, setLoading] = useState(true); + const [error, setError] = useState(null); + const [selectedGpu, setSelectedGpu] = useState(null); + + const refresh = useCallback(async () => { + if (!taskId) return; + try { + const [gpu, bench] = await Promise.all([ + fetchJSON(`${BASE(taskId)}/gpu-status`).catch(() => null), + fetchJSON(`${BASE(taskId)}/aggregate-bench`).catch(() => null), + ]); + if (gpu) setGpuData(gpu); + if (bench) setBenchData(bench); + setError(null); + } catch (e) { setError(e.message); } + finally { setLoading(false); } + }, [taskId]); + + useEffect(() => { refresh(); }, [refresh]); + useEffect(() => { + if (!taskId) return; + const id = setInterval(refresh, 3000); + return () => clearInterval(id); + }, [taskId, refresh]); + + if (loading && !gpuData) { + return html`

Multi-GPU Optimization

Loading…

`; + } + if (error && !gpuData) { + return html`

Multi-GPU Optimization

${error}

`; + } + + const workers = gpuData ? gpuData.workers : []; + const runningCount = workers.filter(w => w.running).length; + + return html` +
+

Multi-GPU Optimization + (${runningCount}/${workers.length} GPUs active) +

+ <${GpuCards} workers=${workers} selectedGpu=${selectedGpu} onSelectGpu=${setSelectedGpu} /> +
+ + ${selectedGpu != null ? html` + <${GpuDetailPanel} taskId=${taskId} gpuIdx=${selectedGpu} onClose=${() => setSelectedGpu(null)} /> + ` : null} + + <${PerGpuAgents} workers=${workers} /> + +
+

Aggregated Shape Benchmarks

+ <${AggregatedBenchTable} benchData=${benchData} /> +
+ + <${BestKernelsSummary} benchData=${benchData} /> + `; +} diff --git a/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js b/metainfer/tasks/evolve_kernel/static/ok-evolve-runtime-api.js similarity index 69% rename from metainfer/tasks/evolve_kernel/static/ok-runtime-api.js rename to metainfer/tasks/evolve_kernel/static/ok-evolve-runtime-api.js index 8f572de0..5395ea99 100644 --- a/metainfer/tasks/evolve_kernel/static/ok-runtime-api.js +++ b/metainfer/tasks/evolve_kernel/static/ok-evolve-runtime-api.js @@ -56,3 +56,27 @@ export async function getRetrospective(taskId, n) { if (!res.ok) throw new Error(`retrospective ${n}: ${res.status}`); return res.json(); } + +export async function getKernelLineage(taskId, kernelId) { + const res = await fetch(`${BASE(taskId)}/kernel-library/${kernelId}/lineage`); + if (!res.ok) throw new Error(`kernel lineage: ${res.status}`); + return res.json(); +} + +export async function getFailures(taskId) { + const res = await fetch(`${BASE(taskId)}/failures`); + if (!res.ok) throw new Error(`failures: ${res.status}`); + return res.json(); +} + +export async function getShapeBenchmark(taskId) { + const res = await fetch(`${BASE(taskId)}/shape-benchmark`); + if (!res.ok) throw new Error(`shape-benchmark: ${res.status}`); + return res.json(); +} + +export async function refreshShapeBenchmark(taskId) { + const res = await fetch(`${BASE(taskId)}/shape-benchmark/refresh`, { method: "POST" }); + if (!res.ok) throw new Error(`shape-benchmark refresh: ${res.status}`); + return res.json(); +} diff --git a/metainfer/tasks/evolve_kernel/static/ok-state-graph.js b/metainfer/tasks/evolve_kernel/static/ok-evolve-state-graph.js similarity index 100% rename from metainfer/tasks/evolve_kernel/static/ok-state-graph.js rename to metainfer/tasks/evolve_kernel/static/ok-evolve-state-graph.js diff --git a/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js b/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js deleted file mode 100644 index 6bfc58ac..00000000 --- a/metainfer/tasks/evolve_kernel/static/ok-kernel-library.js +++ /dev/null @@ -1,61 +0,0 @@ -// Kernel library table component for evolve-kernel. -// Shows ranked kernels with exec_time, complexity, combined_score. - -import { html } from "htm/preact"; - -function scoreClass(val) { - if (val == null || val === 0) return ""; - if (val > 0.8) return "ok-score-good"; - if (val > 0.4) return "ok-score-ok"; - return "ok-score-bad"; -} - -function formatMs(ms) { - if (ms == null || ms === 0) return "—"; - if (ms < 1) return (ms * 1000).toFixed(1) + " μs"; - if (ms < 1000) return ms.toFixed(3) + " ms"; - return (ms / 1000).toFixed(2) + " s"; -} - -export function KernelLibrary({ library, selectedKernelId, onSelectKernel }) { - if (!library || !library.kernels || library.kernels.length === 0) { - return html`

Kernel library is empty — waiting for the first optimization result.

`; - } - - const kernels = library.kernels; - - return html` - - - - - - - - - - - - - - ${kernels.map((k, i) => html` - onSelectKernel && onSelectKernel(k)} - style=${{ cursor: onSelectKernel ? "pointer" : "default" }} - > - - - - - - - - - `)} - -
#IDExec TimeComplexityCombinedIterLines
${i + 1}${k.id.slice(0, 8)}${formatMs(k.exec_time_ms)}${(k.complexity_score * 100).toFixed(0)}%${k.combined_score.toFixed(4)}${k.iteration_added || 0}${k.code_lines || 0}
-

${library.size} kernel(s) in library (max 10) — ranked by combined score

- `; -} diff --git a/metainfer/tasks/evolve_kernel/static/ok.css b/metainfer/tasks/evolve_kernel/static/ok.css index 8b4988c2..80eff7ab 100644 --- a/metainfer/tasks/evolve_kernel/static/ok.css +++ b/metainfer/tasks/evolve_kernel/static/ok.css @@ -11,65 +11,166 @@ grid-column: 1 / -1; } +/* ---- Kernel Table ---- */ + +.ok-kernel-table-wrap { + max-height: 340px; + overflow-y: auto; + border: 1px solid var(--border-color, #30363d); + border-radius: 6px; +} + .ok-kernel-table { width: 100%; border-collapse: collapse; - font-size: 0.85rem; + font-size: 0.82rem; +} + +.ok-kernel-table thead { + position: sticky; + top: 0; + z-index: 1; } .ok-kernel-table th, .ok-kernel-table td { - padding: 0.4rem 0.6rem; + padding: 0.35rem 0.5rem; text-align: left; - border-bottom: 1px solid var(--border-color, #30363d); + border-bottom: 1px solid rgba(48, 54, 61, 0.6); } .ok-kernel-table th { font-weight: 600; - color: var(--muted, #8b949e); - font-size: 0.8rem; + color: #8b949e; + font-size: 0.7rem; text-transform: uppercase; + letter-spacing: 0.3px; + background: var(--panel-bg, #161b22); +} + +.ok-kernel-table tr { + transition: background 0.12s; } .ok-kernel-table tr:hover td { - background: var(--hover-bg, rgba(177, 186, 196, 0.08)); + background: rgba(177, 186, 196, 0.06); } .ok-kernel-table tr.selected td { - background: var(--accent-bg, rgba(88, 166, 255, 0.12)); - border-left: 2px solid var(--accent, #58a6ff); + background: rgba(88, 166, 255, 0.08); + border-left: none; } -.ok-score-good { - color: #3fb950; - font-weight: 600; +.ok-kernel-table tr.selected td:first-child { + box-shadow: inset 3px 0 0 var(--accent, #58a6ff); } -.ok-score-ok { - color: #d29922; +/* Row color coding by speedup */ +.ok-kernel-table tr.speedup-good td { + background: rgba(63, 185, 80, 0.04); } +.ok-kernel-table tr.speedup-bad td { + background: rgba(248, 81, 73, 0.04); +} +.ok-kernel-table tr.speedup-good:hover td, +.ok-kernel-table tr.speedup-bad:hover td { + background: rgba(177, 186, 196, 0.08); +} +.ok-kernel-table tr.speedup-good.selected td, +.ok-kernel-table tr.speedup-bad.selected td { + background: rgba(88, 166, 255, 0.08) !important; +} + +/* ---- Combined score bar ---- */ -.ok-score-bad { - color: #f85149; +.ok-combined-bar { + display: inline-block; + width: 60px; + height: 6px; + background: rgba(255,255,255,0.08); + border-radius: 3px; + overflow: hidden; + vertical-align: middle; + margin-right: 0.25rem; +} +.ok-combined-bar-fill { + height: 100%; + border-radius: 3px; + background: linear-gradient(90deg, #3fb950, #58a6ff); + transition: width 0.3s; } +/* ---- Headroom badge (inline) ---- */ + +.ok-hr-badge { + display: inline-block; + padding: 0.1rem 0.35rem; + border-radius: 8px; + font-size: 0.6rem; + font-weight: 700; + text-transform: uppercase; + letter-spacing: 0.2px; + white-space: nowrap; +} +.ok-hr-badge.near_optimal { background: rgba(63,185,80,0.15); color: #3fb950; } +.ok-hr-badge.memory_bound { background: rgba(88,166,255,0.15); color: #58a6ff; } +.ok-hr-badge.compute_bound { background: rgba(210,153,34,0.15); color: #d29922; } +.ok-hr-badge.inefficient { background: rgba(139,148,158,0.12); color: #8b949e; } + +/* ---- Score colors ---- */ + +.ok-score-good { color: #3fb950; font-weight: 600; } +.ok-score-ok { color: #d29922; } +.ok-score-bad { color: #f85149; } + +/* ---- Code Preview (dark editor theme) ---- */ + .ok-code-preview { max-height: 300px; overflow: auto; - background: var(--code-bg, #0d1117); + background: #1a1b1e; border: 1px solid var(--border-color, #30363d); border-radius: 6px; - padding: 0.75rem; - font-family: monospace; - font-size: 0.8rem; + padding: 0.75rem 0.75rem 0.75rem 3.5rem; + font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, monospace; + font-size: 0.78rem; white-space: pre; - line-height: 1.4; + line-height: 1.45; + tab-size: 4; + counter-reset: line; + position: relative; } .ok-code-preview.large { max-height: 600px; } +.ok-code-preview::before { + content: ""; + position: absolute; + left: 2.8rem; + top: 0; + bottom: 0; + width: 1px; + background: rgba(255,255,255,0.06); +} + +.ok-code-preview .ok-code-line { + counter-increment: line; +} +.ok-code-preview .ok-code-line::before { + content: counter(line); + position: absolute; + left: 0.5rem; + width: 2rem; + text-align: right; + color: rgba(255,255,255,0.2); + font-size: 0.7rem; + user-select: none; +} + +/* ---- Harness status ---- */ + .ok-harness-status { display: inline-flex; align-items: center; @@ -79,21 +180,11 @@ font-size: 0.8rem; font-weight: 600; } +.ok-harness-status.pass { background: rgba(63,185,80,0.15); color: #3fb950; } +.ok-harness-status.fail { background: rgba(248,81,73,0.15); color: #f85149; } +.ok-harness-status.pending { background: rgba(210,153,34,0.15); color: #d29922; } -.ok-harness-status.pass { - background: rgba(63, 185, 80, 0.15); - color: #3fb950; -} - -.ok-harness-status.fail { - background: rgba(248, 81, 73, 0.15); - color: #f85149; -} - -.ok-harness-status.pending { - background: rgba(210, 153, 34, 0.15); - color: #d29922; -} +/* ---- Badges ---- */ .ok-badge { display: inline-block; @@ -101,19 +192,14 @@ border-radius: 10px; font-size: 0.75rem; font-weight: 600; - background: var(--accent-bg, rgba(88, 166, 255, 0.15)); + background: var(--accent-bg, rgba(88,166,255,0.15)); color: var(--accent, #58a6ff); } +.ok-badge.green { background: rgba(63,185,80,0.15); color: #3fb950; } +.ok-badge.yellow { background: rgba(210,153,34,0.15); color: #d29922; } +.ok-badge.red { background: rgba(248,81,73,0.15); color: #f85149; } -.ok-badge.green { - background: rgba(63, 185, 80, 0.15); - color: #3fb950; -} - -.ok-badge.yellow { - background: rgba(210, 153, 34, 0.15); - color: #d29922; -} +/* ---- Flow indicator ---- */ .ok-flow-indicator { display: flex; @@ -124,25 +210,25 @@ border: 1px solid var(--border-color, #30363d); border-radius: 6px; margin-bottom: 0.75rem; + flex-wrap: wrap; } - .ok-flow-indicator .ok-phase-label { font-weight: 600; padding: 0.25rem 0.5rem; border-radius: 4px; - background: var(--accent-bg, rgba(88, 166, 255, 0.15)); + background: var(--accent-bg, rgba(88,166,255,0.15)); } - .ok-flow-indicator .ok-phase-divider { color: var(--muted, #8b949e); } +/* ---- Meta cards ---- */ + .ok-meta-grid { display: grid; - grid-template-columns: repeat(auto-fill, minmax(180px, 1fr)); + grid-template-columns: repeat(auto-fill, minmax(160px, 1fr)); gap: 0.75rem; } - .ok-meta-card { padding: 0.75rem; background: var(--panel-bg, #161b22); @@ -150,15 +236,541 @@ border-radius: 6px; text-align: center; } - .ok-meta-card .ok-meta-value { font-size: 1.5rem; font-weight: 700; font-variant-numeric: tabular-nums; } - .ok-meta-card .ok-meta-label { font-size: 0.75rem; color: var(--muted, #8b949e); margin-top: 0.25rem; } + +/* ---- Stats summary bar ---- */ + +.ok-stats-bar { + display: flex; + align-items: center; + gap: 1.2rem; + padding: 0.5rem 0.75rem; + background: var(--panel-bg, #161b22); + border: 1px solid var(--border-color, #30363d); + border-radius: 6px; + margin-bottom: 0.75rem; + flex-wrap: wrap; + font-size: 0.82rem; +} +.ok-stats-item { + display: flex; + align-items: center; + gap: 0.3rem; +} +.ok-stats-item .ok-stats-label { + color: #8b949e; + font-size: 0.7rem; + text-transform: uppercase; + letter-spacing: 0.3px; +} +.ok-stats-item .ok-stats-value { + font-weight: 700; + font-variant-numeric: tabular-nums; +} +.ok-stats-sep { + color: rgba(255,255,255,0.1); + font-size: 0.8rem; +} + +/* Convergence circle */ +.ok-conv-circle { + display: inline-flex; + align-items: center; + justify-content: center; + width: 22px; + height: 22px; + border-radius: 50%; + font-size: 0.6rem; + font-weight: 700; + background: rgba(255,255,255,0.06); + border: 2px solid rgba(255,255,255,0.12); +} +.ok-conv-circle.warn { border-color: #d29922; color: #d29922; } +.ok-conv-circle.done { border-color: #3fb950; color: #3fb950; background: rgba(63,185,80,0.08); } +.ok-conv-circle.clear { border-color: rgba(255,255,255,0.08); } + +/* ---- Kernel Lineage ---- */ + +.ok-lineage-panel { + background: var(--panel-bg, #161b22); + border: 1px solid var(--border-color, #30363d); + border-radius: 6px; + padding: 0.75rem 1rem; + margin-bottom: 1rem; +} +.ok-lineage-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 0.5rem; + padding-bottom: 0.5rem; + border-bottom: 1px solid var(--border-color, #30363d); +} +.ok-lineage-title { font-weight: 700; font-size: 0.9rem; } +.ok-lineage-id { + font-size: 0.7rem; + background: var(--code-bg, #0d1117); + padding: 0.1rem 0.3rem; + border-radius: 3px; + margin: 0 0.1rem; + cursor: default; +} +.ok-lineage-body { + display: flex; + flex-direction: column; + gap: 0.5rem; +} +.ok-lineage-metrics { + display: flex; + gap: 1.5rem; + flex-wrap: wrap; +} +.ok-lineage-metric { + display: flex; + flex-direction: column; + gap: 0.15rem; +} +.ok-lineage-metric-label { + font-size: 0.7rem; + color: var(--muted, #8b949e); + text-transform: uppercase; +} +.ok-lineage-metric-value { + font-size: 0.9rem; + font-weight: 600; + font-variant-numeric: tabular-nums; +} +.ok-lineage-link { color: var(--accent, #58a6ff); cursor: pointer; text-decoration: underline; text-underline-offset: 2px; } +.ok-lineage-link:hover { color: #79c0ff; } +.ok-lineage-summary { + background: var(--code-bg, #0d1117); + border: 1px solid var(--border-color, #30363d); + border-radius: 4px; + padding: 0.5rem 0.75rem; +} +.ok-lineage-summary-label { + font-size: 0.7rem; + color: var(--muted, #8b949e); + text-transform: uppercase; + margin-bottom: 0.35rem; + font-weight: 700; +} +.ok-lineage-summary-text { font-size: 0.85rem; line-height: 1.5; } +.ok-lineage-summary-items { margin: 0.25rem 0; padding-left: 1.2rem; font-size: 0.82rem; line-height: 1.45; } +.ok-lineage-summary-items li { margin-bottom: 0.3rem; color: #c9d1d9; } +.ok-lineage-summary-items li::marker { color: #58a6ff; } +.ok-lineage-retrospective { margin-top: 0.25rem; } +.ok-lineage-retrospective summary { cursor: pointer; font-size: 0.8rem; color: var(--accent, #58a6ff); font-weight: 600; } + +/* ---- Failure Log ---- */ + +.ok-failure-log { display: flex; flex-direction: column; gap: 0.25rem; } +.ok-failure-entry { border: 1px solid var(--border-color, #30363d); border-radius: 4px; overflow: hidden; } +.ok-failure-header { + display: flex; align-items: center; gap: 0.5rem; + padding: 0.4rem 0.6rem; cursor: pointer; font-size: 0.82rem; transition: background 0.15s; +} +.ok-failure-header:hover { filter: brightness(1.1); } +.ok-failure-header.severity-error { background: rgba(248,81,73,0.08); border-left: 3px solid #f85149; } +.ok-failure-header.severity-warning { background: rgba(210,153,34,0.06); border-left: 3px solid #d29922; } +.ok-failure-severity { font-size: 1rem; width: 18px; text-align: center; flex-shrink: 0; } +.severity-error .ok-failure-severity { color: #f85149; } +.severity-warning .ok-failure-severity { color: #d29922; } +.ok-failure-iteration { font-weight: 600; font-size: 0.78rem; color: var(--accent, #58a6ff); white-space: nowrap; flex-shrink: 0; } +.ok-failure-phase { font-size: 0.78rem; color: var(--muted, #8b949e); white-space: nowrap; flex-shrink: 0; } +.ok-failure-summary { flex: 1; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 0.82rem; } +.ok-failure-attempts { font-size: 0.7rem; color: var(--muted, #8b949e); background: var(--code-bg, #0d1117); padding: 0.1rem 0.35rem; border-radius: 8px; white-space: nowrap; flex-shrink: 0; } +.ok-failure-toggle { font-size: 0.6rem; color: var(--muted, #8b949e); flex-shrink: 0; } +.ok-failure-detail { padding: 0.5rem 0.6rem 0.6rem 2.2rem; border-top: 1px solid var(--border-color, #30363d); font-size: 0.82rem; } +.ok-failure-meta { display: flex; gap: 1rem; margin-bottom: 0.4rem; font-size: 0.78rem; color: var(--muted, #8b949e); } + +/* ---- Headroom / Roofline ---- */ + +.ok-headroom-card { + background: var(--code-bg, #0d1117); + border: 1px solid var(--border-color, #30363d); + border-radius: 6px; + padding: 0.6rem 0.75rem; +} +.ok-headroom-header { + display: flex; justify-content: space-between; align-items: center; + margin-bottom: 0.5rem; padding-bottom: 0.4rem; + border-bottom: 1px solid var(--border-color, #30363d); +} +.ok-headroom-title { font-weight: 700; font-size: 0.85rem; } +.ok-headroom-bottleneck { + display: inline-block; padding: 0.15rem 0.5rem; border-radius: 10px; + font-size: 0.7rem; font-weight: 700; text-transform: uppercase; letter-spacing: 0.3px; +} +.ok-headroom-metrics { display: flex; flex-direction: column; gap: 0.5rem; margin-bottom: 0.5rem; } +.ok-headroom-bar-group { display: flex; align-items: center; gap: 0.5rem; } +.ok-headroom-bar-label { display: flex; flex-direction: column; min-width: 110px; font-size: 0.78rem; font-weight: 600; } +.ok-headroom-bar-track { flex: 1; height: 12px; background: rgba(255,255,255,0.06); border-radius: 6px; overflow: hidden; position: relative; } +.ok-headroom-bar-fill { height: 100%; border-radius: 6px; transition: width 0.4s ease; } +.ok-headroom-bar-bw { background: linear-gradient(90deg, #1f6feb, #3fb950); } +.ok-headroom-bar-comp { background: linear-gradient(90deg, #a371f7, #d29922); } +.ok-headroom-bar-fill .ok-bar-value { + position: absolute; right: 6px; top: 50%; transform: translateY(-50%); + font-size: 0.6rem; font-weight: 700; color: #fff; text-shadow: 0 1px 2px rgba(0,0,0,0.5); + white-space: nowrap; +} +.ok-headroom-summary-line { + display: flex; justify-content: space-between; align-items: center; + font-size: 0.8rem; margin-bottom: 0.4rem; +} +.ok-headroom-advice { + font-size: 0.78rem; line-height: 1.5; color: var(--muted, #8b949e); + padding: 0.4rem 0.5rem; background: rgba(255,255,255,0.03); border-radius: 4px; margin-top: 0.25rem; +} +.ok-headroom-suggestions { + margin: 0.4rem 0 0 1rem; padding: 0; font-size: 0.78rem; line-height: 1.5; color: var(--muted, #8b949e); +} +.ok-headroom-suggestions li { margin-bottom: 0.25rem; } + +/* ---- Roofline Callout (clean, no plot) ---- */ + +.ok-roofline-callout { + background: var(--code-bg, #0d1117); + border: 1px solid var(--border-color, #30363d); + border-left: 3px solid #58a6ff; + border-radius: 6px; + padding: 0.7rem 0.85rem; +} +.ok-roofline-callout-header { + display: flex; align-items: center; gap: 0.35rem; margin-bottom: 0.35rem; +} +.ok-roofline-icon { font-size: 1.1rem; } +.ok-roofline-callout-desc { + font-size: 0.78rem; color: #8b949e; margin: 0; line-height: 1.4; +} +.ok-roofline-eq-row { + display: flex; align-items: baseline; gap: 0.3rem; + font-size: 0.8rem; padding: 2px 0; +} +.ok-roofline-eff-bar { + display: flex; align-items: center; gap: 0.5rem; +} + +/* ---- Shape Target Cards ---- */ + +.ok-shape-targets { + display: flex; flex-wrap: wrap; gap: 0.5rem; margin-bottom: 0.75rem; +} +.ok-shape-target { + background: rgba(255,255,255,0.02); + border: 1px solid rgba(255,255,255,0.06); + border-radius: 6px; + padding: 0.4rem 0.6rem; + min-width: 200px; +} +.ok-shape-target-header { + display: flex; align-items: center; gap: 0.3rem; margin-bottom: 0.2rem; +} +.ok-shape-target-name { + font-weight: 700; font-size: 0.78rem; font-family: monospace; +} +.ok-shape-target-m { + display: flex; gap: 0.4rem; flex-wrap: wrap; +} +.ok-shape-m-tag { + display: inline-block; + padding: 0.1rem 0.4rem; + border-radius: 3px; + font-size: 0.65rem; + font-weight: 600; + font-family: monospace; +} +.ok-shape-m-tag.decode { background: rgba(88,166,255,0.12); color: #58a6ff; } +.ok-shape-m-tag.large { background: rgba(63,185,80,0.1); color: #3fb950; } + +.muted-badge { + background: rgba(255,255,255,0.04) !important; + color: #8b949e !important; +} + +/* ---- Compact Lineage Metrics ---- */ + +.ok-lineage-metrics-compact { + display: grid; + grid-template-columns: repeat(3, 1fr); + gap: 0.5rem; + margin-bottom: 0.25rem; +} +.ok-lm-item { + text-align: center; + padding: 0.35rem 0.25rem; + background: rgba(255,255,255,0.02); + border-radius: 4px; +} +.ok-lm-val { display: block; font-size: 0.95rem; font-weight: 700; } +.ok-lm-lbl { display: block; font-size: 0.6rem; color: #8b949e; text-transform: uppercase; margin-top: 0.1rem; } + +/* ---- Shape benchmark row emphasis ---- */ + +.ok-shape-table tr.row-decode td { background: rgba(88,166,255,0.04); } +.ok-shape-table tr.row-large td { background: rgba(63,185,80,0.04); } +.ok-shape-table tr.row-decode td:first-child::after { content: " ⚡"; font-size: 0.6rem; } +.ok-shape-table tr.row-large td:first-child::after { content: " 📦"; font-size: 0.6rem; } + +/* ---- Mini Roofline Scatter Plot (deprecated, replaced by callout) ---- */ + +.ok-roofline-plot { + width: 100%; + height: 140px; + background: rgba(0,0,0,0.2); + border-radius: 4px; + border: 1px solid rgba(255,255,255,0.06); + position: relative; + overflow: hidden; + margin-bottom: 0.5rem; +} +.ok-roofline-plot .axis-label { + position: absolute; + font-size: 0.6rem; + color: rgba(255,255,255,0.3); + font-family: monospace; +} +.ok-roofline-plot .axis-label.x { bottom: 4px; left: 50%; transform: translateX(-50%); } +.ok-roofline-plot .axis-label.y { left: 4px; top: 50%; transform: translateY(-50%) rotate(-90deg); } +.ok-roofline-plot .ridge-line { + position: absolute; + left: 0; top: 0; + width: 100%; height: 100%; + pointer-events: none; +} +.ok-roofline-dot { + position: absolute; + width: 10px; height: 10px; + border-radius: 50%; + background: #58a6ff; + border: 2px solid #fff; + transform: translate(-50%, -50%); + box-shadow: 0 0 8px rgba(88,166,255,0.5); + transition: all 0.3s; + z-index: 2; +} + +/* ---- Hipprof profiling ---- */ + +.ok-profiling-badges { + display: flex; + gap: 0.5rem; + flex-wrap: wrap; + margin-top: 0.4rem; +} +.ok-profiling-badge { + padding: 0.2rem 0.5rem; + border-radius: 4px; + font-size: 0.7rem; + font-weight: 600; + background: rgba(255,255,255,0.04); + border: 1px solid rgba(255,255,255,0.08); +} +.ok-profiling-badge .label { color: #8b949e; margin-right: 0.3rem; } +.ok-profiling-badge .value { color: #c9d1d9; font-variant-numeric: tabular-nums; } + +/* ---- Shape Benchmark Panel ---- */ + +.ok-shape-bench { display: flex; flex-wrap: wrap; gap: 1rem; } +.ok-shape-group { flex: 1; min-width: 320px; max-width: 520px; } +.ok-shape-group-label { + font-weight: 700; font-size: 0.82rem; color: var(--accent, #58a6ff); + margin-bottom: 0.35rem; padding: 0.2rem 0.5rem; + background: rgba(88,166,255,0.08); border-radius: 4px; +} +.ok-shape-table { font-size: 0.8rem; } +.ok-shape-table th:first-child, .ok-shape-table td:first-child { width: 40px; } +.ok-shape-table th:nth-child(2), .ok-shape-table td:nth-child(2) { width: 100px; } + +/* ---- Multi-GPU Dashboard ---- */ + +.ok-mgpu-cards { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(240px, 1fr)); + gap: 0.75rem; + margin-bottom: 0.5rem; +} + +.ok-mgpu-card { + background: var(--panel-bg, #161b22); + border: 1px solid var(--border-color, #30363d); + border-radius: 8px; + padding: 0.7rem 0.85rem; + transition: border-color 0.2s, box-shadow 0.2s; + position: relative; + overflow: hidden; +} + +/* Phase-colored left bar */ +.ok-mgpu-card::before { + content: ""; + position: absolute; + left: 0; top: 0; bottom: 0; + width: 3px; + background: #30363d; + transition: background 0.3s; +} +.ok-mgpu-card.bootstrap::before { background: #d29922; } +.ok-mgpu-card.optimize::before { background: #58a6ff; } +.ok-mgpu-card.done::before { background: #3fb950; } + +.ok-mgpu-card:hover { border-color: var(--accent, #58a6ff); box-shadow: 0 2px 12px rgba(88,166,255,0.08); } +.ok-mgpu-card.selected { + border-color: var(--accent, #58a6ff); + box-shadow: 0 0 0 2px rgba(88,166,255,0.3), 0 2px 12px rgba(88,166,255,0.1); +} + +.ok-mgpu-card-header { + display: flex; justify-content: space-between; align-items: center; + margin-bottom: 0.4rem; padding-bottom: 0.35rem; + border-bottom: 1px solid var(--border-color, #30363d); +} +.ok-mgpu-card-title { font-weight: 700; font-size: 0.9rem; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } +.ok-mgpu-card-status { + font-size: 0.7rem; font-weight: 600; padding: 0.15rem 0.4rem; border-radius: 10px; white-space: nowrap; +} +.ok-mgpu-card-status.running { background: rgba(63,185,80,0.15); color: #3fb950; } +.ok-mgpu-card-status.stopped { background: rgba(139,148,158,0.15); color: #8b949e; } + +/* Phase dot colors */ +.ok-phase-dot { + display: inline-block; + width: 7px; height: 7px; + border-radius: 50%; + margin-right: 3px; + vertical-align: middle; +} +.ok-phase-dot.bootstrap { background: #d29922; } +.ok-phase-dot.optimize { background: #58a6ff; } +.ok-phase-dot.idle { background: #8b949e; } +.ok-phase-dot.finished { background: #3fb950; } +.ok-phase-dot.starting { background: #a371f7; } + +.ok-mgpu-card-body { + display: flex; flex-direction: column; gap: 0.35rem; +} +.ok-mgpu-card-metrics { + display: flex; gap: 1rem; +} +.ok-mgpu-card-metric { + display: flex; flex-direction: column; gap: 0.1rem; + min-width: 60px; +} +.ok-mgpu-card-label { + font-size: 0.6rem; color: var(--muted, #8b949e); text-transform: uppercase; letter-spacing: 0.3px; +} +.ok-mgpu-card-value { + font-size: 0.85rem; font-weight: 600; font-variant-numeric: tabular-nums; +} +.ok-mgpu-card-shapes { + font-size: 0.6rem; color: #8b949e; margin-top: 0.15rem; + max-height: 2.4em; overflow: hidden; text-overflow: ellipsis; + line-height: 1.2; opacity: 0.7; +} +.ok-mgpu-card-iter-bar { + height: 3px; border-radius: 2px; background: rgba(255,255,255,0.06); + margin-top: 0.2rem; overflow: hidden; +} +.ok-mgpu-card-iter-fill { + height: 100%; border-radius: 2px; + background: linear-gradient(90deg, #58a6ff, #3fb950); + transition: width 1s; +} + +/* ---- Multi-GPU Agent Grid (collapsible) ---- */ + +.ok-mgpu-agents-grid { + display: grid; + grid-template-columns: repeat(auto-fill, minmax(240px, 1fr)); + gap: 0.5rem; + max-height: 250px; + overflow-y: auto; +} +.ok-mgpu-agent-group { + background: rgba(255,255,255,0.02); + border: 1px solid rgba(255,255,255,0.04); + border-radius: 6px; + padding: 0.35rem; +} +.ok-mgpu-agent-group-header { + font-weight: 700; font-size: 0.72rem; color: #58a6ff; + margin-bottom: 0.25rem; padding-bottom: 0.15rem; + border-bottom: 1px solid rgba(255,255,255,0.05); + cursor: pointer; + display: flex; justify-content: space-between; align-items: center; + user-select: none; +} +.ok-mgpu-agent-group-header:hover { color: #79c0ff; } +.ok-mgpu-agent-group-header .collapse-icon { + font-size: 0.6rem; color: #8b949e; transition: transform 0.2s; +} +.ok-mgpu-agent-group-header .collapse-icon.open { transform: rotate(90deg); } +.ok-mgpu-agent-group.collapsed .ok-mgpu-agent-row { display: none; } +.ok-mgpu-agent-group.collapsed .ok-mgpu-agent-row.agent-running { display: flex; } +.ok-mgpu-agent-row { + display: flex; align-items: center; justify-content: space-between; + padding: 2px 3px; font-size: 0.68rem; gap: 0.3rem; + border-radius: 3px; +} +.ok-mgpu-agent-row.agent-running { background: rgba(63,185,80,0.06); } +.ok-mgpu-agent-row.agent-failed { background: rgba(248,81,73,0.06); } +.ok-mgpu-agent-role { flex: 1; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; } +.ok-mgpu-agent-status { min-width: 45px; text-align: right; white-space: nowrap; } +.ok-mgpu-agent-error { + font-size: 0.58rem; max-width: 60px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; + color: #8b949e; +} + +/* ---- Agent labels (per-GPU detail panel) ---- */ + +.ok-agent-list-scroll { + max-height: 200px; + overflow-y: auto; +} + +/* ---- State Machine Animation ---- */ + +@keyframes ok-pulse { + 0%, 100% { opacity: 1; } + 50% { opacity: 0.6; } +} +.state-graph .node-rect.active { + animation: ok-pulse 1.8s ease-in-out infinite; +} +.state-graph .node-rect.visited { + fill: rgba(88,166,255,0.08); + stroke: #58a6ff; + stroke-width: 1.5; +} +.state-graph .node-rect.future { + fill: rgba(255,255,255,0.02); + stroke: rgba(255,255,255,0.1); +} +.state-graph .edge.active-edge { + stroke: #58a6ff !important; + stroke-width: 2; + stroke-dasharray: none; +} +.state-graph .edge.potential { + stroke-dasharray: 4 3; + opacity: 0.4; +} + +/* ---- Combined timeline ---- */ + +.ok-timeline-stream { + font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, monospace; +} +.ok-timeline-row { + display: flex; gap: 0.3rem; align-items: baseline; flex-wrap: wrap; + border-bottom: 1px solid rgba(255,255,255,0.03); +}