Feat/fused matmul tune (#2726)

Author: nathanielsimard

Cargo.lock

@@ -153,8 +153,8 @@ ahash = { version = "0.8.11", default-features = false }
 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 = "dd0c2fa86bc58133cde596e13ddb1d7ad33ac44c" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "dd0c2fa86bc58133cde596e13ddb1d7ad33ac44c" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "2a6dd3e60b686230a8f686aafd246342259f7003" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "2a6dd3e60b686230a8f686aafd246342259f7003" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

Cargo.toml

@@ -1,5 +1,9 @@
 use backend_comparison::persistence::save;
-use burn::tensor::{activation::relu, backend::Backend, Distribution, Shape, Tensor};
+use burn::tensor::{
+    activation::{gelu, relu},
+    backend::Backend,
+    Distribution, Shape, Tensor,
+};
 use burn_common::benchmark::{run_benchmark, Benchmark};
 use derive_new::new;
 
@@ -14,7 +18,7 @@ impl<B: Backend, const D: usize> Benchmark for MatmulBenchmark<B, D> {
     type Args = (Tensor<B, D>, Tensor<B, D>, Tensor<B, 1>);
 
     fn name(&self) -> String {
-        "matmul_bias_relu".into()
+        "matmul_relu_bias_gelu".into()
     }
 
     fn shapes(&self) -> Vec<Vec<usize>> {
@@ -23,7 +27,7 @@ impl<B: Backend, const D: usize> Benchmark for MatmulBenchmark<B, D> {
 
     fn execute(&self, (lhs, rhs, bias): Self::Args) {
         let bias = bias.unsqueeze();
-        relu(lhs.matmul(rhs) + bias);
+        gelu(relu(lhs.matmul(rhs)) + bias);
     }
 
     fn prepare(&self) -> Self::Args {

backend-comparison/benches/matmul_fused.rs

@@ -59,6 +59,84 @@ pub(crate) struct OperationConverter {
     scalar_u8: Vec<u8>,
 }
 
+/// Fork of a [context](Context) which owns its data.
+pub struct ContextOwned<H> {
+    tensors: HashMap<TensorId, TensorDescription>,
+    handles: HandleContainer<H>,
+    scalar_f32: Vec<f32>,
+    scalar_f16: Vec<f16>,
+    scalar_bf16: Vec<bf16>,
+    scalar_i64: Vec<i64>,
+    scalar_i32: Vec<i32>,
+    scalar_i16: Vec<i16>,
+    scalar_i8: Vec<i8>,
+    scalar_u64: Vec<u64>,
+    scalar_u32: Vec<u32>,
+    scalar_u16: Vec<u16>,
+    scalar_u8: Vec<u8>,
+}
+
+impl<H: Clone> ContextOwned<H> {
+    /// Convert into [context](Context).
+    pub fn as_context(&mut self) -> Context<'_, H> {
+        Context {
+            tensors: &mut self.tensors,
+            handles: &mut self.handles,
+            scalar_f32: &self.scalar_f32,
+            scalar_f16: &self.scalar_f16,
+            scalar_bf16: &self.scalar_bf16,
+            scalar_i64: &self.scalar_i64,
+            scalar_i32: &self.scalar_i32,
+            scalar_i16: &self.scalar_i16,
+            scalar_i8: &self.scalar_i8,
+            scalar_u64: &self.scalar_u64,
+            scalar_u32: &self.scalar_u32,
+            scalar_u16: &self.scalar_u16,
+            scalar_u8: &self.scalar_u8,
+        }
+    }
+
+    /// Fork the context again.
+    pub fn fork(&self) -> ContextOwned<H> {
+        ContextOwned {
+            tensors: self.tensors.clone(),
+            handles: self.handles.fork(),
+            scalar_f32: self.scalar_f32.clone(),
+            scalar_f16: self.scalar_f16.clone(),
+            scalar_bf16: self.scalar_bf16.clone(),
+            scalar_i64: self.scalar_i64.clone(),
+            scalar_i32: self.scalar_i32.clone(),
+            scalar_i16: self.scalar_i16.clone(),
+            scalar_i8: self.scalar_i8.clone(),
+            scalar_u64: self.scalar_u64.clone(),
+            scalar_u32: self.scalar_u32.clone(),
+            scalar_u16: self.scalar_u16.clone(),
+            scalar_u8: self.scalar_u8.clone(),
+        }
+    }
+}
+
+impl<H: Clone> Context<'_, H> {
+    /// Fork the context into an [owned context](ContextOwned).
+    pub fn fork(&self) -> ContextOwned<H> {
+        ContextOwned {
+            tensors: self.tensors.clone(),
+            handles: self.handles.fork(),
+            scalar_f32: self.scalar_f32.clone(),
+            scalar_f16: self.scalar_f16.clone(),
+            scalar_bf16: self.scalar_bf16.clone(),
+            scalar_i64: self.scalar_i64.clone(),
+            scalar_i32: self.scalar_i32.clone(),
+            scalar_i16: self.scalar_i16.clone(),
+            scalar_i8: self.scalar_i8.clone(),
+            scalar_u64: self.scalar_u64.clone(),
+            scalar_u32: self.scalar_u32.clone(),
+            scalar_u16: self.scalar_u16.clone(),
+            scalar_u8: self.scalar_u8.clone(),
+        }
+    }
+}
+
 pub(crate) trait RelativeOps {
     /// Convert (usually an [`OperationDescription`]) to a relative form.
     ///

crates/burn-fusion/src/stream/context.rs

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

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

@@ -47,7 +47,13 @@ impl<R: JitRuntime> OptimizationBuilder<JitOptimization<R>> for MatmulBuilder<R>
                 let rhs = self.builder.input_unhandled(&op.rhs);
                 let out = self.builder.output_unhandled(&op.out);
 
-                self.matmul = Some(FusedMatmul::new(lhs, rhs, out, op.clone()));
+                self.matmul = Some(FusedMatmul::new(
+                    lhs,
+                    rhs,
+                    out,
+                    op.clone(),
+                    Default::default(),
+                ));
             } else {
                 self.builder.close();
             }

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

@@ -2,3 +2,4 @@ pub(crate) mod args;
 pub(crate) mod builder;
 pub(crate) mod optimization;
 pub(crate) mod spec;
+pub(crate) mod tune;

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

@@ -12,7 +12,9 @@ use burn_tensor::Shape;
 use cubecl::linalg::matmul::components;
 use cubecl::linalg::matmul::components::tile::accelerated::Accelerated;
 use cubecl::linalg::matmul::components::MatmulProblem;
-use cubecl::linalg::matmul::kernels::matmul::{MatmulSelector, StandardSelector};
+use cubecl::linalg::matmul::kernels::matmul::{
+    MatmulSelector, PipelinedSelector, SpecializedSelector, StandardSelector,
+};
 use cubecl::linalg::matmul::kernels::{MatmulAvailabilityError, MatmulLaunchError};
 use cubecl::linalg::tensor::{matrix_layout, MatrixLayout};
 use cubecl::{client::ComputeClient, prelude::*};
@@ -26,30 +28,65 @@ use crate::fusion::on_write::{
 
 use super::args::FusedMatmulInputLaunch;
 use super::spec::FusedMatmulSpec;
+use super::tune::fused_matmul_autotune;
 
-#[derive(new)]
 /// Fuse matmul operation followed by elemwise operations into a single kernel.
 pub struct MatmulOptimization<R: JitRuntime> {
     trace: FuseOnWriteTrace,
     trace_fallback: FuseOnWriteTrace,
-    client: ComputeClient<R::Server, R::Channel>,
-    device: R::Device,
-    len: usize,
-    matmul: FusedMatmul,
+    pub(crate) client: ComputeClient<R::Server, R::Channel>,
+    pub(crate) device: R::Device,
+    pub(crate) len: usize,
+    pub(crate) matmul_standard: FusedMatmul,
+    pub(crate) matmul_pipelined: FusedMatmul,
+    pub(crate) matmul_specialized: FusedMatmul,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
 /// State for the [matrix optimization](MatmulOptimizationState).
 pub struct MatmulOptimizationState {
     trace: FuseOnWriteTrace,
     trace_fallback: FuseOnWriteTrace,
-    matmul: FusedMatmul,
+    matmul_standard: FusedMatmul,
+    matmul_pipelined: FusedMatmul,
+    matmul_specialized: FusedMatmul,
     len: usize,
 }
 
 impl<R: JitRuntime> MatmulOptimization<R> {
+    pub fn new(
+        trace: FuseOnWriteTrace,
+        trace_fallback: FuseOnWriteTrace,
+        client: ComputeClient<R::Server, R::Channel>,
+        device: R::Device,
+        len: usize,
+        matmul: FusedMatmul,
+    ) -> Self {
+        let mut matmul_standard = matmul.clone();
+        let mut matmul_specialized = matmul.clone();
+        let mut matmul_pipelined = matmul;
+
+        matmul_standard.selector = FusedMatmulSelector::Standard;
+        matmul_specialized.selector = FusedMatmulSelector::Specialized;
+        matmul_pipelined.selector = FusedMatmulSelector::Pipelined;
+
+        Self {
+            trace,
+            trace_fallback,
+            client,
+            device,
+            len,
+            matmul_standard,
+            matmul_pipelined,
+            matmul_specialized,
+        }
+    }
     /// Execute the optimization.
     pub fn execute<BT: BoolElement>(&mut self, context: &mut Context<'_, JitFusionHandle<R>>) {
+        #[cfg(feature = "autotune")]
+        fused_matmul_autotune::<R, BT>(self, context);
+
+        #[cfg(not(feature = "autotune"))]
         if self.execute_fused::<BT>(context).is_err() {
             self.execute_fallback::<BT>(context);
         }
@@ -68,7 +105,9 @@ impl<R: JitRuntime> MatmulOptimization<R> {
             len: state.len,
             client: R::client(device),
             device: device.clone(),
-            matmul: state.matmul.clone(),
+            matmul_standard: state.matmul_standard.clone(),
+            matmul_specialized: state.matmul_specialized.clone(),
+            matmul_pipelined: state.matmul_pipelined.clone(),
         }
     }
 
@@ -77,21 +116,51 @@ impl<R: JitRuntime> MatmulOptimization<R> {
         MatmulOptimizationState {
             trace: self.trace.clone(),
             trace_fallback: self.trace_fallback.clone(),
-            matmul: self.matmul.clone(),
+            matmul_standard: self.matmul_standard.clone(),
+            matmul_specialized: self.matmul_specialized.clone(),
+            matmul_pipelined: self.matmul_pipelined.clone(),
             len: self.len,
         }
     }
 
-    fn execute_fused<BT: BoolElement>(
-        &mut self,
+    pub fn execute_standard_fused<BT: BoolElement>(
+        &self,
         context: &mut Context<'_, JitFusionHandle<R>>,
     ) -> Result<(), FusedMatmulError> {
-        self.trace
-            .run::<R, BT, FusedMatmul>(&self.client, &self.device, context, &self.matmul)
+        self.trace.run::<R, BT, FusedMatmul>(
+            &self.client,
+            &self.device,
+            context,
+            &self.matmul_standard,
+        )
     }
 
-    fn execute_fallback<BT: BoolElement>(&mut self, context: &mut Context<'_, JitFusionHandle<R>>) {
-        match self.matmul.lhs.precision() {
+    pub fn execute_specialized_fused<BT: BoolElement>(
+        &self,
+        context: &mut Context<'_, JitFusionHandle<R>>,
+    ) -> Result<(), FusedMatmulError> {
+        self.trace.run::<R, BT, FusedMatmul>(
+            &self.client,
+            &self.device,
+            context,
+            &self.matmul_specialized,
+        )
+    }
+
+    pub fn execute_pipelined_fused<BT: BoolElement>(
+        &self,
+        context: &mut Context<'_, JitFusionHandle<R>>,
+    ) -> Result<(), FusedMatmulError> {
+        self.trace.run::<R, BT, FusedMatmul>(
+            &self.client,
+            &self.device,
+            context,
+            &self.matmul_pipelined,
+        )
+    }
+
+    pub fn execute_fallback<BT: BoolElement>(&self, context: &mut Context<'_, JitFusionHandle<R>>) {
+        match self.matmul_standard.lhs.precision() {
             ElemwisePrecision::F32 => self.run_fallback::<BT, f32>(context),
             ElemwisePrecision::F16 => self.run_fallback::<BT, f16>(context),
             ElemwisePrecision::BF16 => self.run_fallback::<BT, bf16>(context),
@@ -100,13 +169,25 @@ impl<R: JitRuntime> MatmulOptimization<R> {
     }
 
     fn run_fallback<BT: BoolElement, EG: FloatElement>(
-        &mut self,
+        &self,
         context: &mut Context<'_, JitFusionHandle<R>>,
     ) {
         let (out_tensor, out_desc) = {
-            let lhs = context.tensors.get(&self.matmul.op.lhs.id).unwrap().clone();
-            let rhs = context.tensors.get(&self.matmul.op.rhs.id).unwrap().clone();
-            let out = context.tensors.get(&self.matmul.op.out.id).unwrap().clone();
+            let lhs = context
+                .tensors
+                .get(&self.matmul_standard.op.lhs.id)
+                .unwrap()
+                .clone();
+            let rhs = context
+                .tensors
+                .get(&self.matmul_standard.op.rhs.id)
+                .unwrap()
+                .clone();
+            let out = context
+                .tensors
+                .get(&self.matmul_standard.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);
@@ -136,12 +217,21 @@ impl<R: JitRuntime> MatmulOptimization<R> {
     }
 }
 
+#[derive(Default, Clone, Serialize, Deserialize, Debug)]
+pub enum FusedMatmulSelector {
+    #[default]
+    Standard,
+    Pipelined,
+    Specialized,
+}
+
 #[derive(new, Clone, Serialize, Deserialize, Debug)]
 pub struct FusedMatmul {
     lhs: Arg,
     rhs: Arg,
     out: Arg,
-    op: BinaryOperationDescription,
+    pub(crate) op: BinaryOperationDescription,
+    pub(crate) selector: FusedMatmulSelector,
 }
 
 #[derive(Debug)]
@@ -261,15 +351,43 @@ impl FusedMatmul {
             }
         };
 
-        match matmul_launch_kernel::<R, EG, StandardSelector<Accelerated>>(
-            client,
-            FusedMatmulInputLaunch::new(inputs, config, &self.lhs, &self.rhs, &self.out),
-            outputs,
-            problem,
-            plane_size,
-        ) {
-            Ok(_) => Ok(()),
-            Err(err) => Err(FusedMatmulError::LaunchError(err)),
+        match self.selector {
+            FusedMatmulSelector::Standard => {
+                match matmul_launch_kernel::<R, EG, StandardSelector<Accelerated>>(
+                    client,
+                    FusedMatmulInputLaunch::new(inputs, config, &self.lhs, &self.rhs, &self.out),
+                    outputs,
+                    problem,
+                    plane_size,
+                ) {
+                    Ok(_) => Ok(()),
+                    Err(err) => Err(FusedMatmulError::LaunchError(err)),
+                }
+            }
+            FusedMatmulSelector::Pipelined => {
+                match matmul_launch_kernel::<R, EG, PipelinedSelector<Accelerated>>(
+                    client,
+                    FusedMatmulInputLaunch::new(inputs, config, &self.lhs, &self.rhs, &self.out),
+                    outputs,
+                    problem,
+                    plane_size,
+                ) {
+                    Ok(_) => Ok(()),
+                    Err(err) => Err(FusedMatmulError::LaunchError(err)),
+                }
+            }
+            FusedMatmulSelector::Specialized => {
+                match matmul_launch_kernel::<R, EG, SpecializedSelector<Accelerated>>(
+                    client,
+                    FusedMatmulInputLaunch::new(inputs, config, &self.lhs, &self.rhs, &self.out),
+                    outputs,
+                    problem,
+                    plane_size,
+                ) {
+                    Ok(_) => Ok(()),
+                    Err(err) => Err(FusedMatmulError::LaunchError(err)),
+                }
+            }
         }
     }
 }

crates/burn-jit/src/fusion/matmul/optimization.rs

@@ -0,0 +1,133 @@
+use crate::{
+    fusion::{
+        tune::{TuneContext, TuneInput},
+        JitFusionHandle,
+    },
+    kernel::matmul::MatmulAutotuneKey,
+    BoolElement, JitRuntime, JitTuneId,
+};
+use burn_fusion::stream::Context;
+use cubecl::{
+    tune::{local_tuner, LocalTuner, TunableSet},
+    AutotuneKey,
+};
+use serde::{Deserialize, Serialize};
+
+use super::optimization::MatmulOptimization;
+
+#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
+pub struct FusedMatmulAutotuneKey {
+    matmul_key: MatmulAutotuneKey,
+    #[autotune(anchor)]
+    num_ops_fused: usize,
+}
+
+/// Executes autotune on matmul operations
+pub fn fused_matmul_autotune<R: JitRuntime, BT: BoolElement>(
+    optimization: &MatmulOptimization<R>,
+    context: &mut Context<JitFusionHandle<R>>,
+) {
+    static TUNER: LocalTuner<FusedMatmulAutotuneKey, JitTuneId> = local_tuner!();
+
+    let tunables = TunableSet::new(create_key::<R>, input_gen::<R>)
+        .with_tunable(tune_standard_fused::<R, BT>)
+        .with_tunable(tune_specialized_fused::<R, BT>)
+        .with_tunable(tune_pipelined_fused::<R, BT>)
+        .with_tunable(tune_fallback::<R, BT>);
+
+    TUNER.execute(
+        &JitTuneId::new::<R>(&optimization.device),
+        &optimization.client,
+        &tunables,
+        TuneInput::new(context, optimization),
+    );
+}
+
+pub(crate) fn create_key<R: JitRuntime>(
+    input: &TuneInput<R, MatmulOptimization<R>>,
+) -> FusedMatmulAutotuneKey {
+    let opt = input.optimization();
+    let context = match input.context() {
+        TuneContext::Original(context) => context,
+        TuneContext::Fork(_) => panic!("Not supported when generating key"),
+    };
+
+    let lhs = context.tensors.get(&opt.matmul_standard.op.lhs.id).unwrap();
+    let rhs = context.tensors.get(&opt.matmul_standard.op.rhs.id).unwrap();
+    let out = context.tensors.get(&opt.matmul_standard.op.out.id).unwrap();
+
+    let key = MatmulAutotuneKey::from_shape(
+        &lhs.shape.clone().into(),
+        &rhs.shape.clone().into(),
+        out.dtype,
+    );
+    FusedMatmulAutotuneKey::new(key, opt.len)
+}
+
+fn input_gen<R: JitRuntime>(
+    _key: &FusedMatmulAutotuneKey,
+    input: &TuneInput<R, MatmulOptimization<R>>,
+) -> TuneInput<R, MatmulOptimization<R>> {
+    input.clone()
+}
+
+fn tune_standard_fused<R: JitRuntime, BT: BoolElement>(
+    input: TuneInput<R, MatmulOptimization<R>>,
+) -> Result<(), String> {
+    let optimization = input.optimization();
+    let context = input.context();
+
+    match context {
+        TuneContext::Original(context) => optimization.execute_standard_fused::<BT>(context),
+        TuneContext::Fork(mut context_owned) => {
+            optimization.execute_standard_fused::<BT>(&mut context_owned.as_context())
+        }
+    }
+    .map_err(|e| format!("{e:?}"))
+}
+
+fn tune_specialized_fused<R: JitRuntime, BT: BoolElement>(
+    input: TuneInput<R, MatmulOptimization<R>>,
+) -> Result<(), String> {
+    let optimization = input.optimization();
+    let context = input.context();
+
+    match context {
+        TuneContext::Original(context) => optimization.execute_specialized_fused::<BT>(context),
+        TuneContext::Fork(mut context_owned) => {
+            optimization.execute_specialized_fused::<BT>(&mut context_owned.as_context())
+        }
+    }
+    .map_err(|e| format!("{e:?}"))
+}
+
+fn tune_pipelined_fused<R: JitRuntime, BT: BoolElement>(
+    input: TuneInput<R, MatmulOptimization<R>>,
+) -> Result<(), String> {
+    let optimization = input.optimization();
+    let context = input.context();
+
+    match context {
+        TuneContext::Original(context) => optimization.execute_pipelined_fused::<BT>(context),
+        TuneContext::Fork(mut context_owned) => {
+            optimization.execute_pipelined_fused::<BT>(&mut context_owned.as_context())
+        }
+    }
+    .map_err(|e| format!("{e:?}"))
+}
+
+fn tune_fallback<R: JitRuntime, BT: BoolElement>(
+    input: TuneInput<R, MatmulOptimization<R>>,
+) -> Result<(), String> {
+    let optimization = input.optimization();
+    let context = input.context();
+
+    match context {
+        TuneContext::Original(context) => optimization.execute_fallback::<BT>(context),
+        TuneContext::Fork(mut context_owned) => {
+            optimization.execute_fallback::<BT>(&mut context_owned.as_context())
+        }
+    };
+
+    Ok(())
+}

crates/burn-jit/src/fusion/matmul/tune.rs

@@ -3,5 +3,6 @@ mod base;
 pub(crate) mod elemwise;
 pub(crate) mod matmul;
 pub(crate) mod on_write;
+pub(crate) mod tune;
 
 pub use base::*;

crates/burn-jit/src/fusion/mod.rs

@@ -0,0 +1,108 @@
+use super::JitFusionHandle;
+use crate::JitRuntime;
+use burn_fusion::stream::{Context, ContextOwned};
+
+/// Fusion context used when tuning kernels.
+///
+/// Either the original context is returned or a fork of the original.
+/// The fork is only given when performing autotuning, and not when actually performing the
+/// operation.
+pub enum TuneContext<'a, R: JitRuntime> {
+    Original(&'a mut Context<'a, JitFusionHandle<R>>),
+    Fork(Box<ContextOwned<JitFusionHandle<R>>>),
+}
+
+/// Fusion input wrapper containing the context and the optimization.
+///
+/// # Safety
+///
+/// This should only be used with the [tuner](cubecl::tune::LocalTuner), since safety assumptions
+/// are made based on its behavior.
+pub struct TuneInput<R: JitRuntime, O> {
+    context: UnsafeTuneContext<R>,
+    optimization: *const O,
+}
+
+/// Unsafe wrapper around the context.
+///
+/// # Safety
+///
+/// The wrapper removes the context lifetime.
+///
+/// For it to be correct, the context must not be used after the invocation of the
+/// [cubecl::tune::LocalTuner::execute] function. This is the case, since autotune functions are
+/// tuned using a cloned version of the input; therefore, a fork of the context will be used to find
+/// the best kernel to use, which can be async.
+enum UnsafeTuneContext<R: JitRuntime> {
+    Original(*mut Context<'static, JitFusionHandle<R>>),
+    Fork(Box<ContextOwned<JitFusionHandle<R>>>),
+}
+
+unsafe impl<R: JitRuntime> Send for UnsafeTuneContext<R> {}
+unsafe impl<R: JitRuntime, O> Send for TuneInput<R, O> {}
+
+impl<R: JitRuntime, O> TuneInput<R, O> {
+    /// Create a new autotune input from the [context](Context) and an optimization.
+    pub fn new(context: &mut Context<JitFusionHandle<R>>, optimization: &O) -> Self {
+        let context = UnsafeTuneContext::new(context);
+        // We can erase the lifetime for the same reason we do with the context.
+        let optimization = core::ptr::from_ref(optimization);
+
+        Self {
+            context,
+            optimization,
+        }
+    }
+
+    /// Retrieve the [autotune context](TuneContext) for the current input.
+    pub fn context(&self) -> TuneContext<'static, R> {
+        self.context.get()
+    }
+
+    /// Retrieve the optimization for the current input.
+    pub fn optimization(&self) -> &O {
+        unsafe { self.optimization.as_ref().unwrap() }
+    }
+}
+
+impl<R: JitRuntime> UnsafeTuneContext<R> {
+    fn new(context: &mut Context<'_, JitFusionHandle<R>>) -> Self {
+        let ptr = core::ptr::from_mut(context);
+
+        // It is necessary for the lifetime.
+        #[allow(clippy::unnecessary_cast)]
+        Self::Original(ptr as *mut Context<'static, _>)
+    }
+
+    fn get(&self) -> TuneContext<'static, R> {
+        match self {
+            UnsafeTuneContext::Original(ptr) => {
+                TuneContext::Original(unsafe { ptr.as_mut().unwrap() })
+            }
+            UnsafeTuneContext::Fork(context) => TuneContext::Fork(Box::new(context.fork())),
+        }
+    }
+}
+
+impl<R: JitRuntime, O> Clone for TuneInput<R, O> {
+    fn clone(&self) -> Self {
+        Self {
+            context: self.context.clone(),
+            optimization: self.optimization,
+        }
+    }
+}
+
+impl<R: JitRuntime> Clone for UnsafeTuneContext<R> {
+    fn clone(&self) -> Self {
+        let context = match self {
+            UnsafeTuneContext::Original(ptr) => {
+                let context: &mut Context<'static, JitFusionHandle<R>> =
+                    unsafe { ptr.as_mut().unwrap() };
+                context.fork()
+            }
+            UnsafeTuneContext::Fork(context) => context.fork(),
+        };
+        UnsafeTuneContext::Fork(Box::new(context))
+    }
+}

crates/burn-jit/src/fusion/tune.rs

@@ -22,7 +22,7 @@ pub struct MatmulAutotuneKey {
 }
 
 impl MatmulAutotuneKey {
-    fn from_shape(lhs_shape: &Shape, rhs_shape: &Shape, dtype: DType) -> Self {
+    pub(crate) fn from_shape(lhs_shape: &Shape, rhs_shape: &Shape, dtype: DType) -> Self {
         let ndims = lhs_shape.num_dims();
         let m = lhs_shape.dims[ndims - 2];
         let k = lhs_shape.dims[ndims - 1];

crates/burn-jit/src/kernel/matmul/tune/key.rs

@@ -26,6 +26,23 @@ pub struct HandleContainer<H> {
     pub handles_orphan: Vec<TensorId>,
 }
 
+impl<H: Clone> HandleContainer<H> {
+    /// Fork the container, useful for autotune.
+    pub fn fork(&self) -> Self {
+        let mut handles = HashMap::with_capacity(self.handles.len());
+
+        for (id, handle) in self.handles.iter() {
+            handles.insert(*id, handle.clone());
+        }
+
+        Self {
+            handles,
+            counter: self.counter,
+            handles_orphan: self.handles_orphan.clone(),
+        }
+    }
+}
+
 impl<H> core::fmt::Debug for HandleContainer<H> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         f.debug_struct("HandleContainer")
@@ -37,6 +54,7 @@ impl<H> core::fmt::Debug for HandleContainer<H> {
 }
 
 /// Backend [tensor handle](ReprBackend::Handle) wrapper tracking their creation state
+#[derive(Clone)]
 pub enum Handle<H> {
     /// No [tensor handle](ReprBackend::Handle) has been created yet
     NotInit,

crates/burn-tensor/src/repr/handle.rs

@@ -1,3 +1,5 @@
+#![recursion_limit = "256"]
+
 use burn::{
     nn::transformer::TransformerEncoderConfig,
     optim::{decay::WeightDecayConfig, AdamConfig},