| 169 | |
| 170 | #[inline] |
| 171 | fn dispatch_unary( |
| 172 | s: &candle_core::MetalStorage, l: &Layout, |
| 173 | f32_kernel: &'static str, f16_kernel: &'static str, label: &'static str, |
| 174 | ) -> Result<(candle_core::MetalStorage, Shape)> { |
| 175 | let device = s.device(); |
| 176 | let el = l.shape().elem_count(); |
| 177 | let kernel_name: &'static str = match s.dtype() { |
| 178 | DType::F32 => f32_kernel, |
| 179 | DType::F16 => f16_kernel, |
| 180 | dt => candle_core::bail!("{label} metal: unsupported dtype {dt:?}"), |
| 181 | }; |
| 182 | let pipeline = PIPELINE_CACHE.get_or_create(device, kernel_name)?; |
| 183 | let output = device.new_buffer(el, s.dtype(), label)?; |
| 184 | let encoder = device.command_encoder()?; |
| 185 | encoder.set_compute_pipeline_state(&pipeline); |
| 186 | let offset = l.start_offset() * s.dtype().size_in_bytes(); |
| 187 | candle_metal_kernels::utils::set_param(&encoder, 0, (s.buffer(), offset)); |
| 188 | candle_metal_kernels::utils::set_param(&encoder, 1, (&*output, 0usize)); |
| 189 | candle_metal_kernels::utils::set_param(&encoder, 2, el as u32); |
| 190 | let grid = objc2_metal::MTLSize { width: el, height: 1, depth: 1 }; |
| 191 | let group = candle_metal_kernels::utils::get_block_dims(el, 1, 1); |
| 192 | encoder.dispatch_threads(grid, group); |
| 193 | Ok((candle_core::MetalStorage::new(output, device.clone(), el, s.dtype()), l.shape().clone())) |
| 194 | } |
| 195 | |
| 196 | // ─── CustomOp structs ─────────────────────────────────────────────── |
| 197 | |