Rewrite legacy cpa kernels (#2455)

tracel-ai · Nov 7, 2024 · 8f83b2c · 8f83b2c
1 parent c3549fb
commit 8f83b2c
Show file tree

Hide file tree

Showing 54 changed files with 2,176 additions and 4,260 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -153,8 +153,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 = "9460a3244aa2b42e1d6c36bd25b65f814f81ecd0" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9460a3244aa2b42e1d6c36bd25b65f814f81ecd0" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "4198b192c60bffff7ec51920bcd7191560d6c98b" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "4198b192c60bffff7ec51920bcd7191560d6c98b" }
 ### 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-jit/src/kernel/conv/conv2d/implicit_gemm.rs b/crates/burn-jit/src/kernel/conv/conv2d/implicit_gemm.rs
@@ -455,7 +455,7 @@ fn execute_gemm<F: Float, FMat: Float>(
 
         // Run CMMA
         cmma::load(&matrices.a, input_tile.as_slice(), cmma_k);
-        cmma::load(&matrices.b, weight_tile.as_slice(), cmma_n);
+        cmma::load(&matrices.b, weight_tile.as_slice(), cmma_k);
 
         cmma::execute::<FMat, FMat, F, F>(&matrices.a, &matrices.b, &matrices.acc, &matrices.acc);
     }