himada-dispatch 0.1.1

use crate::HwdnaError;
#[cfg(feature = "vulkan")]
use crate::vulkan::VulkanRuntime;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum GpuBackend {
    None,
    Metal,
    Cuda,
    Vulkan,
}

pub fn detect_gpu() -> GpuBackend {
    #[cfg(all(target_os = "macos", feature = "metal"))]
    {
        if is_metal_available() { return GpuBackend::Metal; }
    }
    #[cfg(feature = "cuda")]
    {
        if is_cuda_available() { return GpuBackend::Cuda; }
    }
    #[cfg(feature = "vulkan")]
    {
        if crate::vulkan::VulkanRuntime::is_available() {
            return GpuBackend::Vulkan;
        }
    }
    GpuBackend::None
}

#[cfg(all(target_os = "macos", feature = "metal"))]
fn is_metal_available() -> bool {
    metal_available_impl()
}

#[cfg(all(target_os = "macos", feature = "metal"))]
fn metal_available_impl() -> bool {
    unsafe {
        let lib = libc::dlopen(b"Metal.framework/Metal\0".as_ptr() as *const i8, libc::RTLD_LAZY);
        if lib.is_null() { return false; }
        let sym = libc::dlsym(lib, b"MTLCreateSystemDefaultDevice\0".as_ptr() as *const i8);
        libc::dlclose(lib);
        !sym.is_null()
    }
}

#[cfg(feature = "cuda")]
fn is_cuda_available() -> bool {
    unsafe {
        let lib = libc::dlopen(b"libcuda.so.1\0".as_ptr() as *const i8, libc::RTLD_LAZY);
        if lib.is_null() {
            let lib = libc::dlopen(b"libcuda.dylib\0".as_ptr() as *const i8, libc::RTLD_LAZY);
            if lib.is_null() { return false; }
        }
        libc::dlclose(lib);
        true
    }
}

/// Metal compute shaders compiled at runtime.
#[cfg(all(target_os = "macos", feature = "metal"))]
mod metal_kernels {
use crate::HwdnaError;
    use metal::*;

    const DOT_MSL: &str = "
        kernel void dot_f64(
            device const double *a [[buffer(0)]],
            device const double *b [[buffer(1)]],
            device double *result [[buffer(2)]],
            uint id [[thread_position_in_grid]],
            uint grid_size [[threads_per_grid]])
        {
            double sum = 0.0;
            uint len = grid_size;
            for (uint i = id; i < len; i += grid_size) {
                sum += a[i] * b[i];
            }
            result[id] = sum;
        }
    ";

    const MATMUL_MSL: &str = "
        kernel void matmul_f64(
            device const double *a [[buffer(0)]],
            device const double *b [[buffer(1)]],
            device double *c [[buffer(2)]],
            constant uint &n [[buffer(3)]],
            uint2 gid [[thread_position_in_grid]])
        {
            uint row = gid.y;
            uint col = gid.x;
            if (row >= n || col >= n) return;
            double sum = 0.0;
            for (uint i = 0; i < n; ++i) {
                sum += a[row * n + i] * b[i * n + col];
            }
            c[row * n + col] = sum;
        }
    ";

    const REDUCE_SUM_MSL: &str = "
        kernel void reduce_sum_f64(
            device const double *data [[buffer(0)]],
            device double *result [[buffer(1)]],
            uint id [[thread_position_in_grid]],
            uint grid_size [[threads_per_grid]])
        {
            double sum = 0.0;
            uint len = grid_size;
            for (uint i = id; i < len; i += grid_size) {
                sum += data[i];
            }
            result[id] = sum;
        }
    ";

    pub struct MetalCompute {
        device: Device,
        queue: CommandQueue,
        dot_pso: ComputePipelineState,
        matmul_pso: ComputePipelineState,
        reduce_pso: ComputePipelineState,
    }

    impl MetalCompute {
        pub fn new() -> Result<Self, HwdnaError> {
            let device = Device::system_default().ok_or(HwdnaError::NoSupportedKernels)?;
            let queue = device.new_command_queue();

            let dot_pso = compile_kernel(&device, DOT_MSL, "dot_f64")?;
            let matmul_pso = compile_kernel(&device, MATMUL_MSL, "matmul_f64")?;
            let reduce_pso = compile_kernel(&device, REDUCE_SUM_MSL, "reduce_sum_f64")?;

            Ok(Self { device, queue, dot_pso, matmul_pso, reduce_pso })
        }

        pub fn dot(&self, a: &[f64], b: &[f64]) -> Result<f64, HwdnaError> {
            let len = a.len().min(b.len());
            if len == 0 { return Ok(0.0); }

            let a_buf = self.device.new_buffer_with_data(
                a.as_ptr() as *const std::ffi::c_void,
                (len * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );
            let b_buf = self.device.new_buffer_with_data(
                b.as_ptr() as *const std::ffi::c_void,
                (len * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );

            let thread_count = 256u64.min(len as u64);
            let result_buf = self.device.new_buffer(
                (thread_count * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );

            let cmd_buf = self.queue.new_command_buffer();
            let encoder = cmd_buf.new_compute_command_encoder();
            encoder.set_compute_pipeline_state(&self.dot_pso);
            encoder.set_buffer(0, Some(&a_buf), 0);
            encoder.set_buffer(1, Some(&b_buf), 0);
            encoder.set_buffer(2, Some(&result_buf), 0);

            let grid_size = MTLSize::new(thread_count, 1, 1);
            let group_size = MTLSize::new(thread_count.min(64), 1, 1);
            encoder.dispatch_threads(grid_size, group_size);
            encoder.end_encoding();
            cmd_buf.commit();
            cmd_buf.wait_until_completed();

            let ptr = result_buf.contents() as *const f64;
            let mut sum = 0.0;
            for i in 0..thread_count as usize {
                sum += unsafe { *ptr.add(i) };
            }
            Ok(sum)
        }

        pub fn matmul(&self, a: &[f64], b: &[f64], c: &mut [f64], n: usize) -> Result<(), HwdnaError> {
            let a_buf = self.device.new_buffer_with_data(
                a.as_ptr() as *const std::ffi::c_void,
                (n * n * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );
            let b_buf = self.device.new_buffer_with_data(
                b.as_ptr() as *const std::ffi::c_void,
                (n * n * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );
            let c_buf = self.device.new_buffer_with_data(
                c.as_ptr() as *const std::ffi::c_void,
                (n * n * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );

            let n_val = n as u32;
            let n_buf = self.device.new_buffer_with_data(
                &n_val as *const u32 as *const std::ffi::c_void,
                4,
                MTLResourceOptions::StorageModeShared,
            );

            let cmd_buf = self.queue.new_command_buffer();
            let encoder = cmd_buf.new_compute_command_encoder();
            encoder.set_compute_pipeline_state(&self.matmul_pso);
            encoder.set_buffer(0, Some(&a_buf), 0);
            encoder.set_buffer(1, Some(&b_buf), 0);
            encoder.set_buffer(2, Some(&c_buf), 0);
            encoder.set_buffer(3, Some(&n_buf), 0);

            let grid_size = MTLSize::new(n as u64, n as u64, 1);
            let group_size = MTLSize::new(16.min(n as u64), 16.min(n as u64), 1);
            encoder.dispatch_threads(grid_size, group_size);
            encoder.end_encoding();
            cmd_buf.commit();
            cmd_buf.wait_until_completed();

            let ptr = c_buf.contents() as *const f64;
            for i in 0..(n * n) {
                c[i] = unsafe { *ptr.add(i) };
            }
            Ok(())
        }

        pub fn reduce_sum(&self, data: &[f64]) -> Result<f64, HwdnaError> {
            let len = data.len();
            if len == 0 { return Ok(0.0); }

            let data_buf = self.device.new_buffer_with_data(
                data.as_ptr() as *const std::ffi::c_void,
                (len * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );

            let thread_count = 256u64.min(len as u64);
            let result_buf = self.device.new_buffer(
                (thread_count * 8) as u64,
                MTLResourceOptions::StorageModeShared,
            );

            let cmd_buf = self.queue.new_command_buffer();
            let encoder = cmd_buf.new_compute_command_encoder();
            encoder.set_compute_pipeline_state(&self.reduce_pso);
            encoder.set_buffer(0, Some(&data_buf), 0);
            encoder.set_buffer(1, Some(&result_buf), 0);

            let grid_size = MTLSize::new(thread_count, 1, 1);
            let group_size = MTLSize::new(thread_count.min(64), 1, 1);
            encoder.dispatch_threads(grid_size, group_size);
            encoder.end_encoding();
            cmd_buf.commit();
            cmd_buf.wait_until_completed();

            let ptr = result_buf.contents() as *const f64;
            let mut sum = 0.0;
            for i in 0..thread_count as usize {
                sum += unsafe { *ptr.add(i) };
            }
            Ok(sum)
        }
    }

    fn compile_kernel(device: &Device, source: &str, name: &str) -> Result<ComputePipelineState, HwdnaError> {
        let lib = device.new_library_with_source(source, &CompileOptions::new())
            .map_err(|_| HwdnaError::NoSupportedKernels)?;
        let func = lib.get_function(name, None)
            .map_err(|_| HwdnaError::NoSupportedKernels)?;
        device.new_compute_pipeline_state_with_function(&func)
            .map_err(|_| HwdnaError::NoSupportedKernels)
    }
}

/// GPU runtime wrapping Metal or other backends.
pub struct GpuRuntime {
    backend: GpuBackend,
    initialized: bool,
    #[cfg(all(target_os = "macos", feature = "metal"))]
    metal: Option<metal_kernels::MetalCompute>,
    #[cfg(feature = "vulkan")]
    vulkan: Option<VulkanRuntime>,
}

impl GpuRuntime {
    pub fn new() -> Self {
        let backend = detect_gpu();
        Self {
            backend,
            initialized: false,
            #[cfg(all(target_os = "macos", feature = "metal"))]
            metal: None,
            #[cfg(feature = "vulkan")]
            vulkan: None,
        }
    }

    pub fn backend(&self) -> GpuBackend { self.backend }
    pub fn is_available(&self) -> bool { self.backend != GpuBackend::None }

    pub fn init(&mut self) -> Result<(), HwdnaError> {
        if self.initialized { return Ok(()); }
        match self.backend {
            GpuBackend::None => Err(HwdnaError::NoSupportedKernels),
            GpuBackend::Metal => {
                #[cfg(all(target_os = "macos", feature = "metal"))]
                {
                    let mc = metal_kernels::MetalCompute::new()?;
                    self.metal = Some(mc);
                    self.initialized = true;
                    return Ok(());
                }
                #[cfg(not(all(target_os = "macos", feature = "metal")))]
                Err(HwdnaError::NoSupportedKernels)
            }
            GpuBackend::Cuda => {
                Err(HwdnaError::NoSupportedKernels)
            }
            GpuBackend::Vulkan => {
                #[cfg(feature = "vulkan")]
                {
                    let mut vk = VulkanRuntime::new();
                    vk.init()?;
                    self.vulkan = Some(vk);
                    self.initialized = true;
                    return Ok(());
                }
                #[cfg(not(feature = "vulkan"))]
                Err(HwdnaError::NoSupportedKernels)
            }
        }
    }

    #[cfg_attr(not(any(all(target_os = "macos", feature = "metal"), feature = "vulkan")), allow(unused_variables))]
    pub fn dot_f64(&self, a: &[f64], b: &[f64]) -> Result<f64, HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(all(target_os = "macos", feature = "metal"))]
        if let Some(ref m) = self.metal { return m.dot(a, b); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.dot_f64(a, b); }
        Err(HwdnaError::NoSupportedKernels)
    }

    #[cfg_attr(not(any(all(target_os = "macos", feature = "metal"), feature = "vulkan")), allow(unused_variables))]
    pub fn matmul_f64(&self, a: &[f64], b: &[f64], c: &mut [f64], n: usize) -> Result<(), HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(all(target_os = "macos", feature = "metal"))]
        if let Some(ref mc) = self.metal { return mc.matmul(a, b, c, n); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.matmul_f64(a, b, c, n); }
        Err(HwdnaError::NoSupportedKernels)
    }

    #[cfg_attr(not(any(all(target_os = "macos", feature = "metal"), feature = "vulkan")), allow(unused_variables))]
    pub fn reduce_sum_f64(&self, data: &[f64]) -> Result<f64, HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(all(target_os = "macos", feature = "metal"))]
        if let Some(ref m) = self.metal { return m.reduce_sum(data); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.reduce_sum_f64(data); }
        Err(HwdnaError::NoSupportedKernels)
    }

    #[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
    pub fn dot_f32(&self, a: &[f32], b: &[f32]) -> Result<f32, HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.dot_f32(a, b); }
        Err(HwdnaError::NoSupportedKernels)
    }

    #[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
    pub fn matmul_f32(&self, a: &[f32], b: &[f32], c: &mut [f32], n: usize) -> Result<(), HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.matmul_f32(a, b, c, n); }
        Err(HwdnaError::NoSupportedKernels)
    }

    #[cfg_attr(not(feature = "vulkan"), allow(unused_variables))]
    pub fn reduce_sum_f32(&self, data: &[f32]) -> Result<f32, HwdnaError> {
        if !self.initialized { return Err(HwdnaError::NoSupportedKernels); }
        #[cfg(feature = "vulkan")]
        if let Some(ref vk) = self.vulkan { return vk.reduce_sum_f32(data); }
        Err(HwdnaError::NoSupportedKernels)
    }
}