v2: Initial clean branch with unit tests + CI/CD pipeline
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,289 @@
|
||||
@preconcurrency import Foundation
|
||||
@preconcurrency import Metal
|
||||
|
||||
public enum E4BError: Error, LocalizedError {
|
||||
case noMetalDevice
|
||||
case libraryNotFound
|
||||
case sourceCompilationFailed(String)
|
||||
case pipelineCreationFailed(String)
|
||||
case pipelineNotFound(String)
|
||||
case bufferCreationFailed
|
||||
case commandBufferCreationFailed
|
||||
case commandEncoderCreationFailed
|
||||
|
||||
public var errorDescription: String? {
|
||||
switch self {
|
||||
case .noMetalDevice: return "No Metal-capable GPU found"
|
||||
case .libraryNotFound: return "Metal library not found"
|
||||
case .sourceCompilationFailed(let detail): return "Metal source compilation failed: \(detail)"
|
||||
case .pipelineCreationFailed(let detail): return "Pipeline creation failed: \(detail)"
|
||||
case .pipelineNotFound(let name): return "Kernel '\(name)' not found in library"
|
||||
case .bufferCreationFailed: return "Failed to allocate Metal buffer"
|
||||
case .commandBufferCreationFailed: return "Failed to create command buffer"
|
||||
case .commandEncoderCreationFailed: return "Failed to create command encoder"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pure-Swift Metal engine with pipeline cache, proper threadgroup sizing,
|
||||
/// buffer pool for reuse, and support for both runtime source compilation and pre-compiled metallib.
|
||||
public final class MarkBaseEngine: @unchecked Sendable {
|
||||
public let device: MTLDevice
|
||||
public let commandQueue: MTLCommandQueue
|
||||
public let bufferPool: BufferPool
|
||||
public private(set) var library: MTLLibrary?
|
||||
private var pipelineCache: [String: MTLComputePipelineState] = [:]
|
||||
|
||||
public init() throws {
|
||||
guard let device = MTLCreateSystemDefaultDevice() else {
|
||||
throw E4BError.noMetalDevice
|
||||
}
|
||||
self.device = device
|
||||
guard let queue = device.makeCommandQueue() else {
|
||||
throw E4BError.bufferCreationFailed
|
||||
}
|
||||
self.commandQueue = queue
|
||||
self.bufferPool = BufferPool(device: device)
|
||||
}
|
||||
|
||||
/// Initialize with Metal kernels auto-compiled from source.
|
||||
/// Loads original, optimized, and fusion kernels.
|
||||
public convenience init(autoCompile: Bool = true) throws {
|
||||
try self.init()
|
||||
if autoCompile {
|
||||
try compileSource(MetalKernels.fullOptimizedSourceWithFusion)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Library ──────────────────────────────────────
|
||||
|
||||
/// Compile Metal source at runtime.
|
||||
public func compileSource(_ source: String) throws {
|
||||
pipelineCache.removeAll()
|
||||
library = try device.makeLibrary(source: source, options: nil)
|
||||
}
|
||||
|
||||
/// Load a pre-compiled .metallib from a file path.
|
||||
public func loadMetallib(path: String) throws {
|
||||
let data = try Data(contentsOf: URL(fileURLWithPath: path))
|
||||
try loadMetallib(data: data)
|
||||
}
|
||||
|
||||
/// Load a pre-compiled .metallib from Data.
|
||||
public func loadMetallib(data: Data) throws {
|
||||
pipelineCache.removeAll()
|
||||
library = try data.withUnsafeBytes { rawPtr in
|
||||
try device.makeLibrary(data: DispatchData(bytes: rawPtr))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Pipeline (with cache) ─────────────────────────
|
||||
|
||||
/// Return a cached or newly-created compute pipeline state.
|
||||
public func pipeline(named name: String) throws -> MTLComputePipelineState {
|
||||
if let cached = pipelineCache[name] { return cached }
|
||||
guard let lib = library else { throw E4BError.libraryNotFound }
|
||||
guard let fn = lib.makeFunction(name: name) else {
|
||||
throw E4BError.pipelineNotFound(name)
|
||||
}
|
||||
let pso = try device.makeComputePipelineState(function: fn)
|
||||
pipelineCache[name] = pso
|
||||
return pso
|
||||
}
|
||||
|
||||
/// Clear the pipeline cache (e.g. after recompiling the library).
|
||||
public func clearPipelineCache() { pipelineCache.removeAll() }
|
||||
|
||||
// ── Buffers ───────────────────────────────────────
|
||||
|
||||
/// Acquire a buffer from the pool (or create new if none available)
|
||||
public func acquireBuffer(length: Int) -> MTLBuffer {
|
||||
return bufferPool.acquire(length: length)
|
||||
}
|
||||
|
||||
/// Release a buffer back to the pool for reuse
|
||||
public func releaseBuffer(_ buffer: MTLBuffer) {
|
||||
bufferPool.release(buffer)
|
||||
}
|
||||
|
||||
public func makeBuffer<T>(_ values: [T]) throws -> MTLBuffer {
|
||||
let count = values.count * MemoryLayout<T>.stride
|
||||
guard let buf = values.withUnsafeBytes({ rawPtr in
|
||||
device.makeBuffer(bytes: rawPtr.baseAddress!, length: count,
|
||||
options: .storageModeShared)
|
||||
}) else { throw E4BError.bufferCreationFailed }
|
||||
return buf
|
||||
}
|
||||
|
||||
public func makeBuffer(length: Int) throws -> MTLBuffer {
|
||||
guard let buf = device.makeBuffer(length: length, options: .storageModeShared) else {
|
||||
throw E4BError.bufferCreationFailed
|
||||
}
|
||||
return buf
|
||||
}
|
||||
|
||||
// ── Threadgroup sizing ────────────────────────────
|
||||
|
||||
/// Optimal 1D threadgroup size for a given pipeline.
|
||||
public func threadgroupSize1D(_ pipeline: MTLComputePipelineState,
|
||||
count: Int) -> MTLSize {
|
||||
let w = min(pipeline.maxTotalThreadsPerThreadgroup, count)
|
||||
return MTLSize(width: w, height: 1, depth: 1)
|
||||
}
|
||||
|
||||
/// Optimal 2D threadgroup size for a given pipeline.
|
||||
public func threadgroupSize2D(_ pipeline: MTLComputePipelineState,
|
||||
grid: (width: Int, height: Int)) -> MTLSize {
|
||||
let w = pipeline.threadExecutionWidth
|
||||
let h = pipeline.maxTotalThreadsPerThreadgroup / w
|
||||
return MTLSize(
|
||||
width: min(w, grid.width),
|
||||
height: min(h, grid.height),
|
||||
depth: 1
|
||||
)
|
||||
}
|
||||
|
||||
// ── Synchronous dispatch ──────────────────────────
|
||||
|
||||
/// Synchronous 1D dispatch with proper threadgroup sizing.
|
||||
@discardableResult
|
||||
public func dispatch1D(
|
||||
_ pipeline: MTLComputePipelineState,
|
||||
buffers: [MTLBuffer],
|
||||
count: Int
|
||||
) throws -> MTLCommandBuffer {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
let tg = threadgroupSize1D(pipeline, count: count)
|
||||
enc.setComputePipelineState(pipeline)
|
||||
for (i, buf) in buffers.enumerated() { enc.setBuffer(buf, offset: 0, index: i) }
|
||||
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
||||
threadsPerThreadgroup: tg)
|
||||
enc.endEncoding()
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted()
|
||||
return cmdBuf
|
||||
}
|
||||
|
||||
/// Synchronous 2D dispatch with proper threadgroup sizing.
|
||||
@discardableResult
|
||||
public func dispatch2D(
|
||||
_ pipeline: MTLComputePipelineState,
|
||||
buffers: [MTLBuffer],
|
||||
grid: (width: Int, height: Int)
|
||||
) throws -> MTLCommandBuffer {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
let tg = threadgroupSize2D(pipeline, grid: grid)
|
||||
enc.setComputePipelineState(pipeline)
|
||||
for (i, buf) in buffers.enumerated() { enc.setBuffer(buf, offset: 0, index: i) }
|
||||
enc.dispatchThreads(MTLSize(width: grid.width, height: grid.height, depth: 1),
|
||||
threadsPerThreadgroup: tg)
|
||||
enc.endEncoding()
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted()
|
||||
return cmdBuf
|
||||
}
|
||||
|
||||
// ── Async dispatch ────────────────────────────────
|
||||
|
||||
/// Batch dispatch: execute multiple kernel operations in a single command buffer
|
||||
/// Returns after all GPU work completes
|
||||
public func batchDispatch(_ operations: [(MTLComputePipelineState, [MTLBuffer], MTLSize)]) throws {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
|
||||
for (pso, buffers, gridSize) in operations {
|
||||
enc.setComputePipelineState(pso)
|
||||
for (i, buf) in buffers.enumerated() {
|
||||
enc.setBuffer(buf, offset: 0, index: i)
|
||||
}
|
||||
enc.dispatchThreads(gridSize, threadsPerThreadgroup: threadgroupSize1D(pso, count: Int(gridSize.width)))
|
||||
}
|
||||
|
||||
enc.endEncoding()
|
||||
cmdBuf.commit()
|
||||
cmdBuf.waitUntilCompleted()
|
||||
}
|
||||
|
||||
/// Create a batch encoder for manual use
|
||||
/// Caller is responsible for calling endEncoding() and commit/wait on returned command buffer
|
||||
public func makeBatchEncoder() throws -> (MTLCommandBuffer, MTLComputeCommandEncoder) {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
return (cmdBuf, enc)
|
||||
}
|
||||
|
||||
/// Asynchronous 1D dispatch with proper threadgroup sizing.
|
||||
/// Returns when GPU work completes; buffer contents are safe to read after.
|
||||
public func dispatch1DAsync(
|
||||
_ pipeline: MTLComputePipelineState,
|
||||
buffers: [MTLBuffer],
|
||||
count: Int
|
||||
) async throws {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
let tg = threadgroupSize1D(pipeline, count: count)
|
||||
enc.setComputePipelineState(pipeline)
|
||||
for (i, buf) in buffers.enumerated() { enc.setBuffer(buf, offset: 0, index: i) }
|
||||
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
|
||||
threadsPerThreadgroup: tg)
|
||||
enc.endEncoding()
|
||||
return await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
|
||||
cmdBuf.addCompletedHandler { _ in continuation.resume() }
|
||||
cmdBuf.commit()
|
||||
}
|
||||
}
|
||||
|
||||
/// Asynchronous 2D dispatch with proper threadgroup sizing.
|
||||
/// Returns when GPU work completes; buffer contents are safe to read after.
|
||||
public func dispatch2DAsync(
|
||||
_ pipeline: MTLComputePipelineState,
|
||||
buffers: [MTLBuffer],
|
||||
grid: (width: Int, height: Int)
|
||||
) async throws {
|
||||
guard let cmdBuf = commandQueue.makeCommandBuffer() else {
|
||||
throw E4BError.commandBufferCreationFailed
|
||||
}
|
||||
guard let enc = cmdBuf.makeComputeCommandEncoder() else {
|
||||
throw E4BError.commandEncoderCreationFailed
|
||||
}
|
||||
let tg = threadgroupSize2D(pipeline, grid: grid)
|
||||
enc.setComputePipelineState(pipeline)
|
||||
for (i, buf) in buffers.enumerated() { enc.setBuffer(buf, offset: 0, index: i) }
|
||||
enc.dispatchThreads(MTLSize(width: grid.width, height: grid.height, depth: 1),
|
||||
threadsPerThreadgroup: tg)
|
||||
enc.endEncoding()
|
||||
return await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
|
||||
cmdBuf.addCompletedHandler { _ in continuation.resume() }
|
||||
cmdBuf.commit()
|
||||
}
|
||||
}
|
||||
|
||||
// ── Read-back ─────────────────────────────────────
|
||||
|
||||
public func readFloats(from buffer: MTLBuffer, offset: Int = 0, count: Int) -> [Float] {
|
||||
let ptr = buffer.contents().assumingMemoryBound(to: Float.self)
|
||||
return Array(UnsafeBufferPointer(start: ptr + offset, count: count))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user