Remove some unsafe code (#3385)

Author: nathanielsimard

crates/burn-cubecl/src/fusion.rs

@@ -16,6 +16,7 @@ use burn_ir::{BackendIr, TensorHandle};
 use burn_tensor::{DType, Shape};
 use core::marker::PhantomData;
 use half::{bf16, f16};
+use std::sync::Arc;
 
 impl<R, BT> burn_fusion::Optimization<FusionCubeRuntime<R, BT>> for CubeOptimization<R>
 where
@@ -34,11 +35,11 @@ where
             Self::ElementWise(op) => op.execute::<BT>(context),
             Self::Matmul(op) => op.execute::<BT>(context, |index| {
                 let operation = execution.operation_within_optimization(index);
-                Box::new(FallbackOperationUnsafe::new(operation))
+                Box::new(FallbackOperationWrapper::new(operation))
             }),
             Self::Reduce(op) => op.execute::<BT>(context, |index| {
                 let operation = execution.operation_within_optimization(index);
-                Box::new(FallbackOperationUnsafe::new(operation))
+                Box::new(FallbackOperationWrapper::new(operation))
             }),
         }
     }
@@ -62,33 +63,21 @@ where
     }
 }
 
-/// This is only safe because we know the fallback must be executed before the cubecl context is dropped.
-///
-/// The safer alternatives would require fused operation to be cloned, so that it could
-/// escape the lifetime of the context's execution, which doesn't make sense since
-/// its only goal is to modify the context it operates on.
-struct FallbackOperationUnsafe<O> {
-    operation: *const O,
+struct FallbackOperationWrapper<O: Clone> {
+    operation: O,
 }
 
-unsafe impl<O> Send for FallbackOperationUnsafe<O> {}
-unsafe impl<O> Sync for FallbackOperationUnsafe<O> {}
-
-impl<O> FallbackOperationUnsafe<O> {
-    fn new(op: &O) -> Self {
-        let ptr = core::ptr::from_ref(op);
-
-        Self { operation: ptr }
+impl<O: Clone> FallbackOperationWrapper<O> {
+    fn new(op: O) -> Self {
+        Self { operation: op }
     }
 }
 
 impl<R: CubeRuntime, BT: BoolElement> FallbackOperation<R>
-    for FallbackOperationUnsafe<Box<dyn Operation<FusionCubeRuntime<R, BT>>>>
+    for FallbackOperationWrapper<Arc<dyn Operation<FusionCubeRuntime<R, BT>>>>
 {
     fn run(&self, context: &mut burn_fusion::stream::Context<'_, CubeFusionHandle<R>>) {
-        unsafe {
-            self.operation.as_ref().unwrap().execute(context.handles);
-        }
+        self.operation.as_ref().execute(context.handles);
     }
 }
 

crates/burn-fusion/src/client/mutex.rs

@@ -43,7 +43,7 @@ where
     {
         self.server
             .lock()
-            .register(streams, repr, Box::new(operation))
+            .register(streams, repr, Arc::new(operation))
     }
 
     fn drain(&self) {

crates/burn-fusion/src/server.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::{
     FusionBackend, FusionRuntime,
     stream::{MultiStream, OperationStreams, StreamId, execution::Operation},
@@ -25,7 +27,7 @@ where
         &mut self,
         streams: OperationStreams,
         repr: OperationIr,
-        operation: Box<dyn Operation<R>>,
+        operation: Arc<dyn Operation<R>>,
     ) {
         self.streams
             .register(streams, repr, operation, &mut self.handles)

crates/burn-fusion/src/stream/execution/ordering.rs

@@ -8,7 +8,7 @@ use super::Operation;
 
 /// Manage the execution of potentially multiple optimizations and operations out of order.
 pub struct OrderedExecution<R: FusionRuntime> {
-    operations: Vec<Box<dyn Operation<R>>>,
+    operations: Vec<Arc<dyn Operation<R>>>,
     num_executed: usize,
     ordering: Option<Arc<Vec<usize>>>,
 }
@@ -18,25 +18,25 @@ impl<R: FusionRuntime> OrderedExecution<R> {
     ///
     /// This is useful to implement fallback for optimizations.
     #[allow(clippy::borrowed_box)]
-    pub fn operation_within_optimization(&self, index: usize) -> &Box<dyn Operation<R>> {
+    pub fn operation_within_optimization(&self, index: usize) -> Arc<dyn Operation<R>> {
         match &self.ordering {
             Some(val) => {
                 let index = val[index];
-                &self.operations[index]
+                self.operations[index].clone()
             }
             None => panic!("No ordering provided"),
         }
     }
 
-    pub(crate) fn new(operations: Vec<Box<dyn Operation<R>>>) -> Self {
+    pub(crate) fn new(operations: Vec<Arc<dyn Operation<R>>>) -> Self {
         Self {
             operations,
             num_executed: 0,
             ordering: None,
         }
     }
 
-    pub(crate) fn finish(mut self) -> (Vec<Box<dyn Operation<R>>>, usize) {
+    pub(crate) fn finish(mut self) -> (Vec<Arc<dyn Operation<R>>>, usize) {
         self.operations.drain(0..self.num_executed);
         (self.operations, self.num_executed)
     }

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

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use burn_ir::{HandleContainer, OperationIr, TensorId, TensorIr, TensorStatus};
 use hashbrown::{HashMap, HashSet};
 
@@ -47,7 +49,7 @@ impl<R: FusionRuntime> MultiStream<R> {
         &mut self,
         streams: OperationStreams,
         mut repr: OperationIr,
-        operation: Box<dyn Operation<R>>,
+        operation: Arc<dyn Operation<R>>,
         handles: &mut HandleContainer<R::FusionHandle>,
     ) {
         let id = self.resolve_streams(&streams, handles, &mut repr);
@@ -133,7 +135,7 @@ impl<R: FusionRuntime> MultiStream<R> {
         id: StreamId,
         repr: OperationIr,
         streams: &OperationStreams,
-        operation: Box<dyn Operation<R>>,
+        operation: Arc<dyn Operation<R>>,
         handles: &mut HandleContainer<R::FusionHandle>,
     ) -> usize {
         let stream = match self.streams.get_mut(&id) {
@@ -355,7 +357,7 @@ impl<R: FusionRuntime> MultiStream<R> {
                 current,
             };
 
-            let op = Box::new(DropOp { id: tensor.id });
+            let op = Arc::new(DropOp { id: tensor.id });
             self.register(streams, OperationIr::Drop(tensor), op, handles);
         }
     }

crates/burn-fusion/src/stream/queue/base.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::FusionRuntime;
 use crate::stream::{OperationConverter, OperationStreams, RelativeOps, execution::Operation};
 use burn_common::id::StreamId;
@@ -17,7 +19,7 @@ pub struct OperationQueue<R: FusionRuntime> {
     /// determine which operations can be fused.
     pub(crate) relative: Vec<OperationIr>,
     pub(crate) converter: OperationConverter,
-    pub(crate) operations: Vec<Box<dyn Operation<R>>>,
+    pub(crate) operations: Vec<Arc<dyn Operation<R>>>,
     pub(crate) variables: HashMap<TensorId, (StreamId, TensorStatus)>,
 }
 
@@ -47,7 +49,7 @@ impl<R: FusionRuntime> OperationQueue<R> {
     pub fn add(
         &mut self,
         global: OperationIr,
-        operation: Box<dyn Operation<R>>,
+        operation: Arc<dyn Operation<R>>,
         streams: &OperationStreams,
         current: StreamId,
     ) {

crates/burn-fusion/src/stream/queue/execution.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use burn_ir::{HandleContainer, TensorStatus};
 
 use crate::{
@@ -84,8 +86,8 @@ impl<'a, R: FusionRuntime> QueueExecution<'a, R> {
         optimization: &mut BlockOptimization<R::Optimization>,
         converter: &'a mut OperationConverter,
         handles: &'a mut HandleContainer<R::FusionHandle>,
-        operations: Vec<Box<dyn Operation<R>>>,
-    ) -> (Vec<Box<dyn Operation<R>>>, usize) {
+        operations: Vec<Arc<dyn Operation<R>>>,
+    ) -> (Vec<Arc<dyn Operation<R>>>, usize) {
         let execution = OrderedExecution::new(operations);
 
         if matches!(&optimization.strategy, ExecutionStrategy::Composed(..)) {