| 45 | struct S; |
| 46 | impl Map2 for S { |
| 47 | fn f<T: DeviceRepr + WithDType>( |
| 48 | &self, |
| 49 | gate: &CudaSlice<T>, |
| 50 | l1: &Layout, |
| 51 | up: &CudaSlice<T>, |
| 52 | l2: &Layout, |
| 53 | dev: &CudaDevice, |
| 54 | ) -> Result<CudaSlice<T>> { |
| 55 | let gate = match l1.contiguous_offsets() { |
| 56 | Some((o1, o2)) => gate.slice(o1..o2), |
| 57 | None => candle_core::bail!("silu_mul: gate must be contiguous"), |
| 58 | }; |
| 59 | let up = match l2.contiguous_offsets() { |
| 60 | Some((o1, o2)) => up.slice(o1..o2), |
| 61 | None => candle_core::bail!("silu_mul: up must be contiguous"), |
| 62 | }; |
| 63 | let el = l1.shape().elem_count(); |
| 64 | let cfg = LaunchConfig::for_num_elems(el as u32); |
| 65 | let func = dev.get_or_load_custom_func( |
| 66 | &kernel_name::<T>("silu_mul"), |
| 67 | "cake_fused_ops", |
| 68 | FUSED_OPS_PTX, |
| 69 | )?; |
| 70 | let out = unsafe { dev.alloc::<T>(el)? }; |
| 71 | let mut builder = func.builder(); |
| 72 | builder.arg(&el); |
| 73 | builder.arg(&gate); |
| 74 | builder.arg(&up); |
| 75 | builder.arg(&out); |
| 76 | unsafe { builder.launch(cfg) }.w()?; |
| 77 | Ok(out) |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | use candle_core::backend::BackendStorage; |