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

Tmp commit for the heap experiment (heap is indeed decreasing).

Browse Source
This commit is contained in:
Nicolas Patry
2023-11-14 17:04:23 +01:00
parent 51f05e997d
commit e8c1c31245
4 changed files with 176 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
Cargo.toml
View File

@ -61,7 +61,8 @@ tracing-subscriber = "0.3.7"
wav = "1.0.0" wav = "1.0.0"
yoke = { version = "0.7.2", features = ["derive"] } yoke = { version = "0.7.2", features = ["derive"] }
zip = { version = "0.6.6", default-features = false } zip = { version = "0.6.6", default-features = false }
metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] } # metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
metal = { path = "../metal-rs", features = ["mps"] }
[profile.release-with-debug] [profile.release-with-debug]
inherits = "release" inherits = "release"

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

@ -67,12 +67,28 @@ impl MetalDevice {
self.command_buffer.read().unwrap() self.command_buffer.read().unwrap()
} }
pub fn wait_until_completed(&self) { pub fn commit_wait_until_completed(&self) {
let mut old = self.command_buffer.write().unwrap(); let mut old = self.command_buffer.try_write().unwrap();
let status = old.status();
use metal::MTLCommandBufferStatus::{
Committed, Completed, Enqueued, Error, NotEnqueued, Scheduled,
};
// match old.status() {}
if old.status() == metal::MTLCommandBufferStatus::Completed {
return;
}
old.commit(); old.commit();
old.wait_until_completed(); old.wait_until_completed();
let command_buffer = self.command_queue.new_owned_command_buffer(); // let count = old.retain_count();
// println!("Count {count:?}");
let command_buffer = self.command_queue.new_command_buffer().to_owned();
*old = command_buffer; *old = command_buffer;
// let count = old.retain_count();
// // println!("Count after {count:?}");
// old.release();
// let count = old.retain_count();
// println!("Count after release {count:?}");
// self.command_buffer.replace_with(|_| command_buffer) // self.command_buffer.replace_with(|_| command_buffer)
} }
@ -86,18 +102,21 @@ impl MetalDevice {
pub fn new_buffer(&self, element_count: usize, dtype: DType) -> Buffer { pub fn new_buffer(&self, element_count: usize, dtype: DType) -> Buffer {
let size = (element_count * dtype.size_in_bytes()) as NSUInteger; let size = (element_count * dtype.size_in_bytes()) as NSUInteger;
self.heap // println!("Creating buffer {size}");
let buffer = self
.heap
.new_buffer(size, MTLResourceOptions::StorageModeShared) .new_buffer(size, MTLResourceOptions::StorageModeShared)
.expect(" New buffer") .expect("New buffer");
// println!("{:?}", self.heap.used_size());
buffer
} }
pub fn new_buffer_with_data<T>(&self, data: &[T]) -> Buffer { pub fn new_buffer_with_data<T>(&self, data: &[T]) -> Buffer {
let size = core::mem::size_of_val(data) as NSUInteger;
let option = metal::MTLResourceOptions::StorageModeShared; let option = metal::MTLResourceOptions::StorageModeShared;
self.device.new_buffer_with_data( // println!("Creating data buffer {size}");
data.as_ptr() as *const core::ffi::c_void, self.device
core::mem::size_of_val(data) as NSUInteger, .new_buffer_with_data(data.as_ptr() as *const core::ffi::c_void, size, option)
option,
)
} }
} }
@ -124,7 +143,7 @@ impl BackendStorage for MetalStorage {
} }
fn to_cpu_storage(&self) -> Result<CpuStorage> { fn to_cpu_storage(&self) -> Result<CpuStorage> {
self.device.wait_until_completed(); self.device.commit_wait_until_completed();
match self.dtype { match self.dtype {
DType::U8 => Ok(CpuStorage::U8( DType::U8 => Ok(CpuStorage::U8(
@ -335,93 +354,95 @@ impl BackendStorage for MetalStorage {
let shape = layout.shape(); let shape = layout.shape();
let el_count = shape.elem_count(); let el_count = shape.elem_count();
let mut buffer = device.new_buffer(el_count, dtype); let mut buffer = device.new_buffer(el_count, dtype);
let command_buffer = device.command_buffer(); {
if layout.is_contiguous() && layout.start_offset() == 0 { let command_buffer = device.command_buffer();
use candle_metal_kernels::unary::contiguous; if layout.is_contiguous() && layout.start_offset() == 0 {
use candle_metal_kernels::unary::contiguous;
let kernel_name = match (B::KERNEL, dtype) { let kernel_name = match (B::KERNEL, dtype) {
("ucos", DType::F32) => contiguous::cos::FLOAT, ("ucos", DType::F32) => contiguous::cos::FLOAT,
("usin", DType::F32) => contiguous::sin::FLOAT, ("usin", DType::F32) => contiguous::sin::FLOAT,
("usqr", DType::F32) => contiguous::sqr::FLOAT, ("usqr", DType::F32) => contiguous::sqr::FLOAT,
("usqrt", DType::F32) => contiguous::sqrt::FLOAT, ("usqrt", DType::F32) => contiguous::sqrt::FLOAT,
("uneg", DType::F32) => contiguous::neg::FLOAT, ("uneg", DType::F32) => contiguous::neg::FLOAT,
("uexp", DType::F32) => contiguous::exp::FLOAT, ("uexp", DType::F32) => contiguous::exp::FLOAT,
("ulog", DType::F32) => contiguous::log::FLOAT, ("ulog", DType::F32) => contiguous::log::FLOAT,
("ugelu", DType::F32) => contiguous::gelu::FLOAT, ("ugelu", DType::F32) => contiguous::gelu::FLOAT,
("ugelu_erf", DType::F32) => contiguous::gelu_erf::FLOAT, ("ugelu_erf", DType::F32) => contiguous::gelu_erf::FLOAT,
("uerf", DType::F32) => contiguous::erf::FLOAT, ("uerf", DType::F32) => contiguous::erf::FLOAT,
("uceil", DType::F32) => contiguous::ceil::FLOAT, ("uceil", DType::F32) => contiguous::ceil::FLOAT,
("ufloor", DType::F32) => contiguous::floor::FLOAT, ("ufloor", DType::F32) => contiguous::floor::FLOAT,
("uround", DType::F32) => contiguous::round::FLOAT, ("uround", DType::F32) => contiguous::round::FLOAT,
("ucos", DType::F16) => contiguous::cos::HALF, ("ucos", DType::F16) => contiguous::cos::HALF,
("usin", DType::F16) => contiguous::sin::HALF, ("usin", DType::F16) => contiguous::sin::HALF,
("usqr", DType::F16) => contiguous::sqr::HALF, ("usqr", DType::F16) => contiguous::sqr::HALF,
("usqrt", DType::F16) => contiguous::sqrt::HALF, ("usqrt", DType::F16) => contiguous::sqrt::HALF,
("uneg", DType::F16) => contiguous::neg::HALF, ("uneg", DType::F16) => contiguous::neg::HALF,
("uexp", DType::F16) => contiguous::exp::HALF, ("uexp", DType::F16) => contiguous::exp::HALF,
("ulog", DType::F16) => contiguous::log::HALF, ("ulog", DType::F16) => contiguous::log::HALF,
("ugelu", DType::F16) => contiguous::gelu::HALF, ("ugelu", DType::F16) => contiguous::gelu::HALF,
("ugelu_erf", DType::F16) => contiguous::gelu_erf::HALF, ("ugelu_erf", DType::F16) => contiguous::gelu_erf::HALF,
("uerf", DType::F16) => contiguous::erf::HALF, ("uerf", DType::F16) => contiguous::erf::HALF,
("uceil", DType::F16) => contiguous::ceil::HALF, ("uceil", DType::F16) => contiguous::ceil::HALF,
("ufloor", DType::F16) => contiguous::floor::HALF, ("ufloor", DType::F16) => contiguous::floor::HALF,
("uround", DType::F16) => contiguous::round::HALF, ("uround", DType::F16) => contiguous::round::HALF,
(name, dtype) => todo!("Match {name} - {dtype:?}"), (name, dtype) => todo!("Match {name} - {dtype:?}"),
}; };
candle_metal_kernels::call_unary_contiguous( candle_metal_kernels::call_unary_contiguous(
&device.device, &device.device,
&command_buffer, &command_buffer,
&device.kernels, &device.kernels,
kernel_name, kernel_name,
el_count, el_count,
&self.buffer, &self.buffer,
&mut buffer, &mut buffer,
) )
.map_err(MetalError::from)?; .map_err(MetalError::from)?;
} else { } else {
use candle_metal_kernels::unary::strided; use candle_metal_kernels::unary::strided;
let kernel_name = match (B::KERNEL, dtype) { let kernel_name = match (B::KERNEL, dtype) {
("ucos", DType::F32) => strided::cos::FLOAT, ("ucos", DType::F32) => strided::cos::FLOAT,
("usin", DType::F32) => strided::sin::FLOAT, ("usin", DType::F32) => strided::sin::FLOAT,
("usqr", DType::F32) => strided::sqr::FLOAT, ("usqr", DType::F32) => strided::sqr::FLOAT,
("usqrt", DType::F32) => strided::sqrt::FLOAT, ("usqrt", DType::F32) => strided::sqrt::FLOAT,
("uneg", DType::F32) => strided::neg::FLOAT, ("uneg", DType::F32) => strided::neg::FLOAT,
("uexp", DType::F32) => strided::exp::FLOAT, ("uexp", DType::F32) => strided::exp::FLOAT,
("ulog", DType::F32) => strided::log::FLOAT, ("ulog", DType::F32) => strided::log::FLOAT,
("ugelu", DType::F32) => strided::gelu::FLOAT, ("ugelu", DType::F32) => strided::gelu::FLOAT,
("ugelu_erf", DType::F32) => strided::gelu_erf::FLOAT, ("ugelu_erf", DType::F32) => strided::gelu_erf::FLOAT,
("uerf", DType::F32) => strided::erf::FLOAT, ("uerf", DType::F32) => strided::erf::FLOAT,
("uceil", DType::F32) => strided::ceil::FLOAT, ("uceil", DType::F32) => strided::ceil::FLOAT,
("ufloor", DType::F32) => strided::floor::FLOAT, ("ufloor", DType::F32) => strided::floor::FLOAT,
("uround", DType::F32) => strided::round::FLOAT, ("uround", DType::F32) => strided::round::FLOAT,
("ucos", DType::F16) => strided::cos::HALF, ("ucos", DType::F16) => strided::cos::HALF,
("usin", DType::F16) => strided::sin::HALF, ("usin", DType::F16) => strided::sin::HALF,
("usqr", DType::F16) => strided::sqr::HALF, ("usqr", DType::F16) => strided::sqr::HALF,
("usqrt", DType::F16) => strided::sqrt::HALF, ("usqrt", DType::F16) => strided::sqrt::HALF,
("uneg", DType::F16) => strided::neg::HALF, ("uneg", DType::F16) => strided::neg::HALF,
("uexp", DType::F16) => strided::exp::HALF, ("uexp", DType::F16) => strided::exp::HALF,
("ulog", DType::F16) => strided::log::HALF, ("ulog", DType::F16) => strided::log::HALF,
("ugelu", DType::F16) => strided::gelu::HALF, ("ugelu", DType::F16) => strided::gelu::HALF,
("ugelu_erf", DType::F16) => strided::gelu_erf::HALF, ("ugelu_erf", DType::F16) => strided::gelu_erf::HALF,
("uerf", DType::F16) => strided::erf::HALF, ("uerf", DType::F16) => strided::erf::HALF,
("uceil", DType::F16) => strided::ceil::HALF, ("uceil", DType::F16) => strided::ceil::HALF,
("ufloor", DType::F16) => strided::floor::HALF, ("ufloor", DType::F16) => strided::floor::HALF,
("uround", DType::F16) => strided::round::HALF, ("uround", DType::F16) => strided::round::HALF,
(name, dtype) => todo!("Match {name} - {dtype:?}"), (name, dtype) => todo!("Match {name} - {dtype:?}"),
}; };
candle_metal_kernels::call_unary_strided( candle_metal_kernels::call_unary_strided(
&device.device, &device.device,
&command_buffer, &command_buffer,
&device.kernels, &device.kernels,
kernel_name, kernel_name,
layout.dims(), layout.dims(),
&self.buffer, &self.buffer,
layout.stride(), layout.stride(),
layout.start_offset() * self.dtype.size_in_bytes(), layout.start_offset() * self.dtype.size_in_bytes(),
&mut buffer, &mut buffer,
0, 0,
) )
.map_err(MetalError::from)?; .map_err(MetalError::from)?;
}
} }
Ok(Self { Ok(Self {
buffer, buffer,
@ -769,59 +790,61 @@ impl BackendStorage for MetalStorage {
let out_buffer = self.device.new_buffer(elem_count, self.dtype); let out_buffer = self.device.new_buffer(elem_count, self.dtype);
let command_buffer = self.device.command_buffer(); {
for bi in 0..b { let command_buffer = self.device.command_buffer();
// Create matrix objects for bi in 0..b {
let left_matrix = Matrix::init_with_buffer_descriptor( // Create matrix objects
&self.buffer, let left_matrix = Matrix::init_with_buffer_descriptor(
(bi * stride_left + lhs_l.start_offset() as u64) * size, &self.buffer,
&left_descriptor, (bi * stride_left + lhs_l.start_offset() as u64) * size,
) &left_descriptor,
.ok_or_else(|| { )
MetalError::from("Failed to create matrix multiplication kernel".to_string()) .ok_or_else(|| {
})?; MetalError::from("Failed to create matrix multiplication kernel".to_string())
let right_matrix = Matrix::init_with_buffer_descriptor( })?;
&rhs.buffer, let right_matrix = Matrix::init_with_buffer_descriptor(
(bi * stride_right + rhs_l.start_offset() as u64) * size, &rhs.buffer,
&right_descriptor, (bi * stride_right + rhs_l.start_offset() as u64) * size,
) &right_descriptor,
.ok_or_else(|| { )
MetalError::from("Failed to create matrix multiplication kernel".to_string()) .ok_or_else(|| {
})?; MetalError::from("Failed to create matrix multiplication kernel".to_string())
})?;
let result_matrix = Matrix::init_with_buffer_descriptor( let result_matrix = Matrix::init_with_buffer_descriptor(
&out_buffer, &out_buffer,
bi * m * n * size, bi * m * n * size,
&result_descriptor, &result_descriptor,
) )
.ok_or_else(|| { .ok_or_else(|| {
MetalError::from("Failed to create matrix multiplication kernel".to_string()) MetalError::from("Failed to create matrix multiplication kernel".to_string())
})?; })?;
let alpha = 1.0f64; let alpha = 1.0f64;
let beta = 0.0f64; let beta = 0.0f64;
// Create kernel // Create kernel
let matrix_multiplication = MatrixMultiplication::init( let matrix_multiplication = MatrixMultiplication::init(
&self.device, &self.device,
transpose_left, transpose_left,
transpose_right, transpose_right,
m, m,
n, n,
k, k,
alpha, alpha,
beta, beta,
) )
.ok_or_else(|| { .ok_or_else(|| {
MetalError::from("Failed to create matrix multiplication kernel".to_string()) MetalError::from("Failed to create matrix multiplication kernel".to_string())
})?; })?;
// Encode kernel to command buffer // Encode kernel to command buffer
matrix_multiplication.encode_to_command_buffer( matrix_multiplication.encode_to_command_buffer(
&command_buffer, &command_buffer,
&left_matrix, &left_matrix,
&right_matrix, &right_matrix,
&result_matrix, &result_matrix,
); );
}
} }
Ok(Self { Ok(Self {
@ -891,7 +914,7 @@ impl BackendDevice for MetalDevice {
descriptor.set_size(size.size); descriptor.set_size(size.size);
descriptor.set_storage_mode(metal::MTLStorageMode::Shared); descriptor.set_storage_mode(metal::MTLStorageMode::Shared);
let heap = device.new_heap(&descriptor); let heap = device.new_heap(&descriptor);
let command_buffer = Arc::new(RwLock::new(command_queue.new_owned_command_buffer())); let command_buffer = Arc::new(RwLock::new(command_queue.new_command_buffer().to_owned()));
let kernels = Arc::new(Kernels::new()); let kernels = Arc::new(Kernels::new());
Ok(Self { Ok(Self {
device, device,

3
candle-metal-kernels/Cargo.toml
View File

@ -10,7 +10,8 @@ categories = ["science"]
license = "MIT OR Apache-2.0" license = "MIT OR Apache-2.0"
[dependencies] [dependencies]
metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] } # metal = { git = "https://github.com/ivarflakstad/metal-rs.git", features = ["mps"] }
metal = { path = "../../metal-rs", features = ["mps"] }
once_cell = "1.18.0" once_cell = "1.18.0"
thiserror = "1" thiserror = "1"
tracing = "0.1.37" tracing = "0.1.37"

1
candle-metal-kernels/src/lib.rs
View File

@ -1133,6 +1133,7 @@ mod tests {
let device = Device::system_default().expect("no device found"); let device = Device::system_default().expect("no device found");
let options = CompileOptions::new(); let options = CompileOptions::new();
options.set_fast_math_enabled(true);
let library = device.new_library_with_source(INDEXING, &options).unwrap(); let library = device.new_library_with_source(INDEXING, &options).unwrap();
let left = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; let left = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];