wafer-core 0.1.38__py3-none-any.whl → 0.1.39__py3-none-any.whl

wafer_core/tools/compile/compiler.py

@@ -0,0 +1,314 @@
+"""CUDA compiler wrapper.
+
+This module provides the core compilation logic that can be used locally
+or through Modal. It handles request/response conversion and provides
+a clean interface for the API and CLI.
+"""
+
+import subprocess
+from dataclasses import asdict
+
+from wafer_core.tools.compile.types import (
+    CompileRequest,
+    CompileResponse,
+    OutputFormat,
+)
+
+
+def request_to_dict(request: CompileRequest) -> dict:
+    """Convert a CompileRequest to a dict for Modal invocation.
+
+    Args:
+        request: The compile request
+
+    Returns:
+        Dict suitable for passing to Modal function
+    """
+    return {
+        "files": dict(request.files),
+        "arch": request.arch,
+        "flags": list(request.flags),
+        "output": [fmt.value for fmt in request.output],
+    }
+
+
+def response_from_dict(data: dict) -> CompileResponse:
+    """Convert a Modal response dict to a CompileResponse.
+
+    Args:
+        data: Dict returned from Modal function
+
+    Returns:
+        CompileResponse object
+    """
+    return CompileResponse(
+        success=data.get("success", False),
+        ptx=data.get("ptx"),
+        sass=data.get("sass"),
+        stderr=data.get("stderr", ""),
+        compilation_time_ms=data.get("compilation_time_ms", 0),
+    )
+
+
+async def compile_cuda_remote(
+    request: CompileRequest,
+    *,
+    modal_token_id: str | None = None,
+    modal_token_secret: str | None = None,
+) -> CompileResponse:
+    """Compile CUDA code using Modal (remote execution).
+
+    This function spawns a subprocess to call Modal, avoiding event loop
+    conflicts between the caller's event loop and Modal's asyncio.
+
+    Args:
+        request: The compile request
+        modal_token_id: Optional Modal token ID (uses env var if not provided)
+        modal_token_secret: Optional Modal token secret
+
+    Returns:
+        CompileResponse with PTX/SASS or error
+    """
+    import asyncio
+    import json
+    import os
+    import tempfile
+    import time
+    from pathlib import Path
+
+    start_time = time.time()
+
+    # Write request to temp file
+    request_dict = request_to_dict(request)
+
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".json", delete=False
+    ) as request_file:
+        json.dump(request_dict, request_file)
+        request_path = request_file.name
+
+    try:
+        # Create a Python script that calls Modal using Function.lookup
+        # This calls the deployed function without needing to rebuild the image
+        script = f'''
+import json
+import modal
+
+# Load request
+with open("{request_path}") as f:
+    request = json.load(f)
+
+# Look up the deployed function
+compile_fn = modal.Function.from_name("cuda-compile", "compile_cuda")
+
+# Call the function remotely
+result = compile_fn.remote(request)
+
+# Output result as JSON
+print(json.dumps(result))
+'''
+
+        # Run in subprocess to avoid event loop conflicts
+        env = os.environ.copy()
+        if modal_token_id:
+            env["MODAL_TOKEN_ID"] = modal_token_id
+        if modal_token_secret:
+            env["MODAL_TOKEN_SECRET"] = modal_token_secret
+
+        # Use the same Python interpreter that's running this code
+        import sys
+        python_executable = sys.executable
+
+        # Use asyncio.create_subprocess_exec for async subprocess execution
+        proc = await asyncio.create_subprocess_exec(
+            python_executable, "-c", script,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=env,
+        )
+        stdout_bytes, stderr_bytes = await proc.communicate()
+
+        if proc.returncode != 0:
+            stderr = stderr_bytes.decode() if stderr_bytes else "Unknown error"
+            # Check for common Modal auth errors
+            if "MODAL_TOKEN" in stderr or "AuthError" in stderr or "not authenticated" in stderr.lower():
+                return CompileResponse.error(
+                    "Modal not configured. Set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET environment variables, "
+                    "or run 'modal token new' to authenticate.",
+                    compilation_time_ms=int((time.time() - start_time) * 1000),
+                )
+            return CompileResponse.error(
+                f"Compilation failed: {stderr}",
+                compilation_time_ms=int((time.time() - start_time) * 1000),
+            )
+
+        # Parse result
+        stdout = stdout_bytes.decode() if stdout_bytes else "{}"
+        try:
+            response_dict = json.loads(stdout)
+            return response_from_dict(response_dict)
+        except json.JSONDecodeError as e:
+            return CompileResponse.error(
+                f"Failed to parse Modal response: {e}\nOutput: {stdout[:500]}",
+                compilation_time_ms=int((time.time() - start_time) * 1000),
+            )
+
+    finally:
+        # Clean up temp file
+        Path(request_path).unlink(missing_ok=True)
+
+
+def compile_cuda_local(request: CompileRequest) -> CompileResponse:
+    """Compile CUDA code locally using nvcc.
+
+    This function requires nvcc to be installed locally.
+    Primarily useful for testing without Modal.
+
+    Args:
+        request: The compile request
+
+    Returns:
+        CompileResponse with PTX/SASS or error
+    """
+    import os
+    import tempfile
+    import time
+    from pathlib import Path
+
+    start_time = time.time()
+
+    # Check if nvcc is available
+    try:
+        subprocess.run(
+            ["nvcc", "--version"],
+            capture_output=True,
+            check=True,
+            timeout=10,
+        )
+    except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
+        return CompileResponse.error(
+            "nvcc not found. Install CUDA toolkit or use Modal for remote compilation.",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+
+            # Write all files to temp directory
+            for filename, content in request.files.items():
+                file_path = tmp_path / filename
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                file_path.write_text(content)
+
+            main_cu_path = tmp_path / request.main_cu_file
+            include_dir = main_cu_path.parent
+
+            results: dict[str, str | None] = {"ptx": None, "sass": None}
+
+            # Build base nvcc command
+            base_cmd = [
+                "nvcc",
+                "-arch",
+                request.arch,
+                f"-I{include_dir}",
+            ]
+            base_cmd.extend(request.flags)
+
+            # Generate PTX if requested
+            if OutputFormat.PTX in request.output:
+                ptx_output = tmp_path / "output.ptx"
+                ptx_cmd = base_cmd + [
+                    "--ptx",
+                    "-o",
+                    str(ptx_output),
+                    str(main_cu_path),
+                ]
+
+                ptx_result = subprocess.run(
+                    ptx_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                )
+
+                if ptx_result.returncode != 0:
+                    return CompileResponse.error(
+                        ptx_result.stderr or ptx_result.stdout,
+                        compilation_time_ms=int((time.time() - start_time) * 1000),
+                    )
+
+                if ptx_output.exists():
+                    results["ptx"] = ptx_output.read_text()
+
+            # Generate SASS if requested
+            if OutputFormat.SASS in request.output:
+                cubin_output = tmp_path / "output.cubin"
+                cubin_cmd = base_cmd + [
+                    "--cubin",
+                    "-o",
+                    str(cubin_output),
+                    str(main_cu_path),
+                ]
+
+                cubin_result = subprocess.run(
+                    cubin_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                )
+
+                if cubin_result.returncode != 0:
+                    if results["ptx"]:
+                        return CompileResponse(
+                            success=True,
+                            ptx=results["ptx"],
+                            sass=None,
+                            stderr=f"SASS generation failed: {cubin_result.stderr}",
+                            compilation_time_ms=int((time.time() - start_time) * 1000),
+                        )
+                    return CompileResponse.error(
+                        cubin_result.stderr or cubin_result.stdout,
+                        compilation_time_ms=int((time.time() - start_time) * 1000),
+                    )
+
+                if cubin_output.exists():
+                    sass_result = subprocess.run(
+                        ["cuobjdump", "--dump-sass", str(cubin_output)],
+                        capture_output=True,
+                        text=True,
+                        timeout=30,
+                        cwd=tmpdir,
+                    )
+
+                    if sass_result.returncode == 0:
+                        results["sass"] = sass_result.stdout
+
+            if not results["ptx"] and not results["sass"]:
+                return CompileResponse.error(
+                    "No output generated",
+                    compilation_time_ms=int((time.time() - start_time) * 1000),
+                )
+
+            return CompileResponse(
+                success=True,
+                ptx=results["ptx"],
+                sass=results["sass"],
+                stderr="",
+                compilation_time_ms=int((time.time() - start_time) * 1000),
+            )
+
+    except subprocess.TimeoutExpired:
+        return CompileResponse.error(
+            "Compilation timed out",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )
+    except Exception as e:
+        import traceback
+
+        return CompileResponse.error(
+            f"Internal error: {e}\n{traceback.format_exc()}",
+            compilation_time_ms=int((time.time() - start_time) * 1000),
+        )

wafer_core/tools/compile/modal_compile.py

@@ -0,0 +1,359 @@
+"""Modal function for cloud CUDA compilation.
+
+This module provides a CPU-only Modal function that compiles CUDA code
+and returns PTX/SASS assembly. No GPU is required - nvcc can generate
+PTX/SASS for any architecture without a physical GPU.
+
+Usage:
+    # From Python
+    from wafer_core.tools.compile.modal_compile import compile_cuda
+    result = compile_cuda.remote(request_dict)
+
+    # Deploy Modal app
+    modal deploy modal_compile.py
+"""
+
+from enum import Enum
+
+import modal
+
+
+# Inline OutputFormat to avoid import chain issues with Modal deployment
+class OutputFormat(str, Enum):
+    """Supported output formats."""
+
+    PTX = "ptx"
+    SASS = "sass"
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Modal Image Configuration
+# ══════════════════════════════════════════════════════════════════════════════
+
+# Build a CPU-only image with nvcc for CUDA compilation
+# No GPU driver needed - nvcc can generate PTX/SASS without a GPU
+compile_image = (
+    # Start with NVIDIA's CUDA development image (includes nvcc)
+    # Use debian-slim + manual CUDA install for smaller image
+    modal.Image.from_registry(
+        "nvidia/cuda:13.0.1-devel-ubuntu22.04",
+        add_python="3.12",
+    )
+    # System dependencies
+    .apt_install("git", "build-essential", "wget")
+    # Install PyTorch headers (CPU-only, we just need the headers)
+    .pip_install(
+        "torch>=2.5.0",
+        index_url="https://download.pytorch.org/whl/cpu",
+        extra_index_url="https://pypi.org/simple",
+    )
+    # Install CUTLASS headers (v4.3.5)
+    .run_commands(
+        "git clone --depth 1 --branch v4.3.5 https://github.com/NVIDIA/cutlass.git /usr/local/cutlass",
+    )
+    # Set environment variables for CUDA and libraries
+    .env(
+        {
+            "CUDA_HOME": "/usr/local/cuda",
+            "PATH": "/usr/local/cuda/bin:$PATH",
+            "CUTLASS_PATH": "/usr/local/cutlass/include",
+        }
+    )
+    # Verify nvcc is working
+    .run_commands(
+        "nvcc --version",
+        "ls -la /usr/local/cuda/bin/nvcc",
+    )
+)
+
+# Create Modal app
+app = modal.App(name="cuda-compile", image=compile_image)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Modal Function - CUDA Compilation
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+@app.function(
+    # CPU only - no GPU needed for compilation!
+    cpu=4,
+    memory=8192,  # 8GB RAM
+    timeout=120,  # 2 minute timeout
+)
+def compile_cuda(request: dict) -> dict:
+    """Compile CUDA code and return PTX/SASS.
+
+    This function runs on a CPU-only container with nvcc installed.
+    No GPU is required because nvcc can cross-compile for any target architecture.
+
+    Args:
+        request: Dict with the following fields:
+            - files: dict[str, str] - Mapping of filename to content
+            - arch: str - Target architecture (e.g., "sm_90a")
+            - flags: list[str] - Additional nvcc flags
+            - output: list[str] - Output formats ("ptx", "sass")
+
+    Returns:
+        Dict with the following fields:
+            - success: bool - Whether compilation succeeded
+            - ptx: str | None - Generated PTX code
+            - sass: str | None - Generated SASS code
+            - stderr: str - Compiler warnings/errors
+            - compilation_time_ms: int - Time taken in ms
+    """
+    import os
+    import subprocess
+    import tempfile
+    import time
+    from pathlib import Path
+
+    start_time = time.time()
+
+    # Extract request fields
+    files: dict[str, str] = request.get("files", {})
+    arch: str = request.get("arch", "sm_90a")
+    flags: list[str] = request.get("flags", [])
+    output_formats: list[str] = request.get("output", ["ptx", "sass"])
+
+    # Validate request
+    if not files:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "No files provided",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+    # Find the main .cu file
+    cu_files = [f for f in files.keys() if f.endswith(".cu")]
+    if not cu_files:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "No .cu file found",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+    main_cu_file = cu_files[0]
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+
+            # Write all files to temp directory, preserving subdirectory structure
+            for filename, content in files.items():
+                file_path = tmp_path / filename
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                file_path.write_text(content)
+
+            # Determine the directory of the main .cu file for include paths
+            main_cu_path = tmp_path / main_cu_file
+            include_dir = main_cu_path.parent
+
+            results: dict[str, str | None] = {"ptx": None, "sass": None}
+
+            # Build base nvcc command with common flags
+            base_cmd = [
+                "nvcc",
+                "-arch",
+                arch,
+                # Include the temp directory for user headers
+                f"-I{include_dir}",
+                # Include PyTorch headers
+                "-I/usr/local/lib/python3.12/site-packages/torch/include",
+                "-I/usr/local/lib/python3.12/site-packages/torch/include/torch/csrc/api/include",
+                # Include CUTLASS headers
+                "-I/usr/local/cutlass/include",
+                # Standard CUDA headers are already in the default path
+            ]
+
+            # Add user-specified flags
+            base_cmd.extend(flags)
+
+            # Generate PTX if requested
+            if OutputFormat.PTX.value in output_formats:
+                ptx_output = tmp_path / "output.ptx"
+                ptx_cmd = base_cmd + [
+                    "--ptx",  # Generate PTX
+                    "-o",
+                    str(ptx_output),
+                    str(main_cu_path),
+                ]
+
+                ptx_result = subprocess.run(
+                    ptx_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                    env={
+                        **os.environ,
+                        "CUDA_HOME": "/usr/local/cuda",
+                        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
+                    },
+                )
+
+                if ptx_result.returncode != 0:
+                    return {
+                        "success": False,
+                        "ptx": None,
+                        "sass": None,
+                        "stderr": ptx_result.stderr or ptx_result.stdout,
+                        "compilation_time_ms": int((time.time() - start_time) * 1000),
+                    }
+
+                if ptx_output.exists():
+                    results["ptx"] = ptx_output.read_text()
+
+            # Generate SASS if requested
+            if OutputFormat.SASS.value in output_formats:
+                # First compile to cubin, then disassemble to SASS
+                cubin_output = tmp_path / "output.cubin"
+                cubin_cmd = base_cmd + [
+                    "--cubin",  # Generate cubin (binary)
+                    "-o",
+                    str(cubin_output),
+                    str(main_cu_path),
+                ]
+
+                cubin_result = subprocess.run(
+                    cubin_cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=60,
+                    cwd=tmpdir,
+                    env={
+                        **os.environ,
+                        "CUDA_HOME": "/usr/local/cuda",
+                        "PATH": f"/usr/local/cuda/bin:{os.environ.get('PATH', '')}",
+                    },
+                )
+
+                if cubin_result.returncode != 0:
+                    # If we already have PTX, that's a partial success
+                    if results["ptx"]:
+                        return {
+                            "success": True,
+                            "ptx": results["ptx"],
+                            "sass": None,
+                            "stderr": f"SASS generation failed: {cubin_result.stderr}",
+                            "compilation_time_ms": int(
+                                (time.time() - start_time) * 1000
+                            ),
+                        }
+                    return {
+                        "success": False,
+                        "ptx": None,
+                        "sass": None,
+                        "stderr": cubin_result.stderr or cubin_result.stdout,
+                        "compilation_time_ms": int((time.time() - start_time) * 1000),
+                    }
+
+                # Disassemble cubin to SASS using cuobjdump
+                if cubin_output.exists():
+                    sass_cmd = [
+                        "cuobjdump",
+                        "--dump-sass",
+                        str(cubin_output),
+                    ]
+
+                    sass_result = subprocess.run(
+                        sass_cmd,
+                        capture_output=True,
+                        text=True,
+                        timeout=30,
+                        cwd=tmpdir,
+                    )
+
+                    if sass_result.returncode == 0:
+                        results["sass"] = sass_result.stdout
+                    else:
+                        # SASS generation failed but we might have PTX
+                        if results["ptx"]:
+                            return {
+                                "success": True,
+                                "ptx": results["ptx"],
+                                "sass": None,
+                                "stderr": f"SASS disassembly failed: {sass_result.stderr}",
+                                "compilation_time_ms": int(
+                                    (time.time() - start_time) * 1000
+                                ),
+                            }
+
+            # Check if we got any output
+            if not results["ptx"] and not results["sass"]:
+                return {
+                    "success": False,
+                    "ptx": None,
+                    "sass": None,
+                    "stderr": "No output generated",
+                    "compilation_time_ms": int((time.time() - start_time) * 1000),
+                }
+
+            return {
+                "success": True,
+                "ptx": results["ptx"],
+                "sass": results["sass"],
+                "stderr": "",
+                "compilation_time_ms": int((time.time() - start_time) * 1000),
+            }
+
+    except subprocess.TimeoutExpired:
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": "Compilation timed out",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+    except Exception as e:
+        import traceback
+
+        return {
+            "success": False,
+            "ptx": None,
+            "sass": None,
+            "stderr": f"Internal error: {e}\n{traceback.format_exc()}",
+            "compilation_time_ms": int((time.time() - start_time) * 1000),
+        }
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Health Check
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+@app.function(cpu=1, memory=1024, timeout=30)
+def health_check() -> dict:
+    """Verify the compilation environment is working.
+
+    Returns:
+        Dict with status and version info
+    """
+    import subprocess
+
+    try:
+        # Get nvcc version
+        nvcc_result = subprocess.run(
+            ["nvcc", "--version"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
+        # Parse version from output
+        version_line = nvcc_result.stdout.strip().split("\n")[-1]
+
+        return {
+            "status": "ok",
+            "nvcc_version": version_line,
+            "nvcc_available": nvcc_result.returncode == 0,
+        }
+    except Exception as e:
+        return {
+            "status": "error",
+            "error": str(e),
+            "nvcc_available": False,
+        }

wafer_core/tools/compile/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for the cloud CUDA compiler."""