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

wafer_core/tools/compile/benchmark.py

@@ -0,0 +1,636 @@
+"""Benchmark script for CUDA compilation performance.
+
+This script measures compilation time for kernels of different sizes
+to track performance improvements from optimizations.
+
+Usage:
+    python -m wafer_core.tools.compile.benchmark
+    
+    # Or with specific test:
+    python -m wafer_core.tools.compile.benchmark --kernel simple
+    python -m wafer_core.tools.compile.benchmark --kernel medium
+    python -m wafer_core.tools.compile.benchmark --kernel complex
+"""
+
+import argparse
+import statistics
+import time
+from typing import NamedTuple
+
+# ============================================================================
+# Test Kernels
+# ============================================================================
+
+SIMPLE_KERNEL = """\
+// Simple vector addition kernel (~20 lines)
+__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];
+    }
+}
+"""
+
+MEDIUM_KERNEL = """\
+// Medium complexity kernel with shared memory (~100 lines)
+#include <cuda_runtime.h>
+
+#define TILE_SIZE 16
+
+__global__ void tiled_matmul(
+    const float* __restrict__ A,
+    const float* __restrict__ B,
+    float* __restrict__ C,
+    int M, int N, int K
+) {
+    __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 < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        // Load tile from A
+        if (row < M && t * TILE_SIZE + tx < K) {
+            As[ty][tx] = A[row * K + t * TILE_SIZE + tx];
+        } else {
+            As[ty][tx] = 0.0f;
+        }
+
+        // Load tile from B
+        if (t * TILE_SIZE + ty < K && col < N) {
+            Bs[ty][tx] = B[(t * TILE_SIZE + ty) * N + col];
+        } else {
+            Bs[ty][tx] = 0.0f;
+        }
+
+        __syncthreads();
+
+        // Compute partial dot product
+        #pragma unroll
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum = fmaf(As[ty][k], Bs[k][tx], sum);
+        }
+
+        __syncthreads();
+    }
+
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+// Reduction kernel
+__global__ void reduce_sum(const float* input, float* output, int n) {
+    extern __shared__ float sdata[];
+
+    unsigned int tid = threadIdx.x;
+    unsigned int i = blockIdx.x * blockDim.x * 2 + threadIdx.x;
+
+    float mySum = (i < n) ? input[i] : 0.0f;
+    if (i + blockDim.x < n) {
+        mySum += input[i + blockDim.x];
+    }
+    sdata[tid] = mySum;
+    __syncthreads();
+
+    for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {
+        if (tid < s) {
+            sdata[tid] = mySum = mySum + sdata[tid + s];
+        }
+        __syncthreads();
+    }
+
+    if (tid < 32) {
+        volatile float* smem = sdata;
+        smem[tid] = mySum = mySum + smem[tid + 32];
+        smem[tid] = mySum = mySum + smem[tid + 16];
+        smem[tid] = mySum = mySum + smem[tid + 8];
+        smem[tid] = mySum = mySum + smem[tid + 4];
+        smem[tid] = mySum = mySum + smem[tid + 2];
+        smem[tid] = mySum = mySum + smem[tid + 1];
+    }
+
+    if (tid == 0) {
+        output[blockIdx.x] = sdata[0];
+    }
+}
+"""
+
+COMPLEX_KERNEL = """\
+// Complex kernel with multiple features (~500 lines)
+#include <cuda_runtime.h>
+#include <cooperative_groups.h>
+
+namespace cg = cooperative_groups;
+
+// Constants
+constexpr int BLOCK_SIZE = 256;
+constexpr int TILE_SIZE = 16;
+
+// ============================================================================
+// Kernel 1: Vector operations with shared memory and reduction
+// ============================================================================
+template <typename T, int BlockSize>
+__global__ void reduceSum(const T* __restrict__ input, T* __restrict__ output, int N) {
+    __shared__ T sdata[BlockSize];
+    
+    unsigned int tid = threadIdx.x;
+    unsigned int i = blockIdx.x * (blockDim.x * 2) + threadIdx.x;
+    
+    T mySum = (i < N) ? input[i] : T(0);
+    if (i + blockDim.x < N) {
+        mySum += input[i + blockDim.x];
+    }
+    sdata[tid] = mySum;
+    __syncthreads();
+    
+    #pragma unroll
+    for (unsigned int s = blockDim.x / 2; s > 32; s >>= 1) {
+        if (tid < s) {
+            sdata[tid] = mySum = mySum + sdata[tid + s];
+        }
+        __syncthreads();
+    }
+    
+    if (tid < 32) {
+        volatile T* smem = sdata;
+        if (BlockSize >= 64) mySum += smem[tid + 32];
+        smem[tid] = mySum;
+        if (BlockSize >= 32) mySum += smem[tid + 16];
+        smem[tid] = mySum;
+        if (BlockSize >= 16) mySum += smem[tid + 8];
+        smem[tid] = mySum;
+        if (BlockSize >= 8) mySum += smem[tid + 4];
+        smem[tid] = mySum;
+        if (BlockSize >= 4) mySum += smem[tid + 2];
+        smem[tid] = mySum;
+        if (BlockSize >= 2) mySum += smem[tid + 1];
+        smem[tid] = mySum;
+    }
+    
+    if (tid == 0) {
+        output[blockIdx.x] = sdata[0];
+    }
+}
+
+// ============================================================================
+// Kernel 2: Matrix transpose with shared memory
+// ============================================================================
+__global__ void matrixTranspose(const float* __restrict__ input, 
+                                float* __restrict__ output,
+                                int width, int height) {
+    __shared__ float tile[TILE_SIZE][TILE_SIZE + 1];
+    
+    int xIndex = blockIdx.x * TILE_SIZE + threadIdx.x;
+    int yIndex = blockIdx.y * TILE_SIZE + threadIdx.y;
+    
+    if (xIndex < width && yIndex < height) {
+        tile[threadIdx.y][threadIdx.x] = input[yIndex * width + xIndex];
+    }
+    
+    __syncthreads();
+    
+    xIndex = blockIdx.y * TILE_SIZE + threadIdx.x;
+    yIndex = blockIdx.x * TILE_SIZE + threadIdx.y;
+    
+    if (xIndex < height && yIndex < width) {
+        output[yIndex * height + xIndex] = tile[threadIdx.x][threadIdx.y];
+    }
+}
+
+// ============================================================================
+// Kernel 3: Softmax with cooperative groups
+// ============================================================================
+__global__ void softmaxKernel(const float* __restrict__ input,
+                              float* __restrict__ output,
+                              int N, int stride) {
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
+    
+    extern __shared__ float shared[];
+    
+    int row = blockIdx.x;
+    const float* rowInput = input + row * stride;
+    float* rowOutput = output + row * stride;
+    
+    float maxVal = -INFINITY;
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        maxVal = fmaxf(maxVal, rowInput[i]);
+    }
+    
+    for (int offset = warp.size() / 2; offset > 0; offset /= 2) {
+        maxVal = fmaxf(maxVal, warp.shfl_down(maxVal, offset));
+    }
+    
+    if (warp.thread_rank() == 0) {
+        shared[threadIdx.x / 32] = maxVal;
+    }
+    block.sync();
+    
+    if (threadIdx.x < blockDim.x / 32) {
+        maxVal = shared[threadIdx.x];
+    } else {
+        maxVal = -INFINITY;
+    }
+    
+    for (int offset = 16; offset > 0; offset /= 2) {
+        maxVal = fmaxf(maxVal, __shfl_down_sync(0xffffffff, maxVal, offset));
+    }
+    maxVal = __shfl_sync(0xffffffff, maxVal, 0);
+    
+    float sum = 0.0f;
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        float val = expf(rowInput[i] - maxVal);
+        rowOutput[i] = val;
+        sum += val;
+    }
+    
+    for (int offset = warp.size() / 2; offset > 0; offset /= 2) {
+        sum += warp.shfl_down(sum, offset);
+    }
+    
+    if (warp.thread_rank() == 0) {
+        shared[threadIdx.x / 32] = sum;
+    }
+    block.sync();
+    
+    if (threadIdx.x < blockDim.x / 32) {
+        sum = shared[threadIdx.x];
+    } else {
+        sum = 0.0f;
+    }
+    
+    for (int offset = 16; offset > 0; offset /= 2) {
+        sum += __shfl_down_sync(0xffffffff, sum, offset);
+    }
+    sum = __shfl_sync(0xffffffff, sum, 0);
+    
+    float invSum = 1.0f / sum;
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        rowOutput[i] *= invSum;
+    }
+}
+
+// ============================================================================
+// Kernel 4: Fused multiply-add with vectorized loads
+// ============================================================================
+__global__ void fusedMulAddVec4(const float4* __restrict__ A,
+                                 const float4* __restrict__ B,
+                                 float4* __restrict__ C,
+                                 float alpha, float beta, int N) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    
+    if (idx < N) {
+        float4 a = A[idx];
+        float4 b = B[idx];
+        float4 c;
+        
+        c.x = fmaf(alpha, a.x, beta * b.x);
+        c.y = fmaf(alpha, a.y, beta * b.y);
+        c.z = fmaf(alpha, a.z, beta * b.z);
+        c.w = fmaf(alpha, a.w, beta * b.w);
+        
+        C[idx] = c;
+    }
+}
+
+// ============================================================================
+// Kernel 5: Simple GEMM with shared memory tiling
+// ============================================================================
+__global__ void matmulTiled(const float* __restrict__ A,
+                            const float* __restrict__ B,
+                            float* __restrict__ C,
+                            int M, int N, int K) {
+    __shared__ float As[TILE_SIZE][TILE_SIZE];
+    __shared__ float Bs[TILE_SIZE][TILE_SIZE];
+    
+    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
+    int col = blockIdx.x * TILE_SIZE + threadIdx.x;
+    
+    float sum = 0.0f;
+    
+    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
+        int tiledCol = t * TILE_SIZE + threadIdx.x;
+        int tiledRow = t * TILE_SIZE + threadIdx.y;
+        
+        As[threadIdx.y][threadIdx.x] = (row < M && tiledCol < K) ? 
+                                        A[row * K + tiledCol] : 0.0f;
+        Bs[threadIdx.y][threadIdx.x] = (tiledRow < K && col < N) ? 
+                                        B[tiledRow * N + col] : 0.0f;
+        
+        __syncthreads();
+        
+        #pragma unroll
+        for (int k = 0; k < TILE_SIZE; k++) {
+            sum = fmaf(As[threadIdx.y][k], Bs[k][threadIdx.x], sum);
+        }
+        
+        __syncthreads();
+    }
+    
+    if (row < M && col < N) {
+        C[row * N + col] = sum;
+    }
+}
+
+// ============================================================================
+// Device helper functions
+// ============================================================================
+__device__ __forceinline__ float warpReduceSum(float val) {
+    for (int offset = 16; offset > 0; offset /= 2) {
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    }
+    return val;
+}
+
+__device__ __forceinline__ float blockReduceSum(float val) {
+    __shared__ float shared[32];
+    
+    int lane = threadIdx.x % 32;
+    int wid = threadIdx.x / 32;
+    
+    val = warpReduceSum(val);
+    
+    if (lane == 0) shared[wid] = val;
+    __syncthreads();
+    
+    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
+    
+    if (wid == 0) val = warpReduceSum(val);
+    
+    return val;
+}
+
+// ============================================================================
+// Kernel 6: Layer normalization
+// ============================================================================
+__global__ void layerNorm(const float* __restrict__ input,
+                          const float* __restrict__ gamma,
+                          const float* __restrict__ beta,
+                          float* __restrict__ output,
+                          int N, float eps) {
+    int row = blockIdx.x;
+    const float* rowInput = input + row * N;
+    float* rowOutput = output + row * N;
+    
+    float sum = 0.0f;
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        sum += rowInput[i];
+    }
+    sum = blockReduceSum(sum);
+    
+    __shared__ float s_mean, s_var;
+    if (threadIdx.x == 0) {
+        s_mean = sum / N;
+    }
+    __syncthreads();
+    
+    float var = 0.0f;
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        float diff = rowInput[i] - s_mean;
+        var += diff * diff;
+    }
+    var = blockReduceSum(var);
+    
+    if (threadIdx.x == 0) {
+        s_var = rsqrtf(var / N + eps);
+    }
+    __syncthreads();
+    
+    for (int i = threadIdx.x; i < N; i += blockDim.x) {
+        float normalized = (rowInput[i] - s_mean) * s_var;
+        rowOutput[i] = fmaf(normalized, gamma[i], beta[i]);
+    }
+}
+"""
+
+
+# ============================================================================
+# Benchmark Results
+# ============================================================================
+
+class BenchmarkResult(NamedTuple):
+    """Result of a single benchmark run."""
+    kernel_name: str
+    kernel_lines: int
+    compile_time_ms: int
+    success: bool
+    ptx_lines: int | None
+    sass_lines: int | None
+    error: str | None
+
+
+def count_lines(code: str) -> int:
+    """Count non-empty lines in code."""
+    return len([line for line in code.split('\n') if line.strip()])
+
+
+def run_benchmark(
+    kernel_name: str,
+    kernel_code: str,
+    arch: str = "sm_90a",
+    output_formats: list[str] | None = None,
+    num_runs: int = 3,
+) -> list[BenchmarkResult]:
+    """Run benchmark for a kernel.
+    
+    Args:
+        kernel_name: Name of the kernel for reporting
+        kernel_code: CUDA source code
+        arch: Target architecture
+        output_formats: Output formats to request (default: ["ptx", "sass"])
+        num_runs: Number of benchmark runs
+        
+    Returns:
+        List of BenchmarkResult for each run
+    """
+    import modal
+    
+    if output_formats is None:
+        output_formats = ["ptx", "sass"]
+    
+    # Get the deployed function
+    compile_fn = modal.Function.from_name("cuda-compile", "compile_cuda")
+    
+    kernel_lines = count_lines(kernel_code)
+    results: list[BenchmarkResult] = []
+    
+    for run in range(num_runs):
+        print(f"  Run {run + 1}/{num_runs}...", end=" ", flush=True)
+        
+        start_time = time.time()
+        
+        try:
+            result = compile_fn.remote({
+                "files": {"kernel.cu": kernel_code},
+                "arch": arch,
+                "flags": ["-O3", "-lineinfo"],
+                "output": output_formats,
+            })
+            
+            elapsed_ms = int((time.time() - start_time) * 1000)
+            
+            if result["success"]:
+                ptx_lines = count_lines(result["ptx"]) if result.get("ptx") else None
+                sass_lines = count_lines(result["sass"]) if result.get("sass") else None
+                
+                results.append(BenchmarkResult(
+                    kernel_name=kernel_name,
+                    kernel_lines=kernel_lines,
+                    compile_time_ms=elapsed_ms,
+                    success=True,
+                    ptx_lines=ptx_lines,
+                    sass_lines=sass_lines,
+                    error=None,
+                ))
+                print(f"{elapsed_ms}ms")
+            else:
+                results.append(BenchmarkResult(
+                    kernel_name=kernel_name,
+                    kernel_lines=kernel_lines,
+                    compile_time_ms=elapsed_ms,
+                    success=False,
+                    ptx_lines=None,
+                    sass_lines=None,
+                    error=result.get("stderr", "Unknown error"),
+                ))
+                print(f"FAILED ({elapsed_ms}ms)")
+                
+        except Exception as e:
+            elapsed_ms = int((time.time() - start_time) * 1000)
+            results.append(BenchmarkResult(
+                kernel_name=kernel_name,
+                kernel_lines=kernel_lines,
+                compile_time_ms=elapsed_ms,
+                success=False,
+                ptx_lines=None,
+                sass_lines=None,
+                error=str(e),
+            ))
+            print(f"ERROR: {e}")
+    
+    return results
+
+
+def print_summary(results: list[BenchmarkResult]) -> None:
+    """Print benchmark summary."""
+    successful = [r for r in results if r.success]
+    
+    if not successful:
+        print("\n  No successful runs!")
+        if results:
+            print(f"  Error: {results[0].error}")
+        return
+    
+    times = [r.compile_time_ms for r in successful]
+    mean_time = statistics.mean(times)
+    
+    if len(times) > 1:
+        stdev = statistics.stdev(times)
+        min_time = min(times)
+        max_time = max(times)
+        print(f"\n  Results: {mean_time:.0f}ms avg (min: {min_time}ms, max: {max_time}ms, stdev: {stdev:.0f}ms)")
+    else:
+        print(f"\n  Results: {mean_time:.0f}ms")
+    
+    # Show output sizes
+    if successful[0].ptx_lines:
+        print(f"  PTX output: {successful[0].ptx_lines} lines")
+    if successful[0].sass_lines:
+        print(f"  SASS output: {successful[0].sass_lines} lines")
+
+
+def run_all_benchmarks(num_runs: int = 3) -> dict[str, list[BenchmarkResult]]:
+    """Run benchmarks for all kernel sizes."""
+    print("=" * 60)
+    print("CUDA Compilation Benchmark")
+    print("=" * 60)
+    
+    kernels = [
+        ("simple", SIMPLE_KERNEL),
+        ("medium", MEDIUM_KERNEL),
+        ("complex", COMPLEX_KERNEL),
+    ]
+    
+    all_results: dict[str, list[BenchmarkResult]] = {}
+    
+    for name, code in kernels:
+        lines = count_lines(code)
+        print(f"\n{name.upper()} KERNEL ({lines} lines)")
+        print("-" * 40)
+        
+        results = run_benchmark(name, code, num_runs=num_runs)
+        all_results[name] = results
+        print_summary(results)
+    
+    # Print final summary
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    
+    for name in ["simple", "medium", "complex"]:
+        results = all_results.get(name, [])
+        successful = [r for r in results if r.success]
+        if successful:
+            avg_time = statistics.mean([r.compile_time_ms for r in successful])
+            print(f"  {name:10s}: {avg_time:6.0f}ms ({results[0].kernel_lines} lines)")
+        else:
+            print(f"  {name:10s}: FAILED")
+    
+    return all_results
+
+
+def main() -> None:
+    """Main entry point."""
+    parser = argparse.ArgumentParser(description="Benchmark CUDA compilation")
+    parser.add_argument(
+        "--kernel",
+        choices=["simple", "medium", "complex", "all"],
+        default="all",
+        help="Which kernel to benchmark",
+    )
+    parser.add_argument(
+        "--runs",
+        type=int,
+        default=3,
+        help="Number of benchmark runs per kernel",
+    )
+    parser.add_argument(
+        "--arch",
+        default="sm_90a",
+        help="Target GPU architecture",
+    )
+    
+    args = parser.parse_args()
+    
+    if args.kernel == "all":
+        run_all_benchmarks(num_runs=args.runs)
+    else:
+        kernel_map = {
+            "simple": SIMPLE_KERNEL,
+            "medium": MEDIUM_KERNEL,
+            "complex": COMPLEX_KERNEL,
+        }
+        code = kernel_map[args.kernel]
+        lines = count_lines(code)
+        
+        print(f"\n{args.kernel.upper()} KERNEL ({lines} lines)")
+        print("-" * 40)
+        
+        results = run_benchmark(
+            args.kernel,
+            code,
+            arch=args.arch,
+            num_runs=args.runs,
+        )
+        print_summary(results)
+
+
+if __name__ == "__main__":
+    main()