From 16f1991d6ba333e71b9401a4c3a9607e534fe652 Mon Sep 17 00:00:00 2001 From: Nathaniel Simard Date: Thu, 24 Oct 2024 16:41:53 -0400 Subject: [PATCH] Chore/update cubecl (#2410) --- Cargo.lock | 37 +++++++++++++------ Cargo.toml | 4 +- .../burn-import/onnx-tests/tests/test_onnx.rs | 2 +- crates/burn-jit/src/kernel/index/slice.rs | 16 ++++---- .../burn-jit/src/kernel/matmul/tune/base.rs | 2 +- crates/burn-jit/src/kernel/reduce/base.rs | 4 +- .../burn-jit/src/kernel/reduce/tune/base.rs | 4 +- 7 files changed, 43 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4172180cbb..5db706d772 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -224,6 +224,17 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-lock" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff6e472cdea888a4bd64f342f09b3f50e1886d32afe8df3d663c01140b811b18" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "atoi" version = "2.0.0" @@ -1448,7 +1459,7 @@ dependencies = [ [[package]] name = "cubecl" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "cubecl-core", "cubecl-cuda", @@ -1461,12 +1472,13 @@ dependencies = [ [[package]] name = "cubecl-common" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "derive-new", "embassy-futures", "futures-lite", "getrandom", + "log", "portable-atomic", "rand", "serde", @@ -1477,7 +1489,7 @@ dependencies = [ [[package]] name = "cubecl-core" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "bytemuck", "cubecl-common", @@ -1494,7 +1506,7 @@ dependencies = [ [[package]] name = "cubecl-cpp" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "bytemuck", "cubecl-common", @@ -1508,7 +1520,7 @@ dependencies = [ [[package]] name = "cubecl-cuda" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "bytemuck", "cubecl-common", @@ -1524,7 +1536,7 @@ dependencies = [ [[package]] name = "cubecl-hip" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "bytemuck", "cubecl-common", @@ -1549,7 +1561,7 @@ dependencies = [ [[package]] name = "cubecl-linalg" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "bytemuck", "cubecl-core", @@ -1560,7 +1572,7 @@ dependencies = [ [[package]] name = "cubecl-macros" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "cubecl-common", "darling", @@ -1575,7 +1587,7 @@ dependencies = [ [[package]] name = "cubecl-opt" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "cubecl-common", "cubecl-core", @@ -1590,9 +1602,10 @@ dependencies = [ [[package]] name = "cubecl-runtime" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "async-channel", + "async-lock", "cfg_aliases 0.2.1", "cubecl-common", "derive-new", @@ -1610,7 +1623,7 @@ dependencies = [ [[package]] name = "cubecl-spirv" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "cubecl-common", "cubecl-core", @@ -1623,7 +1636,7 @@ dependencies = [ [[package]] name = "cubecl-wgpu" version = "0.2.0" -source = "git+https://github.com/tracel-ai/cubecl?rev=fb2a5c87802c9d01801f16d93cada5fe924989f3#fb2a5c87802c9d01801f16d93cada5fe924989f3" +source = "git+https://github.com/tracel-ai/cubecl?rev=44be0489c04c1c88bef61bde4c386e6772693b94#44be0489c04c1c88bef61bde4c386e6772693b94" dependencies = [ "ash", "async-channel", diff --git a/Cargo.toml b/Cargo.toml index c6eaf0999b..bcb5ee930d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -152,8 +152,8 @@ ahash = { version = "0.8.11", default-features = false } portable-atomic-util = { version = "0.2.2", features = ["alloc"] } ### For the main burn branch. ### -cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" } -cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "fb2a5c87802c9d01801f16d93cada5fe924989f3" } +cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "44be0489c04c1c88bef61bde4c386e6772693b94" } +cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "44be0489c04c1c88bef61bde4c386e6772693b94" } ### For local development. ### # cubecl = { path = "../cubecl/crates/cubecl", default-features = false } # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false } diff --git a/crates/burn-import/onnx-tests/tests/test_onnx.rs b/crates/burn-import/onnx-tests/tests/test_onnx.rs index 24ddaeca45..fae79b134d 100644 --- a/crates/burn-import/onnx-tests/tests/test_onnx.rs +++ b/crates/burn-import/onnx-tests/tests/test_onnx.rs @@ -1526,7 +1526,7 @@ mod tests { // because the output tensor is too large to compare with the expected tensor. let output_sum = output.sum().into_scalar(); - let expected_sum = 33.810329; // example result running the corresponding PyTorch model (conv_transpose1d.py) + let expected_sum = 33.810_33; // example result running the corresponding PyTorch model (conv_transpose1d.py) assert!(expected_sum.approx_eq(output_sum, (1.0e-4, 2))); } diff --git a/crates/burn-jit/src/kernel/index/slice.rs b/crates/burn-jit/src/kernel/index/slice.rs index e0820fead9..30aa0aca00 100644 --- a/crates/burn-jit/src/kernel/index/slice.rs +++ b/crates/burn-jit/src/kernel/index/slice.rs @@ -110,21 +110,23 @@ pub(crate) fn slice( indices: &[Range], ) -> JitTensor { let mut dims = tensor.shape.dims.clone(); - let mut offset_start = 0; - let mut offset_end = 0; + let mut offset_start = 0u64; + let mut offset_end = 0u64; for i in 0..indices.len() { - offset_start += tensor.strides[i] * indices[i].start; - offset_end += tensor.strides[i] * (dims[i] - indices[i].end); + offset_start += (tensor.strides[i] * indices[i].start) as u64; + offset_end += (tensor.strides[i] * (dims[i] - indices[i].end)) as u64; dims[i] = indices[i].end - indices[i].start; } - let offset_start = offset_start * E::cube_elem().size(); - let offset_end = offset_end * E::cube_elem().size(); + let offset_start = offset_start * E::cube_elem().size() as u64; + let offset_end = offset_end * E::cube_elem().size() as u64; let memory_offset_alignment = tensor.client.properties().memory_properties().alignment; - if offset_start % memory_offset_alignment == 0 && offset_end % memory_offset_alignment == 0 { + if offset_start % memory_offset_alignment == 0u64 + && offset_end % memory_offset_alignment == 0u64 + { JitTensor::new( tensor.client, tensor diff --git a/crates/burn-jit/src/kernel/matmul/tune/base.rs b/crates/burn-jit/src/kernel/matmul/tune/base.rs index c09fca0f91..1bd7187c5f 100644 --- a/crates/burn-jit/src/kernel/matmul/tune/base.rs +++ b/crates/burn-jit/src/kernel/matmul/tune/base.rs @@ -92,7 +92,7 @@ pub fn matmul_autotune( macro_rules! matmul_tune_ops { ($name:ident, $func:expr) => { - #[derive(new)] + #[derive(new, Debug)] pub(crate) struct $name { lhs: JitTensor, rhs: JitTensor, diff --git a/crates/burn-jit/src/kernel/reduce/base.rs b/crates/burn-jit/src/kernel/reduce/base.rs index eaa72e65e3..8b78a4cb16 100644 --- a/crates/burn-jit/src/kernel/reduce/base.rs +++ b/crates/burn-jit/src/kernel/reduce/base.rs @@ -9,7 +9,7 @@ use super::{ #[allow(dead_code)] pub(crate) trait ReduceDimAlgorithm: - ReduceDimNaive + ReduceDimShared + core::fmt::Debug + ReduceDimNaive + ReduceDimShared { } @@ -64,7 +64,9 @@ impl Default for ReduceStrategy { macro_rules! reduce_operation { ($name:ident, $ops:ident) => { + #[derive(Debug)] pub(crate) struct $ops; + impl ReduceDimAlgorithm for $ops {} /// Executes the reduce operation with the given strategy. diff --git a/crates/burn-jit/src/kernel/reduce/tune/base.rs b/crates/burn-jit/src/kernel/reduce/tune/base.rs index 2f6902e0e1..e12bd1ac58 100644 --- a/crates/burn-jit/src/kernel/reduce/tune/base.rs +++ b/crates/burn-jit/src/kernel/reduce/tune/base.rs @@ -133,7 +133,7 @@ pub(crate) fn reduce_dim_autotune< output } -#[derive(new)] +#[derive(new, Debug)] // Probably better on balanced tensor shapes pub(crate) struct ReduceDimNaiveAutotune< RD: ReduceDimAlgorithm, @@ -169,7 +169,7 @@ where } } -#[derive(new)] +#[derive(new, Debug)] // Probably better on tensors large along reduce dim pub(crate) struct ReduceDimSharedAutotune< RD: ReduceDimAlgorithm,