Fix/fusion compilation speed (#3155)

Author: nathanielsimard

Cargo.lock

@@ -156,9 +156,9 @@ portable-atomic = { version = "1.11.0" }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "150829d0876e5ced8e937f18abbc7e3c757e11c7" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "150829d0876e5ced8e937f18abbc7e3c757e11c7" }
-cubecl-std = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "150829d0876e5ced8e937f18abbc7e3c757e11c7" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "48939d2ca47473f0f30526962c3b6fb8c9b558e0" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "48939d2ca47473f0f30526962c3b6fb8c9b558e0" }
+cubecl-std = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "48939d2ca47473f0f30526962c3b6fb8c9b558e0" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

Cargo.toml

@@ -1,10 +1,10 @@
-use std::marker::PhantomData;
-
 use crate::reduce::optimization::{ReduceOptimization, ReduceOptimizationState};
+use std::marker::PhantomData;
 
 use super::elemwise::optimization::{ElemwiseOptimization, ElemwiseOptimizationState};
 use super::matmul::optimization::{MatmulOptimization, MatmulOptimizationState};
 
+use burn_fusion::stream::Context;
 use burn_tensor::DType;
 use cubecl::client::ComputeClient;
 use cubecl::ir::Elem;
@@ -35,6 +35,10 @@ pub enum CubeOptimizationState {
     Reduce(ReduceOptimizationState),
 }
 
+pub trait FallbackOperation<R: Runtime>: Send + Sync {
+    fn run(&self, context: &mut Context<'_, CubeFusionHandle<R>>);
+}
+
 pub(crate) fn strides_dyn_rank(shape: &[usize]) -> Vec<usize> {
     let mut strides = vec![0; shape.len()];
 

crates/burn-cubecl-fusion/src/base.rs

@@ -1,6 +1,3 @@
-use std::sync::Arc;
-
-use super::MatmulFallbackFn;
 use burn_fusion::{OptimizationBuilder, OptimizationStatus};
 use burn_ir::{FloatOperationIr, OperationIr};
 use cubecl::Runtime;
@@ -19,15 +16,10 @@ pub struct MatmulBuilder<R: Runtime> {
     builder_fallback: FuseOptimizationBuilder,
     device: R::Device,
     matmul: Option<FusedMatmul>,
-    fallback: Arc<dyn MatmulFallbackFn<R>>,
 }
 
 impl<R: Runtime> MatmulBuilder<R> {
-    pub fn new(
-        device: R::Device,
-        bool_precision: FusePrecision,
-        fallback: Arc<dyn MatmulFallbackFn<R>>,
-    ) -> Self {
+    pub fn new(device: R::Device, bool_precision: FusePrecision) -> Self {
         let client = R::client(&device);
         let props = client.properties();
         let max_bindings = props.hardware_properties().max_bindings;
@@ -43,7 +35,6 @@ impl<R: Runtime> MatmulBuilder<R> {
             builder_fallback: FuseOptimizationBuilder::new(max_bindings, bool_precision, settings),
             device,
             matmul: None,
-            fallback,
         }
     }
 }
@@ -94,7 +85,6 @@ impl<R: Runtime> OptimizationBuilder<CubeOptimization<R>> for MatmulBuilder<R> {
             self.device.clone(),
             self.len(),
             self.matmul.as_ref().unwrap().clone(),
-            self.fallback.clone(),
         );
 
         CubeOptimization::Matmul(matmul)

crates/burn-cubecl-fusion/src/matmul/builder.rs

@@ -4,5 +4,3 @@ pub mod optimization;
 pub(crate) mod args;
 pub(crate) mod spec;
 pub(crate) mod tune;
-
-pub use optimization::MatmulFallbackFn;

crates/burn-cubecl-fusion/src/matmul/mod.rs

@@ -1,7 +1,7 @@
 use std::any::TypeId;
-use std::sync::Arc;
 
 use crate::CubeFusionHandle;
+use crate::FallbackOperation;
 use crate::elemwise::optimization::ElemwiseRunner;
 use crate::shared::ir::FusePrecision;
 use crate::shared::ir::RefLayout;
@@ -10,7 +10,7 @@ use crate::shared::trace::TuneOutput;
 use crate::shared::trace::Vectorization;
 
 use burn_fusion::stream::Context;
-use burn_ir::{BinaryOpIr, TensorStatus};
+use burn_ir::BinaryOpIr;
 use cubecl::linalg::matmul::components;
 use cubecl::linalg::matmul::components::MatmulPrecision;
 use cubecl::linalg::matmul::components::MatmulProblem;
@@ -44,15 +44,7 @@ pub struct MatmulOptimization<R: Runtime> {
     pub(crate) len: usize,
     pub(crate) matmul_simple: FusedMatmul,
     pub(crate) matmul_double_buffering: FusedMatmul,
-    fallback: Arc<dyn MatmulFallbackFn<R>>,
-}
-
-pub trait MatmulFallbackFn<R: Runtime>: Send + Sync {
-    fn run(
-        &self,
-        lhs: (CubeFusionHandle<R>, &[usize]),
-        rhs: (CubeFusionHandle<R>, &[usize]),
-    ) -> CubeFusionHandle<R>;
+    fallback: Option<Box<dyn FallbackOperation<R>>>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -73,7 +65,6 @@ impl<R: Runtime> MatmulOptimization<R> {
         device: R::Device,
         len: usize,
         matmul: FusedMatmul,
-        fallback: Arc<dyn MatmulFallbackFn<R>>,
     ) -> Self {
         let mut matmul_simple = matmul.clone();
         let mut matmul_double_buffering = matmul;
@@ -89,11 +80,18 @@ impl<R: Runtime> MatmulOptimization<R> {
             len,
             matmul_simple,
             matmul_double_buffering,
-            fallback,
+            fallback: None,
         }
     }
     /// Execute the optimization.
-    pub fn execute<BT: CubeElement>(&mut self, context: &mut Context<'_, CubeFusionHandle<R>>) {
+    pub fn execute<BT: CubeElement>(
+        &mut self,
+        context: &mut Context<'_, CubeFusionHandle<R>>,
+        fallback: impl FnOnce(usize) -> Box<dyn FallbackOperation<R>>,
+    ) {
+        // The index of the fallback matmul is always 0.
+        self.fallback = Some(fallback(0));
+
         #[cfg(feature = "autotune")]
         fused_matmul_autotune::<R, BT>(self, context);
 
@@ -109,11 +107,7 @@ impl<R: Runtime> MatmulOptimization<R> {
     }
 
     /// Create an optimization from its [state](MatmulOptimizationState).
-    pub fn from_state(
-        device: &R::Device,
-        state: MatmulOptimizationState,
-        fallback: Arc<dyn MatmulFallbackFn<R>>,
-    ) -> Self {
+    pub fn from_state(device: &R::Device, state: MatmulOptimizationState) -> Self {
         Self {
             trace: state.trace,
             trace_fallback: state.trace_fallback,
@@ -122,7 +116,7 @@ impl<R: Runtime> MatmulOptimization<R> {
             device: device.clone(),
             matmul_simple: state.matmul_simple.clone(),
             matmul_double_buffering: state.matmul_double_buffering.clone(),
-            fallback,
+            fallback: None,
         }
     }
 
@@ -170,31 +164,11 @@ impl<R: Runtime> MatmulOptimization<R> {
         &self,
         context: &mut Context<'_, CubeFusionHandle<R>>,
     ) -> TuneOutput<R> {
-        let (out_tensor, out_desc) = {
-            let lhs = context
-                .tensors
-                .get(&self.matmul_simple.op.lhs.id)
-                .unwrap()
-                .clone();
-            let rhs = context
-                .tensors
-                .get(&self.matmul_simple.op.rhs.id)
-                .unwrap()
-                .clone();
-            let out = context
-                .tensors
-                .get(&self.matmul_simple.op.out.id)
-                .unwrap()
-                .clone();
-
-            let lhs_handle = context.handles.get_handle(&lhs.id, &TensorStatus::ReadOnly);
-            let rhs_handle = context.handles.get_handle(&rhs.id, &TensorStatus::ReadOnly);
-            let out_handle = self
-                .fallback
-                .run((lhs_handle, &lhs.shape), (rhs_handle, &rhs.shape));
-
-            (out_handle, out)
-        };
+        self.fallback
+            .as_ref()
+            .expect("A fallback operation should be available")
+            .run(context);
+
         #[cfg(feature = "autotune-checks")]
         let mut output = TuneOutput::Checked {
             handles: Default::default(),
@@ -204,12 +178,16 @@ impl<R: Runtime> MatmulOptimization<R> {
 
         #[cfg(feature = "autotune-checks")]
         if let TuneOutput::Checked { handles } = &mut output {
+            let out_desc = context.tensors.get(&self.matmul_simple.op.out.id).unwrap();
+            let handle_out = context
+                .handles
+                .get_handle(&out_desc.id, &burn_ir::TensorStatus::ReadOnly);
+
             handles.insert(
                 self.matmul_simple.op.out.id,
-                (out_desc.shape.clone(), out_tensor.clone()),
+                (out_desc.shape.clone(), handle_out.clone()),
             );
         }
-        context.handles.register_handle(out_desc.id, out_tensor);
 
         let output_write = self
             .trace_fallback