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

wafer_core/tools/compile/tests/test_compiler.py

@@ -0,0 +1,675 @@
+"""Tests for the cloud CUDA compiler.
+
+These tests verify request validation, serialization, and compilation logic.
+Integration tests that actually invoke nvcc are in test_modal_integration.py.
+"""
+
+import pytest
+
+from wafer_core.tools.compile.types import (
+    CompileRequest,
+    CompileResponse,
+    OutputFormat,
+    VALID_ARCHITECTURES,
+)
+
+
+class TestCompileRequest:
+    """Test request validation and serialization."""
+
+    def test_single_file_request(self) -> None:
+        """Test creating a request with a single file."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+        )
+        assert request.files == {"kernel.cu": "__global__ void test() {}"}
+        assert request.arch == "sm_90a"  # default
+        assert request.flags == ()
+        assert request.output == (OutputFormat.PTX, OutputFormat.SASS)
+        assert request.main_cu_file == "kernel.cu"
+
+    def test_multi_file_request(self) -> None:
+        """Test creating a request with multiple files."""
+        request = CompileRequest(
+            files={
+                "main.cu": '#include "utils.cuh"\n__global__ void test() {}',
+                "utils.cuh": "__device__ float square(float x) { return x * x; }",
+            },
+        )
+        assert len(request.files) == 2
+        assert "main.cu" in request.files
+        assert "utils.cuh" in request.files
+
+    def test_custom_arch(self) -> None:
+        """Test specifying a custom architecture."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            arch="sm_80",
+        )
+        assert request.arch == "sm_80"
+
+    def test_custom_flags(self) -> None:
+        """Test specifying custom compiler flags."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            flags=("-O3", "--maxrregcount=64", "-lineinfo"),
+        )
+        assert request.flags == ("-O3", "--maxrregcount=64", "-lineinfo")
+
+    def test_ptx_only_output(self) -> None:
+        """Test requesting only PTX output."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            output=(OutputFormat.PTX,),
+        )
+        assert request.output == (OutputFormat.PTX,)
+
+    def test_sass_only_output(self) -> None:
+        """Test requesting only SASS output."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            output=(OutputFormat.SASS,),
+        )
+        assert request.output == (OutputFormat.SASS,)
+
+    def test_invalid_arch_rejected(self) -> None:
+        """Test that invalid architecture is rejected."""
+        with pytest.raises(ValueError, match="Invalid architecture"):
+            CompileRequest(
+                files={"kernel.cu": "__global__ void test() {}"},
+                arch="sm_999",
+            )
+
+    def test_empty_files_rejected(self) -> None:
+        """Test that empty files dict is rejected."""
+        with pytest.raises(ValueError, match="At least one file is required"):
+            CompileRequest(files={})
+
+    def test_no_cu_file_rejected(self) -> None:
+        """Test that missing .cu file is rejected."""
+        with pytest.raises(ValueError, match="At least one .cu file is required"):
+            CompileRequest(files={"utils.cuh": "// header only"})
+
+    def test_empty_output_rejected(self) -> None:
+        """Test that empty output list is rejected."""
+        with pytest.raises(ValueError, match="At least one output format is required"):
+            CompileRequest(
+                files={"kernel.cu": "__global__ void test() {}"},
+                output=(),
+            )
+
+    def test_request_is_frozen(self) -> None:
+        """Test that request is immutable."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+        )
+        with pytest.raises(AttributeError):
+            request.arch = "sm_80"  # type: ignore[misc]
+
+
+class TestCompileResponse:
+    """Test response creation and helper methods."""
+
+    def test_success_response(self) -> None:
+        """Test creating a successful response."""
+        response = CompileResponse(
+            success=True,
+            ptx=".version 8.0\n.target sm_90a",
+            sass="MOV R0, R1;",
+            compilation_time_ms=150,
+        )
+        assert response.success is True
+        assert response.ptx is not None
+        assert response.sass is not None
+        assert response.stderr == ""
+        assert response.compilation_time_ms == 150
+
+    def test_error_response_helper(self) -> None:
+        """Test the error response helper method."""
+        response = CompileResponse.error(
+            "kernel.cu(10): error: expected a ';'",
+            compilation_time_ms=50,
+        )
+        assert response.success is False
+        assert response.ptx is None
+        assert response.sass is None
+        assert "expected a ';'" in response.stderr
+        assert response.compilation_time_ms == 50
+
+    def test_response_is_frozen(self) -> None:
+        """Test that response is immutable."""
+        response = CompileResponse(success=True)
+        with pytest.raises(AttributeError):
+            response.success = False  # type: ignore[misc]
+
+
+class TestValidArchitectures:
+    """Test the set of valid architectures."""
+
+    def test_common_architectures_valid(self) -> None:
+        """Test that common architectures are in the valid set."""
+        common_archs = ["sm_80", "sm_86", "sm_89", "sm_90", "sm_90a"]
+        for arch in common_archs:
+            assert arch in VALID_ARCHITECTURES, f"{arch} should be valid"
+
+    def test_hopper_architectures(self) -> None:
+        """Test Hopper architectures."""
+        assert "sm_90" in VALID_ARCHITECTURES
+        assert "sm_90a" in VALID_ARCHITECTURES
+
+    def test_blackwell_architectures(self) -> None:
+        """Test Blackwell architectures."""
+        assert "sm_100" in VALID_ARCHITECTURES
+        assert "sm_100a" in VALID_ARCHITECTURES
+
+
+class TestOutputFormat:
+    """Test output format enum."""
+
+    def test_ptx_value(self) -> None:
+        """Test PTX enum value."""
+        assert OutputFormat.PTX.value == "ptx"
+        assert OutputFormat.PTX == "ptx"  # str enum comparison
+
+    def test_sass_value(self) -> None:
+        """Test SASS enum value."""
+        assert OutputFormat.SASS.value == "sass"
+        assert OutputFormat.SASS == "sass"  # str enum comparison
+
+
+# =============================================================================
+# Test CUDA Code Examples
+# =============================================================================
+
+# These are example CUDA kernels used in tests.
+# They are validated for correct syntax in integration tests.
+
+VECTOR_ADD_KERNEL = """\
+__global__ void vector_add(float* a, float* b, float* c, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        c[idx] = a[idx] + b[idx];
+    }
+}
+"""
+
+MATRIX_MULTIPLY_KERNEL = """\
+__global__ void matmul(float* A, float* B, float* C, int M, int N, int K) {
+    int row = blockIdx.y * blockDim.y + threadIdx.y;
+    int col = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (row < M && col < N) {
+        float sum = 0.0f;
+        for (int k = 0; k < K; k++) {
+            sum += A[row * K + k] * B[k * N + col];
+        }
+        C[row * N + col] = sum;
+    }
+}
+"""
+
+REDUCTION_KERNEL = """\
+__global__ void reduce_sum(float* input, float* output, int n) {
+    extern __shared__ float sdata[];
+
+    int tid = threadIdx.x;
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    sdata[tid] = (i < n) ? input[i] : 0.0f;
+    __syncthreads();
+
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            sdata[tid] += sdata[tid + s];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0) {
+        output[blockIdx.x] = sdata[0];
+    }
+}
+"""
+
+SHARED_MEMORY_KERNEL = """\
+#define TILE_SIZE 16
+
+__global__ void tiled_matmul(float* A, float* B, float* C, int N) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+
+    int bx = blockIdx.x, by = blockIdx.y;
+    int tx = threadIdx.x, ty = threadIdx.y;
+
+    int row = by * TILE_SIZE + ty;
+    int col = bx * TILE_SIZE + tx;
+
+    float sum = 0.0f;
+
+    for (int t = 0; t < (N + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        if (row < N && t * TILE_SIZE + tx < N)
+            As[ty][tx] = A[row * N + t * TILE_SIZE + tx];
+        else
+            As[ty][tx] = 0.0f;
+
+        if (t * TILE_SIZE + ty < N && col < N)
+            Bs[ty][tx] = B[(t * TILE_SIZE + ty) * N + col];
+        else
+            Bs[ty][tx] = 0.0f;
+
+        __syncthreads();
+
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum += As[ty][k] * Bs[k][tx];
+        }
+        __syncthreads();
+    }
+
+    if (row < N && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+"""
+
+
+class TestSimpleKernels:
+    """Test that simple kernel code can be used in requests."""
+
+    def test_vector_add_kernel(self) -> None:
+        """Test vector add kernel can be packaged in a request."""
+        request = CompileRequest(
+            files={"vector_add.cu": VECTOR_ADD_KERNEL},
+        )
+        assert "vector_add" in request.files["vector_add.cu"]
+
+    def test_matrix_multiply_kernel(self) -> None:
+        """Test matrix multiply kernel can be packaged in a request."""
+        request = CompileRequest(
+            files={"matmul.cu": MATRIX_MULTIPLY_KERNEL},
+        )
+        assert "matmul" in request.files["matmul.cu"]
+
+    def test_reduction_kernel(self) -> None:
+        """Test reduction kernel can be packaged in a request."""
+        request = CompileRequest(
+            files={"reduce.cu": REDUCTION_KERNEL},
+        )
+        assert "__shared__" in request.files["reduce.cu"]
+
+    def test_shared_memory_kernel(self) -> None:
+        """Test shared memory kernel can be packaged in a request."""
+        request = CompileRequest(
+            files={"tiled_matmul.cu": SHARED_MEMORY_KERNEL},
+        )
+        assert "TILE_SIZE" in request.files["tiled_matmul.cu"]
+
+
+# Multi-file test examples
+
+HEADER_FILE = """\
+#pragma once
+
+__device__ float square(float x) {
+    return x * x;
+}
+
+__device__ float cube(float x) {
+    return x * x * x;
+}
+"""
+
+MAIN_FILE_WITH_HEADER = """\
+#include "utils.cuh"
+
+__global__ void apply_square(float* data, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        data[idx] = square(data[idx]);
+    }
+}
+"""
+
+NESTED_HEADER_A = """\
+#pragma once
+
+#define MY_CONSTANT 42
+"""
+
+NESTED_HEADER_B = """\
+#pragma once
+
+#include "constants.cuh"
+
+__device__ int get_constant() {
+    return MY_CONSTANT;
+}
+"""
+
+MAIN_WITH_NESTED = """\
+#include "helper.cuh"
+
+__global__ void use_constant(int* out) {
+    *out = get_constant();
+}
+"""
+
+
+class TestMultiFile:
+    """Test multi-file compilation requests."""
+
+    def test_kernel_with_header(self) -> None:
+        """Test kernel that includes a header."""
+        request = CompileRequest(
+            files={
+                "main.cu": MAIN_FILE_WITH_HEADER,
+                "utils.cuh": HEADER_FILE,
+            },
+        )
+        assert len(request.files) == 2
+        assert request.main_cu_file == "main.cu"
+
+    def test_kernel_with_multiple_headers(self) -> None:
+        """Test kernel with multiple headers."""
+        extra_header = "__device__ float add(float a, float b) { return a + b; }"
+        request = CompileRequest(
+            files={
+                "main.cu": '#include "utils.cuh"\n#include "math.cuh"\n'
+                + MAIN_FILE_WITH_HEADER.split("\n", 1)[1],
+                "utils.cuh": HEADER_FILE,
+                "math.cuh": extra_header,
+            },
+        )
+        assert len(request.files) == 3
+
+    def test_nested_includes(self) -> None:
+        """Test nested header includes."""
+        request = CompileRequest(
+            files={
+                "main.cu": MAIN_WITH_NESTED,
+                "helper.cuh": NESTED_HEADER_B,
+                "constants.cuh": NESTED_HEADER_A,
+            },
+        )
+        assert len(request.files) == 3
+        assert "MY_CONSTANT" in request.files["constants.cuh"]
+
+    def test_relative_include_paths(self) -> None:
+        """Test that relative include paths are preserved in files dict."""
+        request = CompileRequest(
+            files={
+                "src/main.cu": '#include "../include/utils.cuh"\n__global__ void test() {}',
+                "include/utils.cuh": "// utils",
+            },
+        )
+        assert "src/main.cu" in request.files
+        assert "include/utils.cuh" in request.files
+
+
+# PyTorch extension examples
+
+PYTORCH_ACCESSOR_KERNEL = """\
+#include <torch/types.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+__global__ void fused_add_kernel(
+    const float* __restrict__ a,
+    const float* __restrict__ b,
+    float* __restrict__ out,
+    int64_t size
+) {
+    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < size) {
+        out[idx] = a[idx] + b[idx];
+    }
+}
+"""
+
+PYTORCH_PACKED_ACCESSOR_KERNEL = """\
+#include <torch/types.h>
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+template <typename scalar_t>
+__global__ void accessor_kernel(
+    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> input,
+    torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> output
+) {
+    int row = blockIdx.x * blockDim.x + threadIdx.x;
+    int col = blockIdx.y * blockDim.y + threadIdx.y;
+
+    if (row < input.size(0) && col < input.size(1)) {
+        output[row][col] = input[row][col] * 2;
+    }
+}
+"""
+
+
+class TestPyTorchHeaders:
+    """Test code using PyTorch CUDA types."""
+
+    def test_pytorch_accessor(self) -> None:
+        """Test PyTorch accessor kernel can be packaged."""
+        request = CompileRequest(
+            files={"pytorch_kernel.cu": PYTORCH_ACCESSOR_KERNEL},
+        )
+        assert "torch/types.h" in request.files["pytorch_kernel.cu"]
+
+    def test_pytorch_packed_accessor(self) -> None:
+        """Test PyTorch packed accessor kernel can be packaged."""
+        request = CompileRequest(
+            files={"accessor_kernel.cu": PYTORCH_PACKED_ACCESSOR_KERNEL},
+        )
+        assert "PackedTensorAccessor" in request.files["accessor_kernel.cu"]
+
+
+# CUTLASS examples
+
+CUTLASS_TYPES_KERNEL = """\
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_types.h>
+
+__global__ void use_cutlass_types() {
+    cutlass::half_t a = cutlass::half_t(1.0f);
+    cutlass::half_t b = cutlass::half_t(2.0f);
+    cutlass::half_t c = a + b;
+}
+"""
+
+CUTLASS_GEMM_INCLUDE = """\
+#include <cutlass/cutlass.h>
+#include <cutlass/gemm/device/gemm.h>
+
+// Just testing include paths work
+using ColumnMajor = cutlass::layout::ColumnMajor;
+"""
+
+
+class TestCUTLASSHeaders:
+    """Test code using CUTLASS headers."""
+
+    def test_cutlass_types(self) -> None:
+        """Test CUTLASS types kernel can be packaged."""
+        request = CompileRequest(
+            files={"cutlass_kernel.cu": CUTLASS_TYPES_KERNEL},
+        )
+        assert "cutlass/cutlass.h" in request.files["cutlass_kernel.cu"]
+
+    def test_cutlass_gemm_include(self) -> None:
+        """Test CUTLASS GEMM include can be packaged."""
+        request = CompileRequest(
+            files={"gemm_kernel.cu": CUTLASS_GEMM_INCLUDE},
+        )
+        assert "cutlass/gemm" in request.files["gemm_kernel.cu"]
+
+
+class TestOutputFormats:
+    """Test output format expectations for responses."""
+
+    def test_ptx_output_structure(self) -> None:
+        """Test PTX output has expected structure."""
+        # PTX typically starts with version/target info
+        mock_ptx = """\
+.version 8.0
+.target sm_90a
+.address_size 64
+
+.visible .entry vector_add(
+    .param .u64 a,
+    .param .u64 b,
+    .param .u64 c,
+    .param .u32 n
+)
+{
+    .reg .pred %p<2>;
+    .reg .f32 %f<4>;
+    // ...
+}
+"""
+        response = CompileResponse(
+            success=True,
+            ptx=mock_ptx,
+            compilation_time_ms=100,
+        )
+        assert response.ptx is not None
+        assert ".version" in response.ptx
+        assert ".target" in response.ptx
+        assert ".entry" in response.ptx
+
+    def test_sass_output_structure(self) -> None:
+        """Test SASS output has expected structure."""
+        # SASS is assembly with register operations
+        mock_sass = """\
+        code for sm_90a
+                Function : vector_add
+        .headerflags    @"EF_CUDA_SM90 EF_CUDA_PTX_SM(8.0)"
+        /*0000*/   MOV R1, c[0x0][0x28] ;
+        /*0010*/   S2R R0, SR_CTAID.X ;
+        /*0020*/   IMAD.U32 R0, R0, c[0x0][0x0], R0 ;
+"""
+        response = CompileResponse(
+            success=True,
+            sass=mock_sass,
+            compilation_time_ms=100,
+        )
+        assert response.sass is not None
+        assert "MOV" in response.sass or "code for" in response.sass
+
+    def test_both_outputs(self) -> None:
+        """Test response can have both PTX and SASS."""
+        response = CompileResponse(
+            success=True,
+            ptx=".version 8.0",
+            sass="MOV R0, R1",
+            compilation_time_ms=150,
+        )
+        assert response.ptx is not None
+        assert response.sass is not None
+
+
+class TestArchitectures:
+    """Test different GPU architecture specifications."""
+
+    def test_sm_90a_hopper(self) -> None:
+        """Test Hopper architecture request."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            arch="sm_90a",
+        )
+        assert request.arch == "sm_90a"
+
+    def test_sm_89_ada(self) -> None:
+        """Test Ada Lovelace architecture request."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            arch="sm_89",
+        )
+        assert request.arch == "sm_89"
+
+    def test_sm_80_ampere(self) -> None:
+        """Test Ampere architecture request."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            arch="sm_80",
+        )
+        assert request.arch == "sm_80"
+
+    def test_sm_100_blackwell(self) -> None:
+        """Test Blackwell architecture request."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            arch="sm_100",
+        )
+        assert request.arch == "sm_100"
+
+
+class TestCompilerFlags:
+    """Test custom nvcc compiler flags."""
+
+    def test_optimization_O3(self) -> None:
+        """Test O3 optimization flag."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            flags=("-O3",),
+        )
+        assert "-O3" in request.flags
+
+    def test_lineinfo_flag(self) -> None:
+        """Test lineinfo debug flag."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            flags=("-lineinfo",),
+        )
+        assert "-lineinfo" in request.flags
+
+    def test_maxrregcount(self) -> None:
+        """Test maxrregcount register limit flag."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            flags=("--maxrregcount=64",),
+        )
+        assert "--maxrregcount=64" in request.flags
+
+    def test_multiple_flags(self) -> None:
+        """Test multiple compiler flags."""
+        request = CompileRequest(
+            files={"kernel.cu": "__global__ void test() {}"},
+            flags=("-O3", "-lineinfo", "--maxrregcount=64"),
+        )
+        assert len(request.flags) == 3
+
+
+class TestErrors:
+    """Test compilation error reporting."""
+
+    def test_syntax_error_reported(self) -> None:
+        """Test syntax error is included in response."""
+        response = CompileResponse.error(
+            "kernel.cu(10): error: expected a ';'"
+        )
+        assert response.success is False
+        assert "expected a ';'" in response.stderr
+
+    def test_missing_include_error(self) -> None:
+        """Test missing include error."""
+        response = CompileResponse.error(
+            "kernel.cu(1): fatal error: cannot open include file: missing.h"
+        )
+        assert response.success is False
+        assert "cannot open include file" in response.stderr
+
+    def test_undefined_symbol_error(self) -> None:
+        """Test undefined symbol error."""
+        response = CompileResponse.error(
+            "kernel.cu(5): error: identifier 'undefined_func' is undefined"
+        )
+        assert response.success is False
+        assert "undefined" in response.stderr
+
+    def test_error_line_number_included(self) -> None:
+        """Test that error messages include line numbers."""
+        response = CompileResponse.error(
+            "kernel.cu(42): error: some error message"
+        )
+        assert "(42)" in response.stderr

wafer_core/tools/compile/tests/test_data/utils.cuh

@@ -0,0 +1,10 @@
+// Utility functions header
+#pragma once
+
+__device__ float square(float x) {
+    return x * x;
+}
+
+__device__ float cube(float x) {
+    return x * x * x;
+}

wafer_core/tools/compile/tests/test_data/vector_add.cu

@@ -0,0 +1,7 @@
+// Simple vector addition kernel
+__global__ void vector_add(float* a, float* b, float* c, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        c[idx] = a[idx] + b[idx];
+    }
+}

wafer_core/tools/compile/tests/test_data/with_header.cu

@@ -0,0 +1,9 @@
+// Kernel using header
+#include "utils.cuh"
+
+__global__ void apply_square(float* data, int n) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < n) {
+        data[idx] = square(data[idx]);
+    }
+}