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

wafer_core/tools/dispatch_baseline/__init__.py

@@ -0,0 +1,73 @@
+"""Kernel trace tool for discovering SOTA kernels.
+
+Given a PyTorch operation, traces what kernel PyTorch dispatches to on target hardware.
+
+Usage:
+    from wafer_core.tools.dispatch_baseline import quick_trace
+
+    result = quick_trace(
+        "torch.matmul(A, B)",
+        {"A": (4096, 4096), "B": (4096, 4096)},
+        hardware="H100",
+    )
+    print(result.summary())
+"""
+
+from wafer_core.tools.dispatch_baseline.client import (
+    lookup_baseline,
+    store_baseline,
+)
+from wafer_core.tools.dispatch_baseline.codegen import (
+    generate_trace_script,
+    parse_op_string,
+    update_dtypes,
+    update_shapes,
+)
+from wafer_core.tools.dispatch_baseline.dtypes import (
+    HardwareSpec,
+    KernelInfo,
+    KernelTraceConfig,
+    KernelTraceResult,
+    OpSpec,
+    RooflineAnalysis,
+    TensorSpec,
+)
+from wafer_core.tools.dispatch_baseline.executor import (
+    TraceExecutionResult,
+    quick_trace,
+    trace_kernel_local,
+    trace_kernel_remote,
+)
+from wafer_core.tools.dispatch_baseline.roofline import (
+    HARDWARE_SPECS,
+    compute_roofline,
+    get_hardware_spec,
+)
+
+__all__ = [
+    # Data types
+    "HardwareSpec",
+    "KernelInfo",
+    "KernelTraceConfig",
+    "KernelTraceResult",
+    "OpSpec",
+    "RooflineAnalysis",
+    "TensorSpec",
+    "TraceExecutionResult",
+    # Codegen
+    "generate_trace_script",
+    "parse_op_string",
+    "update_dtypes",
+    "update_shapes",
+    # Execution
+    "quick_trace",
+    "trace_kernel_local",
+    "trace_kernel_remote",
+    # Database
+    "lookup_baseline",
+    "store_baseline",
+    # Roofline
+    "HARDWARE_SPECS",
+    "compute_roofline",
+    "get_hardware_spec",
+]

wafer_core/tools/dispatch_baseline/analyzer.py

@@ -0,0 +1,174 @@
+"""Analyzer for kernel trace output.
+
+Parses profiler output to extract kernel information and identify the primary kernel.
+"""
+
+import json
+import re
+from dataclasses import dataclass
+from typing import Any
+
+from wafer_core.tools.dispatch_baseline.dtypes import KernelInfo, KernelTraceResult, OpSpec
+
+
+@dataclass(frozen=True)
+class ParsedTraceResult:
+    """Parsed trace output with environment info for caching."""
+
+    result: KernelTraceResult
+    pytorch_version: str
+    runtime_version: str
+    gpu_arch: str
+
+
+def parse_trace_output(output: str, op_spec: OpSpec, hardware: str) -> ParsedTraceResult:
+    """Parse kernel trace script output and extract kernel information.
+
+    Args:
+        output: Raw output from the trace script (stdout)
+        op_spec: The operation that was traced
+        hardware: Target hardware name
+
+    Returns:
+        ParsedTraceResult with extracted kernel information and environment info
+    """
+    # Look for our JSON marker in the output
+    json_match = re.search(r"KERNEL_TRACE_RESULT_JSON:(.+)$", output, re.MULTILINE)
+
+    if not json_match:
+        return ParsedTraceResult(
+            result=KernelTraceResult(
+                op_spec=op_spec,
+                hardware=hardware,
+                kernels=[],
+                primary_kernel=None,
+                raw_output=output,
+                error="Could not find KERNEL_TRACE_RESULT_JSON marker in output",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    try:
+        result_json = json.loads(json_match.group(1))
+    except json.JSONDecodeError as e:
+        return ParsedTraceResult(
+            result=KernelTraceResult(
+                op_spec=op_spec,
+                hardware=hardware,
+                kernels=[],
+                primary_kernel=None,
+                raw_output=output,
+                error=f"Failed to parse JSON: {e}",
+            ),
+            pytorch_version="unknown",
+            runtime_version="unknown",
+            gpu_arch="unknown",
+        )
+
+    # Extract environment info
+    env_info = result_json.get("environment", {})
+    pytorch_version = env_info.get("pytorch_version", "unknown")
+    runtime_version = env_info.get("runtime_version", "unknown")
+    gpu_arch = env_info.get("gpu_arch", "unknown")
+
+    # Extract kernels from JSON
+    kernels = []
+    for k in result_json.get("kernels", []):
+        kernel_info = _parse_kernel_dict(k)
+        if kernel_info:
+            kernels.append(kernel_info)
+
+    # Identify primary kernel (longest duration)
+    primary_kernel = kernels[0] if kernels else None
+
+    return ParsedTraceResult(
+        result=KernelTraceResult(
+            op_spec=op_spec,
+            hardware=hardware,
+            kernels=kernels,
+            primary_kernel=primary_kernel,
+            raw_output=output,
+        ),
+        pytorch_version=pytorch_version,
+        runtime_version=runtime_version,
+        gpu_arch=gpu_arch,
+    )
+
+
+def _parse_kernel_dict(kernel_dict: dict[str, Any]) -> KernelInfo | None:
+    """Parse a kernel dictionary into KernelInfo."""
+    name = kernel_dict.get("name", "")
+    if not name:
+        return None
+
+    return KernelInfo(
+        name=name,
+        duration_us=kernel_dict.get("duration_us", 0.0),
+    )
+
+
+def extract_kernels_from_nsys(nsys_output: str) -> list[dict[str, Any]]:
+    """Extract kernel information from nsys output.
+
+    This is a fallback for when torch.profiler doesn't capture all kernels.
+
+    Args:
+        nsys_output: Output from `nsys stats` or similar
+
+    Returns:
+        List of kernel dictionaries
+    """
+    kernels = []
+
+    # Look for GPU kernel summary lines
+    # Format varies, but typically: name, duration, count, etc.
+    lines = nsys_output.split("\n")
+
+    for line in lines:
+        # Skip headers and empty lines
+        if not line.strip() or "---" in line or "Name" in line:
+            continue
+
+        # Try to parse as nsys kernel line
+        # This is a simplified parser - real nsys output varies
+        parts = line.split()
+        if len(parts) >= 2:
+            # Heuristic: last column is often the kernel name
+            name = parts[-1]
+            if any(c.isalpha() for c in name) and "_" in name:
+                kernels.append({"name": name, "duration_us": 0.0})
+
+    return kernels
+
+
+def merge_kernel_infos(
+    profiler_kernels: list[KernelInfo], nsys_kernels: list[dict[str, Any]]
+) -> list[KernelInfo]:
+    """Merge kernel info from multiple sources.
+
+    Prefers profiler data when available, but adds any missing kernels from nsys.
+
+    Args:
+        profiler_kernels: Kernels from torch.profiler
+        nsys_kernels: Kernels from nsys
+
+    Returns:
+        Merged list of KernelInfo objects
+    """
+    seen_names = {k.name for k in profiler_kernels}
+    result = list(profiler_kernels)
+
+    for nsys_kernel in nsys_kernels:
+        name = nsys_kernel.get("name", "")
+        if name and name not in seen_names:
+            result.append(
+                KernelInfo(
+                    name=name,
+                    duration_us=nsys_kernel.get("duration_us", 0.0),
+                )
+            )
+            seen_names.add(name)
+
+    return result

wafer_core/tools/dispatch_baseline/client.py

@@ -0,0 +1,196 @@
+"""API client for baseline kernel trace results.
+
+Interacts with the Supabase baselines table to store and retrieve kernel dispatch info.
+Results are keyed by (op, shape, dtype, pytorch_version, runtime_version, gpu_arch).
+Since kernel dispatch is deterministic for a given environment, this serves as a
+shared lookup table across all users.
+"""
+
+import logging
+
+import httpx
+
+from wafer_core.tools.dispatch_baseline.dtypes import (
+    KernelInfo,
+    KernelTraceResult,
+    OpSpec,
+)
+
+logger = logging.getLogger(__name__)
+
+# API timeout for baseline lookups
+API_TIMEOUT = 30.0
+
+
+def _get_api_config() -> tuple[str, dict[str, str]]:
+    """Get API URL and auth headers.
+
+    Returns:
+        Tuple of (api_url, headers)
+    """
+    # Import here to avoid circular imports and allow use without CLI
+    try:
+        from wafer.api_client import get_api_url
+        from wafer.auth import get_auth_headers
+
+        return get_api_url(), get_auth_headers()
+    except ImportError:
+        # If CLI not installed, return defaults (won't work without auth)
+        return "https://api.wafer.ai", {}
+
+
+def lookup_baseline(
+    op_spec: OpSpec,
+    hardware: str,
+    pytorch_version: str,
+    runtime_version: str,
+    gpu_arch: str,
+) -> KernelTraceResult | None:
+    """Look up baseline result from the database.
+
+    Args:
+        op_spec: Operation specification
+        hardware: Hardware name (for display, not part of cache key)
+        pytorch_version: PyTorch version string
+        runtime_version: CUDA version string
+        gpu_arch: GPU architecture (e.g., "sm_90")
+
+    Returns:
+        KernelTraceResult if found, None otherwise
+    """
+    api_url, headers = _get_api_config()
+
+    if not headers.get("Authorization"):
+        logger.debug("No auth headers, skipping baseline lookup")
+        return None
+
+    # Build request
+    request_data = {
+        "op": op_spec.op,
+        "inputs": [
+            {"name": t.name, "shape": list(t.shape), "dtype": t.dtype}
+            for t in op_spec.inputs
+        ],
+        "kwargs": op_spec.kwargs,
+        "pytorch_version": pytorch_version,
+        "runtime_version": runtime_version,
+        "gpu_arch": gpu_arch,
+    }
+
+    try:
+        with httpx.Client(timeout=API_TIMEOUT, headers=headers) as client:
+            response = client.post(f"{api_url}/v1/baselines/lookup", json=request_data)
+            response.raise_for_status()
+            data = response.json()
+
+        if not data.get("found"):
+            return None
+
+        # Reconstruct result
+        kernels = [
+            KernelInfo(
+                name=k["name"],
+                duration_us=k["duration_us"],
+            )
+            for k in data.get("kernels", [])
+        ]
+
+        primary_data = data.get("primary_kernel")
+        primary_kernel = KernelInfo(
+            name=primary_data["name"],
+            duration_us=primary_data["duration_us"],
+        ) if primary_data else (kernels[0] if kernels else None)
+
+        return KernelTraceResult(
+            op_spec=op_spec,
+            hardware=hardware,
+            kernels=kernels,
+            primary_kernel=primary_kernel,
+            # Note: roofline will be recomputed by the caller
+        )
+
+    except httpx.HTTPError as e:
+        logger.warning(f"Baseline lookup failed: {e}")
+        return None
+    except Exception as e:
+        logger.warning(f"Baseline lookup error: {e}")
+        return None
+
+
+def store_baseline(
+    result: KernelTraceResult,
+    pytorch_version: str,
+    runtime_version: str,
+    gpu_arch: str,
+) -> bool:
+    """Store baseline result in the database.
+
+    Args:
+        result: Trace result to cache
+        pytorch_version: PyTorch version string
+        runtime_version: CUDA version string
+        gpu_arch: GPU architecture
+
+    Returns:
+        True if stored successfully, False otherwise
+    """
+    if result.error:
+        # Don't cache errors
+        return False
+
+    if not result.primary_kernel:
+        # Don't cache empty results
+        return False
+
+    api_url, headers = _get_api_config()
+
+    if not headers.get("Authorization"):
+        logger.debug("No auth headers, skipping baseline store")
+        return False
+
+    # Build request
+    request_data = {
+        "op": result.op_spec.op,
+        "inputs": [
+            {"name": t.name, "shape": list(t.shape), "dtype": t.dtype}
+            for t in result.op_spec.inputs
+        ],
+        "kwargs": result.op_spec.kwargs,
+        "pytorch_version": pytorch_version,
+        "runtime_version": runtime_version,
+        "gpu_arch": gpu_arch,
+        "hardware_name": result.hardware,
+        "primary_kernel": {
+            "name": result.primary_kernel.name,
+            "duration_us": result.primary_kernel.duration_us,
+        },
+        "kernels": [
+            {
+                "name": k.name,
+                "duration_us": k.duration_us,
+            }
+            for k in result.kernels
+        ],
+    }
+
+    try:
+        with httpx.Client(timeout=API_TIMEOUT, headers=headers) as client:
+            response = client.post(f"{api_url}/v1/baselines/store", json=request_data)
+            response.raise_for_status()
+            data = response.json()
+
+        if data.get("created"):
+            logger.info(f"Stored baseline: {result.op_spec.op} ({gpu_arch})")
+        else:
+            logger.debug(f"Baseline already exists: {result.op_spec.op} ({gpu_arch})")
+
+        return True
+
+    except httpx.HTTPError as e:
+        logger.warning(f"Baseline store failed: {e}")
+        return False
+    except Exception as e:
+        logger.warning(f"Baseline store error: {e}")
+        return False
+
+

wafer_core/tools/dispatch_baseline/codegen.py

@@ -0,0 +1,246 @@
+"""Code generator for kernel trace profiling scripts.
+
+Generates minimal Python scripts that:
+1. Create tensors with specified shapes/dtypes
+2. Run the PyTorch op with warmup
+3. Profile using torch.profiler
+4. Output structured JSON with kernel info
+"""
+
+from wafer_core.tools.dispatch_baseline.dtypes import KernelTraceConfig, OpSpec, TensorSpec
+
+
+def generate_trace_script(config: KernelTraceConfig) -> str:
+    """Generate a profiling script for the given operation.
+
+    Args:
+        config: Kernel trace configuration with op spec and settings
+
+    Returns:
+        Python script as a string
+    """
+    op_spec = config.op_spec
+    tensor_setup = _generate_tensor_setup(op_spec.inputs)
+    op_call = _generate_op_call(op_spec)
+
+    script = f'''"""Auto-generated kernel trace script."""
+import json
+import sys
+import torch
+
+# Ensure CUDA is available
+assert torch.cuda.is_available(), "CUDA not available"
+
+# Setup tensors
+{tensor_setup}
+
+# Warmup
+print("Warming up...", file=sys.stderr)
+for _ in range({config.num_warmup}):
+    _ = {op_call}
+torch.cuda.synchronize()
+
+# Profile with torch.profiler
+print("Profiling...", file=sys.stderr)
+with torch.profiler.profile(
+    activities=[
+        torch.profiler.ProfilerActivity.CPU,
+        torch.profiler.ProfilerActivity.CUDA,
+    ],
+    record_shapes=True,
+    with_stack=False,
+) as prof:
+    for _ in range({config.num_runs}):
+        _ = {op_call}
+    torch.cuda.synchronize()
+
+# Extract kernel information
+kernels = []
+for event in prof.key_averages():
+    if event.device_type == torch.profiler.DeviceType.CUDA:
+        # Use device_time (new) or cuda_time (old) for average time in us
+        duration = getattr(event, 'device_time', None) or getattr(event, 'cuda_time', 0)
+        kernels.append({{
+            "name": event.key,
+            "duration_us": duration,
+            "count": event.count,
+        }})
+
+# Sort by duration (longest first)
+kernels.sort(key=lambda k: k["duration_us"], reverse=True)
+
+# Get environment info for caching
+props = torch.cuda.get_device_properties(0)
+
+# Detect runtime version (CUDA or ROCm)
+if hasattr(torch.version, 'hip') and torch.version.hip:
+    runtime_version = torch.version.hip
+    # ROCm uses gcnArchName for architecture (e.g., "gfx942")
+    gpu_arch = getattr(props, 'gcnArchName', f"gfx{{props.major}}{{props.minor}}")
+else:
+    runtime_version = torch.version.cuda or "unknown"
+    gpu_arch = f"sm_{{props.major}}{{props.minor}}"
+
+env_info = {{
+    "pytorch_version": torch.__version__,
+    "runtime_version": runtime_version,
+    "gpu_arch": gpu_arch,
+    "gpu_name": props.name,
+}}
+
+# Build result
+result = {{
+    "op": "{op_spec.op}",
+    "inputs": {_serialize_inputs(op_spec.inputs)},
+    "num_runs": {config.num_runs},
+    "kernels": kernels,
+    "total_cuda_time_us": sum(k["duration_us"] for k in kernels),
+    "environment": env_info,
+}}
+
+# Output as JSON (marker for parsing)
+print("KERNEL_TRACE_RESULT_JSON:" + json.dumps(result))
+'''
+    return script
+
+
+def _generate_tensor_setup(inputs: list[TensorSpec]) -> str:
+    """Generate tensor creation code."""
+    lines = []
+    for tensor in inputs:
+        shape_str = ", ".join(str(d) for d in tensor.shape)
+        dtype_map = {
+            "float16": "torch.float16",
+            "float32": "torch.float32",
+            "bfloat16": "torch.bfloat16",
+            "float64": "torch.float64",
+            "int8": "torch.int8",
+            "int16": "torch.int16",
+            "int32": "torch.int32",
+            "int64": "torch.int64",
+        }
+        dtype = dtype_map.get(tensor.dtype, f"torch.{tensor.dtype}")
+
+        # Use randn for float types, randint for int types
+        if "int" in tensor.dtype:
+            lines.append(
+                f'{tensor.name} = torch.randint(-128, 127, ({shape_str},), '
+                f'dtype={dtype}, device="{tensor.device}")'
+            )
+        else:
+            lines.append(
+                f'{tensor.name} = torch.randn({shape_str}, dtype={dtype}, device="{tensor.device}")'
+            )
+
+    return "\n".join(lines)
+
+
+def _generate_op_call(op_spec: OpSpec) -> str:
+    """Generate the operation call code."""
+    args = [t.name for t in op_spec.inputs]
+    kwargs = [f"{k}={v}" for k, v in op_spec.kwargs.items()]
+    all_args = ", ".join(args + kwargs)
+    return f"{op_spec.op}({all_args})"
+
+
+def _serialize_inputs(inputs: list[TensorSpec]) -> str:
+    """Serialize inputs for JSON output."""
+    items = []
+    for t in inputs:
+        items.append(f'{{"name": "{t.name}", "shape": {list(t.shape)}, "dtype": "{t.dtype}"}}')
+    return "[" + ", ".join(items) + "]"
+
+
+def parse_op_string(op_string: str) -> OpSpec:
+    """Parse a simple operation string into an OpSpec.
+
+    Supports formats like:
+        - "torch.matmul(A, B)"
+        - "torch.nn.functional.softmax(x, dim=-1)"
+
+    For shapes, use --shape flags separately.
+
+    Args:
+        op_string: Operation string like "torch.matmul(A, B)"
+
+    Returns:
+        OpSpec with parsed operation and placeholder inputs
+    """
+    # Extract op name and arguments
+    if "(" not in op_string or ")" not in op_string:
+        raise ValueError(f"Invalid op format: {op_string}. Expected: op(args)")
+
+    op_name = op_string[: op_string.index("(")].strip()
+    args_str = op_string[op_string.index("(") + 1 : op_string.rindex(")")].strip()
+
+    # Parse arguments (simple comma split, doesn't handle nested calls)
+    args = [a.strip() for a in args_str.split(",") if a.strip()]
+
+    # Separate positional args (tensor names) from kwargs
+    inputs = []
+    kwargs = {}
+
+    for arg in args:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+            kwargs[key.strip()] = value.strip()
+        else:
+            # Assume it's a tensor input with default shape
+            inputs.append(
+                TensorSpec(
+                    name=arg,
+                    shape=(1024, 1024),  # Default shape, override with --shape
+                    dtype="float16",
+                )
+            )
+
+    return OpSpec(op=op_name, inputs=inputs, kwargs=kwargs)
+
+
+def update_shapes(op_spec: OpSpec, shapes: dict[str, tuple[int, ...]]) -> OpSpec:
+    """Update tensor shapes in an OpSpec.
+
+    Args:
+        op_spec: Original OpSpec
+        shapes: Dict mapping tensor name to new shape
+
+    Returns:
+        New OpSpec with updated shapes
+    """
+    new_inputs = []
+    for tensor in op_spec.inputs:
+        if tensor.name in shapes:
+            new_inputs.append(
+                TensorSpec(
+                    name=tensor.name,
+                    shape=shapes[tensor.name],
+                    dtype=tensor.dtype,
+                    device=tensor.device,
+                )
+            )
+        else:
+            new_inputs.append(tensor)
+
+    return OpSpec(op=op_spec.op, inputs=new_inputs, kwargs=op_spec.kwargs)
+
+
+def update_dtypes(op_spec: OpSpec, dtype: str) -> OpSpec:
+    """Update all tensor dtypes in an OpSpec.
+
+    Args:
+        op_spec: Original OpSpec
+        dtype: New dtype for all tensors
+
+    Returns:
+        New OpSpec with updated dtypes
+    """
+    new_inputs = [
+        TensorSpec(
+            name=t.name,
+            shape=t.shape,
+            dtype=dtype,
+            device=t.device,
+        )
+        for t in op_spec.inputs
+    ]
+    return OpSpec(op=op_spec.op, inputs=new_inputs, kwargs=op_spec.kwargs)