ringkernel_metal/

lib.rs

//! Metal Backend for RingKernel
//!
//! This crate provides Apple Metal GPU support for RingKernel.
//! Supports macOS, iOS, and Apple Silicon.
//!
//! # Features
//!
//! - Event-driven kernel execution (Metal compute shaders)
//! - MSL (Metal Shading Language) support
//! - Apple Silicon optimization
//! - Unified memory architecture support
//!
//! # Limitations
//!
//! - No true persistent kernels (Metal doesn't support cooperative groups)
//! - macOS/iOS only
//!
//! # Example
//!
//! ```ignore
//! use ringkernel_metal::MetalRuntime;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let runtime = MetalRuntime::new().await?;
//!     let kernel = runtime.launch("compute", Default::default()).await?;
//!     kernel.activate().await?;
//!     Ok(())
//! }
//! ```

#![warn(missing_docs)]

#[cfg(all(target_os = "macos", feature = "metal"))]
mod device;
#[cfg(all(target_os = "macos", feature = "metal"))]
mod kernel;
#[cfg(all(target_os = "macos", feature = "metal"))]
mod memory;
#[cfg(all(target_os = "macos", feature = "metal"))]
mod runtime;

#[cfg(all(target_os = "macos", feature = "metal"))]
pub use device::MetalDevice;
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use kernel::{
    HaloExchangeConfig, HaloExchangeStats, MetalHaloExchange, MetalHaloMessage,
    MetalK2KInboxHeader, MetalK2KRouteEntry, MetalK2KRoutingTable, MetalKernel,
};
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use memory::MetalBuffer;
#[cfg(all(target_os = "macos", feature = "metal"))]
pub use runtime::MetalRuntime;

// Stub implementation when Metal is not available
#[cfg(not(all(target_os = "macos", feature = "metal")))]
mod stub {
    ringkernel_core::unavailable_backend!(
        MetalRuntime,
        ringkernel_core::runtime::Backend::Metal,
        "Metal"
    );
}

#[cfg(not(all(target_os = "macos", feature = "metal")))]
pub use stub::MetalRuntime;

/// Check if Metal is available at runtime.
pub fn is_metal_available() -> bool {
    #[cfg(all(target_os = "macos", feature = "metal"))]
    {
        metal::Device::system_default().is_some()
    }
    #[cfg(not(all(target_os = "macos", feature = "metal")))]
    {
        false
    }
}

/// MSL (Metal Shading Language) kernel template.
pub const RING_KERNEL_MSL_TEMPLATE: &str = r#"
//
// RingKernel Metal Shading Language Template
// Generated by ringkernel-metal
//

#include <metal_stdlib>
using namespace metal;

// Control block structure (128 bytes)
struct ControlBlock {
    atomic_uint is_active;
    atomic_uint should_terminate;
    atomic_uint has_terminated;
    uint _pad1;

    atomic_ulong messages_processed;
    atomic_ulong messages_in_flight;

    atomic_ulong input_head;
    atomic_ulong input_tail;
    atomic_ulong output_head;
    atomic_ulong output_tail;

    uint input_capacity;
    uint output_capacity;
    uint input_mask;
    uint output_mask;

    // HLC state
    atomic_ulong hlc_physical;
    atomic_ulong hlc_logical;

    atomic_uint last_error;
    atomic_uint error_count;

    uchar _reserved[16];
};

// Message header structure (256 bytes)
struct MessageHeader {
    ulong magic;
    uint version;
    uint flags;
    ulong message_id;
    ulong correlation_id;
    ulong source_kernel;
    ulong dest_kernel;
    ulong message_type;
    uchar priority;
    uchar _reserved1[7];
    ulong payload_size;
    uint checksum;
    uint _reserved2;
    // HLC timestamp (24 bytes)
    ulong ts_physical;
    ulong ts_logical;
    ulong ts_node_id;
    // Deadline
    ulong deadline_physical;
    ulong deadline_logical;
    ulong deadline_node_id;
    uchar _reserved3[104];
};

// Kernel entry point
kernel void ring_kernel_main(
    device ControlBlock* control [[buffer(0)]],
    device uchar* input_queue [[buffer(1)]],
    device uchar* output_queue [[buffer(2)]],
    device uchar* shared_state [[buffer(3)]],
    uint thread_id [[thread_position_in_threadgroup]],
    uint threadgroup_id [[threadgroup_position_in_grid]],
    uint threads_per_group [[threads_per_threadgroup]]
) {
    // Check if kernel should process
    uint is_active = atomic_load_explicit(&control->is_active, memory_order_acquire);
    if (is_active == 0) {
        return;
    }

    // Check termination
    uint should_term = atomic_load_explicit(&control->should_terminate, memory_order_acquire);
    if (should_term != 0) {
        if (thread_id == 0 && threadgroup_id == 0) {
            atomic_store_explicit(&control->has_terminated, 1, memory_order_release);
        }
        return;
    }

    // User kernel code will be inserted here
    // USER_KERNEL_CODE

    // Update message counter
    if (thread_id == 0 && threadgroup_id == 0) {
        atomic_fetch_add_explicit(&control->messages_processed, 1, memory_order_relaxed);
    }
}
"#;

/// MSL (Metal Shading Language) K2K Halo Exchange template.
///
/// This template provides kernel-to-kernel communication for stencil computations.
/// Each threadgroup can exchange halo data with its neighbors.
pub const K2K_HALO_EXCHANGE_MSL_TEMPLATE: &str = r#"
//
// RingKernel Metal K2K Halo Exchange Template
// Generated by ringkernel-metal
//

#include <metal_stdlib>
using namespace metal;

// K2K Inbox Header (64 bytes)
struct K2KInboxHeader {
    atomic_uint message_count;
    uint max_messages;
    atomic_uint head;
    atomic_uint tail;
    uint last_source;
    atomic_uint lock;
    atomic_uint sequence;
    uint _reserved[9];
};

// K2K Route Entry (32 bytes)
struct K2KRouteEntry {
    uint dest_threadgroup;
    uint inbox_offset;
    uint is_active;
    uint hops;
    uint bandwidth_hint;
    uint priority;
    uint _reserved[2];
};

// K2K Routing Table
struct K2KRoutingTable {
    uint self_id;
    uint route_count;
    uint grid_dim_x;
    uint grid_dim_y;
    uint grid_dim_z;
    uint _reserved[3];
    K2KRouteEntry routes[26];  // Max neighbors for 3D Moore neighborhood
};

// Halo Message Header (32 bytes)
struct HaloMessageHeader {
    uint source;
    uint direction;
    uint width;
    uint height;
    uint depth;
    uint element_size;
    uint sequence;
    uint flags;
};

// Try to acquire inbox lock
bool k2k_try_lock(device K2KInboxHeader* inbox) {
    uint expected = 0;
    return atomic_compare_exchange_weak_explicit(
        &inbox->lock, &expected, 1,
        memory_order_acquire, memory_order_relaxed
    );
}

// Release inbox lock
void k2k_unlock(device K2KInboxHeader* inbox) {
    atomic_store_explicit(&inbox->lock, 0, memory_order_release);
}

// Send halo data to neighbor
bool k2k_send_halo(
    device K2KRoutingTable* routing,
    device uchar* inbox_buffer,
    uint dest_id,
    device float* halo_data,
    uint width,
    uint height,
    uint depth,
    uint direction,
    uint thread_id
) {
    // Only thread 0 performs the send
    if (thread_id != 0) return true;

    // Find route to destination
    for (uint i = 0; i < routing->route_count; i++) {
        if (routing->routes[i].dest_threadgroup == dest_id &&
            routing->routes[i].is_active != 0) {

            uint offset = routing->routes[i].inbox_offset;
            device K2KInboxHeader* inbox = (device K2KInboxHeader*)(inbox_buffer + offset);

            // Try to acquire lock
            if (!k2k_try_lock(inbox)) {
                return false;  // Inbox busy
            }

            // Check if inbox has space
            uint count = atomic_load_explicit(&inbox->message_count, memory_order_acquire);
            if (count >= inbox->max_messages) {
                k2k_unlock(inbox);
                return false;  // Inbox full
            }

            // Write message header
            uint msg_offset = offset + 64 + count * (32 + width * height * depth * 4);
            device HaloMessageHeader* msg = (device HaloMessageHeader*)(inbox_buffer + msg_offset);
            msg->source = routing->self_id;
            msg->direction = direction;
            msg->width = width;
            msg->height = height;
            msg->depth = depth;
            msg->element_size = 4;
            msg->sequence = atomic_fetch_add_explicit(&inbox->sequence, 1, memory_order_relaxed);
            msg->flags = 0;

            // Copy halo data
            device float* payload = (device float*)(inbox_buffer + msg_offset + 32);
            uint payload_size = width * height * depth;
            for (uint j = 0; j < payload_size; j++) {
                payload[j] = halo_data[j];
            }

            // Update message count
            atomic_fetch_add_explicit(&inbox->message_count, 1, memory_order_release);
            inbox->last_source = routing->self_id;

            k2k_unlock(inbox);
            return true;
        }
    }

    return false;  // No route found
}

// Receive halo data from neighbors
bool k2k_recv_halo(
    device K2KRoutingTable* routing,
    device uchar* inbox_buffer,
    device float* dest_buffer,
    uint* source_out,
    uint* direction_out,
    uint thread_id
) {
    // Only thread 0 performs the receive
    if (thread_id != 0) return false;

    uint offset = routing->self_id * 4096;  // Assume 4KB per inbox
    device K2KInboxHeader* inbox = (device K2KInboxHeader*)(inbox_buffer + offset);

    // Try to acquire lock
    if (!k2k_try_lock(inbox)) {
        return false;
    }

    // Check if inbox has messages
    uint count = atomic_load_explicit(&inbox->message_count, memory_order_acquire);
    if (count == 0) {
        k2k_unlock(inbox);
        return false;
    }

    // Read oldest message (FIFO)
    uint head = atomic_load_explicit(&inbox->head, memory_order_acquire);
    uint msg_offset = offset + 64 + head * 4064;  // 32 header + max 4032 payload
    device HaloMessageHeader* msg = (device HaloMessageHeader*)(inbox_buffer + msg_offset);

    *source_out = msg->source;
    *direction_out = msg->direction;

    // Copy halo data
    device float* payload = (device float*)(inbox_buffer + msg_offset + 32);
    uint payload_size = msg->width * msg->height * msg->depth;
    for (uint j = 0; j < payload_size; j++) {
        dest_buffer[j] = payload[j];
    }

    // Update head and count
    atomic_fetch_add_explicit(&inbox->head, 1, memory_order_relaxed);
    atomic_fetch_sub_explicit(&inbox->message_count, 1, memory_order_release);

    k2k_unlock(inbox);
    return true;
}

// Halo exchange kernel - sends halo data to all neighbors
kernel void k2k_halo_exchange(
    device K2KRoutingTable* routing [[buffer(0)]],
    device uchar* inbox_buffer [[buffer(1)]],
    device float* local_data [[buffer(2)]],
    constant uint& tile_width [[buffer(3)]],
    constant uint& tile_height [[buffer(4)]],
    constant uint& halo_size [[buffer(5)]],
    uint thread_id [[thread_position_in_threadgroup]],
    uint threadgroup_id [[threadgroup_position_in_grid]]
) {
    // Extract halos from local data and send to neighbors
    // Direction: 0=North, 1=South, 2=West, 3=East, 4=Up, 5=Down
    // Note: Only thread 0 performs sends; other threads synchronize

    uint tw = tile_width;
    uint th = tile_height;
    uint self_id = routing->self_id;

    // Calculate neighbor IDs based on grid position
    uint gx = routing->grid_dim_x;
    uint gy = routing->grid_dim_y;
    uint x = self_id % gx;
    uint y = self_id / gx;

    // Threadgroup-local halo buffer for column gather
    threadgroup float column_halo[256];  // Max halo height

    // Send North halo (top row) to neighbor above (y-1)
    if (y > 0) {
        uint north_neighbor = self_id - gx;
        device float* north_halo = local_data + halo_size * tw;  // First interior row
        k2k_send_halo(routing, inbox_buffer, north_neighbor,
                      north_halo, tw, halo_size, 1, 0, thread_id);  // dir=0 (North)
    }

    // Send South halo (bottom row) to neighbor below (y+1)
    if (y < gy - 1) {
        uint south_neighbor = self_id + gx;
        device float* south_halo = local_data + (th - 2 * halo_size) * tw;  // Last interior row
        k2k_send_halo(routing, inbox_buffer, south_neighbor,
                      south_halo, tw, halo_size, 1, 1, thread_id);  // dir=1 (South)
    }

    // Send West halo (left column) to neighbor left (x-1)
    // Gather column data to threadgroup memory first
    if (x > 0) {
        uint west_neighbor = self_id - 1;
        // Gather left interior column
        if (thread_id < th && thread_id < 256) {
            column_halo[thread_id] = local_data[thread_id * tw + halo_size];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Thread 0 sends the gathered column
        k2k_send_halo(routing, inbox_buffer, west_neighbor,
                      column_halo, halo_size, th, 1, 2, thread_id);  // dir=2 (West)
    }

    // Send East halo (right column) to neighbor right (x+1)
    if (x < gx - 1) {
        uint east_neighbor = self_id + 1;
        // Gather right interior column
        if (thread_id < th && thread_id < 256) {
            column_halo[thread_id] = local_data[thread_id * tw + (tw - 2 * halo_size)];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        // Thread 0 sends the gathered column
        k2k_send_halo(routing, inbox_buffer, east_neighbor,
                      column_halo, halo_size, th, 1, 3, thread_id);  // dir=3 (East)
    }

    threadgroup_barrier(mem_flags::mem_device);
}

// Halo apply kernel - receives halo data and applies to ghost cells
kernel void k2k_halo_apply(
    device K2KRoutingTable* routing [[buffer(0)]],
    device uchar* inbox_buffer [[buffer(1)]],
    device float* local_data [[buffer(2)]],
    constant uint& tile_width [[buffer(3)]],
    constant uint& tile_height [[buffer(4)]],
    constant uint& halo_size [[buffer(5)]],
    uint thread_id [[thread_position_in_threadgroup]],
    uint threadgroup_id [[threadgroup_position_in_grid]]
) {
    // Receive halo data from neighbors and apply to local ghost cells
    // Thread 0 receives messages, then all threads cooperate to apply them

    uint tw = tile_width;
    uint th = tile_height;

    // Threadgroup-shared receive buffer
    threadgroup float recv_buffer[256];  // Max halo size
    threadgroup uint msg_source;
    threadgroup uint msg_direction;
    threadgroup bool has_message;

    // Keep receiving until inbox is empty
    while (true) {
        // Thread 0 attempts to receive
        if (thread_id == 0) {
            has_message = k2k_recv_halo(routing, inbox_buffer, recv_buffer, &msg_source, &msg_direction, 0);
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);

        if (!has_message) break;

        // All threads cooperate to apply the received halo
        switch (msg_direction) {
            case 0: {
                // From North - apply to top ghost row (row 0)
                // Received data is a row of width tw
                if (thread_id < tw) {
                    for (uint h = 0; h < halo_size; h++) {
                        local_data[h * tw + thread_id] = recv_buffer[h * tw + thread_id];
                    }
                }
                break;
            }
            case 1: {
                // From South - apply to bottom ghost row (row th-halo_size to th-1)
                if (thread_id < tw) {
                    for (uint h = 0; h < halo_size; h++) {
                        uint row = th - halo_size + h;
                        local_data[row * tw + thread_id] = recv_buffer[h * tw + thread_id];
                    }
                }
                break;
            }
            case 2: {
                // From West - apply to left ghost column (col 0)
                if (thread_id < th) {
                    for (uint h = 0; h < halo_size; h++) {
                        local_data[thread_id * tw + h] = recv_buffer[thread_id];
                    }
                }
                break;
            }
            case 3: {
                // From East - apply to right ghost column (col tw-halo_size to tw-1)
                if (thread_id < th) {
                    for (uint h = 0; h < halo_size; h++) {
                        uint col = tw - halo_size + h;
                        local_data[thread_id * tw + col] = recv_buffer[thread_id];
                    }
                }
                break;
            }
            case 4: {
                // From Up - apply to top ghost plane (3D)
                // Would need depth dimension; placeholder for 3D support
                break;
            }
            case 5: {
                // From Down - apply to bottom ghost plane (3D)
                // Would need depth dimension; placeholder for 3D support
                break;
            }
        }

        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    threadgroup_barrier(mem_flags::mem_device);
}
"#;