Fusion experimental feature flag (#2690)

* Fusion experimental feature flag * Minor change to vec init --------- Co-authored-by: Guillaume Lagrange <[email protected]>
tracel-ai · Jan 13, 2025 · 28f2b67 · 28f2b67
1 parent ad89dcc
commit 28f2b67
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/crates/burn-jit/Cargo.toml b/crates/burn-jit/Cargo.toml
@@ -25,6 +25,8 @@ export_tests = [
     "paste",
 ]
 fusion = ["burn-fusion"]
+fusion-experimental = ["fusion"]
+
 std = ["cubecl/std"]
 template = []
 

diff --git a/crates/burn-jit/src/fusion/base.rs b/crates/burn-jit/src/fusion/base.rs
@@ -125,16 +125,20 @@ impl<R: JitRuntime, BT: BoolElement> FusionRuntime for FusionJitRuntime<R, BT> {
     fn optimizations(
         device: R::Device,
     ) -> Vec<Box<dyn burn_fusion::OptimizationBuilder<Self::Optimization>>> {
-        vec![
-            Box::new(ElementWiseBuilder::<R>::new(
+        let mut optimizations: Vec<Box<dyn burn_fusion::OptimizationBuilder<Self::Optimization>>> =
+            vec![Box::new(ElementWiseBuilder::<R>::new(
                 device.clone(),
                 BT::as_elem_native_unchecked().into(),
-            )),
-            Box::new(MatmulBuilder::<R>::new(
+            ))];
+
+        if cfg!(feature = "fusion-experimental") {
+            optimizations.push(Box::new(MatmulBuilder::<R>::new(
                 device.clone(),
                 BT::as_elem_native_unchecked().into(),
-            )),
-        ]
+            )));
+        }
+
+        optimizations
     }
 }