locomp 0.1.0__py3-none-any.whl

locomp/__init__.py

@@ -0,0 +1,158 @@
+"""
+Locomp — A GPU compute compiler for Apple Metal.
+Write GPU kernels in Python, compile to native Metal shaders.
+"""
+
+from locomp.autotune import Config, autotune, clear_cache
+from locomp.api import (
+    kernel,
+    program_id,
+    thread_id,
+    local_id,
+    group_size,
+    num_groups,
+    barrier,
+    shared_memory,
+    shared_load,
+    shared_store,
+    simd_sum,
+    simd_max,
+    simd_min,
+    simd_broadcast,
+    simd_shuffle_down,
+    simd_lane_id,
+    simd_group_id,
+    simdgroup_matrix_load,
+    simdgroup_matrix_load_device,
+    simdgroup_matrix_store,
+    simdgroup_matrix_store_device,
+    simdgroup_mac,
+    simdgroup_matrix,
+    atomic_add,
+    atomic_max,
+    atomic_min,
+    arange,
+    load,
+    store,
+    exp,
+    log,
+    sqrt,
+    abs,
+    tanh,
+    sin,
+    cos,
+    asin,
+    acos,
+    atan,
+    atan2,
+    sinh,
+    cosh,
+    exp2,
+    log2,
+    log10,
+    rsqrt,
+    ceil,
+    floor,
+    round,
+    sigmoid,
+    fma,
+    pow,
+    clamp,
+    copysign,
+    fmod,
+    step,
+    where,
+    tensor,
+    empty,
+    zeros,
+    ones,
+    hardware_info,
+    set_device,
+    constexpr,
+    Tensor,
+    Float16,
+    UInt8,
+    Int8,
+    Int32,
+    Bool,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "Config",
+    "autotune",
+    "clear_cache",
+    "kernel",
+    "program_id",
+    "thread_id",
+    "local_id",
+    "group_size",
+    "num_groups",
+    "barrier",
+    "shared_memory",
+    "shared_load",
+    "shared_store",
+    "simd_sum",
+    "simd_max",
+    "simd_min",
+    "simd_broadcast",
+    "simd_shuffle_down",
+    "simd_lane_id",
+    "simd_group_id",
+    "simdgroup_matrix_load",
+    "simdgroup_matrix_load_device",
+    "simdgroup_matrix_store",
+    "simdgroup_matrix_store_device",
+    "simdgroup_mac",
+    "simdgroup_matrix",
+    "atomic_add",
+    "atomic_max",
+    "atomic_min",
+    "arange",
+    "load",
+    "store",
+    "exp",
+    "log",
+    "sqrt",
+    "abs",
+    "tanh",
+    "sin",
+    "cos",
+    "asin",
+    "acos",
+    "atan",
+    "atan2",
+    "sinh",
+    "cosh",
+    "exp2",
+    "log2",
+    "log10",
+    "rsqrt",
+    "ceil",
+    "floor",
+    "round",
+    "sigmoid",
+    "fma",
+    "pow",
+    "clamp",
+    "copysign",
+    "fmod",
+    "step",
+    "where",
+    "tensor",
+    "empty",
+    "zeros",
+    "ones",
+    "atomic_add",
+    "atomic_max",
+    "atomic_min",
+    "hardware_info",
+    "set_device",
+    "constexpr",
+    "Tensor",
+    "Float16",
+    "UInt8",
+    "Int8",
+    "Int32",
+    "Bool",
+]

locomp/_native/fast_dispatch.m

@@ -0,0 +1,162 @@
+// locust/_native/fast_dispatch.m
+// Minimal C bridge for Metal kernel dispatch.
+// Replaces 8+ PyObjC round-trips with a single ctypes call.
+
+#import <Metal/Metal.h>
+
+// Pending command buffer for async dispatch
+static id<MTLCommandBuffer> _pending_cmd_buf = nil;
+
+// Dispatch a compute kernel synchronously. Returns GPU time in milliseconds.
+double locust_dispatch(void *queue_ptr, void *pipeline_ptr,
+                       void **buffer_ptrs, int num_buffers,
+                       int gx, int gy, int gz,
+                       int tx, int ty, int tz) {
+    @autoreleasepool {
+        id<MTLCommandQueue> queue = (__bridge id<MTLCommandQueue>)queue_ptr;
+        id<MTLComputePipelineState> pipeline = (__bridge id<MTLComputePipelineState>)pipeline_ptr;
+
+        id<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
+        id<MTLComputeCommandEncoder> encoder = [cmdBuf computeCommandEncoder];
+        [encoder setComputePipelineState:pipeline];
+
+        for (int i = 0; i < num_buffers; i++) {
+            id<MTLBuffer> buf = (__bridge id<MTLBuffer>)buffer_ptrs[i];
+            [encoder setBuffer:buf offset:0 atIndex:i];
+        }
+
+        MTLSize grid = MTLSizeMake(gx, gy, gz);
+        MTLSize tgSize = MTLSizeMake(tx, ty, tz);
+        [encoder dispatchThreadgroups:grid threadsPerThreadgroup:tgSize];
+        [encoder endEncoding];
+
+        [cmdBuf commit];
+        [cmdBuf waitUntilCompleted];
+
+        return (cmdBuf.GPUEndTime - cmdBuf.GPUStartTime) * 1000.0;
+    }
+}
+
+// Dispatch async — commit but don't wait. Returns immediately.
+void locust_dispatch_async(void *queue_ptr, void *pipeline_ptr,
+                           void **buffer_ptrs, int num_buffers,
+                           int gx, int gy, int gz,
+                           int tx, int ty, int tz) {
+    // Wait for any previous pending work first
+    if (_pending_cmd_buf != nil) {
+        [_pending_cmd_buf waitUntilCompleted];
+        _pending_cmd_buf = nil;
+    }
+
+    id<MTLCommandQueue> queue = (__bridge id<MTLCommandQueue>)queue_ptr;
+    id<MTLComputePipelineState> pipeline = (__bridge id<MTLComputePipelineState>)pipeline_ptr;
+
+    id<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
+    id<MTLComputeCommandEncoder> encoder = [cmdBuf computeCommandEncoder];
+    [encoder setComputePipelineState:pipeline];
+
+    for (int i = 0; i < num_buffers; i++) {
+        id<MTLBuffer> buf = (__bridge id<MTLBuffer>)buffer_ptrs[i];
+        [encoder setBuffer:buf offset:0 atIndex:i];
+    }
+
+    MTLSize grid = MTLSizeMake(gx, gy, gz);
+    MTLSize tgSize = MTLSizeMake(tx, ty, tz);
+    [encoder dispatchThreadgroups:grid threadsPerThreadgroup:tgSize];
+    [encoder endEncoding];
+
+    [cmdBuf commit];
+    _pending_cmd_buf = cmdBuf;
+}
+
+// Wait for pending async work. Returns 0 if nothing pending.
+double locust_sync(void) {
+    if (_pending_cmd_buf == nil) return 0.0;
+    [_pending_cmd_buf waitUntilCompleted];
+    double gpu_ms = (_pending_cmd_buf.GPUEndTime - _pending_cmd_buf.GPUStartTime) * 1000.0;
+    _pending_cmd_buf = nil;
+    return gpu_ms;
+}
+
+// Dispatch same kernel N times in one command buffer. Returns avg GPU time in ms.
+double locust_dispatch_repeat(void *queue_ptr, void *pipeline_ptr,
+                              void **buffer_ptrs, int num_buffers,
+                              int gx, int gy, int gz,
+                              int tx, int ty, int tz,
+                              int repeat) {
+    @autoreleasepool {
+        id<MTLCommandQueue> queue = (__bridge id<MTLCommandQueue>)queue_ptr;
+        id<MTLComputePipelineState> pipeline = (__bridge id<MTLComputePipelineState>)pipeline_ptr;
+
+        id<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
+
+        MTLSize grid = MTLSizeMake(gx, gy, gz);
+        MTLSize tgSize = MTLSizeMake(tx, ty, tz);
+
+        for (int r = 0; r < repeat; r++) {
+            id<MTLComputeCommandEncoder> encoder = [cmdBuf computeCommandEncoder];
+            [encoder setComputePipelineState:pipeline];
+            for (int i = 0; i < num_buffers; i++) {
+                id<MTLBuffer> buf = (__bridge id<MTLBuffer>)buffer_ptrs[i];
+                [encoder setBuffer:buf offset:0 atIndex:i];
+            }
+            [encoder dispatchThreadgroups:grid threadsPerThreadgroup:tgSize];
+            [encoder endEncoding];
+        }
+
+        [cmdBuf commit];
+        [cmdBuf waitUntilCompleted];
+
+        double total_ms = (cmdBuf.GPUEndTime - cmdBuf.GPUStartTime) * 1000.0;
+        return total_ms / repeat;
+    }
+}
+
+// --- Batch mode ---
+// Multiple different kernels in one command buffer.
+static id<MTLCommandBuffer> _batch_cmd_buf = nil;
+
+void locust_batch_begin(void *queue_ptr) {
+    if (_batch_cmd_buf != nil) {
+        [_batch_cmd_buf waitUntilCompleted];
+        _batch_cmd_buf = nil;
+    }
+    // Also drain any pending async work
+    if (_pending_cmd_buf != nil) {
+        [_pending_cmd_buf waitUntilCompleted];
+        _pending_cmd_buf = nil;
+    }
+    id<MTLCommandQueue> queue = (__bridge id<MTLCommandQueue>)queue_ptr;
+    _batch_cmd_buf = [queue commandBuffer];
+}
+
+void locust_batch_dispatch(void *pipeline_ptr,
+                           void **buffer_ptrs, int num_buffers,
+                           int gx, int gy, int gz,
+                           int tx, int ty, int tz) {
+    if (_batch_cmd_buf == nil) return;
+
+    id<MTLComputePipelineState> pipeline = (__bridge id<MTLComputePipelineState>)pipeline_ptr;
+
+    id<MTLComputeCommandEncoder> encoder = [_batch_cmd_buf computeCommandEncoder];
+    [encoder setComputePipelineState:pipeline];
+
+    for (int i = 0; i < num_buffers; i++) {
+        id<MTLBuffer> buf = (__bridge id<MTLBuffer>)buffer_ptrs[i];
+        [encoder setBuffer:buf offset:0 atIndex:i];
+    }
+
+    MTLSize grid = MTLSizeMake(gx, gy, gz);
+    MTLSize tgSize = MTLSizeMake(tx, ty, tz);
+    [encoder dispatchThreadgroups:grid threadsPerThreadgroup:tgSize];
+    [encoder endEncoding];
+}
+
+double locust_batch_end(void) {
+    if (_batch_cmd_buf == nil) return 0.0;
+    [_batch_cmd_buf commit];
+    [_batch_cmd_buf waitUntilCompleted];
+    double gpu_ms = (_batch_cmd_buf.GPUEndTime - _batch_cmd_buf.GPUStartTime) * 1000.0;
+    _batch_cmd_buf = nil;
+    return gpu_ms;
+}