diff options
Diffstat (limited to 'candle-core/src/metal_backend.rs')
| -rw-r--r-- | candle-core/src/metal_backend.rs | 178 |
1 files changed, 144 insertions, 34 deletions
diff --git a/candle-core/src/metal_backend.rs b/candle-core/src/metal_backend.rs index 73a532e6..24beeb7a 100644 --- a/candle-core/src/metal_backend.rs +++ b/candle-core/src/metal_backend.rs @@ -106,7 +106,7 @@ pub struct MetalDevice { /// Whenever we actually allocate a new buffer, we make a full sweep to cleanup unused buffers /// (strong_count = 1). buffers: AllocatedBuffers, - + /// Seed for random number generation. seed: Arc<Mutex<u64>>, } @@ -355,7 +355,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "affine_f32", DType::F16 => "affine_f16", - dtype => crate::bail!("Affine {dtype:?}"), + dtype => crate::bail!("Metal contiguous affine {dtype:?} not implemented"), }; candle_metal_kernels::call_affine( &device.device, @@ -373,7 +373,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "affine_f32_strided", DType::F16 => "affine_f16_strided", - dtype => crate::bail!("Affine {dtype:?}"), + dtype => crate::bail!("Metal strided affine {dtype:?} not implemented"), }; candle_metal_kernels::call_affine_strided( &device.device, @@ -406,7 +406,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "powf_f32", DType::F16 => "powf_f16", - dtype => crate::bail!("Powf {dtype:?}"), + dtype => crate::bail!("Metal contiguous powf {dtype:?} not implemented"), }; candle_metal_kernels::call_powf( &device.device, @@ -423,7 +423,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "powf_f32_strided", DType::F16 => "powf_f16_strided", - dtype => crate::bail!("Powf {dtype:?}"), + dtype => crate::bail!("Metal strided powf {dtype:?} not implemented"), }; candle_metal_kernels::call_powf_strided( &device.device, @@ -455,7 +455,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "elu_f32", DType::F16 => "elu_f16", - dtype => crate::bail!("Powf {dtype:?}"), + dtype => crate::bail!("Metal contiguous elu {dtype:?} not implemented"), }; candle_metal_kernels::call_elu( &device.device, @@ -472,7 +472,7 @@ impl BackendStorage for MetalStorage { let name = match self.dtype { DType::F32 => "elu_f32_strided", DType::F16 => "elu_f16_strided", - dtype => crate::bail!("Powf {dtype:?}"), + dtype => crate::bail!("Metal strided elu {dtype:?} not implemented"), }; candle_metal_kernels::call_elu_strided( &device.device, @@ -534,7 +534,17 @@ impl BackendStorage for MetalStorage { (ReduceOp::Max, DType::BF16) => ("fast_max_bf16_strided", true, false), (ReduceOp::ArgMin, DType::BF16) => ("fast_argmin_bf16_strided", true, true), (ReduceOp::ArgMax, DType::BF16) => ("fast_argmax_bf16_strided", true, true), - (k, dtype) => crate::bail!("Reduce op for non float {k:?} {dtype:?}"), + (ReduceOp::Sum, DType::I64) => ("fast_sum_i64_strided", false, false), + (ReduceOp::Min, DType::I64) => ("fast_min_i64_strided", true, false), + (ReduceOp::Max, DType::I64) => ("fast_max_i64_strided", true, false), + (ReduceOp::ArgMin, DType::I64) => ("fast_argmin_i64_strided", true, true), + (ReduceOp::ArgMax, DType::I64) => ("fast_argmax_i64_strided", true, true), + (ReduceOp::Sum, DType::U8) => ("fast_sum_u8_strided", false, false), + (ReduceOp::Min, DType::U8) => ("fast_min_u8_strided", true, false), + (ReduceOp::Max, DType::U8) => ("fast_max_u8_strided", true, false), + (ReduceOp::ArgMin, DType::U8) => ("fast_argmin_u8_strided", true, true), + (ReduceOp::ArgMax, DType::U8) => ("fast_argmax_u8_strided", true, true), + (k, dtype) => crate::bail!("Metal reduce op {k:?} {dtype:?} not implemented"), }; if check_empty && layout.shape().elem_count() == 0 { Err(crate::Error::EmptyTensor { op: "reduce" }.bt())? @@ -581,11 +591,18 @@ impl BackendStorage for MetalStorage { let kernel_name = match (self.dtype, dtype) { (DType::U32, DType::F32) => "cast_u32_f32", (DType::U32, DType::U8) => "cast_u32_u8", + (DType::U32, DType::I64) => "cast_u32_i64", (DType::U8, DType::U32) => "cast_u8_u32", (DType::U8, DType::F32) => "cast_u8_f32", + (DType::U8, DType::I64) => "cast_u8_i64", (DType::F32, DType::F16) => "cast_f32_f16", (DType::F16, DType::F32) => "cast_f16_f32", - (left, right) => crate::bail!("to dtype {left:?} - {right:?}"), + (DType::I64, DType::F32) => "cast_i64_f32", + (DType::F32, DType::BF16) => "cast_f32_bf16", + (DType::BF16, DType::F32) => "cast_bf16_f32", + (left, right) => { + crate::bail!("Metal contiguous to_dtype {left:?} {right:?} not implemented") + } }; candle_metal_kernels::call_cast_contiguous( &device.device, @@ -602,11 +619,18 @@ impl BackendStorage for MetalStorage { let kernel_name = match (self.dtype, dtype) { (DType::U32, DType::F32) => "cast_u32_f32_strided", (DType::U32, DType::U8) => "cast_u32_u8_strided", + (DType::U32, DType::I64) => "cast_u32_i64_strided", (DType::U8, DType::U32) => "cast_u8_u32_strided", (DType::U8, DType::F32) => "cast_u8_f32_strided", + (DType::U8, DType::I64) => "cast_u8_i64_strided", (DType::F32, DType::F16) => "cast_f32_f16_strided", (DType::F16, DType::F32) => "cast_f16_f32_strided", - (left, right) => crate::bail!("to dtype {left:?} - {right:?}"), + (DType::I64, DType::F32) => "cast_i64_f32_strided", + (DType::F32, DType::BF16) => "cast_f32_bf16_strided", + (DType::BF16, DType::F32) => "cast_bf16_f32_strided", + (left, right) => { + crate::bail!("Metal strided to_dtype {left:?} {right:?} not implemented") + } }; candle_metal_kernels::call_cast_strided( &device.device, @@ -647,9 +671,11 @@ impl BackendStorage for MetalStorage { ("ugelu", DType::F32) => contiguous::gelu::FLOAT, ("ugelu_erf", DType::F32) => contiguous::gelu_erf::FLOAT, ("uerf", DType::F32) => contiguous::erf::FLOAT, + ("uabs", DType::F32) => contiguous::abs::FLOAT, ("uceil", DType::F32) => contiguous::ceil::FLOAT, ("ufloor", DType::F32) => contiguous::floor::FLOAT, ("uround", DType::F32) => contiguous::round::FLOAT, + ("urecip", DType::F32) => contiguous::recip::FLOAT, ("utanh", DType::F32) => contiguous::tanh::FLOAT, ("ucos", DType::F16) => contiguous::cos::HALF, ("usin", DType::F16) => contiguous::sin::HALF, @@ -661,11 +687,15 @@ impl BackendStorage for MetalStorage { ("ugelu", DType::F16) => contiguous::gelu::HALF, ("ugelu_erf", DType::F16) => contiguous::gelu_erf::HALF, ("uerf", DType::F16) => contiguous::erf::HALF, + ("uabs", DType::F16) => contiguous::abs::HALF, ("uceil", DType::F16) => contiguous::ceil::HALF, ("ufloor", DType::F16) => contiguous::floor::HALF, ("uround", DType::F16) => contiguous::round::HALF, + ("urecip", DType::F16) => contiguous::recip::HALF, ("utanh", DType::F16) => contiguous::tanh::HALF, - (name, dtype) => crate::bail!("Match {name} - {dtype:?}"), + (name, dtype) => { + crate::bail!("Metal contiguous unary {name} {dtype:?} not implemented") + } }; candle_metal_kernels::call_unary_contiguous( &device.device, @@ -690,6 +720,7 @@ impl BackendStorage for MetalStorage { ("ugelu", DType::F32) => strided::gelu::FLOAT, ("ugelu_erf", DType::F32) => strided::gelu_erf::FLOAT, ("uerf", DType::F32) => strided::erf::FLOAT, + ("uabs", DType::F32) => strided::abs::FLOAT, ("uceil", DType::F32) => strided::ceil::FLOAT, ("ufloor", DType::F32) => strided::floor::FLOAT, ("uround", DType::F32) => strided::round::FLOAT, @@ -703,10 +734,13 @@ impl BackendStorage for MetalStorage { ("ugelu", DType::F16) => strided::gelu::HALF, ("ugelu_erf", DType::F16) => strided::gelu_erf::HALF, ("uerf", DType::F16) => strided::erf::HALF, + ("uabs", DType::F16) => strided::abs::HALF, ("uceil", DType::F16) => strided::ceil::HALF, ("ufloor", DType::F16) => strided::floor::HALF, ("uround", DType::F16) => strided::round::HALF, - (name, dtype) => crate::bail!("Match {name} - {dtype:?}"), + (name, dtype) => { + crate::bail!("Metal strided unary {name} {dtype:?} not implemented") + } }; candle_metal_kernels::call_unary_strided( &device.device, @@ -759,7 +793,10 @@ impl BackendStorage for MetalStorage { let name = match (self.dtype, t.dtype()) { (DType::U8, DType::F32) => "where_u8_f32", (DType::U8, DType::F16) => "where_u8_f16", - (left, right) => crate::bail!("where {left:?} - {right:?} not implemented"), + (DType::U8, DType::I64) => "where_u8_i64", + (DType::U8, DType::U32) => "where_u8_u32", + (DType::U8, DType::U8) => "where_u8_u8", + (left, right) => crate::bail!("Metal where_cond {left:?} {right:?} not implemented"), }; candle_metal_kernels::call_where_cond_strided( &device.device, @@ -806,7 +843,7 @@ impl BackendStorage for MetalStorage { let command_buffer = self.device.command_buffer()?; let name = match self.dtype { DType::F32 => "im2col1d_f32", - dtype => crate::bail!("conv1d metal {dtype:?} not implemented"), + dtype => crate::bail!("Metal conv1d {dtype:?} not implemented"), }; candle_metal_kernels::call_im2col1d_strided( &self.device.device, @@ -859,7 +896,7 @@ impl BackendStorage for MetalStorage { _kernel_l: &Layout, _params: &ParamsConvTranspose1D, ) -> Result<Self> { - crate::bail!("conv_transpose1d metal") + crate::bail!("Metal conv_transpose1d not implemented") } fn conv2d( @@ -890,7 +927,7 @@ impl BackendStorage for MetalStorage { let command_buffer = self.device.command_buffer()?; let name = match self.dtype { DType::F32 => "im2col_f32", - dtype => crate::bail!("conv1d metal {dtype:?} not implemented"), + dtype => crate::bail!("Metal conv2d {dtype:?} not implemented"), }; candle_metal_kernels::call_im2col_strided( &self.device.device, @@ -946,19 +983,19 @@ impl BackendStorage for MetalStorage { _kernel_l: &Layout, _params: &ParamsConvTranspose2D, ) -> Result<Self> { - crate::bail!("conv_tranpose2d metal") + crate::bail!("Metal conv_tranpose2d not implemented") } fn avg_pool2d(&self, _: &Layout, _: (usize, usize), _: (usize, usize)) -> Result<Self> { - crate::bail!("avg_pool2d metal") + crate::bail!("Metal avg_pool2d not implemented") } fn max_pool2d(&self, _: &Layout, _: (usize, usize), _: (usize, usize)) -> Result<Self> { - crate::bail!("max_pool2d metal") + crate::bail!("Metal max_pool2d not implemented") } fn upsample_nearest1d(&self, _: &Layout, _: usize) -> Result<Self> { - crate::bail!("upsample_nearest1d metal") + crate::bail!("Metal upsample_nearest1d not implemented") } fn upsample_nearest2d(&self, inp_l: &Layout, out_w: usize, out_h: usize) -> Result<Self> { @@ -971,7 +1008,7 @@ impl BackendStorage for MetalStorage { } let name = match self.dtype { DType::F32 => "upsample_nearest2d_f32", - dtype => crate::bail!("Not implemented {dtype:?} for upsample_nearest2d, metal"), + dtype => crate::bail!("Metal upsample_nearest2d {dtype:?} not implemented"), }; let dst_el = out_w * out_h * dims[0] * dims[1]; @@ -1009,7 +1046,7 @@ impl BackendStorage for MetalStorage { let name = match (ids.dtype, self.dtype) { (DType::U32, DType::F32) => "gather_u32_f32", (DType::U32, DType::F16) => "gather_u32_f16", - (left, right) => crate::bail!("gather metal {left:?} {right:?} not implemented"), + (left, right) => crate::bail!("Metal gather {left:?} {right:?} not implemented"), }; let command_buffer = self.device.command_buffer()?; candle_metal_kernels::call_gather( @@ -1082,7 +1119,7 @@ impl BackendStorage for MetalStorage { && ids_l.is_contiguous() && ids_l.start_offset() == 0) { - crate::bail!("Non contiguous index select not implemented"); + crate::bail!("Metal strided index_select not implemented"); } let left_size: usize = src_l.dims()[..dim].iter().product(); let right_size: usize = src_l.dims()[dim + 1..].iter().product(); @@ -1094,7 +1131,9 @@ impl BackendStorage for MetalStorage { let name = match (ids.dtype, self.dtype) { (DType::U32, DType::F32) => "is_u32_f32", (DType::U32, DType::F16) => "is_u32_f16", - (left, right) => crate::bail!("index select metal {left:?} {right:?}"), + (left, right) => { + crate::bail!("Metal contiguous index_select {left:?} {right:?} not implemented") + } }; let command_buffer = self.device.command_buffer()?; candle_metal_kernels::call_index_select( @@ -1135,7 +1174,7 @@ impl BackendStorage for MetalStorage { let name = match (ids.dtype, self.dtype) { (DType::U32, DType::F32) => "ia_u32_f32", _ => Err(MetalError::UnexpectedDType { - msg: "index-add ids should be u8/u32/i64", + msg: "index-add ids should be u32", expected: DType::U32, got: ids.dtype(), })?, @@ -1216,9 +1255,10 @@ impl BackendStorage for MetalStorage { DType::F32 => candle_metal_kernels::unary::strided::copy::FLOAT, DType::F16 => candle_metal_kernels::unary::strided::copy::HALF, DType::BF16 => candle_metal_kernels::unary::strided::copy::BFLOAT, + DType::I64 => candle_metal_kernels::unary::strided::copy::I64, DType::U32 => candle_metal_kernels::unary::strided::copy::U32, DType::U8 => candle_metal_kernels::unary::strided::copy::U8, - dtype => crate::bail!("copy_strided not implemented for {dtype:?}"), + dtype => crate::bail!("Metal copy_strided {dtype:?} not implemented"), }; candle_metal_kernels::call_unary_strided( &self.device.device, @@ -1290,7 +1330,39 @@ impl MetalStorage { ("lt", DType::F16) => (contiguous::lt::HALF, DType::U8), ("ge", DType::F16) => (contiguous::ge::HALF, DType::U8), ("gt", DType::F16) => (contiguous::gt::HALF, DType::U8), - (name, dtype) => crate::bail!("Binary {name} - {dtype:?} not implemented"), + ("add", DType::I64) => (contiguous::add::I64, self.dtype), + ("sub", DType::I64) => (contiguous::sub::I64, self.dtype), + ("mul", DType::I64) => (contiguous::mul::I64, self.dtype), + ("div", DType::I64) => (contiguous::div::I64, self.dtype), + ("eq", DType::I64) => (contiguous::eq::I64, DType::U8), + ("ne", DType::I64) => (contiguous::ne::I64, DType::U8), + ("le", DType::I64) => (contiguous::le::I64, DType::U8), + ("lt", DType::I64) => (contiguous::lt::I64, DType::U8), + ("ge", DType::I64) => (contiguous::ge::I64, DType::U8), + ("gt", DType::I64) => (contiguous::gt::I64, DType::U8), + ("add", DType::U32) => (contiguous::add::U32, self.dtype), + ("sub", DType::U32) => (contiguous::sub::U32, self.dtype), + ("mul", DType::U32) => (contiguous::mul::U32, self.dtype), + ("div", DType::U32) => (contiguous::div::U32, self.dtype), + ("eq", DType::U32) => (contiguous::eq::U32, DType::U8), + ("ne", DType::U32) => (contiguous::ne::U32, DType::U8), + ("le", DType::U32) => (contiguous::le::U32, DType::U8), + ("lt", DType::U32) => (contiguous::lt::U32, DType::U8), + ("ge", DType::U32) => (contiguous::ge::U32, DType::U8), + ("gt", DType::U32) => (contiguous::gt::U32, DType::U8), + ("add", DType::U8) => (contiguous::add::U8, self.dtype), + ("sub", DType::U8) => (contiguous::sub::U8, self.dtype), + ("mul", DType::U8) => (contiguous::mul::U8, self.dtype), + ("div", DType::U8) => (contiguous::div::U8, self.dtype), + ("eq", DType::U8) => (contiguous::eq::U8, DType::U8), + ("ne", DType::U8) => (contiguous::ne::U8, DType::U8), + ("le", DType::U8) => (contiguous::le::U8, DType::U8), + ("lt", DType::U8) => (contiguous::lt::U8, DType::U8), + ("ge", DType::U8) => (contiguous::ge::U8, DType::U8), + ("gt", DType::U8) => (contiguous::gt::U8, DType::U8), + (name, dtype) => { + crate::bail!("Metal contiguous binary {name} {dtype:?} not implemented") + } }; let buffer = device.new_buffer(el_count, dtype, op)?; candle_metal_kernels::call_binary_contiguous( @@ -1333,7 +1405,45 @@ impl MetalStorage { ("lt", DType::F16) => (strided::lt::HALF, DType::U8), ("ge", DType::F16) => (strided::ge::HALF, DType::U8), ("gt", DType::F16) => (strided::gt::HALF, DType::U8), - (name, dtype) => crate::bail!("Binary strided {name} - {dtype:?} not implemented"), + ("badd", DType::I64) => (strided::add::I64, self.dtype), + ("bsub", DType::I64) => (strided::sub::I64, self.dtype), + ("bmul", DType::I64) => (strided::mul::I64, self.dtype), + ("bdiv", DType::I64) => (strided::div::I64, self.dtype), + ("bminimum", DType::I64) => (strided::min::I64, self.dtype), + ("bmaximum", DType::I64) => (strided::max::I64, self.dtype), + ("eq", DType::I64) => (strided::eq::I64, DType::U8), + ("ne", DType::I64) => (strided::ne::I64, DType::U8), + ("le", DType::I64) => (strided::le::I64, DType::U8), + ("lt", DType::I64) => (strided::lt::I64, DType::U8), + ("ge", DType::I64) => (strided::ge::I64, DType::U8), + ("gt", DType::I64) => (strided::gt::I64, DType::U8), + ("badd", DType::U32) => (strided::add::U32, self.dtype), + ("bsub", DType::U32) => (strided::sub::U32, self.dtype), + ("bmul", DType::U32) => (strided::mul::U32, self.dtype), + ("bdiv", DType::U32) => (strided::div::U32, self.dtype), + ("bminimum", DType::U32) => (strided::min::U32, self.dtype), + ("bmaximum", DType::U32) => (strided::max::U32, self.dtype), + ("eq", DType::U32) => (strided::eq::U32, DType::U8), + ("ne", DType::U32) => (strided::ne::U32, DType::U8), + ("le", DType::U32) => (strided::le::U32, DType::U8), + ("lt", DType::U32) => (strided::lt::U32, DType::U8), + ("ge", DType::U32) => (strided::ge::U32, DType::U8), + ("gt", DType::U32) => (strided::gt::U32, DType::U8), + ("badd", DType::U8) => (strided::add::U8, self.dtype), + ("bsub", DType::U8) => (strided::sub::U8, self.dtype), + ("bmul", DType::U8) => (strided::mul::U8, self.dtype), + ("bdiv", DType::U8) => (strided::div::U8, self.dtype), + ("bminimum", DType::U8) => (strided::min::U8, self.dtype), + ("bmaximum", DType::U8) => (strided::max::U8, self.dtype), + ("eq", DType::U8) => (strided::eq::U8, DType::U8), + ("ne", DType::U8) => (strided::ne::U8, DType::U8), + ("le", DType::U8) => (strided::le::U8, DType::U8), + ("lt", DType::U8) => (strided::lt::U8, DType::U8), + ("ge", DType::U8) => (strided::ge::U8, DType::U8), + ("gt", DType::U8) => (strided::gt::U8, DType::U8), + (name, dtype) => { + crate::bail!("Metal strided binary {name} {dtype:?} not implemented") + } }; let buffer = device.new_buffer(el_count, dtype, op)?; candle_metal_kernels::call_binary_strided( @@ -1389,12 +1499,6 @@ impl BackendDevice for MetalDevice { }) } - fn set_seed(&self, seed: u64) -> Result<()> { - let mut s = self.seed.try_lock().map_err(MetalError::from)?; - *s = seed; - Ok(()) - } - fn location(&self) -> crate::DeviceLocation { crate::DeviceLocation::Metal { gpu_id: self.registry_id() as usize, @@ -1504,6 +1608,12 @@ impl BackendDevice for MetalDevice { Ok(Self::Storage::new(buffer, self.clone(), dtype)) } + + fn set_seed(&self, seed: u64) -> Result<()> { + let mut s = self.seed.try_lock().map_err(MetalError::from)?; + *s = seed; + Ok(()) + } } fn read_to_vec<T: Clone>(buffer: &Buffer, n: usize) -> Vec<T> { |