markbaseengine/Sources/MarkBase/MetalKernels.swift

import Foundation

/// Single source of truth for Metal kernel source code.
/// Tests use these constants instead of duplicating inline strings.
public enum MetalKernels {
    public static let vectorAdd = """
    #include <metal_stdlib>
    using namespace metal;
    kernel void vector_add(
        device const float *a [[buffer(0)]],
        device const float *b [[buffer(1)]],
        device float *c [[buffer(2)]],
        constant uint &n [[buffer(3)]],
        uint id [[thread_position_in_grid]]
    ) {
        if (id < n) c[id] = a[id] + b[id];
    }
    """

    public static let matmulF32 = """
    #include <metal_stdlib>
    using namespace metal;
    kernel void matmul_f32(
        device const float *A [[buffer(0)]],
        device const float *B [[buffer(1)]],
        device float *C [[buffer(2)]],
        constant uint &M [[buffer(3)]],
        constant uint &N [[buffer(4)]],
        constant uint &K [[buffer(5)]],
        uint2 gid [[thread_position_in_grid]]
    ) {
        if (gid.x >= N || gid.y >= M) return;
        float sum = 0.0;
        for (uint k = 0; k < K; k++)
            sum += A[gid.y * K + k] * B[k * N + gid.x];
        C[gid.y * N + gid.x] = sum;
    }
    """

    /// Full E4B inference kernel source (reads from .metal file at runtime).
    /// Use for JIT compilation in tests.
    public static var e4bSource: String {
        let url = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent()
            .appendingPathComponent("Metal/MetalKernels.metal")
        return try! String(contentsOf: url, encoding: .utf8)
    }
    
    /// Optimized SIMD kernel source for Phase 1.
    /// Includes attention, matmul, and norm optimizations.
    public static var optimizedSource: String {
        let url = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent()
            .appendingPathComponent("Metal/OptimizedKernels.metal")
        return try! String(contentsOf: url, encoding: .utf8)
    }
    
    /// Combined source: original + optimized kernels.
    /// Use for production deployment.
    public static var combinedSource: String {
        return e4bSource + "\n" + optimizedSource
    }
    
    /// Fusion kernel source for Phase 1.3.
    /// Includes kernel fusion optimizations.
    public static var fusionSource: String {
        let url = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent()
            .appendingPathComponent("Metal/FusionKernels.metal")
        return try! String(contentsOf: url, encoding: .utf8)
    }
    
    /// Fused kernel source for Phase 2.
    /// Includes advanced kernel fusion for embedding, norm, and matmul.
    public static var fusedKernelsSource: String {
        let url = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent()
            .appendingPathComponent("Metal/FusedKernels.metal")
        return try! String(contentsOf: url, encoding: .utf8)
    }
    
    /// Full optimized source: original + SIMD + fusion.
    /// Maximum optimization without MPS.
    public static var fullOptimizedSource: String {
        return combinedSource + "\n" + fusionSource + "\n" + fusedKernelsSource
    }
    
    /// Full optimized source with all kernels.
    /// Strips duplicate #include and using namespace from subsequent files.
    public static var fullOptimizedSourceWithFusion: String {
        // Start with first file (includes its #include and using namespace)
        var result = e4bSource
        
        // Strip preamble from optimized source
        let optStripped = optimizedSource
            .replacingOccurrences(of: "#include <metal_stdlib>\n", with: "")
            .replacingOccurrences(of: "using namespace metal;\n", with: "")
        result += "\n" + optStripped
        
        // Strip preamble from fusion source
        let fusionStripped = fusionSource
            .replacingOccurrences(of: "#include <metal_stdlib>\n", with: "")
            .replacingOccurrences(of: "using namespace metal;\n", with: "")
        result += "\n" + fusionStripped
        
        // Strip preamble from fused kernels source
        let fusedStripped = fusedKernelsSource
            .replacingOccurrences(of: "#include <metal_stdlib>\n", with: "")
            .replacingOccurrences(of: "using namespace metal;\n", with: "")
        result += "\n" + fusedStripped
        
        // Strip preamble from embedding kernels source
        let embStripped = embeddingKernelsSource
            .replacingOccurrences(of: "#include <metal_stdlib>\n", with: "")
            .replacingOccurrences(of: "using namespace metal;\n", with: "")
        result += "\n" + embStripped
        
        return result
    }

    /// Embedding kernel source for EmbeddingGemma.
    /// Includes RoPE, bidirectional sliding window attention, Q/K norm, and GELU.
    public static var embeddingKernelsSource: String {
        let url = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent()
            .appendingPathComponent("Metal/EmbeddingKernels.metal")
        return try! String(contentsOf: url, encoding: .utf8)
    }
}