huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 02:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Tmp state.

Browse Source
This commit is contained in:
Nicolas Patry
2023-11-10 15:35:46 +01:00
committed by Nicolas Patry
parent f710fab02e
commit d46670f7c0
14 changed files with 699 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-core/src/device.rs
View File

@ -8,7 +8,7 @@ use crate::{CpuStorage, DType, Result, Shape, Storage, WithDType};
pub enum DeviceLocation {
Cpu,
Cuda { gpu_id: usize },
Metal,
Metal { gpu_id: usize },
}
#[derive(Debug, Clone)]

8
candle-core/src/display.rs
View File

@ -14,7 +14,9 @@ impl Tensor {
crate::DeviceLocation::Cuda { gpu_id } => {
format!(", cuda:{}", gpu_id)
}
_ => todo!(),
crate::DeviceLocation::Metal { gpu_id } => {
format!(", metal:{}", gpu_id)
}
};
write!(f, "Tensor[")?;
@ -477,7 +479,9 @@ impl std::fmt::Display for Tensor {
crate::DeviceLocation::Cuda { gpu_id } => {
format!(", cuda:{}", gpu_id)
}
crate::DeviceLocation::Metal => todo!(),
crate::DeviceLocation::Metal { gpu_id } => {
format!(", metal:{}", gpu_id)
}
};
write!(

157
candle-core/src/metal_backend.rs
View File

@ -100,11 +100,30 @@ impl BackendStorage for MetalStorage {
}
fn to_cpu_storage(&self) -> Result<CpuStorage> {
// TODO Is this necessary
// self.buffer.synchronize();
match self.dtype {
DType::U8 => Ok(CpuStorage::U8(
self.buffer.read_to_vec(self.buffer.length() as usize / 1),
)),
DType::U32 => Ok(CpuStorage::U32(
self.buffer.read_to_vec(self.buffer.length() as usize / 4),
)),
DType::I64 => Ok(CpuStorage::I64(
self.buffer.read_to_vec(self.buffer.length() as usize / 8),
)),
DType::F16 => Ok(CpuStorage::F16(
self.buffer.read_to_vec(self.buffer.length() as usize / 2),
)),
DType::BF16 => Ok(CpuStorage::BF16(
self.buffer.read_to_vec(self.buffer.length() as usize / 2),
)),
DType::F32 => Ok(CpuStorage::F32(
self.buffer.read_to_vec(self.buffer.length() as usize / 4),
)),
dtype => todo!("Unsupported dtype {dtype:?}"),
DType::F64 => Ok(CpuStorage::F64(
self.buffer.read_to_vec(self.buffer.length() as usize / 8),
)),
}
}
@ -132,6 +151,7 @@ impl BackendStorage for MetalStorage {
)
.unwrap();
command_buffer.commit();
command_buffer.wait_until_completed();
return Ok(Self {
buffer,
device: device.clone(),
@ -200,6 +220,7 @@ impl BackendStorage for MetalStorage {
)
.map_err(MetalError::from)?;
command_buffer.commit();
command_buffer.wait_until_completed();
Ok(Self {
buffer,
@ -242,6 +263,7 @@ impl BackendStorage for MetalStorage {
}
command_buffer.commit();
command_buffer.wait_until_completed();
// command_buffer.wait_until_scheduled();
// debug!(
// "cast {:?} - {:?} - {:?}",
@ -289,6 +311,7 @@ impl BackendStorage for MetalStorage {
todo!("TODO Implement the kernel calling {}", B::KERNEL);
}
command_buffer.commit();
command_buffer.wait_until_completed();
Ok(Self {
buffer,
@ -361,6 +384,7 @@ impl BackendStorage for MetalStorage {
.map_err(MetalError::from)?;
}
command_buffer.commit();
command_buffer.wait_until_completed();
Ok(Self {
buffer,
@ -400,6 +424,7 @@ impl BackendStorage for MetalStorage {
)
.map_err(MetalError::from)?;
command_buffer.commit();
command_buffer.wait_until_completed();
Ok(Self {
buffer,
device,
@ -489,6 +514,7 @@ impl BackendStorage for MetalStorage {
let dtype = self.dtype;
let device = self.device();
let mut buffer = device.new_buffer(dst_el, dtype);
let out = self.to_cpu_storage().unwrap();
let name = match (ids.dtype, self.dtype) {
(DType::U32, DType::F32) => "is_u32_f32",
(left, right) => todo!("index select metal {left:?} {right:?}"),
@ -508,6 +534,7 @@ impl BackendStorage for MetalStorage {
)
.map_err(MetalError::from)?;
command_buffer.commit();
command_buffer.wait_until_completed();
Ok(Self {
buffer,
device: device.clone(),
@ -556,39 +583,42 @@ impl BackendStorage for MetalStorage {
if el_count == 0 {
return Ok(());
}
if src_l.is_contiguous() {
let command_buffer = self.device.command_queue.new_command_buffer();
let blip = command_buffer.new_blit_command_encoder();
blip.copy_from_buffer(
&self.buffer,
src_l.start_offset() as u64,
&dst.buffer,
dst_offset as u64,
self.buffer.length(),
);
} else {
let command_buffer = self.device.command_queue.new_command_buffer();
let kernel_name = match self.dtype {
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 => todo!("copy_strided not implemented for {dtype:?}"),
};
candle_metal_kernels::call_unary_strided(
&self.device.device,
&command_buffer,
&self.device.kernels,
kernel_name,
src_l.dims(),
&self.buffer,
&src_l.stride(),
src_l.start_offset(),
&mut dst.buffer,
dst_offset,
)
.map_err(MetalError::from)?;
command_buffer.commit();
}
// todo!("Copy strided {:?}", src_l.is_contiguous());
// if src_l.is_contiguous() {
// let command_buffer = self.device.command_queue.new_command_buffer();
// let blip = command_buffer.new_blit_command_encoder();
// blip.copy_from_buffer(
// &self.buffer,
// src_l.start_offset() as u64,
// &dst.buffer,
// dst_offset as u64,
// self.buffer.length(),
// );
// } else {
let command_buffer = self.device.command_queue.new_command_buffer();
let kernel_name = match self.dtype {
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 => todo!("copy_strided not implemented for {dtype:?}"),
};
candle_metal_kernels::call_unary_strided(
&self.device.device,
&command_buffer,
&self.device.kernels,
kernel_name,
src_l.dims(),
&self.buffer,
&src_l.stride(),
src_l.start_offset(),
&mut dst.buffer,
dst_offset,
)
.map_err(MetalError::from)?;
command_buffer.commit();
command_buffer.wait_until_completed();
// todo!("Output {:?}", dst.buffer.read_to_vec::<f32>(10));
// }
Ok(())
}
}
@ -616,28 +646,29 @@ impl MetalStorage {
match (self.dtype, rhs.dtype) {
(DType::F32, DType::F32) => {
let mut out_buffer = self.device.new_buffer(elem_count, self.dtype);
if b != 1 {
// debug!("TODO implement batched matmul for B={b}");
// bail!("Didn't implemented strided matmul yet");
return Ok(Self {
buffer: out_buffer,
device: self.device.clone(),
dtype: self.dtype(),
});
}
if !lhs_l.is_contiguous() || !rhs_l.is_contiguous() {
// debug!(
// "TODO non contiguous matmul yet {:?} {:?} - {:?} - {transpose_right}",
// lhs_l.is_contiguous(),
// rhs_l.is_contiguous(),
// rhs_l
// );
return Ok(Self {
buffer: out_buffer,
device: self.device.clone(),
dtype: self.dtype(),
});
}
// if b != 1 {
// // debug!("TODO implement batched matmul for B={b}");
// crate::bail!("Didn't implemented strided matmul yet");
// return Ok(Self {
// buffer: out_buffer,
// device: self.device.clone(),
// dtype: self.dtype(),
// });
//}
// if !lhs_l.is_contiguous() || !rhs_l.is_contiguous() {
// // debug!(
// // "TODO non contiguous matmul yet {:?} {:?} - {:?} - {transpose_right}",
// // lhs_l.is_contiguous(),
// // rhs_l.is_contiguous(),
// // rhs_l
// // );
// crate::bail!("No not contiguous matmul");
// return Ok(Self {
// buffer: out_buffer,
// device: self.device.clone(),
// dtype: self.dtype(),
// });
// }
// debug!("TODO GEMM");
let command_buffer = self.device.command_queue.new_command_buffer();
@ -659,7 +690,15 @@ impl MetalStorage {
.map_err(MetalError::from)?;
command_buffer.commit();
command_buffer.wait_until_completed();
// command_buffer.wait_until_scheduled();
//
let left = self.buffer.read_to_vec::<f32>(10);
let right = rhs.buffer.read_to_vec::<f32>(10);
let out = out_buffer.read_to_vec::<f32>(10);
println!("{b} {m} {n} {k} ");
println!("{left:?} {right:?} {out:?}");
Ok(Self {
buffer: out_buffer,
@ -709,7 +748,9 @@ impl BackendDevice for MetalDevice {
}
fn location(&self) -> crate::DeviceLocation {
crate::DeviceLocation::Metal
crate::DeviceLocation::Metal {
gpu_id: self.registry_id() as usize,
}
}
fn same_device(&self, rhs: &Self) -> bool {
@ -767,6 +808,8 @@ impl BackendDevice for MetalDevice {
option,
),
};
// TODO is that necessary ?
// buffer.did_modify_range(metal::NSRange::new(0, buffer.length()));
// debug!("Allocate 2 - buffer size {}", buffer.length());
Ok(Self::Storage {
buffer,

8
candle-core/src/tensor.rs
View File

@ -157,6 +157,8 @@ pub(crate) fn from_storage<S: Into<Shape>>(
) -> Tensor {
let dtype = storage.dtype();
let device = storage.device();
let shape = shape.into();
// println!("{:?} {storage:?}", shape);
let tensor_ = Tensor_ {
id: TensorId::new(),
storage: Arc::new(RwLock::new(storage)),
@ -166,7 +168,11 @@ pub(crate) fn from_storage<S: Into<Shape>>(
dtype,
device,
};
Tensor(Arc::new(tensor_))
let result = Tensor(Arc::new(tensor_));
// todo!(" from_storage");
// let result = result.to_device(&Device::Cpu).unwrap();
// todo!(" {result}");
result
}
impl Tensor {