benchmark

backend-comparison/benches/matmul.rs

@@ -53,9 +53,9 @@ fn bench<B: Backend>(
 ) {
     const D: usize = 3;
     let batch_size = 32;
-    let m = 256;
-    let k = 1024;
-    let n = 256;
+    let m = 1024;
+    let k = 256;
+    let n = 1024;
     let shape_lhs = [batch_size, m, k].into();
     let shape_rhs = [batch_size, k, n].into();
 

crates/burn-jit/src/kernel/matmul/base.rs

@@ -31,10 +31,7 @@ pub enum MatmulStrategy {
 #[cfg(feature = "autotune")]
 impl Default for MatmulStrategy {
     fn default() -> Self {
-        MatmulStrategy::Simple {
-            grid_x: 16,
-            grid_y: 16,
-        }
+        MatmulStrategy::Autotune
     }
 }
 

crates/burn-jit/src/kernel/matmul/tiling2d_cube/base.rs

@@ -193,6 +193,8 @@ pub fn matmul_tiling_2d_cube<R: JitRuntime, E: FloatElement, const D: usize>(
         } => (tensor, transposed),
         MemoryLayout::HighlyPermuted => (into_contiguous(tensor), false),
     };
+
+    // let check_layout = |tensor: JitTensor<R, E, D>| (into_contiguous(tensor), false);
     let (lhs, lhs_transposed) = check_layout(lhs);
     let (rhs, rhs_transposed) = check_layout(rhs);