626 lines
18 KiB
Rust
626 lines
18 KiB
Rust
//! Metal Compute Pipeline and Kernel Management
|
|
//!
|
|
//! This module provides compute pipeline creation, shader compilation,
|
|
//! and kernel execution for Metal compute operations.
|
|
|
|
use crate::device::MetalDevice;
|
|
use crate::error::{MetalError, Result};
|
|
use crate::memory::MetalBuffer;
|
|
use dashmap::DashMap;
|
|
use std::sync::Arc;
|
|
use tracing::{debug, info};
|
|
|
|
#[cfg(target_os = "macos")]
|
|
use objc2::rc::Retained;
|
|
#[cfg(target_os = "macos")]
|
|
use objc2::runtime::ProtocolObject;
|
|
#[cfg(target_os = "macos")]
|
|
use objc2_foundation::NSString;
|
|
#[cfg(target_os = "macos")]
|
|
use objc2_metal::{
|
|
MTLCommandBuffer, MTLCommandQueue, MTLComputeCommandEncoder,
|
|
MTLComputePipelineState, MTLDevice, MTLLibrary, MTLSize,
|
|
MTLCommandEncoder,
|
|
};
|
|
|
|
/// Metal kernel launch configuration
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct MetalLaunchConfig {
|
|
/// Total threads in grid (width, height, depth)
|
|
pub grid_size: (usize, usize, usize),
|
|
/// Threads per threadgroup (width, height, depth)
|
|
pub threadgroup_size: (usize, usize, usize),
|
|
}
|
|
|
|
impl MetalLaunchConfig {
|
|
/// Create a 1D launch configuration
|
|
pub fn new_1d(total_threads: usize, threads_per_group: usize) -> Self {
|
|
Self {
|
|
grid_size: (total_threads, 1, 1),
|
|
threadgroup_size: (threads_per_group.min(1024), 1, 1),
|
|
}
|
|
}
|
|
|
|
/// Create a 2D launch configuration
|
|
pub fn new_2d(width: usize, height: usize, tile_size: usize) -> Self {
|
|
let tile = tile_size.min(32);
|
|
Self {
|
|
grid_size: (width, height, 1),
|
|
threadgroup_size: (tile, tile, 1),
|
|
}
|
|
}
|
|
|
|
/// Validate the launch configuration
|
|
pub fn validate(&self) -> Result<()> {
|
|
let (gx, gy, gz) = self.grid_size;
|
|
let (tx, ty, tz) = self.threadgroup_size;
|
|
|
|
if gx == 0 || gy == 0 || gz == 0 {
|
|
return Err(MetalError::InvalidLaunchConfig(
|
|
"Grid size dimensions must be positive".into()
|
|
));
|
|
}
|
|
|
|
if tx == 0 || ty == 0 || tz == 0 {
|
|
return Err(MetalError::InvalidLaunchConfig(
|
|
"Threadgroup size dimensions must be positive".into()
|
|
));
|
|
}
|
|
|
|
let total_threads = tx * ty * tz;
|
|
if total_threads > 1024 {
|
|
return Err(MetalError::InvalidLaunchConfig(format!(
|
|
"Too many threads per threadgroup: {total_threads} (max 1024)"
|
|
)));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Metal command queue for submitting work to the GPU
|
|
#[cfg(target_os = "macos")]
|
|
pub struct MetalCommandQueue {
|
|
queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub struct MetalCommandQueue {
|
|
_private: (),
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
impl MetalCommandQueue {
|
|
/// Create a new command queue from a device
|
|
pub fn new(device: &MetalDevice) -> Result<Self> {
|
|
let queue = device.new_command_queue()?;
|
|
Ok(Self { queue })
|
|
}
|
|
|
|
/// Get the underlying MTLCommandQueue
|
|
pub fn mtl_queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
|
|
&self.queue
|
|
}
|
|
|
|
/// Create a new command buffer
|
|
pub fn command_buffer(&self) -> Result<Retained<ProtocolObject<dyn MTLCommandBuffer>>> {
|
|
self.queue.commandBuffer()
|
|
.ok_or_else(|| MetalError::CommandBuffer("Failed to create command buffer".into()))
|
|
}
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
impl MetalCommandQueue {
|
|
/// Create a new command queue (stub)
|
|
pub fn new(_device: &MetalDevice) -> Result<Self> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
}
|
|
|
|
/// Compiled Metal compute kernel
|
|
#[cfg(target_os = "macos")]
|
|
pub struct MetalKernel {
|
|
/// Kernel function name
|
|
name: String,
|
|
/// Compiled pipeline state
|
|
pipeline: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
|
|
/// Maximum total threads per threadgroup
|
|
max_threads_per_threadgroup: usize,
|
|
/// Thread execution width (SIMD group size)
|
|
thread_execution_width: usize,
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub struct MetalKernel {
|
|
name: String,
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
impl MetalKernel {
|
|
/// Get kernel name
|
|
pub fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
/// Get maximum threads per threadgroup for this kernel
|
|
pub fn max_threads_per_threadgroup(&self) -> usize {
|
|
self.max_threads_per_threadgroup
|
|
}
|
|
|
|
/// Get thread execution width (SIMD group size)
|
|
pub fn thread_execution_width(&self) -> usize {
|
|
self.thread_execution_width
|
|
}
|
|
|
|
/// Get the pipeline state
|
|
pub fn pipeline(&self) -> &ProtocolObject<dyn MTLComputePipelineState> {
|
|
&self.pipeline
|
|
}
|
|
|
|
/// Calculate optimal threadgroup size for 1D dispatch
|
|
pub fn optimal_threadgroup_size_1d(&self, total_elements: usize) -> usize {
|
|
let max = self.max_threads_per_threadgroup;
|
|
let width = self.thread_execution_width;
|
|
|
|
// Prefer multiples of thread execution width
|
|
if total_elements >= max {
|
|
max
|
|
} else if total_elements >= width {
|
|
(total_elements / width) * width
|
|
} else {
|
|
total_elements
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
impl MetalKernel {
|
|
/// Get kernel name
|
|
pub fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
/// Get maximum threads per threadgroup
|
|
pub fn max_threads_per_threadgroup(&self) -> usize {
|
|
0
|
|
}
|
|
|
|
/// Get thread execution width
|
|
pub fn thread_execution_width(&self) -> usize {
|
|
0
|
|
}
|
|
}
|
|
|
|
/// Metal compute pipeline manager
|
|
///
|
|
/// Handles shader compilation, kernel caching, and execution
|
|
#[cfg(target_os = "macos")]
|
|
pub struct MetalComputePipeline {
|
|
device: Arc<MetalDevice>,
|
|
/// Cached compiled kernels
|
|
kernel_cache: DashMap<String, Arc<MetalKernel>>,
|
|
/// Compiled shader library
|
|
library: Option<Retained<ProtocolObject<dyn MTLLibrary>>>,
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub struct MetalComputePipeline {
|
|
_private: (),
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
impl MetalComputePipeline {
|
|
/// Create a new compute pipeline manager
|
|
pub fn new(device: Arc<MetalDevice>) -> Self {
|
|
Self {
|
|
device,
|
|
kernel_cache: DashMap::new(),
|
|
library: None,
|
|
}
|
|
}
|
|
|
|
/// Compile MSL shader source code into a library
|
|
pub fn compile_library(&mut self, source: &str) -> Result<()> {
|
|
debug!("Compiling Metal shader library");
|
|
|
|
let ns_source = NSString::from_str(source);
|
|
let library = self.device.mtl_device()
|
|
.newLibraryWithSource_options_error(&ns_source, None)
|
|
.map_err(|e| MetalError::ShaderCompilation(format!("{e:?}")))?;
|
|
|
|
self.library = Some(library);
|
|
info!("Metal shader library compiled successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Get or compile a kernel by name
|
|
pub fn get_kernel(&self, name: &str) -> Result<Arc<MetalKernel>> {
|
|
// Check cache first
|
|
if let Some(kernel) = self.kernel_cache.get(name) {
|
|
return Ok(kernel.value().clone());
|
|
}
|
|
|
|
// Need library to compile kernel
|
|
let library = self.library.as_ref()
|
|
.ok_or_else(|| MetalError::ShaderCompilation("No shader library compiled".into()))?;
|
|
|
|
// Get function from library
|
|
let ns_name = NSString::from_str(name);
|
|
let function = library.newFunctionWithName(&ns_name)
|
|
.ok_or_else(|| MetalError::KernelNotFound(name.to_string()))?;
|
|
|
|
// Create pipeline state
|
|
let pipeline = self.device.mtl_device()
|
|
.newComputePipelineStateWithFunction_error(&function)
|
|
.map_err(|e| MetalError::PipelineCreation(format!("{e:?}")))?;
|
|
|
|
let kernel = Arc::new(MetalKernel {
|
|
name: name.to_string(),
|
|
max_threads_per_threadgroup: pipeline.maxTotalThreadsPerThreadgroup(),
|
|
thread_execution_width: pipeline.threadExecutionWidth(),
|
|
pipeline,
|
|
});
|
|
|
|
self.kernel_cache.insert(name.to_string(), kernel.clone());
|
|
debug!("Cached Metal kernel: {}", name);
|
|
|
|
Ok(kernel)
|
|
}
|
|
|
|
/// Execute a kernel with the given buffers
|
|
pub fn execute<T: Copy + 'static>(
|
|
&self,
|
|
kernel: &MetalKernel,
|
|
config: MetalLaunchConfig,
|
|
buffers: &[&MetalBuffer<T>],
|
|
) -> Result<()> {
|
|
config.validate()?;
|
|
|
|
// Create command buffer
|
|
let cmd_buffer = self.device.command_queue_retained().commandBuffer()
|
|
.ok_or_else(|| MetalError::CommandBuffer("Failed to create command buffer".into()))?;
|
|
|
|
// Create compute encoder
|
|
let encoder = cmd_buffer.computeCommandEncoder()
|
|
.ok_or_else(|| MetalError::Execution("Failed to create compute encoder".into()))?;
|
|
|
|
unsafe {
|
|
// Set pipeline state
|
|
encoder.setComputePipelineState(kernel.pipeline());
|
|
|
|
// Bind buffers
|
|
for (index, buffer) in buffers.iter().enumerate() {
|
|
encoder.setBuffer_offset_atIndex(
|
|
Some(buffer.mtl_buffer()),
|
|
0,
|
|
index
|
|
);
|
|
}
|
|
|
|
// Dispatch threads
|
|
let grid_size = MTLSize {
|
|
width: config.grid_size.0,
|
|
height: config.grid_size.1,
|
|
depth: config.grid_size.2,
|
|
};
|
|
|
|
let threadgroup_size = MTLSize {
|
|
width: config.threadgroup_size.0,
|
|
height: config.threadgroup_size.1,
|
|
depth: config.threadgroup_size.2,
|
|
};
|
|
|
|
encoder.dispatchThreads_threadsPerThreadgroup(grid_size, threadgroup_size);
|
|
}
|
|
|
|
// End encoding and execute
|
|
encoder.endEncoding();
|
|
cmd_buffer.commit();
|
|
cmd_buffer.waitUntilCompleted();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Execute a kernel with data buffers and scalar constant buffers
|
|
///
|
|
/// This is useful for kernels that need both data buffers (A, B, C matrices)
|
|
/// and scalar parameters (M, N, K dimensions) passed as separate buffers.
|
|
///
|
|
/// Buffers are bound at indices 0..buffers.len(),
|
|
/// scalars are bound at indices buffers.len()..
|
|
pub fn execute_with_scalars<T: Copy + 'static>(
|
|
&self,
|
|
kernel: &MetalKernel,
|
|
config: MetalLaunchConfig,
|
|
buffers: &[&MetalBuffer<T>],
|
|
scalars: &[&dyn crate::memory::AsRawMetalBuffer],
|
|
) -> Result<()> {
|
|
config.validate()?;
|
|
|
|
// Create command buffer
|
|
let cmd_buffer = self.device.command_queue_retained().commandBuffer()
|
|
.ok_or_else(|| MetalError::CommandBuffer("Failed to create command buffer".into()))?;
|
|
|
|
// Create compute encoder
|
|
let encoder = cmd_buffer.computeCommandEncoder()
|
|
.ok_or_else(|| MetalError::Execution("Failed to create compute encoder".into()))?;
|
|
|
|
unsafe {
|
|
// Set pipeline state
|
|
encoder.setComputePipelineState(kernel.pipeline());
|
|
|
|
// Bind data buffers at indices 0..buffers.len()
|
|
for (index, buffer) in buffers.iter().enumerate() {
|
|
encoder.setBuffer_offset_atIndex(
|
|
Some(buffer.mtl_buffer()),
|
|
0,
|
|
index
|
|
);
|
|
}
|
|
|
|
// Bind scalar buffers at indices buffers.len()..
|
|
let scalar_offset = buffers.len();
|
|
for (index, scalar) in scalars.iter().enumerate() {
|
|
encoder.setBuffer_offset_atIndex(
|
|
Some(scalar.as_raw_buffer()),
|
|
0,
|
|
scalar_offset + index
|
|
);
|
|
}
|
|
|
|
// Dispatch threads
|
|
let grid_size = MTLSize {
|
|
width: config.grid_size.0,
|
|
height: config.grid_size.1,
|
|
depth: config.grid_size.2,
|
|
};
|
|
|
|
let threadgroup_size = MTLSize {
|
|
width: config.threadgroup_size.0,
|
|
height: config.threadgroup_size.1,
|
|
depth: config.threadgroup_size.2,
|
|
};
|
|
|
|
encoder.dispatchThreads_threadsPerThreadgroup(grid_size, threadgroup_size);
|
|
}
|
|
|
|
// End encoding and execute
|
|
encoder.endEncoding();
|
|
cmd_buffer.commit();
|
|
cmd_buffer.waitUntilCompleted();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Clear kernel cache
|
|
pub fn clear_cache(&self) {
|
|
self.kernel_cache.clear();
|
|
debug!("Metal kernel cache cleared");
|
|
}
|
|
|
|
/// Get number of cached kernels
|
|
pub fn cache_size(&self) -> usize {
|
|
self.kernel_cache.len()
|
|
}
|
|
|
|
/// Dispatch a kernel with a closure for setting up encoder
|
|
pub fn dispatch<F>(
|
|
&self,
|
|
kernel_name: &str,
|
|
grid_size: [u32; 3],
|
|
threadgroup_size: [u32; 3],
|
|
setup: F,
|
|
) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
let kernel = self.get_kernel(kernel_name)?;
|
|
|
|
let cmd_buffer = self
|
|
.device
|
|
.command_queue()
|
|
.commandBuffer()
|
|
.ok_or_else(|| MetalError::CommandBuffer("Failed to create command buffer".into()))?;
|
|
|
|
let encoder = cmd_buffer
|
|
.computeCommandEncoder()
|
|
.ok_or_else(|| MetalError::Execution("Failed to create compute encoder".into()))?;
|
|
|
|
unsafe {
|
|
encoder.setComputePipelineState(kernel.pipeline());
|
|
}
|
|
|
|
// Wrap encoder and call setup
|
|
let wrapper = ComputeEncoder { encoder: &encoder };
|
|
setup(&wrapper);
|
|
|
|
unsafe {
|
|
let grid = MTLSize {
|
|
width: grid_size[0] as usize,
|
|
height: grid_size[1] as usize,
|
|
depth: grid_size[2] as usize,
|
|
};
|
|
let threads = MTLSize {
|
|
width: threadgroup_size[0] as usize,
|
|
height: threadgroup_size[1] as usize,
|
|
depth: threadgroup_size[2] as usize,
|
|
};
|
|
encoder.dispatchThreads_threadsPerThreadgroup(grid, threads);
|
|
}
|
|
|
|
encoder.endEncoding();
|
|
cmd_buffer.commit();
|
|
cmd_buffer.waitUntilCompleted();
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Dispatch a 1D kernel with automatic threadgroup sizing
|
|
pub fn dispatch_1d<F>(&self, kernel_name: &str, total_threads: u32, setup: F) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
let threads_per_group = 256.min(total_threads);
|
|
self.dispatch(
|
|
kernel_name,
|
|
[total_threads, 1, 1],
|
|
[threads_per_group, 1, 1],
|
|
setup,
|
|
)
|
|
}
|
|
|
|
/// Dispatch a 3D kernel with automatic threadgroup sizing
|
|
pub fn dispatch_3d<F>(&self, kernel_name: &str, grid_size: [u32; 3], setup: F) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
let tx = 8.min(grid_size[0]);
|
|
let ty = 8.min(grid_size[1]);
|
|
let tz = 1.min(grid_size[2]);
|
|
self.dispatch(kernel_name, grid_size, [tx, ty, tz], setup)
|
|
}
|
|
}
|
|
|
|
/// Wrapper for compute encoder that provides safe buffer/bytes setting
|
|
#[cfg(target_os = "macos")]
|
|
pub struct ComputeEncoder<'a> {
|
|
encoder: &'a ProtocolObject<dyn MTLComputeCommandEncoder>,
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
impl<'a> ComputeEncoder<'a> {
|
|
/// Set a buffer at the given index
|
|
pub fn set_buffer<T: Copy + 'static>(&self, index: u64, buffer: &MetalBuffer<T>) {
|
|
unsafe {
|
|
self.encoder
|
|
.setBuffer_offset_atIndex(Some(buffer.mtl_buffer()), 0, index as usize);
|
|
}
|
|
}
|
|
|
|
/// Set bytes (scalar value) at the given index
|
|
pub fn set_bytes<T: Copy>(&self, index: u64, value: &T) {
|
|
unsafe {
|
|
self.encoder.setBytes_length_atIndex(
|
|
std::ptr::NonNull::new(value as *const T as *mut std::ffi::c_void).unwrap(),
|
|
std::mem::size_of::<T>(),
|
|
index as usize,
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
impl MetalComputePipeline {
|
|
/// Create a new compute pipeline manager (stub)
|
|
pub fn new(_device: Arc<MetalDevice>) -> Self {
|
|
Self { _private: () }
|
|
}
|
|
|
|
/// Compile shader source (stub)
|
|
pub fn compile_library(&mut self, _source: &str) -> Result<()> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Get kernel (stub)
|
|
pub fn get_kernel(&self, _name: &str) -> Result<Arc<MetalKernel>> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Execute kernel (stub)
|
|
pub fn execute<T: Copy + 'static>(
|
|
&self,
|
|
_kernel: &MetalKernel,
|
|
_config: MetalLaunchConfig,
|
|
_buffers: &[&MetalBuffer<T>],
|
|
) -> Result<()> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Execute kernel with scalars (stub)
|
|
pub fn execute_with_scalars<T: Copy + 'static>(
|
|
&self,
|
|
_kernel: &MetalKernel,
|
|
_config: MetalLaunchConfig,
|
|
_buffers: &[&MetalBuffer<T>],
|
|
_scalars: &[&dyn crate::memory::AsRawMetalBuffer],
|
|
) -> Result<()> {
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Clear cache
|
|
pub fn clear_cache(&self) {}
|
|
|
|
/// Get cache size
|
|
pub fn cache_size(&self) -> usize {
|
|
0
|
|
}
|
|
|
|
/// Dispatch a kernel (stub)
|
|
pub fn dispatch<F>(
|
|
&self,
|
|
_kernel_name: &str,
|
|
_grid_size: [u32; 3],
|
|
_threadgroup_size: [u32; 3],
|
|
_setup: F,
|
|
) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Dispatch a 1D kernel (stub)
|
|
pub fn dispatch_1d<F>(&self, _kernel_name: &str, _total_threads: u32, _setup: F) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
|
|
/// Dispatch a 3D kernel (stub)
|
|
pub fn dispatch_3d<F>(&self, _kernel_name: &str, _grid_size: [u32; 3], _setup: F) -> Result<()>
|
|
where
|
|
F: FnOnce(&ComputeEncoder),
|
|
{
|
|
Err(MetalError::NotAvailable)
|
|
}
|
|
}
|
|
|
|
/// Wrapper for compute encoder (stub)
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub struct ComputeEncoder<'a> {
|
|
_phantom: std::marker::PhantomData<&'a ()>,
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
impl<'a> ComputeEncoder<'a> {
|
|
/// Set a buffer at the given index (stub)
|
|
pub fn set_buffer<T: Copy + 'static>(&self, _index: u64, _buffer: &MetalBuffer<T>) {}
|
|
|
|
/// Set bytes at the given index (stub)
|
|
pub fn set_bytes<T: Copy>(&self, _index: u64, _value: &T) {}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_launch_config_validation() {
|
|
let valid = MetalLaunchConfig::new_1d(1000, 256);
|
|
assert!(valid.validate().is_ok());
|
|
|
|
let invalid = MetalLaunchConfig {
|
|
grid_size: (0, 1, 1),
|
|
threadgroup_size: (256, 1, 1),
|
|
};
|
|
assert!(invalid.validate().is_err());
|
|
|
|
let too_many_threads = MetalLaunchConfig {
|
|
grid_size: (1000, 1, 1),
|
|
threadgroup_size: (1025, 1, 1),
|
|
};
|
|
assert!(too_many_threads.validate().is_err());
|
|
}
|
|
}
|