wafer-core 0.1.31__py3-none-any.whl → 0.1.32__py3-none-any.whl

wafer_core/tools/dispatch_baseline/dtypes.py

@@ -0,0 +1,217 @@
+"""Data types for kernel trace tool."""
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class TensorSpec:
+    """Specification for a tensor operand."""
+
+    name: str
+    shape: tuple[int, ...]
+    dtype: str = "float16"  # e.g., "float16", "float32", "bfloat16", "int8"
+    device: str = "cuda"
+
+    def __str__(self) -> str:
+        shape_str = "x".join(str(d) for d in self.shape)
+        return f"{self.name}[{shape_str}] ({self.dtype})"
+
+
+@dataclass(frozen=True)
+class OpSpec:
+    """Specification for a PyTorch operation to trace.
+
+    Example:
+        OpSpec(
+            op="torch.matmul",
+            inputs=[
+                TensorSpec("A", (4096, 4096), "float16"),
+                TensorSpec("B", (4096, 4096), "float16"),
+            ],
+        )
+    """
+
+    op: str  # e.g., "torch.matmul", "torch.nn.functional.softmax", "torch.conv2d"
+    inputs: list[TensorSpec]
+    kwargs: dict[str, str] = field(default_factory=dict)  # e.g., {"dim": "-1"}
+
+    def __str__(self) -> str:
+        inputs_str = ", ".join(t.name for t in self.inputs)
+        kwargs_str = ", ".join(f"{k}={v}" for k, v in self.kwargs.items())
+        if kwargs_str:
+            return f"{self.op}({inputs_str}, {kwargs_str})"
+        return f"{self.op}({inputs_str})"
+
+
+@dataclass(frozen=True)
+class KernelInfo:
+    """Information about a dispatched kernel."""
+
+    name: str  # e.g., "sm90_xmma_gemm_f16f16f16_f32..."
+    duration_us: float = 0.0  # Duration in microseconds
+    grid_size: tuple[int, int, int] | None = None  # (X, Y, Z)
+    block_size: tuple[int, int, int] | None = None  # (X, Y, Z)
+    registers_per_thread: int | None = None
+    shared_memory_bytes: int | None = None
+    # Performance metrics (from NCU, if available)
+    compute_throughput_tflops: float | None = None
+    memory_throughput_tbps: float | None = None  # TB/s
+    achieved_occupancy_pct: float | None = None
+
+    def __str__(self) -> str:
+        return f"{self.name} {self.duration_us:.1f} µs"
+
+
+@dataclass(frozen=True)
+class HardwareSpec:
+    """Hardware specifications for roofline analysis."""
+
+    name: str  # e.g., "H100", "MI300X"
+    peak_fp16_tflops: float
+    peak_fp32_tflops: float
+    peak_memory_bw_tbps: float  # TB/s
+    # Optional extras
+    peak_fp8_tflops: float | None = None
+    peak_int8_tops: float | None = None
+    shared_memory_per_sm_kb: float | None = None
+
+
+@dataclass(frozen=True)
+class RooflineAnalysis:
+    """Roofline analysis result."""
+
+    # Raw metrics
+    achieved_tflops: float
+    achieved_memory_bw_tbps: float
+    # Percentages of peak
+    compute_pct_of_peak: float
+    memory_bw_pct_of_peak: float
+    # Bottleneck identification
+    bottleneck: Literal["compute", "memory", "balanced"]
+    arithmetic_intensity: float  # FLOPS per byte
+
+
+@dataclass(frozen=True)
+class KernelTraceConfig:
+    """Configuration for kernel tracing."""
+
+    op_spec: OpSpec
+    hardware: str  # e.g., "H100", "MI300X"
+    num_warmup: int = 10
+    num_runs: int = 100
+    use_ncu: bool = False  # Use NCU for detailed metrics (slower)
+    timeout_seconds: int = 120
+
+
+@dataclass(frozen=True)
+class KernelTraceResult:
+    """Result of tracing a PyTorch operation."""
+
+    op_spec: OpSpec
+    hardware: str
+    # Kernels dispatched (may be multiple for fused ops)
+    kernels: list[KernelInfo]
+    # Primary kernel (typically the longest-running one)
+    primary_kernel: KernelInfo | None
+    # Roofline analysis (if metrics available)
+    roofline: RooflineAnalysis | None = None
+    # Raw profiler output (for debugging)
+    raw_output: str | None = None
+    # Error, if any
+    error: str | None = None
+
+    @property
+    def total_duration_us(self) -> float:
+        """Total duration of all kernels."""
+        return sum(k.duration_us for k in self.kernels)
+
+    def summary(self) -> str:
+        """Generate human-readable summary."""
+        if self.error:
+            return f"Error: {self.error}"
+
+        # Build operation description with shapes
+        op_desc = self._format_op_description()
+
+        lines = [
+            "═" * 65,
+            f" BASELINE: {op_desc}",
+            "═" * 65,
+            "",
+        ]
+
+        # Primary kernel section
+        if self.primary_kernel:
+            lines.append(" Primary Kernel:")
+            kernel_name = self._truncate_kernel_name(self.primary_kernel.name, 55)
+            lines.append(f"   → {kernel_name}")
+            lines.append(f"     Duration: {self._format_duration(self.primary_kernel.duration_us)}")
+            lines.append("")
+
+        # Roofline section
+        if self.roofline:
+            lines.append(" Roofline Analysis:")
+            
+            # Format TFLOPS with commas for readability
+            tflops_str = f"{self.roofline.achieved_tflops:,.1f}"
+            bw_str = f"{self.roofline.achieved_memory_bw_tbps:.2f}"
+            
+            lines.append(f"   Achieved:    {tflops_str} TFLOPS  ({self.roofline.compute_pct_of_peak:.1f}% of {self.hardware} peak)")
+            lines.append(f"   Memory BW:   {bw_str} TB/s     ({self.roofline.memory_bw_pct_of_peak:.1f}% of peak)")
+            lines.append(f"   Bottleneck:  {self.roofline.bottleneck.upper()}")
+            lines.append("")
+
+        # All kernels section (if more than one)
+        if len(self.kernels) > 1:
+            lines.append(" All Kernels:")
+            total_dur = self.total_duration_us
+            for i, k in enumerate(self.kernels[:5]):  # Show top 5
+                pct = (k.duration_us / total_dur * 100) if total_dur > 0 else 0
+                name = self._truncate_kernel_name(k.name, 45)
+                dur_str = self._format_duration(k.duration_us)
+                marker = "→" if k == self.primary_kernel else " "
+                lines.append(f"   {marker} {i+1}. {name:<45}  {dur_str:>10}  ({pct:>4.1f}%)")
+            if len(self.kernels) > 5:
+                lines.append(f"      ... and {len(self.kernels) - 5} more kernels")
+            lines.append("")
+
+        lines.append("─" * 65)
+
+        return "\n".join(lines)
+
+    def _format_op_description(self) -> str:
+        """Format operation with shape info."""
+        op_name = self.op_spec.op.split(".")[-1]  # e.g., "matmul" from "torch.matmul"
+        
+        # Format input shapes
+        shape_strs = []
+        for t in self.op_spec.inputs:
+            shape_str = "×".join(str(d) for d in t.shape)
+            shape_strs.append(shape_str)
+        
+        # Get dtype from first input
+        dtype = self.op_spec.inputs[0].dtype if self.op_spec.inputs else "float16"
+        dtype_upper = dtype.upper().replace("FLOAT", "FP").replace("BFLOAT", "BF")
+        
+        if len(shape_strs) == 2:
+            return f"{op_name} ({shape_strs[0]} @ {shape_strs[1]}) {dtype_upper}"
+        elif len(shape_strs) == 1:
+            return f"{op_name} ({shape_strs[0]}) {dtype_upper}"
+        else:
+            return f"{op_name} {dtype_upper}"
+
+    def _truncate_kernel_name(self, name: str, max_len: int) -> str:
+        """Truncate long kernel names."""
+        if len(name) <= max_len:
+            return name
+        return name[:max_len - 3] + "..."
+
+    def _format_duration(self, us: float) -> str:
+        """Format duration nicely."""
+        if us >= 1000:
+            return f"{us / 1000:.2f} ms"
+        elif us >= 1:
+            return f"{us:.1f} µs"
+        else:
+            return f"{us * 1000:.1f} ns"

wafer_core/tools/dispatch_baseline/executor.py

@@ -0,0 +1,360 @@
+"""Executor for kernel trace operations.
+
+Runs profiling scripts on remote GPUs and returns structured results.
+"""
+
+import logging
+import tempfile
+from dataclasses import dataclass, replace
+from pathlib import Path
+
+from wafer_core.ssh import ExecResult, SSHClient
+from wafer_core.tools.dispatch_baseline.analyzer import ParsedTraceResult, parse_trace_output
+from wafer_core.tools.dispatch_baseline.codegen import generate_trace_script
+from wafer_core.tools.dispatch_baseline.dtypes import (
+    KernelTraceConfig,
+    KernelTraceResult,
+    OpSpec,
+    TensorSpec,
+)
+from wafer_core.tools.dispatch_baseline.roofline import (
+    compute_roofline,
+    estimate_attention_bytes,
+    estimate_attention_flops,
+    estimate_matmul_bytes,
+    estimate_matmul_flops,
+    estimate_softmax_bytes,
+    estimate_softmax_flops,
+    get_hardware_spec,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class TraceExecutionResult:
+    """Result of trace execution with environment info for caching."""
+
+    result: KernelTraceResult
+    pytorch_version: str
+    runtime_version: str
+    gpu_arch: str
+    from_cache: bool = False
+
+
+def trace_kernel_local(config: KernelTraceConfig) -> TraceExecutionResult:
+    """Trace a kernel on the local GPU.
+
+    Args:
+        config: Kernel trace configuration
+
+    Returns:
+        TraceExecutionResult with kernel information and environment info
+    """
+    import subprocess
+    import sys
+
+    script = generate_trace_script(config)
+
+    # Write script to temp file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(script)
+        script_path = f.name
+
+    try:
+        # Run script
+        result = subprocess.run(
+            [sys.executable, script_path],
+            capture_output=True,
+            text=True,
+            timeout=config.timeout_seconds,
+        )
+
+        output = result.stdout + result.stderr
+
+        if result.returncode != 0:
+            return TraceExecutionResult(
+                result=KernelTraceResult(
+                    op_spec=config.op_spec,
+                    hardware=config.hardware,
+                    kernels=[],
+                    primary_kernel=None,
+                    raw_output=output,
+                    error=f"Script failed with exit code {result.returncode}: {result.stderr}",
+                ),
+                pytorch_version="unknown",
+                runtime_version="unknown",
+                gpu_arch="unknown",
+            )
+
+        # Parse output (now includes environment info)
+        parsed = parse_trace_output(output, config.op_spec, config.hardware)
+
+        # Add roofline analysis if we can estimate FLOPS
+        trace_result = _add_roofline_analysis(parsed.result, config)
+
+        return TraceExecutionResult(
+            result=trace_result,
+            pytorch_version=parsed.pytorch_version,
+            runtime_version=parsed.runtime_version,
+            gpu_arch=parsed.gpu_arch,
+        )
+
+    except subprocess.TimeoutExpired:
+        return TraceExecutionResult(
+            result=KernelTraceResult(
+                op_spec=config.op_spec,
+                hardware=config.hardware,
+                kernels=[],
+                primary_kernel=None,
+                error=f"Script timed out after {config.timeout_seconds}s",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    finally:
+        # Cleanup
+        Path(script_path).unlink(missing_ok=True)
+
+
+def trace_kernel_remote(
+    config: KernelTraceConfig,
+    ssh_client: SSHClient,
+    workspace_path: str = "/tmp",
+) -> TraceExecutionResult:
+    """Trace a kernel on a remote GPU via SSH.
+
+    Args:
+        config: Kernel trace configuration
+        ssh_client: Connected SSH client
+        workspace_path: Remote directory to use for temporary files
+
+    Returns:
+        TraceExecutionResult with kernel information and environment info
+    """
+    script = generate_trace_script(config)
+    script_filename = "dispatch_baseline_script.py"
+    remote_script_path = f"{workspace_path}/{script_filename}"
+
+    try:
+        # Upload script
+        logger.debug(f"Uploading trace script to {remote_script_path}")
+        ssh_client.upload_content(script, remote_script_path)
+
+        # Run script
+        # TODO: Add timeout to SSHClient.exec() to prevent hanging on remote failures
+        logger.debug("Running trace script...")
+        run_cmd = f"cd {workspace_path} && python {script_filename}"
+        result: ExecResult = ssh_client.exec(run_cmd)
+
+        output = result.stdout + result.stderr
+
+        if result.exit_code != 0:
+            return TraceExecutionResult(
+                result=KernelTraceResult(
+                    op_spec=config.op_spec,
+                    hardware=config.hardware,
+                    kernels=[],
+                    primary_kernel=None,
+                    raw_output=output,
+                    error=f"Script failed with exit code {result.exit_code}: {result.stderr}",
+                ),
+                pytorch_version="unknown",
+                runtime_version="unknown",
+                gpu_arch="unknown",
+            )
+
+        # Parse output (now includes environment info)
+        parsed = parse_trace_output(output, config.op_spec, config.hardware)
+
+        # Add roofline analysis
+        trace_result = _add_roofline_analysis(parsed.result, config)
+
+        return TraceExecutionResult(
+            result=trace_result,
+            pytorch_version=parsed.pytorch_version,
+            runtime_version=parsed.runtime_version,
+            gpu_arch=parsed.gpu_arch,
+        )
+
+    except Exception as e:
+        logger.exception("Failed to trace kernel remotely")
+        return TraceExecutionResult(
+            result=KernelTraceResult(
+                op_spec=config.op_spec,
+                hardware=config.hardware,
+                kernels=[],
+                primary_kernel=None,
+                error=f"Remote execution failed: {e}",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    finally:
+        # Cleanup remote script
+        try:
+            ssh_client.exec(f"rm -f {remote_script_path}")
+        except Exception:
+            pass
+
+
+def _add_roofline_analysis(
+    result: KernelTraceResult, config: KernelTraceConfig
+) -> KernelTraceResult:
+    """Add roofline analysis to trace result if possible.
+    
+    Supports: matmul, softmax, attention, and generic elementwise ops.
+    """
+    if result.primary_kernel is None:
+        return result
+    
+    if config.hardware is None:
+        return result
+
+    # Try to estimate FLOPS based on operation type
+    op_lower = config.op_spec.op.lower()
+    shapes = [t.shape for t in config.op_spec.inputs]
+    dtype_bytes = _get_dtype_bytes(config.op_spec.inputs[0].dtype) if config.op_spec.inputs else 2
+    
+    flops = 0.0
+    bytes_transferred = 0.0
+
+    if "matmul" in op_lower or "mm" in op_lower or "linear" in op_lower:
+        # Matrix multiplication: C[M,N] = A[M,K] @ B[K,N]
+        if len(shapes) >= 2 and len(shapes[0]) >= 2 and len(shapes[1]) >= 2:
+            # Handle batched matmul by taking last 2 dims
+            m, k1 = shapes[0][-2], shapes[0][-1]
+            k2, n = shapes[1][-2], shapes[1][-1]
+            # Validate batch dimensions match (if present)
+            batch_dims_0 = shapes[0][:-2]
+            batch_dims_1 = shapes[1][:-2]
+            if k1 != k2:
+                logger.warning(f"Matmul inner dims mismatch: {k1} vs {k2}, skipping roofline")
+            elif batch_dims_0 != batch_dims_1:
+                logger.warning(f"Matmul batch dims mismatch: {batch_dims_0} vs {batch_dims_1}, skipping roofline")
+            else:
+                # Account for batch dimensions
+                batch_size = 1
+                for dim in batch_dims_0:
+                    batch_size *= dim
+                flops = batch_size * estimate_matmul_flops(m, n, k1)
+                bytes_transferred = batch_size * estimate_matmul_bytes(m, n, k1, dtype_bytes)
+                
+    elif "softmax" in op_lower:
+        # Softmax: read input, write output
+        if shapes:
+            elements = 1
+            for dim in shapes[0]:
+                elements *= dim
+            flops = estimate_softmax_flops(elements)
+            bytes_transferred = estimate_softmax_bytes(elements, dtype_bytes)
+            
+    elif "attention" in op_lower or "sdpa" in op_lower:
+        # Scaled dot-product attention
+        # Expect inputs like Q[B,H,S,D], K[B,H,S,D], V[B,H,S,D]
+        if len(shapes) >= 3 and len(shapes[0]) == 4:
+            batch, heads, seq_len, head_dim = shapes[0]
+            flops = estimate_attention_flops(batch, heads, seq_len, head_dim)
+            bytes_transferred = estimate_attention_bytes(batch, heads, seq_len, head_dim, dtype_bytes)
+            
+    elif any(op in op_lower for op in ["relu", "gelu", "silu", "tanh", "sigmoid", "exp", "log"]):
+        # Elementwise activation: ~1-5 ops per element, read+write
+        if shapes:
+            elements = 1
+            for dim in shapes[0]:
+                elements *= dim
+            flops = 5.0 * elements  # Conservative estimate for transcendentals
+            bytes_transferred = 2.0 * elements * dtype_bytes  # Read + write
+            
+    elif any(op in op_lower for op in ["add", "sub", "mul", "div"]):
+        # Binary elementwise: 1 op per element
+        if shapes:
+            elements = 1
+            for dim in shapes[0]:
+                elements *= dim
+            flops = float(elements)
+            bytes_transferred = 3.0 * elements * dtype_bytes  # Read 2 inputs + write 1 output
+            
+    elif "layernorm" in op_lower or "layer_norm" in op_lower or "rmsnorm" in op_lower:
+        # Normalization: mean, variance, normalize (~10 ops per element)
+        if shapes:
+            elements = 1
+            for dim in shapes[0]:
+                elements *= dim
+            flops = 10.0 * elements
+            bytes_transferred = 2.0 * elements * dtype_bytes
+            
+    elif "conv" in op_lower:
+        # Convolution is complex, skip for now (would need kernel size, stride, etc.)
+        pass
+
+    if flops > 0 and bytes_transferred > 0:
+        roofline = compute_roofline(
+            result.primary_kernel,
+            config.hardware,
+            flops,
+            bytes_transferred,
+        )
+        if roofline:
+            return replace(result, roofline=roofline)
+
+    return result
+
+
+def _get_dtype_bytes(dtype: str) -> int:
+    """Get bytes per element for a dtype."""
+    dtype_map = {
+        "float16": 2,
+        "float32": 4,
+        "float64": 8,
+        "bfloat16": 2,
+        "int8": 1,
+        "int16": 2,
+        "int32": 4,
+        "int64": 8,
+    }
+    return dtype_map.get(dtype, 2)
+
+
+def quick_trace(
+    op: str,
+    shapes: dict[str, tuple[int, ...]],
+    hardware: str = "H100",
+    dtype: str = "float16",
+) -> KernelTraceResult:
+    """Quick helper to trace an operation locally.
+
+    Args:
+        op: Operation string like "torch.matmul(A, B)"
+        shapes: Dict mapping tensor names to shapes
+        hardware: Hardware name (for roofline analysis)
+        dtype: Data type for tensors
+
+    Returns:
+        KernelTraceResult
+
+    Example:
+        result = quick_trace(
+            "torch.matmul(A, B)",
+            {"A": (4096, 4096), "B": (4096, 4096)},
+            hardware="H100",
+        )
+        print(result.summary())
+    """
+    from wafer_core.tools.dispatch_baseline.codegen import parse_op_string, update_dtypes, update_shapes
+
+    op_spec = parse_op_string(op)
+    op_spec = update_shapes(op_spec, shapes)
+    op_spec = update_dtypes(op_spec, dtype)
+
+    config = KernelTraceConfig(
+        op_spec=op_spec,
+        hardware=hardware,
+    )
+
+    execution_result = trace_kernel_local(config)
+    return execution_result.result