| 98 | |
| 99 | #[inline] |
| 100 | fn dispatch_binary( |
| 101 | s1: &candle_core::MetalStorage, l1: &Layout, |
| 102 | s2: &candle_core::MetalStorage, l2: &Layout, |
| 103 | f32_kernel: &'static str, f16_kernel: &'static str, label: &'static str, |
| 104 | ) -> Result<(candle_core::MetalStorage, Shape)> { |
| 105 | let device = s1.device(); |
| 106 | let el = l1.shape().elem_count(); |
| 107 | let kernel_name: &'static str = match s1.dtype() { |
| 108 | DType::F32 => f32_kernel, |
| 109 | DType::F16 => f16_kernel, |
| 110 | dt => candle_core::bail!("{label} metal: unsupported dtype {dt:?}"), |
| 111 | }; |
| 112 | let pipeline = PIPELINE_CACHE.get_or_create(device, kernel_name)?; |
| 113 | let output = device.new_buffer(el, s1.dtype(), label)?; |
| 114 | let encoder = device.command_encoder()?; |
| 115 | encoder.set_compute_pipeline_state(&pipeline); |
| 116 | let off1 = l1.start_offset() * s1.dtype().size_in_bytes(); |
| 117 | let off2 = l2.start_offset() * s2.dtype().size_in_bytes(); |
| 118 | candle_metal_kernels::utils::set_param(&encoder, 0, (s1.buffer(), off1)); |
| 119 | candle_metal_kernels::utils::set_param(&encoder, 1, (s2.buffer(), off2)); |
| 120 | candle_metal_kernels::utils::set_param(&encoder, 2, (&*output, 0usize)); |
| 121 | candle_metal_kernels::utils::set_param(&encoder, 3, el as u32); |
| 122 | let grid = objc2_metal::MTLSize { width: el, height: 1, depth: 1 }; |
| 123 | let group = candle_metal_kernels::utils::get_block_dims(el, 1, 1); |
| 124 | encoder.dispatch_threads(grid, group); |
| 125 | Ok((candle_core::MetalStorage::new(output, device.clone(), el, s1.dtype()), l1.shape().clone())) |
| 126 | } |
| 127 | |
| 128 | // ─── Helper: dispatch an elementwise 3-input kernel ───────────────── |
| 129 | |