diff --git a/pyproject.toml b/pyproject.toml index 21cccb5..76b32d4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,13 +1,13 @@ [project] name = "xe-forge" -version = "0.2.0" +version = "0.3.0" description = "Multi-stage kernel optimization for GPU computing" authors = [{ name = "Intel Labs" }] requires-python = ">=3.11" dependencies = [ "numpy>=1.24.0", "pydantic>=2.0.0", - "dspy-ai>=3.1.3", + "dspy>=3.3.0b1", "litellm>=1.0.0", "python-dotenv>=1.0.0", "PyYAML>=6.0.0", diff --git a/src/xe_forge/__init__.py b/src/xe_forge/__init__.py index bb06945..265a47f 100644 --- a/src/xe_forge/__init__.py +++ b/src/xe_forge/__init__.py @@ -5,7 +5,6 @@ from xe_forge.config import Config, get_config, override_config from xe_forge.models import OptimizationResult, OptimizationStage -from xe_forge.pipeline import XeForgePipeline __version__ = "0.2.0" __all__ = [ @@ -16,3 +15,11 @@ "get_config", "override_config", ] + + +def __getattr__(name: str): + if name == "XeForgePipeline": + from xe_forge.pipeline import XeForgePipeline + + return XeForgePipeline + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/xe_forge/agents/__init__.py b/src/xe_forge/agents/__init__.py index fceea7c..4258a17 100644 --- a/src/xe_forge/agents/__init__.py +++ b/src/xe_forge/agents/__init__.py @@ -1,5 +1,6 @@ from xe_forge.agents.analyzer_agent import AnalyzerAgent from xe_forge.agents.base import Optimizer +from xe_forge.agents.coordinator import CoordinatorAgent from xe_forge.agents.cover import CoVeR from xe_forge.agents.optimizer_agent import ( SUCCESS_MESSAGE, @@ -16,6 +17,7 @@ "AnalyzerAgent", "AutotuneSignature", "CoVeR", + "CoordinatorAgent", "OptimizationReActSignature", "OptimizationSignature", "Optimizer", diff --git a/src/xe_forge/agents/analyzer_agent.py b/src/xe_forge/agents/analyzer_agent.py index 74a9d30..aa45bcf 100644 --- a/src/xe_forge/agents/analyzer_agent.py +++ b/src/xe_forge/agents/analyzer_agent.py @@ -11,6 +11,19 @@ logger = logging.getLogger(__name__) +_DSL_NAMES: dict[str, str] = { + "triton": "Triton", + "gluon": "Gluon", + "sycl": "SYCL/XeTLA", + "cuda": "CUDA C++", +} + +_DEVICE_DESCS: dict[str, str] = { + "xpu": "Intel XPU (Data Center GPU Max / Ponte Vecchio)", + "cuda": "NVIDIA CUDA GPU", + "cpu": "CPU", +} + # --------------------------------------------------------------------------- # Build the issue category section of the prompt dynamically from the enum @@ -216,47 +229,7 @@ def _build_issue_categories(dsl: DSL = DSL.TRITON) -> str: class AnalysisSignature(dspy.Signature): - # The docstring is dynamically constructed so the issue list - # always matches the live IssueType enum. - __doc__ = f"""Analyze Triton kernel for optimization opportunities. - -You are a world-class expert in Triton GPU/XPU kernel optimization, -numerical linear algebra, and high-performance computing. - -Analyze the given Triton kernel code and, if available, the original PyTorch -implementation for higher-level algorithmic context. - -You must identify ALL applicable optimizations across every category below. -Use your deep knowledge of GPU programming, Triton internals, Intel XPU -architecture, and mathematical optimization. - -{_ISSUE_CATEGORIES_BLOCK} -IMPORTANT: -- Return issues as a JSON array of DetectedIssue objects. -- Each issue MUST have: issue_type (exact string from the list above), - severity (1-5), description, suggested_fix, estimated_speedup. -- issue_type MUST be one of the exact strings listed above (e.g. "dtype_float64", - "missing_grf_mode"). Do NOT invent new type names. -- For fused kernels, pay special attention to ALGORITHMIC issues. -- Return empty array [] ONLY if the kernel is already optimal. - -OPEN-ENDED DISCOVERY (issue_type="open_ended"): -After checking all categories above, ask yourself: is there a high-value -optimization that does not fit any existing type? If yes, use issue_type="open_ended" -and populate open_ended_proposal with the full proposal. Requirements: - - Concrete and implementable — not a vague observation - - Mathematically or logically justified - - Includes a before/after code sketch in open_ended_proposal - - Includes estimated speedup with reasoning -Examples that qualify as open_ended: - * sum(x @ W.T, dim=1) rewritten as x @ W.sum(dim=0) — eliminates O(M*N*K) GEMM - * Weight statistic (colsum, norm) recomputed every forward() — cache in __init__ - * Two-kernel pipeline where the HBM intermediate can be eliminated algebraically -Examples that do NOT qualify (use the named type instead): - * "use better tile sizes" → use suboptimal_tile_size - * "add autotuning" → use missing_autotune - * "fuse these kernels" → use unfused_kernels -""" + """Analyze a Triton kernel for optimization opportunities across all categories.""" kernel_code: dspy.Code["python"] = dspy.InputField(desc="Triton kernel source code to analyze.") reference_code: str = dspy.InputField( @@ -278,30 +251,7 @@ class AnalysisSignature(dspy.Signature): class SyclAnalysisSignature(dspy.Signature): - __doc__ = f"""Analyze SYCL/CUTLASS C++ kernel for optimization opportunities on Intel XPU. - -You are a world-class expert in SYCL, CUTLASS/XeTLA, and Intel XPU GPU kernel optimization. - -Analyze the given SYCL C++ kernel code and identify ALL applicable optimizations. - -=== SYCL/CUTLASS OPTIMIZATION KNOBS === -- TileShape: Shape<_M, _N, _K> (e.g. 256x256x32, 128x128x64, 128x256x32) -- PipelineStages: 2, 3, or 4 (more = more prefetching, but more register pressure) -- MMA Atom: XE_DPAS_TT — SubgroupSize 4 or 8 -- Dispatch Policy: MainloopXeL1Staged (L1 cached), MainloopXeL0Staged -- Data types: bfloat16_t/half_t inputs, float/bfloat16_t accumulators -- Memory layout: RowMajor vs ColumnMajor for A, B, C, D matrices -- Epilogue fusion: LinearCombination, bias addition, activation functions via FusionCallbacks -- GmemTiledCopy: void (auto) or explicit copy atoms for fine-grained control - -{_SYCL_ISSUE_CATEGORIES_BLOCK} -IMPORTANT: -- Return issues as a JSON array of DetectedIssue objects. -- Each issue MUST have: issue_type (exact string from the list above), - severity (1-5), description, suggested_fix, estimated_speedup. -- issue_type MUST be one of the exact strings listed above. -- Return empty array [] ONLY if the kernel is already optimal. -""" + """Analyze a SYCL/CUTLASS C++ kernel for optimization opportunities on Intel XPU.""" kernel_code: dspy.Code["cpp"] = dspy.InputField( desc="SYCL/CUTLASS C++ kernel source code to analyze." @@ -329,10 +279,41 @@ class SyclAnalysisSignature(dspy.Signature): class AnalyzerAgent: """LLM-based analyzer for Triton kernels.""" - def __init__(self, knowledge_base=None, dsl: DSL | str = DSL.TRITON): + def __init__( + self, knowledge_base=None, dsl: DSL | str = DSL.TRITON, extra_instructions: str = "" + ): self.knowledge_base = knowledge_base self.dsl = DSL(dsl) if isinstance(dsl, str) else dsl sig = SyclAnalysisSignature if self.dsl == DSL.SYCL else AnalysisSignature + + # Inject the full analysis guidance from template, including the dynamic issue categories block + from xe_forge.config import get_config + from xe_forge.prompts import render_signature_instructions + + try: + cfg = get_config() + issue_block = ( + _SYCL_ISSUE_CATEGORIES_BLOCK if self.dsl == DSL.SYCL else _ISSUE_CATEGORIES_BLOCK + ) + template_instructions = render_signature_instructions( + "analysis_signature", + dsl=cfg.device_config.dsl, + dsl_name=_DSL_NAMES.get( + str(self.dsl.value if hasattr(self.dsl, "value") else self.dsl), "Triton" + ), + device_type=cfg.device_config.device, + device_description=_DEVICE_DESCS.get(cfg.device_config.device, "Intel XPU"), + defaults={}, + issue_categories_block=issue_block, + ) + sig = sig.append_instructions(template_instructions) + except Exception as e: + logger.debug( + "Analysis template render failed, falling back to extra_instructions: %s", e + ) + + if extra_instructions: + sig = sig.append_instructions(extra_instructions) self.predictor = dspy.Predict(sig) def analyze( diff --git a/src/xe_forge/agents/coordinator.py b/src/xe_forge/agents/coordinator.py new file mode 100644 index 0000000..cb34884 --- /dev/null +++ b/src/xe_forge/agents/coordinator.py @@ -0,0 +1,281 @@ +""" +CoordinatorAgent — dspy.ReActV2-driven agentic orchestrator for kernel optimization. + +Replaces the fixed analyze→plan→stage loop with an LLM that decides at runtime: +what to analyze, which stages to apply, whether to profile before/after a stage, +and when to stop. Code never flows through tool arguments — CoordinatorState holds +all code state so the LLM orchestrates at semantic level. +""" + +from __future__ import annotations + +import logging + +import dspy + +from xe_forge.agents.coordinator_tools import ( + CoordinatorState, + make_analyze_tool, + make_apply_stage_tool, + make_benchmark_tool, + make_profile_tool, + make_retrieve_patterns_tool, + make_status_tool, +) +from xe_forge.models import DSL, StageResult + +logger = logging.getLogger(__name__) + + +class CoordinatorSignature(dspy.Signature): + """Orchestrate GPU kernel optimization to maximize speedup using specialist tools.""" + + kernel_specs: str = dspy.InputField( + desc="Kernel specs: name, input shapes, dtype, FLOP count, device, DSL" + ) + achieved_speedup: float = dspy.OutputField( + desc="Final speedup achieved vs original baseline (1.0 = no improvement)" + ) + optimization_summary: str = dspy.OutputField( + desc="Concise summary: which stages were applied, what changed, final speedup" + ) + + +class CoordinatorAgent(dspy.Module): + """Agentic coordinator that uses dspy.ReActV2 to drive multi-stage kernel optimization. + + DSPy >=3.3.0b1 required — dspy.ReActV2 is used unconditionally. + """ + + def __init__( + self, + analyzer, + executor, + knowledge_base=None, + profiler=None, + max_iters: int = 20, + extra_instructions: str = "", + dsl: DSL | str = DSL.TRITON, + ): + super().__init__() + self.analyzer = analyzer + self.executor = executor + self.knowledge_base = knowledge_base + self.profiler = profiler + self.max_iters = max_iters + self.extra_instructions = extra_instructions + self.dsl = DSL(dsl) if isinstance(dsl, str) else dsl + + def run( + self, + kernel_code: str, + kernel_specs: str, + pytorch_code: str = "", + kernel_name: str | None = None, + input_shapes=None, + flop=None, + dtype=None, + spec_dims=None, + init_args=None, + input_dtypes=None, + xpu_config: dict | None = None, + spec_path: str | None = None, + variant_type: str = "bench-gpu", + ) -> tuple[str, float, str, list[StageResult]]: + """Run the coordinator agent. + + Returns (best_code, best_speedup, summary, stage_results). + """ + state = CoordinatorState( + original_code=kernel_code, + current_code=kernel_code, + best_code=kernel_code, + ) + + # Shared mutable ref so profile_tool can update the vtune_report used by apply_stage + vtune_report_ref: list[str] = [""] + + tools = self._build_tools( + state=state, + pytorch_code=pytorch_code, + kernel_name=kernel_name, + input_shapes=input_shapes, + flop=flop, + dtype=dtype, + spec_dims=spec_dims, + init_args=init_args, + input_dtypes=input_dtypes, + xpu_config=xpu_config or {}, + spec_path=spec_path, + variant_type=variant_type, + vtune_report_ref=vtune_report_ref, + ) + + sig = CoordinatorSignature + # Inject coordinator guidance from template + try: + from xe_forge.config import get_config + from xe_forge.prompts import render_signature_instructions + from xe_forge.prompts.device_prompts import ( + _DEVICE_DESCRIPTIONS, + _DEVICE_TUNING_DEFAULTS, + _DSL_NAMES, + ) + + cfg = get_config() + template_text = render_signature_instructions( + "coordinator_signature", + dsl=str(self.dsl.value if hasattr(self.dsl, "value") else self.dsl), + dsl_name=_DSL_NAMES.get( + str(self.dsl.value if hasattr(self.dsl, "value") else self.dsl), "Triton" + ), + device_type=cfg.device_config.device, + device_description=_DEVICE_DESCRIPTIONS.get(cfg.device_config.device, "Intel XPU"), + defaults=_DEVICE_TUNING_DEFAULTS.get(cfg.device_config.device, {}), + ) + sig = sig.append_instructions(template_text) + except Exception as e: + logger.debug("Coordinator template render failed: %s", e) + + if self.extra_instructions: + sig = sig.append_instructions(self.extra_instructions) + + # DSPy >=3.3.0b1 required + agent = dspy.ReActV2( + signature=sig, + tools=tools, + max_iters=self.max_iters, + ) + + logger.info("CoordinatorAgent: starting ReActV2 (max_iters=%d)", self.max_iters) + + try: + result = agent(kernel_specs=kernel_specs) + except Exception as e: + logger.error("CoordinatorAgent: ReActV2 failed: %s", e) + return ( + state.best_code, + state.best_speedup, + f"Coordinator failed: {e}", + state.stage_results, + ) + + termination = getattr(result, "termination_reason", None) + if termination: + logger.info("CoordinatorAgent termination_reason: %s", termination) + if termination in ("max_iters", "context_window_exceeded"): + logger.warning("Coordinator stopped early due to %s", termination) + + llm_speedup = getattr(result, "achieved_speedup", None) + summary = getattr(result, "optimization_summary", "") or "" + + # Use the better of: LLM-reported speedup vs what Python state tracked + final_speedup = max( + float(llm_speedup) if llm_speedup is not None else 0.0, + state.best_speedup, + ) + + logger.info( + "CoordinatorAgent done: speedup=%.3fx, stages_succeeded=%s", + final_speedup, + state.stages_succeeded, + ) + + return state.best_code, final_speedup, summary, state.stage_results + + def _build_tools( + self, + state: CoordinatorState, + pytorch_code: str, + kernel_name, + input_shapes, + flop, + dtype, + spec_dims, + init_args, + input_dtypes, + xpu_config: dict, + spec_path, + variant_type: str, + vtune_report_ref: list[str], + ) -> list: + tools = [ + dspy.Tool( + make_analyze_tool( + state=state, + analyzer=self.analyzer, + pytorch_code=pytorch_code, + kernel_name=kernel_name, + input_shapes=input_shapes, + flop=flop, + dtype=dtype, + ) + ), + dspy.Tool( + make_retrieve_patterns_tool( + state=state, + knowledge_base=self.knowledge_base, + ) + ), + dspy.Tool( + make_apply_stage_tool( + state=state, + optimizer=self._build_stage_optimizer(xpu_config), + xpu_config=xpu_config, + kernel_name=kernel_name, + input_shapes=input_shapes, + spec_dims=spec_dims, + flop=flop, + dtype=dtype, + pytorch_code=pytorch_code, + init_args=init_args, + input_dtypes=input_dtypes, + vtune_report_ref=vtune_report_ref, + ) + ), + dspy.Tool( + make_benchmark_tool( + state=state, + executor=self.executor, + kernel_name=kernel_name, + input_shapes=input_shapes, + flop=flop, + dtype=dtype, + spec_dims=spec_dims, + init_args=init_args, + input_dtypes=input_dtypes, + ) + ), + dspy.Tool(make_status_tool(state=state)), + ] + + if ( + self.profiler is not None + and hasattr(self.profiler, "available") + and self.profiler.available() + ): + tools.append( + dspy.Tool( + make_profile_tool( + state=state, + profiler=self.profiler, + kernel_name=kernel_name, + spec_path=spec_path, + variant_type=variant_type, + vtune_report_ref=vtune_report_ref, + ) + ) + ) + + return tools + + def _build_stage_optimizer(self, xpu_config: dict): + """Build an OptimizerAgent for the apply_stage tool to delegate to.""" + from xe_forge.agents.optimizer_agent import OptimizerAgent + + return OptimizerAgent( + executor=self.executor, + knowledge_base=self.knowledge_base, + dsl=self.dsl, + extra_instructions=self.extra_instructions, + ) diff --git a/src/xe_forge/agents/coordinator_tools.py b/src/xe_forge/agents/coordinator_tools.py new file mode 100644 index 0000000..7a1c411 --- /dev/null +++ b/src/xe_forge/agents/coordinator_tools.py @@ -0,0 +1,407 @@ +""" +CoordinatorState and tool factory functions for CoordinatorAgent. + +Code never flows through tool arguments — all tools are closures over CoordinatorState +so the coordinator LLM operates at semantic level ("apply block_pointers stage") +without ever handling raw kernel code. +""" + +from __future__ import annotations + +import logging +import tempfile +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path + +from xe_forge.models import KernelAnalysis, OptimizationStage, StageResult + +logger = logging.getLogger(__name__) + + +@dataclass +class CoordinatorState: + """Mutable kernel optimization state shared across all coordinator tools.""" + + original_code: str + current_code: str + best_code: str + best_speedup: float = 1.0 + stages_tried: list[str] = field(default_factory=list) + stages_succeeded: list[str] = field(default_factory=list) + analysis: KernelAnalysis | None = None + profile_text: str = "" + attempt_log: list[str] = field(default_factory=list) + stage_results: list[StageResult] = field(default_factory=list) + + +# --------------------------------------------------------------------------- +# Private helpers +# --------------------------------------------------------------------------- + + +def _format_analysis_for_coordinator(analysis: KernelAnalysis) -> str: + """Compact analysis summary for the coordinator — no code.""" + if not analysis.detected_issues: + return "No optimization opportunities detected. Kernel appears optimal." + + from xe_forge.knowledge.patterns import get_stage_for_issue + + by_stage: dict[OptimizationStage, list] = {} + for iss in analysis.detected_issues: + stage = get_stage_for_issue(iss.issue_type) + by_stage.setdefault(stage, []).append(iss) + + lines = [f"Detected {len(analysis.detected_issues)} issue(s):"] + for stage, issues in sorted(by_stage.items(), key=lambda kv: kv[0].value): + lines.append(f"\n Stage: {stage.value}") + for iss in sorted(issues, key=lambda i: -i.severity): + lines.append(f" [{iss.severity}/5] {iss.issue_type.value}: {iss.description[:120]}") + if iss.suggested_fix: + lines.append(f" Fix: {iss.suggested_fix[:100]}") + + if analysis.has_algorithmic_opportunity: + lines.append( + "\nNote: algorithmic opportunities detected — apply ALGORITHMIC/DISCOVERY first." + ) + return "\n".join(lines) + + +def _parse_stage(stage_str: str) -> OptimizationStage: + """Tolerant string → OptimizationStage parser.""" + from xe_forge.knowledge.patterns import get_stage_for_issue_str + + s = stage_str.strip().lower() + # Direct enum value match + try: + return OptimizationStage(s) + except ValueError: + pass + # Keyword inference via patterns module + inferred = get_stage_for_issue_str(s) + if inferred != OptimizationStage.ANALYSIS: + return inferred + raise ValueError( + f"Unknown stage {stage_str!r}. Valid values: " + + ", ".join(s.value for s in OptimizationStage if s != OptimizationStage.ANALYSIS) + ) + + +def _summarize_stage_result(stage: str, result: StageResult, prev_speedup: float) -> str: + """Single-line summary for coordinator tool return value.""" + if result.success and result.speedup and result.speedup > 1.0: + delta = result.speedup - prev_speedup + return f"IMPROVED: {stage} → {result.speedup:.3f}x speedup" + ( + f" (+{delta:.3f}x vs previous best)" if prev_speedup > 0 else "" + ) + if result.success: + return f"APPLIED: {stage} — no measurable speedup (kernel unchanged or within noise)" + return f"FAILED: {stage} — {result.error_message or 'no details'}" + + +# --------------------------------------------------------------------------- +# Tool factory functions +# --------------------------------------------------------------------------- + + +def make_analyze_tool( + state: CoordinatorState, + analyzer, + pytorch_code: str, + kernel_name: str | None, + input_shapes, + flop, + dtype, +) -> Callable: + def analyze_kernel() -> str: + """Analyze the current kernel for optimization opportunities. + + Returns a structured list of detected issues grouped by stage with severities + and suggested fixes. Call this first and again after significant changes. + """ + try: + analysis = analyzer.analyze( + state.current_code, + pytorch_code or "", + kernel_name or "kernel", + input_shapes, + flop, + target_dtype=str(dtype) if dtype is not None else None, + ) + state.analysis = analysis + summary = _format_analysis_for_coordinator(analysis) + logger.info( + "Coordinator: analyze_kernel returned %d issues", len(analysis.detected_issues) + ) + return summary + except Exception as e: + logger.warning("analyze_kernel failed: %s", e) + return f"Analysis failed: {e}" + + return analyze_kernel + + +def make_retrieve_patterns_tool(state: CoordinatorState, knowledge_base) -> Callable: + def retrieve_patterns(stage: str) -> str: + """Retrieve knowledge base patterns and constraints for a given optimization stage. + + Use before applying a stage to understand best practices and known gotchas. + stage: e.g. "block_pointers", "device_specific", "algorithmic" + """ + if knowledge_base is None: + return f"No knowledge base available for stage {stage!r}. Rely on LLM knowledge." + try: + stage_enum = _parse_stage(stage) + result = knowledge_base.format_for_stage(stage_enum) + return result or f"No patterns found for stage {stage!r}." + except ValueError as e: + return f"Error: {e}" + except Exception as e: + logger.warning("retrieve_patterns failed: %s", e) + return f"Pattern retrieval failed: {e}" + + return retrieve_patterns + + +def make_apply_stage_tool( + state: CoordinatorState, + optimizer, + xpu_config: dict, + kernel_name: str | None, + input_shapes, + spec_dims, + flop, + dtype, + pytorch_code: str, + init_args, + input_dtypes, + vtune_report_ref: list[str], +) -> Callable: + def apply_stage(stage: str, hints: str = "") -> str: + """Apply a single optimization stage to the current kernel. + + If the optimization improves performance, the new kernel automatically becomes + the current version. Returns speedup achieved or a failure reason. + stage: e.g. "block_pointers", "device_specific", "dtype_fix" + hints: optional free-text guidance to pass as extra context (e.g. "focus on reduction loop") + """ + try: + stage_enum = _parse_stage(stage) + stage = stage_enum.value # normalize to canonical stage name + except ValueError as e: + return f"Error: {e}" + + if stage in state.stages_tried: + return f"Stage {stage!r} already tried. Use get_status() to see what's been attempted." + + if state.analysis is None: + return "Must call analyze_kernel() first to identify issues before applying stages." + + state.stages_tried.append(stage) + logger.info("Coordinator: applying stage %s", stage) + + # Build a minimal analysis that includes only this stage's issues + prev_speedup = state.best_speedup + + try: + vtune = vtune_report_ref[0] if vtune_report_ref else "" + perf_ctx = { + "original_ms": None, + "speedup_so_far": prev_speedup if prev_speedup > 0 else None, + } + + stage_result = optimizer.optimize_stage( + code=state.current_code, + stage=stage_enum, + analysis=state.analysis, + xpu_config=xpu_config, + kernel_name=kernel_name, + input_shapes=input_shapes, + spec_dims=spec_dims, + flop=flop, + dtype=dtype, + pytorch_code=pytorch_code or "", + init_args=init_args, + vtune_report=vtune, + perf_context=perf_ctx, + input_dtypes=input_dtypes, + ) + except Exception as e: + logger.warning("apply_stage %s failed with exception: %s", stage, e) + state.stage_results.append( + StageResult( + stage=stage_enum, + success=False, + input_code=state.current_code, + output_code=state.current_code, + error_message=str(e), + ) + ) + return f"FAILED: {stage} — exception: {e}" + + state.stage_results.append(stage_result) + summary = _summarize_stage_result(stage, stage_result, prev_speedup) + + if ( + stage_result.success + and stage_result.output_code + and stage_result.output_code != state.current_code + ): + state.current_code = stage_result.output_code + state.stages_succeeded.append(stage) + spd = stage_result.speedup or 0.0 + if spd > state.best_speedup: + state.best_code = stage_result.output_code + state.best_speedup = spd + + state.attempt_log.append(summary) + logger.info("Coordinator: %s", summary) + return summary + + return apply_stage + + +def make_profile_tool( + state: CoordinatorState, + profiler, + kernel_name: str | None, + spec_path: str | None, + variant_type: str, + vtune_report_ref: list[str], +) -> Callable: + def profile_kernel() -> str: + """Profile the current kernel to identify hardware-level bottlenecks. + + Use when you suspect cache misses, low XVE utilization, or memory bandwidth issues + that static analysis cannot reveal. Requires VTune to be configured. + """ + if profiler is None or not profiler.available(): + return ( + "VTune profiler not available. Use static analysis from analyze_kernel() instead." + ) + if spec_path is None: + return "Cannot profile: no spec_path provided (required for VTune benchmark config)." + try: + tmp = Path(tempfile.mkdtemp()) / f"{kernel_name or 'kernel'}_coordinator_profile.py" + tmp.write_text(state.current_code) + profile_result = profiler.profile(str(tmp), spec_path=spec_path, variant=variant_type) + if profile_result.error: + return f"Profile failed: {profile_result.error}" + report = profile_result.format_for_llm() + state.profile_text = report + vtune_report_ref[0] = report + logger.info("Coordinator: profile_kernel completed") + return report + except Exception as e: + logger.warning("profile_kernel failed: %s", e) + return f"Profile failed: {e}" + + return profile_kernel + + +def make_benchmark_tool( + state: CoordinatorState, + executor, + kernel_name: str | None, + input_shapes, + flop, + dtype, + spec_dims, + init_args, + input_dtypes, +) -> Callable: + def benchmark_current() -> str: + """Benchmark the current kernel against the original baseline. + + Returns speedup and correctness verdict. Use to verify cumulative improvement. + """ + if executor is None: + return "No executor available — cannot benchmark." + if not input_shapes and spec_dims is None: + return "No input shapes or spec_dims — cannot benchmark." + try: + is_sycl = spec_dims is not None and not input_shapes + + if is_sycl: + comparison = executor.compare_kernels( + original_code=state.original_code, + optimized_code=state.current_code, + dims=spec_dims, + ) + else: + comparison = executor.compare_kernels( + original_code=state.original_code, + optimized_code=state.current_code, + kernel_name=kernel_name, + input_shapes=input_shapes, + flop=flop, + dtype=dtype, + init_args=init_args, + input_dtypes=input_dtypes, + ) + + if not comparison.optimized_correct: + return ( + f"INCORRECT: {comparison.feedback_message or 'kernel produces wrong results'}" + ) + + speedup = comparison.speedup or 1.0 + orig_ms = getattr(comparison, "original_time_us", None) + opt_ms = getattr(comparison, "optimized_time_us", None) + + lines = [f"Speedup: {speedup:.3f}x"] + if orig_ms and opt_ms: + lines.append(f"Original: {orig_ms:.1f} µs → Optimized: {opt_ms:.1f} µs") + if comparison.is_slower: + lines.append("WARNING: current code is SLOWER than original") + elif speedup > 1.0: + lines.append("Improvement confirmed vs original baseline.") + + if speedup > state.best_speedup: + state.best_speedup = speedup + state.best_code = state.current_code + + return "\n".join(lines) + except Exception as e: + logger.warning("benchmark_current failed: %s", e) + return f"Benchmark failed: {e}" + + return benchmark_current + + +def make_status_tool(state: CoordinatorState) -> Callable: + def get_status() -> str: + """Return a summary of the optimization state: speedup achieved, stages tried/succeeded, pending issues. + + Use to review progress and decide next steps. + """ + lines = ["=== Coordinator Status ==="] + lines.append(f"Best speedup so far: {state.best_speedup:.3f}x") + lines.append( + f"Stages tried ({len(state.stages_tried)}): {', '.join(state.stages_tried) or 'none'}" + ) + lines.append( + f"Stages succeeded ({len(state.stages_succeeded)}): {', '.join(state.stages_succeeded) or 'none'}" + ) + if state.analysis: + from xe_forge.knowledge.patterns import get_stage_for_issue + + tried = {s.strip().lower() for s in state.stages_tried} + remaining = sorted( + { + get_stage_for_issue(iss.issue_type).value + for iss in state.analysis.detected_issues + } + - tried + ) + if remaining: + lines.append(f"Remaining stages to attempt: {', '.join(remaining[:8])}") + else: + lines.append("All detected stages have been attempted.") + if state.attempt_log: + lines.append("\nAttempt history:") + for entry in state.attempt_log[-5:]: + lines.append(f" {entry}") + return "\n".join(lines) + + return get_status diff --git a/src/xe_forge/agents/cover.py b/src/xe_forge/agents/cover.py index bc18076..c64a43e 100644 --- a/src/xe_forge/agents/cover.py +++ b/src/xe_forge/agents/cover.py @@ -105,13 +105,22 @@ def __init__( .append("next_thought", dspy.OutputField(), type_=str) ) - fallback_signature = dspy.Signature( - {**signature.input_fields, **signature.output_fields}, - signature.instructions, - ).append("trajectory", dspy.InputField(), type_=str) + # Add a reasoning field to the fallback signature so the LLM can think through the + # final extraction step. Use dspy.Reasoning if available (DSPy >=3.3) — it transparently + # routes to native blocks on reasoning models and CoT text on others. + # Fall back to a plain str field for older installs where dspy.Reasoning doesn't exist. + _ReasoningType = getattr(dspy, "Reasoning", str) + fallback_signature = ( + dspy.Signature( + {**signature.input_fields, **signature.output_fields}, + signature.instructions, + ) + .append("trajectory", dspy.InputField(), type_=str) + .append("reasoning", dspy.OutputField(), type_=_ReasoningType) + ) self.cover = dspy.Predict(self.cover_signature) - self.extract = dspy.ChainOfThought(fallback_signature) + self.extract = dspy.Predict(fallback_signature) def _format_trajectory(self, trajectory: dict[str, Any]): """Format trajectory for LLM consumption.""" diff --git a/src/xe_forge/agents/optimizer_agent.py b/src/xe_forge/agents/optimizer_agent.py index eaf6564..221317f 100644 --- a/src/xe_forge/agents/optimizer_agent.py +++ b/src/xe_forge/agents/optimizer_agent.py @@ -12,6 +12,7 @@ from xe_forge.agents.base import Optimizer from xe_forge.agents.cover import CoVeR +from xe_forge.agents.utils import SUCCESS_MESSAGE, extract_gemm_dims, verify_sycl from xe_forge.knowledge.loader import KnowledgeBase from xe_forge.models import ( DSL, @@ -22,97 +23,8 @@ logger = logging.getLogger(__name__) -def _extract_gemm_dims( - input_shapes: list[tuple[int, ...]] | None, -) -> tuple[int, int, int]: - """Extract M, N, K from GEMM input shapes [(M, K), (K, N)].""" - if input_shapes and len(input_shapes) >= 2: - a, b = input_shapes[0], input_shapes[1] - if len(a) >= 2 and len(b) >= 2: - return a[-2], b[-1], a[-1] - return 1024, 1024, 1024 - - -def _verify_sycl(code, original_code, executor, input_shapes, spec_dims=None): - """Verify a SYCL C++ kernel: basic structure check + runtime comparison.""" - if "#include" not in code: - return "MISSING: C++ code must contain #include directives." - if "sycl" not in code.lower() and "cutlass" not in code.lower(): - return "MISSING: Code does not appear to be a SYCL/CUTLASS kernel." - - if executor: - try: - _dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(input_shapes), strict=False) - ) - comparison = executor.compare_kernels( - original_code=original_code, - optimized_code=code, - dims=_dims, - ) - if not comparison.optimized_correct: - return comparison.feedback_message or "Optimized kernel failed." - if comparison.is_slower: - sd = 1.0 / comparison.speedup if comparison.speedup > 0 else float("inf") - return ( - f"PERFORMANCE REGRESSION: {sd:.2f}x SLOWER.\n" - f"Original: {comparison.original_time_ms:.4f}ms ({comparison.original_tflops or 0:.3f} TFlop/s)\n" - f"Optimized: {comparison.optimized_time_ms:.4f}ms ({comparison.optimized_tflops or 0:.3f} TFlop/s)" - ) - logger.info( - f"SYCL optimization verified: {comparison.speedup:.2f}x speedup " - f"({comparison.original_tflops or 0:.3f} -> {comparison.optimized_tflops or 0:.3f} TFlop/s)" - ) - return SUCCESS_MESSAGE - except Exception as e: - return f"RUNTIME ERROR: {e!s}" - - logger.warning("No executor - accepting SYCL code based on static checks only") - return SUCCESS_MESSAGE - - -SUCCESS_MESSAGE = "Success! Optimization verified and kernel is faster." - - class OptimizationSignature(dspy.Signature): - """Apply optimization transformation to Triton kernel. - - You are an expert Triton kernel optimizer for Intel XPU with deep knowledge - of GPU programming, numerical linear algebra, and high-performance computing. - - Optimize the kernel for maximum performance while producing numerically - equivalent outputs. You may change the algorithm if outputs are equivalent. - Maintain the same Model class signature including weights shapes and names. - - === STAGE-SPECIFIC GUIDANCE === - ALGORITHMIC: mathematical simplifications, CSE, loop-invariant hoisting, - caching intermediates, reorder associative ops, tree reductions, - exploit GEMM structure (symmetric, triangular, low-rank). - DTYPE_FIX: float64->float32, proper accumulator precision, remove - unnecessary type conversions. - FUSION: fuse kernel launches, elementwise chains, reduction+elementwise. - MEMORY_ACCESS: fix uncoalesced access, remove transposes from inner loops, - add boundary checks, reduce register pressure. - BLOCK_POINTERS: use tl.make_block_ptr(), boundary_check=(0,1) tuple format, - tl.advance() for pointer updates. - XPU_SPECIFIC: BLOCK_M=256, BLOCK_N=256, BLOCK_K=32, num_warps=32, - GROUP_SIZE_M swizzling. - GRF MODE: grf_mode is a compiler option, NOT a triton.Config() kwarg. - Declare it as tl.constexpr in the kernel signature: - grf_mode: tl.constexpr (values: "default", "128", "256", "auto") - Use "auto" — it automatically selects 256-GRF when register spill > 1000 bytes. - 256-GRF requires num_warps <= 32 (halved thread occupancy). - PERSISTENT_KERNEL: persistent kernel pattern, tune NUM_PROGS. - DISCOVERY: apply the open-ended optimization described in the issues field. - This is a novel optimization not covered by standard stages. Follow the - proposal exactly, preserving all numerical equivalences. - - === CODE REQUIREMENTS === - - Include ALL imports, @triton.jit decorator, kernel function, Model class - - num_warps must be power of 2; block sizes must be powers of 2 - - NEVER replace @triton.jit kernels with torch.matmul, torch.mm, torch.bmm, - or any vendor library (oneDNN, cuBLAS, MKL). Keep all original Triton kernels. - """ + """Apply a single optimization stage to a Triton kernel to improve performance.""" original_code: str = dspy.InputField(desc="Original Triton kernel code for reference") current_code: str = dspy.InputField(desc="Current Triton kernel code to optimize") @@ -146,29 +58,7 @@ class OptimizationSignature(dspy.Signature): class AlgorithmicOptimizationSignature(dspy.Signature): - """Apply algorithmic / mathematical optimization to a Triton kernel. - - You are an expert in numerical linear algebra, compiler optimizations, and - high-performance GPU kernel design. - - Transform the kernel to perform FEWER FLOPs and/or FEWER memory accesses - while producing numerically equivalent results. - - Think about: - 1. Matrix structure exploitation (symmetric, triangular, diagonal, low-rank, sparse) - 2. Associative / distributive law rewrites to reduce FLOPs - 3. Common sub-expression elimination - 4. Loop-invariant code hoisting - 5. Caching intermediates in registers vs recomputing - 6. Tree reductions vs serial reductions - 7. Algebraic simplification of fused computations - - Maintain the Model class signature. Produce equivalent outputs. - - === CODE REQUIREMENTS === - - Include ALL imports, @triton.jit decorator, kernel function, Model class - - NEVER replace @triton.jit kernels with torch.matmul, torch.mm, or any vendor library. - """ + """Apply algorithmic/mathematical optimization to reduce FLOPs while preserving output equivalence.""" original_code: str = dspy.InputField(desc="Original Triton kernel code for reference") current_code: str = dspy.InputField(desc="Current Triton kernel code to optimize") @@ -195,50 +85,7 @@ class AlgorithmicOptimizationSignature(dspy.Signature): class AutotuneSignature(dspy.Signature): - """Add or improve @triton.autotune configuration for a Triton kernel. - - You are an expert in Triton kernel autotuning for Intel XPU. - - Your task: Add or improve the @triton.autotune decorator so the kernel - automatically selects the best configuration at runtime. - - You will receive: - - The current kernel code - - Hardware information (compute units, memory, capabilities) - - Problem shapes (M, N, K dimensions) - - A set of suggested autotune configurations generated from hardware analysis - - Your job: - 1. Add @triton.autotune decorator with a good set of configs to search. - 2. Use the suggested configs as a starting point but ADD more configs - based on your knowledge of what works well for this kernel type. - 3. Include the key= argument so configs are re-evaluated when shapes change. - 4. Ensure num_warps and num_stages are included in each config. - 5. Ensure BLOCK sizes are powers of 2 and appropriate for the hardware. - 6. For Intel XPU, always include at least one config with num_warps=32 - and large tile sizes (256x256). - 7. Remove any hardcoded meta-parameters that are now covered by autotune. - 8. Keep the kernel functionally equivalent. - - Tips for good autotune configs: - - Vary BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K across powers of 2 - - Include both small tiles (64x64) for small problems and large tiles - (256x256) for large problems - - Vary num_warps: try 4, 8, 16, 32 - - Vary num_stages: try 2, 3, 4 - - Include GROUP_SIZE_M for L2 cache swizzling - - Use key= with the shape arguments that affect tiling - - Do NOT put grf_mode in triton.Config() — it causes TypeError at runtime. - grf_mode is a compiler option: declare it as tl.constexpr in the kernel - signature. Use grf_mode="auto" (auto-selects 256-GRF if spill > 1000 bytes) - or grf_mode="256" for large register file. Requires num_warps <= 32. - - === CODE REQUIREMENTS === - - Include ALL imports (torch, triton, triton.language as tl) - - Include @triton.autotune with configs list and key - - Include @triton.jit on the kernel - - Include the Model class with forward() method - """ + """Add or improve @triton.autotune configuration so the kernel auto-selects optimal tile/warp parameters.""" original_code: str = dspy.InputField(desc="Original Triton kernel code for reference") current_code: str = dspy.InputField(desc="Current Triton kernel code to add autotune to") @@ -268,43 +115,7 @@ class AutotuneSignature(dspy.Signature): class SyclOptimizationSignature(dspy.Signature): - """Optimize a SYCL/CUTLASS C++ kernel for Intel XPU. - - You are an expert in SYCL, CUTLASS/XeTLA, Intel XPU GPU architecture, - and high-performance C++ kernel optimization. - - Optimize the kernel for maximum performance while producing numerically - equivalent outputs. You may change template parameters, dispatch policies, - data types, and memory layouts. - - === SYCL/CUTLASS OPTIMIZATION KNOBS === - - TileShape: Shape<_M, _N, _K> — try 256x256x32, 128x128x64, 128x256x32 - - PipelineStages: 2, 3, or 4 — more prefetching vs register pressure - - MMA Atom: XE_DPAS_TT — SubgroupSize 4 or 8 - - Dispatch Policy: MainloopXeL1Staged (L1 cached), MainloopXeL0Staged (uncached) - - Data types: bfloat16_t/half_t inputs, float/bfloat16_t accumulators - - Memory layout: RowMajor vs ColumnMajor for A, B, C, D - - Epilogue: LinearCombination, bias, activation via FusionCallbacks - - GmemTiledCopy: void (auto) or explicit copy atoms - - === STAGE-SPECIFIC GUIDANCE === - ALGORITHMIC: mathematical simplifications, CSE, loop-invariant hoisting, - exploit GEMM structure (symmetric, triangular, low-rank). - DTYPE_FIX: use bfloat16_t/half_t inputs, float accumulators, avoid double. - FUSION: fuse into CUTLASS epilogue callbacks — LinearCombination, bias, activation. - MEMORY_ACCESS: fix layout mismatch (RowMajor vs ColumnMajor), increase PipelineStages - for better prefetching, reduce register pressure. - DEVICE_SPECIFIC: TileShape 256x256x32 or 128x128x64, PipelineStages=2-3, - XE_DPAS_TT<8, float, bfloat16_t>, MainloopXeL1Staged dispatch policy. - DISCOVERY: apply the open-ended optimization described in the issues field. - - === CODE REQUIREMENTS === - - Must be complete, valid SYCL C++ with all #include directives - - Must use cutlass namespace and CUTLASS template types - - Must include ExampleRunner template and main() function - - Must compile with icpx -fsycl - - Keep the same output format (Cutlass GEMM Performance line) - """ + """Apply a single optimization stage to a SYCL/CUTLASS C++ kernel to improve performance.""" original_code: str = dspy.InputField(desc="Original SYCL/CUTLASS C++ kernel for reference") current_code: str = dspy.InputField(desc="Current SYCL C++ kernel code to optimize") @@ -329,26 +140,7 @@ class SyclOptimizationSignature(dspy.Signature): class SyclAlgorithmicOptimizationSignature(dspy.Signature): - """Apply algorithmic / mathematical optimization to a SYCL/CUTLASS C++ kernel. - - You are an expert in numerical linear algebra, compiler optimizations, and - high-performance GPU kernel design for Intel XPU. - - Transform the kernel to perform FEWER FLOPs and/or FEWER memory accesses - while producing numerically equivalent results. - - Think about: - 1. Matrix structure exploitation (symmetric, triangular, diagonal, low-rank) - 2. Associative / distributive law rewrites to reduce FLOPs - 3. Common sub-expression elimination in template expressions - 4. Data layout optimization (RowMajor vs ColumnMajor) - 5. Batch dimension exploitation - - === CODE REQUIREMENTS === - - Must be complete, valid SYCL C++ with all #include directives - - Keep CUTLASS GEMM structure (GemmUniversalAdapter, ExampleRunner, main) - - Must compile with icpx -fsycl - """ + """Apply algorithmic/mathematical optimization to a SYCL/CUTLASS kernel to reduce FLOPs.""" original_code: str = dspy.InputField(desc="Original SYCL C++ kernel for reference") current_code: str = dspy.InputField(desc="Current SYCL C++ kernel to optimize") @@ -440,12 +232,14 @@ def __init__( validator=None, max_iterations=5, dsl: DSL | str = DSL.TRITON, + extra_instructions: str = "", ): self.executor = executor self.validator = validator self.max_iterations = max_iterations self.knowledge_base: KnowledgeBase | None = knowledge_base self.dsl = DSL(dsl) if isinstance(dsl, str) else dsl + self.extra_instructions = extra_instructions if not executor: logger.warning("No executor provided - kernels will NOT be verified at runtime!") @@ -477,10 +271,10 @@ def compile_and_verify(optimized_code: dspy.Code["python"]) -> str: ) if dsl == DSL.SYCL: - result = _verify_sycl(code, original_code, executor, input_shapes, spec_dims) + result = verify_sycl(code, original_code, executor, input_shapes, spec_dims) if result == SUCCESS_MESSAGE and executor: _dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(input_shapes), strict=False) + zip(("M", "N", "K"), extract_gemm_dims(input_shapes), strict=False) ) try: c = executor.compare_kernels( @@ -786,6 +580,9 @@ def optimize_stage( "knowledge_base_context": kb_context, } + if self.extra_instructions: + sig = sig.append_instructions(self.extra_instructions) + cover = CoVeR( signature=sig, tools=[verify_tool], @@ -1205,7 +1002,7 @@ def _final_verify( c = cached_comparison elif self.dsl == DSL.SYCL: _dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(shapes), strict=False) + zip(("M", "N", "K"), extract_gemm_dims(shapes), strict=False) ) c = self.executor.compare_kernels( original_code=orig, diff --git a/src/xe_forge/agents/react_agent.py b/src/xe_forge/agents/react_agent.py index ab93ae1..a571c53 100644 --- a/src/xe_forge/agents/react_agent.py +++ b/src/xe_forge/agents/react_agent.py @@ -12,6 +12,7 @@ import dspy from xe_forge.agents.base import Optimizer +from xe_forge.agents.utils import SUCCESS_MESSAGE, extract_gemm_dims, verify_sycl from xe_forge.knowledge.patterns import get_stage_for_issue try: @@ -29,88 +30,8 @@ logger = logging.getLogger(__name__) -SUCCESS_MESSAGE = "Success! Optimization verified and kernel is faster." - - -def _extract_gemm_dims( - input_shapes: list[tuple[int, ...]] | None, -) -> tuple[int, int, int]: - """Extract M, N, K from GEMM input shapes [(M, K), (K, N)].""" - if input_shapes and len(input_shapes) >= 2: - a, b = input_shapes[0], input_shapes[1] - if len(a) >= 2 and len(b) >= 2: - return a[-2], b[-1], a[-1] - return 1024, 1024, 1024 - - -def _verify_sycl(code, original_code, executor, input_shapes, spec_dims=None): - """Verify a SYCL C++ kernel: basic structure check + runtime comparison.""" - if "#include" not in code: - return "MISSING: C++ code must contain #include directives." - if "sycl" not in code.lower() and "cutlass" not in code.lower(): - return "MISSING: Code does not appear to be a SYCL/CUTLASS kernel." - - if executor: - try: - _dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(input_shapes), strict=False) - ) - comparison = executor.compare_kernels( - original_code=original_code, - optimized_code=code, - dims=_dims, - ) - if not comparison.optimized_correct: - return comparison.feedback_message or "Optimized kernel failed." - if comparison.is_slower: - sd = 1.0 / comparison.speedup if comparison.speedup > 0 else float("inf") - return ( - f"PERFORMANCE REGRESSION: {sd:.2f}x SLOWER.\n" - f"Original: {comparison.original_time_ms:.4f}ms ({comparison.original_tflops or 0:.3f} TFlop/s)\n" - f"Optimized: {comparison.optimized_time_ms:.4f}ms ({comparison.optimized_tflops or 0:.3f} TFlop/s)" - ) - logger.info( - f"SYCL optimization verified: {comparison.speedup:.2f}x speedup " - f"({comparison.original_tflops or 0:.3f} -> {comparison.optimized_tflops or 0:.3f} TFlop/s)" - ) - return SUCCESS_MESSAGE - except Exception as e: - return f"RUNTIME ERROR: {e!s}" - - logger.warning("No executor - accepting SYCL code based on static checks only") - return SUCCESS_MESSAGE - - class OptimizationReActSignature(dspy.Signature): - """Apply optimization transformation to Triton kernel. - - You are an expert Triton kernel optimizer for Intel XPU. - - Your task: Optimize the kernel for maximum performance. - You may change the algorithm/computation approach if it produces equivalent outputs. - Maintain the same model signature, including the weights' shapes and names. This is necessary for having identical initialization process for formal verification which is done by a correctness tool. - - === OPTIMIZATION PRIORITIES === - 1. Apply the specific optimization patterns from knowledge_patterns - 2. Use block pointers with tl.make_block_ptr() for better memory access - 3. Use optimal tile sizes: BLOCK_M=256, BLOCK_N=256, BLOCK_K=32 for XPU - 4. Set num_warps=32 for Intel XPU - 5. Add GROUP_SIZE_M swizzling for better L2 cache utilization - 6. Use boundary_check=(0, 1) tuple format, NOT booleans - - === CODE REQUIREMENTS === - - Include ALL imports (torch, triton, triton.language as tl) - - Include the @triton.jit decorator and kernel function - - Include the Model class with forward() method - - num_warps MUST be a power of 2 (1, 2, 4, 8, 16, 32) - - num_stages MUST be a positive integer - - Block sizes (BLOCK_M, BLOCK_N, BLOCK_K) MUST be powers of 2 - - Block sizes should not exceed 256 for most cases - - === WHAT TO CHANGE === - Focus on the issues listed and apply the patterns from knowledge_patterns. - Be aggressive with optimizations - the verification tool will check correctness. - """ + """Apply an optimization transformation to a Triton kernel; a verification tool checks correctness.""" original_code: dspy.Code["python"] = dspy.InputField( # noqa: UP037 desc="Original Triton kernel code for reference" @@ -133,28 +54,7 @@ class OptimizationReActSignature(dspy.Signature): class SyclOptimizationReActSignature(dspy.Signature): - """Optimize a SYCL/CUTLASS C++ kernel for Intel XPU. - - You are an expert SYCL/CUTLASS kernel optimizer for Intel XPU. - - Your task: Optimize the C++ kernel for maximum performance. - You may change template parameters, dispatch policies, and data types - if the outputs remain numerically equivalent. - - === OPTIMIZATION PRIORITIES === - 1. TileShape: try Shape<_256,_256,_32> or Shape<_128,_128,_64> for BMG - 2. PipelineStages: 2-4 (balance prefetching vs register pressure) - 3. MMA Atom: XE_DPAS_TT<8, float, bfloat16_t> for BMG - 4. Dispatch Policy: MainloopXeL1Staged for L1 caching - 5. Data types: bfloat16_t inputs, float accumulators - 6. Memory layout: match RowMajor/ColumnMajor to access patterns - - === CODE REQUIREMENTS === - - Complete, valid SYCL C++ with all #include directives - - CUTLASS template types, ExampleRunner, and main() - - Must compile with icpx -fsycl - - Keep the Cutlass GEMM Performance output format - """ + """Apply an optimization transformation to a SYCL/CUTLASS C++ kernel; a verification tool checks correctness.""" original_code: dspy.Code[cpp] = dspy.InputField( desc="Original SYCL C++ kernel code for reference" @@ -184,12 +84,14 @@ def __init__( validator: Callable | None = None, max_iterations: int = 5, dsl: DSL | str = DSL.TRITON, + extra_instructions: str = "", ): self.knowledge_base = knowledge_base self.executor = executor self.validator = validator self.max_iterations = max_iterations self.dsl = DSL(dsl) if isinstance(dsl, str) else dsl + self.extra_instructions = extra_instructions if not executor: logger.warning("No executor provided - kernels will NOT be verified at runtime!") @@ -218,7 +120,7 @@ def compile_and_verify(optimized_code: dspy.Code["python"]) -> str: # noqa: UP0 code: str = optimized_code.code if dsl == DSL.SYCL: - return _verify_sycl( + return verify_sycl( code, original_code, executor, @@ -391,16 +293,18 @@ def optimize_stage( input_dtypes=input_dtypes, ) - # Create ReAct agent for this optimization sig = SyclOptimizationReActSignature if self.dsl == DSL.SYCL else OptimizationReActSignature - react_agent = dspy.ReAct( + if self.extra_instructions: + sig = sig.append_instructions(self.extra_instructions) + + react_agent = dspy.ReActV2( signature=sig, tools=[verify_tool], max_iters=self.max_iterations, ) try: - logger.info(f"Starting ReAct optimization (max {self.max_iterations} iterations)") + logger.info(f"Starting ReActV2 optimization (max {self.max_iterations} iterations)") result = react_agent( original_code=original_code, @@ -411,6 +315,15 @@ def optimize_stage( xpu_config=xpu_config_text, ) + # Log termination reason when using ReActV2 + termination_reason = getattr(result, "termination_reason", None) + if termination_reason: + logger.info(f"ReActV2 termination_reason: {termination_reason}") + if termination_reason in ("max_iters", "context_window_exceeded"): + logger.warning( + f"Agent stopped due to {termination_reason} — final output may be incomplete" + ) + # Extract optimized code from result if not hasattr(result, "optimized_code") or result.optimized_code is None: logger.error("Agent didn't return code") @@ -424,7 +337,13 @@ def optimize_stage( optimized_code: str = result.optimized_code.code - trajectory = result.trajectory if hasattr(result, "trajectory") else {} + # ReActV2 uses result.history (list of turn dicts); classic ReAct uses result.trajectory (dict) + trajectory = {} + if hasattr(result, "history") and isinstance(result.history, list): + for idx, turn in enumerate(result.history): + trajectory[f"thought_{idx}"] = str(turn.get("thought", "")) + elif hasattr(result, "trajectory"): + trajectory = result.trajectory or {} # Verify the final code directly (don't rely on trajectory) @@ -444,7 +363,7 @@ def optimize_stage( try: if self.dsl == DSL.SYCL: _dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(input_shapes), strict=False) + zip(("M", "N", "K"), extract_gemm_dims(input_shapes), strict=False) ) comparison = self.executor.compare_kernels( original_code=original_code, diff --git a/src/xe_forge/agents/utils.py b/src/xe_forge/agents/utils.py new file mode 100644 index 0000000..41390a2 --- /dev/null +++ b/src/xe_forge/agents/utils.py @@ -0,0 +1,56 @@ +"""Shared utilities for Xe-Forge agents.""" + +import logging + +logger = logging.getLogger(__name__) + +SUCCESS_MESSAGE = "Success! Optimization verified and kernel is faster." + + +def extract_gemm_dims( + input_shapes: list[tuple[int, ...]] | None, +) -> tuple[int, int, int]: + """Extract M, N, K from GEMM input shapes [(M, K), (K, N)].""" + if input_shapes and len(input_shapes) >= 2: + a, b = input_shapes[0], input_shapes[1] + if len(a) >= 2 and len(b) >= 2: + return a[-2], b[-1], a[-1] + return 1024, 1024, 1024 + + +def verify_sycl(code, original_code, executor, input_shapes, spec_dims=None): + """Verify a SYCL C++ kernel: basic structure check + runtime comparison.""" + if "#include" not in code: + return "MISSING: C++ code must contain #include directives." + if "sycl" not in code.lower() and "cutlass" not in code.lower(): + return "MISSING: Code does not appear to be a SYCL/CUTLASS kernel." + + if executor: + try: + _dims = spec_dims or dict( + zip(("M", "N", "K"), extract_gemm_dims(input_shapes), strict=False) + ) + comparison = executor.compare_kernels( + original_code=original_code, + optimized_code=code, + dims=_dims, + ) + if not comparison.optimized_correct: + return comparison.feedback_message or "Optimized kernel failed." + if comparison.is_slower: + sd = 1.0 / comparison.speedup if comparison.speedup > 0 else float("inf") + return ( + f"PERFORMANCE REGRESSION: {sd:.2f}x SLOWER.\n" + f"Original: {comparison.original_time_ms:.4f}ms ({comparison.original_tflops or 0:.3f} TFlop/s)\n" + f"Optimized: {comparison.optimized_time_ms:.4f}ms ({comparison.optimized_tflops or 0:.3f} TFlop/s)" + ) + logger.info( + f"SYCL optimization verified: {comparison.speedup:.2f}x speedup " + f"({comparison.original_tflops or 0:.3f} -> {comparison.optimized_tflops or 0:.3f} TFlop/s)" + ) + return SUCCESS_MESSAGE + except Exception as e: + return f"RUNTIME ERROR: {e!s}" + + logger.warning("No executor - accepting SYCL code based on static checks only") + return SUCCESS_MESSAGE diff --git a/src/xe_forge/cli.py b/src/xe_forge/cli.py index b7aa7f6..d86bf34 100644 --- a/src/xe_forge/cli.py +++ b/src/xe_forge/cli.py @@ -32,7 +32,9 @@ def _setup_dspy(config: Config) -> None: if config.llm.api_key: os.environ["OPENAI_API_KEY"] = config.llm.api_key - litellm.client_session = httpx.Client(verify=False) + litellm.client_session = httpx.Client() + litellm.aclient_session = httpx.AsyncClient() + litellm.ssl_verify = True lm = dspy.LM( model=config.llm.model, api_base=config.llm.api_base, diff --git a/src/xe_forge/config.py b/src/xe_forge/config.py index 9301b22..f1d2bdc 100644 --- a/src/xe_forge/config.py +++ b/src/xe_forge/config.py @@ -29,7 +29,7 @@ class AgentConfig: max_iterations: int = 5 use_cover: bool = True - strategy: str = "cover" # cover, react, hybrid + strategy: str = "cover" # cover, react, coordinator @dataclass diff --git a/src/xe_forge/models.py b/src/xe_forge/models.py index 01586d4..143febb 100644 --- a/src/xe_forge/models.py +++ b/src/xe_forge/models.py @@ -162,6 +162,7 @@ class OptimizationResult(BaseModel): optimized_tflops: float | None = None original_memory_bw: float | None = None optimized_memory_bw: float | None = None + token_usage: dict | None = None class KnowledgeEntry(BaseModel): diff --git a/src/xe_forge/pipeline.py b/src/xe_forge/pipeline.py index 8e29caf..31625a7 100644 --- a/src/xe_forge/pipeline.py +++ b/src/xe_forge/pipeline.py @@ -8,6 +8,7 @@ import litellm from xe_forge.agents import AnalyzerAgent, Optimizer, OptimizerAgent, OptimizerReActAgent +from xe_forge.agents.utils import extract_gemm_dims from xe_forge.config import Config, get_config from xe_forge.core.device_query import get_device_config_for_pipeline from xe_forge.knowledge.loader import KnowledgeBase, load_knowledge_base @@ -19,21 +20,11 @@ ) from xe_forge.planner import DEFAULT_STAGE_ORDER as PLANNER_DEFAULT_STAGE_ORDER from xe_forge.planner import PlannerAgent +from xe_forge.prompts.device_prompts import PromptLibrary logger = logging.getLogger(__name__) -def _extract_gemm_dims( - input_shapes: list[tuple[int, ...]] | None, -) -> tuple[int, int, int]: - """Extract M, N, K from GEMM input shapes [(M, K), (K, N)].""" - if input_shapes and len(input_shapes) >= 2: - a, b = input_shapes[0], input_shapes[1] - if len(a) >= 2 and len(b) >= 2: - return a[-2], b[-1], a[-1] - return 1024, 1024, 1024 - - DEFAULT_STAGE_ORDER: list[OptimizationStage] = [ OptimizationStage.ANALYSIS, OptimizationStage.ALGORITHMIC, @@ -90,27 +81,76 @@ def __init__( else: logger.info(" Knowledge base: disabled (set KNOWLEDGE_BASE_ENABLED=true to enable)") + _dsl_str = ( + self.config.device_config.dsl.value + if hasattr(self.config.device_config.dsl, "value") + else self.config.device_config.dsl + ) + _prompt_lib = PromptLibrary( + dsl=_dsl_str, + device_type=self.config.device_config.device, + ) + + # Render device-specific optimizer guidance from Jinja2 templates. + # Falls back to device_context_addendum() if template not found. + _dsl = self.config.device_config.dsl + _is_sycl_dsl = str(_dsl.value if hasattr(_dsl, "value") else _dsl) == "sycl" + _opt_template = "sycl_optimization_signature" if _is_sycl_dsl else "optimization_signature" + _react_template = ( + "sycl_optimization_signature" if _is_sycl_dsl else "optimization_react_signature" + ) + _opt_instructions = _prompt_lib.render_for_signature(_opt_template) + _react_instructions = _prompt_lib.render_for_signature(_react_template) + self.analyzer = AnalyzerAgent( knowledge_base=self.knowledge_base, dsl=self.config.device_config.dsl, ) self.planner = PlannerAgent() + self.coordinator = None match self.config.agent.strategy: + case "coordinator": + from xe_forge.agents.coordinator import CoordinatorAgent + + self.coordinator = CoordinatorAgent( + analyzer=self.analyzer, + executor=executor, + knowledge_base=self.knowledge_base, + profiler=self.profiler, + max_iters=self.config.agent.max_iterations * 4, + extra_instructions=_opt_instructions, + dsl=self.config.device_config.dsl, + ) + self.optimizer = None case "cover": - Agent = OptimizerAgent + self.optimizer = OptimizerAgent( + executor=executor, + validator=validator, + max_iterations=self.config.agent.max_iterations, + knowledge_base=self.knowledge_base, + dsl=self.config.device_config.dsl, + extra_instructions=_opt_instructions, + ) case "react": - Agent = OptimizerReActAgent + self.optimizer = OptimizerReActAgent( + executor=executor, + validator=validator, + max_iterations=self.config.agent.max_iterations, + knowledge_base=self.knowledge_base, + dsl=self.config.device_config.dsl, + extra_instructions=_react_instructions, + ) case _: - Agent = OptimizerAgent + self.optimizer = OptimizerAgent( + executor=executor, + validator=validator, + max_iterations=self.config.agent.max_iterations, + knowledge_base=self.knowledge_base, + dsl=self.config.device_config.dsl, + extra_instructions=_opt_instructions, + ) - self.optimizer = Agent( - executor=executor, - validator=validator, - max_iterations=self.config.agent.max_iterations, - knowledge_base=self.knowledge_base, - dsl=self.config.device_config.dsl, - ) self.executor = executor self.validator = validator @@ -233,7 +273,7 @@ def optimize( try: if _is_sycl: _sycl_dims = spec_dims or dict( - zip(("M", "N", "K"), _extract_gemm_dims(input_shapes), strict=False) + zip(("M", "N", "K"), extract_gemm_dims(input_shapes), strict=False) ) orig_r = _bench_ex.execute( kernel_code=kernel_code, @@ -270,12 +310,76 @@ def optimize( except Exception as e: logger.warning("Could not initialize trial tree: %s", e) - candidates = [] - best_k = max(1, self.config.optimization.best_k) + _usage_ctx = dspy.track_usage() + _usage_tracker = _usage_ctx.__enter__() + result_out: OptimizationResult | None = None + try: + candidates = [] + + if self.coordinator is not None: + # --- Coordinator (agentic) path --- + result = OptimizationResult( + kernel_name=display_name, original_code=kernel_code, timestamp=datetime.now() + ) + result.original_tflops, result.original_ms = val_orig_tflops, val_orig_ms + + etd = target_dtype or self.config.optimization.target_dtype + if etd is None and dtype is not None: + etd = { + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.float32: "float32", + }.get(dtype) + + device_type = self.config.device_config.device + xpu_config = get_device_config_for_pipeline( + device_type=device_type, + input_shapes=input_shapes, + config=self.config, + dtype=etd or "float16", + ) + _dsl_str = str( + self.config.device_config.dsl.value + if hasattr(self.config.device_config.dsl, "value") + else self.config.device_config.dsl + ) + kernel_specs = _build_kernel_specs( + kernel_name=display_name, + input_shapes=input_shapes, + flop=flop, + dtype=etd, + device=device_type, + dsl=_dsl_str, + ) + + logger.info("=" * 60 + "\nSTRATEGY: COORDINATOR (agentic)\n" + "=" * 60) + best_code, speedup, _summary, stage_results = self.coordinator.run( + kernel_code=kernel_code, + kernel_specs=kernel_specs, + pytorch_code=reference_code or "", + kernel_name=kernel_name, + input_shapes=input_shapes, + flop=flop, + dtype=dtype, + spec_dims=spec_dims, + init_args=init_args, + input_dtypes=input_dtypes, + xpu_config=xpu_config, + spec_path=spec_path, + variant_type=variant_type, + ) + + result.optimized_code = best_code + result.stages_applied = stage_results + result.success = best_code != kernel_code or speedup > 1.0 + if speedup > 1.0: + result.total_speedup = speedup + candidates.append(result) - for attempt in range(best_k): + # Fixed stage loop runs only when coordinator strategy is not active. + best_k = max(1, self.config.optimization.best_k) if self.coordinator is None else 0 if best_k > 1: - logger.info(f"Attempt {attempt + 1}/{best_k}") + logger.info("Attempt 1/%d", best_k) result = OptimizationResult( kernel_name=display_name, original_code=kernel_code, timestamp=datetime.now() @@ -316,7 +420,6 @@ def optimize( if not analysis.detected_issues: result.success, result.optimized_code = True, kernel_code candidates.append(result) - continue logger.info("=" * 60 + "\nSTAGE: PLANNING\n" + "=" * 60) from xe_forge.knowledge.patterns import get_stage_for_issue @@ -362,7 +465,6 @@ def optimize( if not stages_to_apply: result.success, result.optimized_code = True, kernel_code candidates.append(result) - continue current_code = kernel_code current_ms: float | None = val_orig_ms @@ -535,25 +637,37 @@ def optimize( result.optimized_code, result.success = current_code, True candidates.append(result) - if not candidates: - return OptimizationResult( - kernel_name=display_name, original_code=kernel_code, timestamp=datetime.now() - ) + if not candidates: + result_out = OptimizationResult( + kernel_name=display_name, original_code=kernel_code, timestamp=datetime.now() + ) + return result_out - result = max( - candidates, key=lambda r: r.total_speedup if r.total_speedup is not None else -1.0 - ) - self._save_results(result) - - logger.info("=" * 60 + "\nOPTIMIZATION COMPLETE\n" + "=" * 60) - ok = [s for s in result.stages_applied if s.success] - fail = [s for s in result.stages_applied if not s.success] - logger.info(f"Stages: {len(ok)}/{len(result.stages_applied)} succeeded") - if fail: - logger.info(f"Failed: {[s.stage.value for s in fail]}") - if result.total_speedup: - logger.info(f"Speedup: {result.total_speedup:.2f}x") - return result + result = max( + candidates, key=lambda r: r.total_speedup if r.total_speedup is not None else -1.0 + ) + self._save_results(result) + + logger.info("=" * 60 + "\nOPTIMIZATION COMPLETE\n" + "=" * 60) + ok = [s for s in result.stages_applied if s.success] + fail = [s for s in result.stages_applied if not s.success] + logger.info(f"Stages: {len(ok)}/{len(result.stages_applied)} succeeded") + if fail: + logger.info(f"Failed: {[s.stage.value for s in fail]}") + if result.total_speedup: + logger.info(f"Speedup: {result.total_speedup:.2f}x") + + result_out = result + return result_out + finally: + import sys + + _usage_ctx.__exit__(*sys.exc_info()) + if result_out is not None: + token_usage = _usage_tracker.get_total_tokens() + if token_usage: + logger.info("Token usage this run: %s", token_usage) + result_out.token_usage = token_usage or {} def optimize_file( self, @@ -597,3 +711,32 @@ def _save_results(self, result): f"{comment} Stages: {[s.stage.value for s in result.stages_applied if s.success]}\n\n" ) f.write(result.optimized_code) + + +# --------------------------------------------------------------------------- +# Pipeline helpers +# --------------------------------------------------------------------------- + + +def _build_kernel_specs( + kernel_name: str, + input_shapes, + flop, + dtype, + device: str, + dsl: str, +) -> str: + """Format a compact kernel specification string for the CoordinatorAgent.""" + parts = [f"kernel: {kernel_name}", f"device: {device}", f"dsl: {dsl}"] + if input_shapes: + parts.append(f"input_shapes: {input_shapes}") + if dtype: + parts.append(f"dtype: {dtype}") + if flop: + if flop > 1e12: + parts.append(f"flop: {flop / 1e12:.2f} TFLOP") + elif flop > 1e9: + parts.append(f"flop: {flop / 1e9:.2f} GFLOP") + else: + parts.append(f"flop: {flop:.0f}") + return ", ".join(parts) diff --git a/src/xe_forge/planner.py b/src/xe_forge/planner.py index ba17139..514b863 100644 --- a/src/xe_forge/planner.py +++ b/src/xe_forge/planner.py @@ -58,52 +58,7 @@ class PlanningSignature(dspy.Signature): - """Determine the optimal order to apply optimization stages to a GPU kernel. - - You are an expert in GPU kernel optimization. - You have analyzed a kernel and found issues in specific optimization categories. - Your task: decide the OPTIMAL ORDER to apply the available stages. - - === ORDERING PRINCIPLES === - - Mathematical correctness first: - - ALGORITHMIC and DISCOVERY before everything else — structural rewrites may - eliminate entire categories of lower-level issues. A rewrite that removes - a kernel makes BLOCK_POINTERS for that kernel irrelevant. - - Dtype early, but after structure: - - DTYPE_FIX after ALGORITHMIC/DISCOVERY — convert to fp16 after the algorithm - is settled, not before (avoids converting code that gets rewritten anyway). - - Fusion after dtype: - - FUSION after DTYPE_FIX — fuse already-converted kernels. Fusing then - converting can produce suboptimal mixed-precision boundaries. - - Exception: if FUSION is the dominant issue and DTYPE is minor, move FUSION - earlier to reduce the number of kernels before low-level tuning. - - Memory before block pointers: - - MEMORY_ACCESS before BLOCK_POINTERS — fix coalescing and layout issues - before converting to block pointer API (block pointers assume good layout). - - Low-level tuning last: - - DEVICE_SPECIFIC and AUTOTUNING last — tile sizes, warp counts, and autotune - configs should be applied to the final kernel structure, not intermediate forms. - - Persistent kernel placement: - - PERSISTENT_KERNEL before DEVICE_SPECIFIC — persistence changes the kernel - structure; device tuning should happen on the persistent form. - - Skip PERSISTENT_KERNEL entirely if grid size is small (< 512 tiles) or - if the kernel produces a scalar/[M] output — persistence won't help. - - === RULES === - - Only include stages from the available_stages list (stages with detected issues). - - Do NOT add stages that have no detected issues. - - Do NOT include ANALYSIS. - - Your ordered_stages output must be a JSON array of stage value strings, - e.g. ["algorithmic", "dtype_fix", "device_specific"]. - - Every stage in ordered_stages must appear in available_stages. - - Provide clear rationale explaining why you chose this specific order. - """ + """Determine the optimal order to apply optimization stages to a GPU kernel.""" available_stages: str = dspy.InputField( desc="JSON object mapping stage_name → [issue_type, ...] for stages with detected issues." @@ -132,8 +87,24 @@ class PlanningSignature(dspy.Signature): class PlannerAgent: """LLM-based stage ordering. Falls back to default order on any failure.""" - def __init__(self) -> None: - self.predictor = dspy.Predict(PlanningSignature) + def __init__(self, extra_instructions: str = "") -> None: + sig = PlanningSignature + # Inject planning guidance from template + try: + from xe_forge.config import get_config + from xe_forge.prompts import render_signature_instructions + + cfg = get_config() + template_text = render_signature_instructions( + "planning_signature", + device_description=cfg.device_config.device, + ) + sig = sig.append_instructions(template_text) + except Exception as e: + logger.debug("Planning template render failed: %s", e) + if extra_instructions: + sig = sig.append_instructions(extra_instructions) + self.predictor = dspy.Predict(sig) def plan( self, diff --git a/src/xe_forge/prompts/__init__.py b/src/xe_forge/prompts/__init__.py index 1aaabe9..3d53f8d 100644 --- a/src/xe_forge/prompts/__init__.py +++ b/src/xe_forge/prompts/__init__.py @@ -1,3 +1,20 @@ +from pathlib import Path + +from jinja2 import Environment, FileSystemLoader, select_autoescape + from xe_forge.prompts.device_prompts import PromptLibrary -__all__ = ["PromptLibrary"] +_env = Environment( + loader=FileSystemLoader(Path(__file__).parent / "templates"), + trim_blocks=True, + lstrip_blocks=True, + autoescape=select_autoescape(enabled_extensions=()), +) + + +def render_signature_instructions(template_name: str, **context) -> str: + """Render a Signature instruction block from a Jinja2 template.""" + return _env.get_template(f"{template_name}.md.j2").render(**context) + + +__all__ = ["PromptLibrary", "render_signature_instructions"] diff --git a/src/xe_forge/prompts/device_prompts.py b/src/xe_forge/prompts/device_prompts.py index 9bdb14a..e2f0b85 100644 --- a/src/xe_forge/prompts/device_prompts.py +++ b/src/xe_forge/prompts/device_prompts.py @@ -194,3 +194,54 @@ def planner_signature_doc(self) -> str: f"You are an expert in GPU kernel optimization for {self.device_description()}.\n" "Determine the optimal order to apply optimization stages." ) + + def render_for_signature(self, sig_name: str, **extra_context) -> str: + """Render the Jinja2 template for a given Signature class name. + + Returns the rendered instruction text to be injected via append_instructions(). + Falls back to device_context_addendum() if the template is not found. + """ + from xe_forge.prompts import render_signature_instructions + + context = { + "dsl": self.dsl, + "dsl_name": self.dsl_name(), + "device_type": self.device_type, + "device_description": self.device_description(), + "defaults": self.tuning_defaults(), + **extra_context, + } + try: + return render_signature_instructions(sig_name, **context) + except Exception: + return self.device_context_addendum() + + def device_context_addendum(self) -> str: + """Return a short addendum to append to any Signature for the current device/DSL. + + Returns an empty string for the default XPU/Triton target (already hardcoded in + Signature docstrings). For other device targets, returns a correction block so the + LLM uses the right tile sizes and hardware constraints. + """ + defaults = self.tuning_defaults() + + if self.device_type == "xpu": + return "" + + lines = [ + f"=== TARGET DEVICE OVERRIDE: {self.device_description()} ===", + f"DSL: {self.dsl_name()}", + "Use the following hardware-specific defaults (override any XPU values above):", + ] + for k, v in defaults.items(): + lines.append(f" {k}: {v}") + + if self.device_type == "cuda" and self.dsl == "triton": + lines.append("Use shared memory for reductions. Consider warp-level primitives.") + lines.append( + "Do NOT use tl.extra.intel.libdevice — use standard math intrinsics instead." + ) + elif self.dsl == "sycl": + lines.append("Use SYCL/CUTLASS C++ with icpx -fsycl compilation.") + + return "\n".join(lines) diff --git a/src/xe_forge/prompts/templates/algorithmic_signature.md.j2 b/src/xe_forge/prompts/templates/algorithmic_signature.md.j2 new file mode 100644 index 0000000..61eb9a6 --- /dev/null +++ b/src/xe_forge/prompts/templates/algorithmic_signature.md.j2 @@ -0,0 +1,24 @@ +You are an expert in numerical linear algebra, compiler optimizations, and +high-performance {{ dsl_name }} kernel design for {{ device_description }}. + +Transform the kernel to perform FEWER FLOPs and/or FEWER memory accesses +while producing numerically equivalent results. + +Think about: +1. Matrix structure exploitation (symmetric, triangular, diagonal, low-rank, sparse) +2. Associative / distributive law rewrites to reduce FLOPs +3. Common sub-expression elimination +4. Loop-invariant code hoisting +5. Caching intermediates in registers vs recomputing +6. Tree reductions vs serial reductions +7. Algebraic simplification of fused computations + +Maintain the Model class signature. Produce equivalent outputs. + +=== CODE REQUIREMENTS === +- Include ALL imports, @triton.jit decorator, kernel function, Model class +- NEVER replace @triton.jit kernels with torch.matmul, torch.mm, or any vendor library. +{% if device_type == "xpu" %} +- Use tl.extra.intel.libdevice for exp2/sigmoid operations +- Avoid .cpu() returns — output must stay on XPU +{% endif %} diff --git a/src/xe_forge/prompts/templates/analysis_signature.md.j2 b/src/xe_forge/prompts/templates/analysis_signature.md.j2 new file mode 100644 index 0000000..4c42d6a --- /dev/null +++ b/src/xe_forge/prompts/templates/analysis_signature.md.j2 @@ -0,0 +1,37 @@ +You are a world-class expert in {{ dsl_name }} GPU/XPU kernel optimization, +numerical linear algebra, and high-performance computing. + +Analyze the given {{ dsl_name }} kernel code and, if available, the original PyTorch +implementation for higher-level algorithmic context. + +You must identify ALL applicable optimizations across every category below. +Use your deep knowledge of GPU programming, {{ dsl_name }} internals, {{ device_description }} +architecture, and mathematical optimization. + +{{ issue_categories_block }} + +IMPORTANT: +- Return issues as a JSON array of DetectedIssue objects. +- Each issue MUST have: issue_type (exact string from the list above), + severity (1-5), description, suggested_fix, estimated_speedup. +- issue_type MUST be one of the exact strings listed above (e.g. "dtype_float64", + "missing_grf_mode"). Do NOT invent new type names. +- For fused kernels, pay special attention to ALGORITHMIC issues. +- Return empty array [] ONLY if the kernel is already optimal. + +OPEN-ENDED DISCOVERY (issue_type="open_ended"): +After checking all categories above, ask yourself: is there a high-value +optimization that does not fit any existing type? If yes, use issue_type="open_ended" +and populate open_ended_proposal with the full proposal. Requirements: + - Concrete and implementable — not a vague observation + - Mathematically or logically justified + - Includes a before/after code sketch in open_ended_proposal + - Includes estimated speedup with reasoning +Examples that qualify as open_ended: + * sum(x @ W.T, dim=1) rewritten as x @ W.sum(dim=0) — eliminates O(M*N*K) GEMM + * Weight statistic (colsum, norm) recomputed every forward() — cache in __init__ + * Two-kernel pipeline where the HBM intermediate can be eliminated algebraically +Examples that do NOT qualify (use the named type instead): + * "use better tile sizes" → use suboptimal_tile_size + * "add autotuning" → use missing_autotune + * "fuse these kernels" → use unfused_kernels diff --git a/src/xe_forge/prompts/templates/autotune_signature.md.j2 b/src/xe_forge/prompts/templates/autotune_signature.md.j2 new file mode 100644 index 0000000..76f09d5 --- /dev/null +++ b/src/xe_forge/prompts/templates/autotune_signature.md.j2 @@ -0,0 +1,46 @@ +You are an expert in {{ dsl_name }} kernel autotuning for {{ device_description }}. + +Your task: Add or improve the @triton.autotune decorator so the kernel +automatically selects the best configuration at runtime. + +You will receive: +- The current kernel code +- Hardware information (compute units, memory, capabilities) +- Problem shapes (M, N, K dimensions) +- A set of suggested autotune configurations generated from hardware analysis + +Your job: +1. Add @triton.autotune decorator with a good set of configs to search. +2. Use the suggested configs as a starting point but ADD more configs + based on your knowledge of what works well for this kernel type. +3. Include the key= argument so configs are re-evaluated when shapes change. +4. Ensure num_warps and num_stages are included in each config. +5. Ensure BLOCK sizes are powers of 2 and appropriate for the hardware. +{% if device_type == "xpu" %} +6. For Intel XPU, always include at least one config with num_warps={{ defaults.num_warps }} + and large tile sizes ({{ defaults.BLOCK_M }}x{{ defaults.BLOCK_N }}). +7. Do NOT put grf_mode in triton.Config() — it causes TypeError at runtime. + grf_mode is a compiler option: declare it as tl.constexpr in the kernel + signature. Use grf_mode="auto" (auto-selects 256-GRF if spill > 1000 bytes) + or grf_mode="256" for large register file. Requires num_warps <= 32. +{% elif device_type == "cuda" %} +6. For NVIDIA CUDA, include configs with num_warps in [4, 8, 16] and + tile sizes from 64x64 to {{ defaults.BLOCK_M }}x{{ defaults.BLOCK_N }}. +{% endif %} +8. Remove any hardcoded meta-parameters that are now covered by autotune. +9. Keep the kernel functionally equivalent. + +Tips for good autotune configs: +- Vary BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K across powers of 2 +- Include both small tiles (64x64) for small problems and large tiles + ({{ defaults.BLOCK_M }}x{{ defaults.BLOCK_N }}) for large problems +- Vary num_warps: try 4, 8, 16{% if device_type == "xpu" %}, 32{% endif %} +- Vary num_stages: try 2, 3, 4 +- Include GROUP_SIZE_M for L2 cache swizzling +- Use key= with the shape arguments that affect tiling + +=== CODE REQUIREMENTS === +- Include ALL imports (torch, triton, triton.language as tl) +- Include @triton.autotune with configs list and key +- Include @triton.jit on the kernel +- Include the Model class with forward() method diff --git a/src/xe_forge/prompts/templates/coordinator_signature.md.j2 b/src/xe_forge/prompts/templates/coordinator_signature.md.j2 new file mode 100644 index 0000000..6774206 --- /dev/null +++ b/src/xe_forge/prompts/templates/coordinator_signature.md.j2 @@ -0,0 +1,44 @@ +You are the lead optimization engineer for {{ dsl_name }} kernels on {{ device_description }}, +coordinating a team of specialist tools to maximize kernel performance. + +Your job: decide WHAT to optimize, WHEN, and IN WHAT ORDER. + +=== STRATEGY === +1. Always start with analyze_kernel() to understand what needs fixing before writing any code. +2. For each promising stage, call retrieve_patterns(stage) first to get relevant techniques, + then apply_stage(stage) to try them. +3. Call analyze_kernel() again after significant structural changes to get a fresh view. +4. Call profile_kernel() when you suspect hardware bottlenecks (cache misses, low XVE + utilization) that static analysis cannot reveal. +5. Call benchmark_current() to verify cumulative speedup at any point. +6. Call get_status() to review what you've tried and what remains. +7. Stop when: speedup is satisfactory, no more promising stages remain, + or the kernel appears hardware-bound with no further leverage. + +=== STAGE ORDERING GUIDANCE === +Prefer this rough order unless analysis suggests otherwise: + ALGORITHMIC/DISCOVERY → DTYPE_FIX → FUSION → MEMORY_ACCESS → + BLOCK_POINTERS → PERSISTENT_KERNEL → DEVICE_SPECIFIC → AUTOTUNING + +Rationale: structural rewrites first, then memory layout, then hardware tuning. +You can deviate based on severity — if DEVICE_SPECIFIC issues are critical (severity 5) +and ALGORITHMIC issues are minor (severity 2), start with DEVICE_SPECIFIC. + +{% if device_type == "xpu" %} +=== XPU DEFAULTS === +When apply_stage succeeds, the new code automatically uses: + BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }}, num_warps={{ defaults.num_warps }} + grf_mode="auto" (as tl.constexpr), tl.extra.intel.libdevice for math ops +{% elif device_type == "cuda" %} +=== CUDA DEFAULTS === +When apply_stage succeeds, the new code automatically uses: + BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }}, num_warps={{ defaults.num_warps }} +{% endif %} + +=== IMPORTANT RULES === +- Always analyze before applying any stage. +- Prefer high-severity issues first. +- If a stage fails, move on — don't retry the same stage more than once unless + retrieve_patterns revealed new techniques you haven't tried yet. +- The submit tool will record the best code found automatically. +- Never pass raw code through tool arguments — tools operate on internal state. diff --git a/src/xe_forge/prompts/templates/optimization_react_signature.md.j2 b/src/xe_forge/prompts/templates/optimization_react_signature.md.j2 new file mode 100644 index 0000000..09e6299 --- /dev/null +++ b/src/xe_forge/prompts/templates/optimization_react_signature.md.j2 @@ -0,0 +1,38 @@ +You are an expert {{ dsl_name }} kernel optimizer for {{ device_description }}. +Your task: Optimize the kernel for maximum performance. +You may change the algorithm/computation approach if it produces equivalent outputs. +Maintain the same model signature, including the weights' shapes and names. + +=== OPTIMIZATION PRIORITIES === +1. Apply the specific optimization patterns from knowledge_patterns +2. Use block pointers with tl.make_block_ptr() for better memory access +{% if device_type == "xpu" %} +3. Use optimal tile sizes: BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }} for XPU +4. Set num_warps={{ defaults.num_warps }} for Intel XPU +5. Add GROUP_SIZE_M swizzling for better L2 cache utilization +6. Use boundary_check=(0, 1) tuple format, NOT booleans +{% elif device_type == "cuda" %} +3. Use optimal tile sizes: BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }} for CUDA +4. Set num_warps={{ defaults.num_warps }} for NVIDIA GPU +5. Add GROUP_SIZE_M swizzling for better L2 cache utilization +6. Use boundary_check=(0, 1) tuple format, NOT booleans +{% else %} +3. Use tile sizes: BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }} +4. Set num_warps={{ defaults.num_warps }} +{% endif %} + +=== CODE REQUIREMENTS === +- Include ALL imports (torch, triton, triton.language as tl) +- Include the @triton.jit decorator and kernel function +- Include the Model class with forward() method +- num_warps MUST be a power of 2 (1, 2, 4, 8, 16, 32) +- num_stages MUST be a positive integer +- Block sizes (BLOCK_M, BLOCK_N, BLOCK_K) MUST be powers of 2 +{% if device_type == "xpu" %} +- Use tl.extra.intel.libdevice for exp2/sigmoid operations +- Avoid .cpu() returns — output must stay on XPU +{% endif %} + +=== WHAT TO CHANGE === +Focus on the issues listed and apply the patterns from knowledge_patterns. +Be aggressive with optimizations — the verification tool will check correctness. diff --git a/src/xe_forge/prompts/templates/optimization_signature.md.j2 b/src/xe_forge/prompts/templates/optimization_signature.md.j2 new file mode 100644 index 0000000..fa409ee --- /dev/null +++ b/src/xe_forge/prompts/templates/optimization_signature.md.j2 @@ -0,0 +1,45 @@ +You are an expert {{ dsl_name }} kernel optimizer for {{ device_description }}. +Optimize the kernel for maximum performance while producing numerically equivalent outputs. +You may change the algorithm/computation approach if outputs are equivalent. +Maintain the same Model class signature including weights shapes and names. + +=== STAGE-SPECIFIC GUIDANCE === +ALGORITHMIC: mathematical simplifications, CSE, loop-invariant hoisting, + caching intermediates, reorder associative ops, tree reductions, + exploit GEMM structure (symmetric, triangular, low-rank). +DTYPE_FIX: float64->float32, proper accumulator precision, remove + unnecessary type conversions. +FUSION: fuse kernel launches, elementwise chains, reduction+elementwise. +MEMORY_ACCESS: fix uncoalesced access, remove transposes from inner loops, + add boundary checks, reduce register pressure. +BLOCK_POINTERS: use tl.make_block_ptr(), boundary_check=(0,1) tuple format, + tl.advance() for pointer updates. +{% if device_type == "xpu" %} +XPU_SPECIFIC: BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }}, num_warps={{ defaults.num_warps }}, + GROUP_SIZE_M swizzling. + GRF MODE: grf_mode is a compiler option, NOT a triton.Config() kwarg. + Declare it as tl.constexpr in the kernel signature: + grf_mode: tl.constexpr (values: "default", "128", "256", "auto") + Use "auto" — it automatically selects 256-GRF when register spill > 1000 bytes. + 256-GRF requires num_warps <= 32 (halved thread occupancy). +{% elif device_type == "cuda" %} +DEVICE_SPECIFIC (CUDA): BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }}, num_warps={{ defaults.num_warps }}. + Use shared memory for reductions. Consider warp-level primitives. + Do NOT use tl.extra.intel.libdevice. +{% else %} +DEVICE_SPECIFIC: BLOCK_M={{ defaults.BLOCK_M }}, BLOCK_N={{ defaults.BLOCK_N }}, BLOCK_K={{ defaults.BLOCK_K }}, num_warps={{ defaults.num_warps }}. +{% endif %} +PERSISTENT_KERNEL: persistent kernel pattern, tune NUM_PROGS. +DISCOVERY: apply the open-ended optimization described in the issues field. + This is a novel optimization not covered by standard stages. Follow the + proposal exactly, preserving all numerical equivalences. + +=== CODE REQUIREMENTS === +- Include ALL imports, @triton.jit decorator, kernel function, Model class +- num_warps must be power of 2; block sizes must be powers of 2 +- NEVER replace @triton.jit kernels with torch.matmul, torch.mm, torch.bmm, + or any vendor library (oneDNN, cuBLAS, MKL). Keep all original Triton kernels. +{% if device_type == "xpu" %} +- Use tl.extra.intel.libdevice for exp2/sigmoid operations +- Avoid .cpu() returns — output must stay on XPU +{% endif %} diff --git a/src/xe_forge/prompts/templates/planning_signature.md.j2 b/src/xe_forge/prompts/templates/planning_signature.md.j2 new file mode 100644 index 0000000..e56cd69 --- /dev/null +++ b/src/xe_forge/prompts/templates/planning_signature.md.j2 @@ -0,0 +1,43 @@ +You are an expert in GPU kernel optimization for {{ device_description }}. +You have analyzed a kernel and found issues in specific optimization categories. +Your task: decide the OPTIMAL ORDER to apply the available stages. + +=== ORDERING PRINCIPLES === + +Mathematical correctness first: + - ALGORITHMIC and DISCOVERY before everything else — structural rewrites may + eliminate entire categories of lower-level issues. A rewrite that removes + a kernel makes BLOCK_POINTERS for that kernel irrelevant. + +Dtype early, but after structure: + - DTYPE_FIX after ALGORITHMIC/DISCOVERY — convert to fp16 after the algorithm + is settled, not before (avoids converting code that gets rewritten anyway). + +Fusion after dtype: + - FUSION after DTYPE_FIX — fuse already-converted kernels. Fusing then + converting can produce suboptimal mixed-precision boundaries. + - Exception: if FUSION is the dominant issue and DTYPE is minor, move FUSION + earlier to reduce the number of kernels before low-level tuning. + +Memory before block pointers: + - MEMORY_ACCESS before BLOCK_POINTERS — fix coalescing and layout issues + before converting to block pointer API (block pointers assume good layout). + +Low-level tuning last: + - DEVICE_SPECIFIC and AUTOTUNING last — tile sizes, warp counts, and autotune + configs should be applied to the final kernel structure, not intermediate forms. + +Persistent kernel placement: + - PERSISTENT_KERNEL before DEVICE_SPECIFIC — persistence changes the kernel + structure; device tuning should happen on the persistent form. + - Skip PERSISTENT_KERNEL entirely if grid size is small (< 512 tiles) or + if the kernel produces a scalar/[M] output — persistence won't help. + +=== RULES === +- Only include stages from the available_stages list (stages with detected issues). +- Do NOT add stages that have no detected issues. +- Do NOT include ANALYSIS. +- Your ordered_stages output must be a JSON array of stage value strings, + e.g. ["algorithmic", "dtype_fix", "device_specific"]. +- Every stage in ordered_stages must appear in available_stages. +- Provide clear rationale explaining why you chose this specific order. diff --git a/src/xe_forge/prompts/templates/sycl_algorithmic_signature.md.j2 b/src/xe_forge/prompts/templates/sycl_algorithmic_signature.md.j2 new file mode 100644 index 0000000..e3fa8b6 --- /dev/null +++ b/src/xe_forge/prompts/templates/sycl_algorithmic_signature.md.j2 @@ -0,0 +1,17 @@ +You are an expert in numerical linear algebra, compiler optimizations, and +high-performance GPU kernel design for Intel XPU. + +Transform the SYCL/CUTLASS kernel to perform FEWER FLOPs and/or FEWER memory accesses +while producing numerically equivalent results. + +Think about: +1. Matrix structure exploitation (symmetric, triangular, diagonal, low-rank) +2. Associative / distributive law rewrites to reduce FLOPs +3. Common sub-expression elimination in template expressions +4. Data layout optimization (RowMajor vs ColumnMajor) +5. Batch dimension exploitation + +=== CODE REQUIREMENTS === +- Must be complete, valid SYCL C++ with all #include directives +- Keep CUTLASS GEMM structure (GemmUniversalAdapter, ExampleRunner, main) +- Must compile with icpx -fsycl diff --git a/src/xe_forge/prompts/templates/sycl_optimization_signature.md.j2 b/src/xe_forge/prompts/templates/sycl_optimization_signature.md.j2 new file mode 100644 index 0000000..b7a3266 --- /dev/null +++ b/src/xe_forge/prompts/templates/sycl_optimization_signature.md.j2 @@ -0,0 +1,34 @@ +You are an expert in SYCL, CUTLASS/XeTLA, Intel XPU GPU architecture, +and high-performance C++ kernel optimization. + +Optimize the kernel for maximum performance while producing numerically +equivalent outputs. You may change template parameters, dispatch policies, +data types, and memory layouts. + +=== SYCL/CUTLASS OPTIMIZATION KNOBS === +- TileShape: Shape<_M, _N, _K> — try 256x256x32, 128x128x64, 128x256x32 +- PipelineStages: 2, 3, or 4 — more prefetching vs register pressure +- MMA Atom: XE_DPAS_TT — SubgroupSize 4 or 8 +- Dispatch Policy: MainloopXeL1Staged (L1 cached), MainloopXeL0Staged (uncached) +- Data types: bfloat16_t/half_t inputs, float/bfloat16_t accumulators +- Memory layout: RowMajor vs ColumnMajor for A, B, C, D +- Epilogue: LinearCombination, bias, activation via FusionCallbacks +- GmemTiledCopy: void (auto) or explicit copy atoms + +=== STAGE-SPECIFIC GUIDANCE === +ALGORITHMIC: mathematical simplifications, CSE, loop-invariant hoisting, + exploit GEMM structure (symmetric, triangular, low-rank). +DTYPE_FIX: use bfloat16_t/half_t inputs, float accumulators, avoid double. +FUSION: fuse into CUTLASS epilogue callbacks — LinearCombination, bias, activation. +MEMORY_ACCESS: fix layout mismatch (RowMajor vs ColumnMajor), increase PipelineStages + for better prefetching, reduce register pressure. +DEVICE_SPECIFIC: TileShape 256x256x32 or 128x128x64, PipelineStages=2-3, + XE_DPAS_TT<8, float, bfloat16_t>, MainloopXeL1Staged dispatch policy. +DISCOVERY: apply the open-ended optimization described in the issues field. + +=== CODE REQUIREMENTS === +- Must be complete, valid SYCL C++ with all #include directives +- Must use cutlass namespace and CUTLASS template types +- Must include ExampleRunner template and main() function +- Must compile with icpx -fsycl +- Keep the same output format (Cutlass GEMM Performance line) diff --git a/tests/test_pipeline_usage_tracking.py b/tests/test_pipeline_usage_tracking.py new file mode 100644 index 0000000..9a8a1f9 --- /dev/null +++ b/tests/test_pipeline_usage_tracking.py @@ -0,0 +1,129 @@ +from types import SimpleNamespace + +import pytest + +from xe_forge.models import DSL, KernelAnalysis +from xe_forge.pipeline import XeForgePipeline + + +class DummyTracker: + def __init__(self, tokens: dict): + self._tokens = tokens + + def get_total_tokens(self) -> dict: + return self._tokens + + +class DummyUsageContext: + def __init__(self, tokens: dict): + self.tracker = DummyTracker(tokens) + self.exit_calls = 0 + self.exc_type = None + + def __enter__(self): + return self.tracker + + def __exit__(self, exc_type, exc, tb): + self.exit_calls += 1 + self.exc_type = exc_type + return False + + +class StubAnalyzer: + def __init__(self, *, raises: Exception | None = None): + self._raises = raises + + def analyze(self, *args, **kwargs): + if self._raises is not None: + raise self._raises + return KernelAnalysis(kernel_name="k", detected_issues=[]) + + +class StubPlanner: + def plan(self, **kwargs): + return [] + + +class StubCoordinator: + def run(self, **kwargs): + return kwargs["kernel_code"], 1.25, "ok", [] + + +class StubOptimizer: + def optimize_stage(self, **kwargs): + raise AssertionError("optimize_stage should not be called in this test") + + +def _make_pipeline(): + pipeline = object.__new__(XeForgePipeline) + pipeline.config = SimpleNamespace( + optimization=SimpleNamespace(target_dtype=None, best_k=1), + device_config=SimpleNamespace(device="xpu", dsl=DSL.TRITON), + logging=SimpleNamespace(save_intermediate=False), + ) + pipeline.executor = None + pipeline.trial_manager = None + pipeline.profiler = None + pipeline.optimizer = StubOptimizer() + pipeline.analyzer = StubAnalyzer() + pipeline.planner = StubPlanner() + pipeline.coordinator = None + pipeline._save_results = lambda result: None + pipeline._resolve_tolerances = lambda *args, **kwargs: (1e-3, 1e-3) + return pipeline + + +def test_optimize_sets_token_usage_and_exits_tracker(monkeypatch): + usage_ctx = DummyUsageContext({"prompt_tokens": 10, "completion_tokens": 2}) + monkeypatch.setattr("xe_forge.pipeline.dspy.track_usage", lambda: usage_ctx) + monkeypatch.setattr("xe_forge.pipeline.get_device_config_for_pipeline", lambda **kwargs: {}) + + pipeline = _make_pipeline() + + result = pipeline.optimize( + kernel_code="def kernel():\n return 1\n", + reference_code="", + kernel_name="k", + ) + + assert result.token_usage == {"prompt_tokens": 10, "completion_tokens": 2} + assert usage_ctx.exit_calls == 1 + assert usage_ctx.exc_type is None + + +def test_optimize_exits_tracker_when_exception_raised(monkeypatch): + usage_ctx = DummyUsageContext({"prompt_tokens": 0, "completion_tokens": 0}) + monkeypatch.setattr("xe_forge.pipeline.dspy.track_usage", lambda: usage_ctx) + monkeypatch.setattr("xe_forge.pipeline.get_device_config_for_pipeline", lambda **kwargs: {}) + + pipeline = _make_pipeline() + pipeline.analyzer = StubAnalyzer(raises=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + pipeline.optimize( + kernel_code="def kernel():\n return 1\n", + reference_code="", + kernel_name="k", + ) + + assert usage_ctx.exit_calls == 1 + assert usage_ctx.exc_type is RuntimeError + + +def test_optimize_coordinator_path_also_exits_tracker(monkeypatch): + usage_ctx = DummyUsageContext({"prompt_tokens": 7, "completion_tokens": 3}) + monkeypatch.setattr("xe_forge.pipeline.dspy.track_usage", lambda: usage_ctx) + monkeypatch.setattr("xe_forge.pipeline.get_device_config_for_pipeline", lambda **kwargs: {}) + + pipeline = _make_pipeline() + pipeline.coordinator = StubCoordinator() + + result = pipeline.optimize( + kernel_code="def kernel():\n return 1\n", + reference_code="", + kernel_name="k", + ) + + assert result.token_usage == {"prompt_tokens": 7, "completion_tokens": 3} + assert usage_ctx.exit_calls == 1 + assert usage_ctx.exc_type is None