Fix clippy

Author: nathanielsimard

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

@@ -29,7 +29,7 @@ pub fn fused_matmul_autotune<R: JitRuntime, BT: BoolElement>(
 ) {
     static TUNER: LocalTuner<FusedMatmulAutotuneKey, JitTuneId> = local_tuner!();
 
-    let tunables = TunableSet::new(create_key::<R, BT>, input_gen::<R, BT>)
+    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>)
@@ -43,7 +43,7 @@ pub fn fused_matmul_autotune<R: JitRuntime, BT: BoolElement>(
     );
 }
 
-pub(crate) fn create_key<R: JitRuntime, BT: BoolElement>(
+pub(crate) fn create_key<R: JitRuntime>(
     input: &TuneInput<R, MatmulOptimization<R>>,
 ) -> FusedMatmulAutotuneKey {
     let opt = input.optimization();
@@ -64,7 +64,7 @@ pub(crate) fn create_key<R: JitRuntime, BT: BoolElement>(
     FusedMatmulAutotuneKey::new(key, opt.len)
 }
 
-fn input_gen<R: JitRuntime, BT: BoolElement>(
+fn input_gen<R: JitRuntime>(
     _key: &FusedMatmulAutotuneKey,
     input: &TuneInput<R, MatmulOptimization<R>>,
 ) -> TuneInput<R, MatmulOptimization<R>> {

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

@@ -9,7 +9,7 @@ use burn_fusion::stream::{Context, ContextOwned};
 /// operation.
 pub enum TuneContext<'a, R: JitRuntime> {
     Original(&'a mut Context<'a, JitFusionHandle<R>>),
-    Fork(ContextOwned<JitFusionHandle<R>>),
+    Fork(Box<ContextOwned<JitFusionHandle<R>>>),
 }
 
 /// Fusion input wrapper containing the context and the optimization.
@@ -35,7 +35,7 @@ pub struct TuneInput<R: JitRuntime, O> {
 /// the best kernel to use, which can be async.
 enum UnsafeTuneContext<R: JitRuntime> {
     Original(*mut Context<'static, JitFusionHandle<R>>),
-    Fork(ContextOwned<JitFusionHandle<R>>),
+    Fork(Box<ContextOwned<JitFusionHandle<R>>>),
 }
 
 unsafe impl<R: JitRuntime> Send for UnsafeTuneContext<R> {}
@@ -67,15 +67,19 @@ impl<R: JitRuntime, O> TuneInput<R, O> {
 
 impl<R: JitRuntime> UnsafeTuneContext<R> {
     fn new(context: &mut Context<'_, JitFusionHandle<R>>) -> Self {
-        Self::Original(core::ptr::from_mut(context) as *mut Context<'static, JitFusionHandle<R>>)
+        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(context.fork()),
+            UnsafeTuneContext::Fork(context) => TuneContext::Fork(Box::new(context.fork())),
         }
     }
 }
@@ -84,7 +88,7 @@ impl<R: JitRuntime, O> Clone for TuneInput<R, O> {
     fn clone(&self) -> Self {
         Self {
             context: self.context.clone(),
-            optimization: self.optimization.clone(),
+            optimization: self.optimization,
         }
     }
 }
@@ -99,6 +103,6 @@ impl<R: JitRuntime> Clone for UnsafeTuneContext<R> {
             }
             UnsafeTuneContext::Fork(context) => context.fork(),
         };
-        UnsafeTuneContext::Fork(context)
+        UnsafeTuneContext::Fork(Box::new(context))
     }
 }