Big refactoring

Author: nathanielsimard

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

@@ -2,10 +2,7 @@ use burn_fusion::OptimizationBuilder;
 
 use crate::{
     fusion::{
-        on_write::{
-            builder::{FuseOnWriteBuilder, FuseSettings},
-            ir::ElemwisePrecision,
-        },
+        on_write::{builder::FuseOnWriteBuilder, ir::ElemwisePrecision, settings::FuseSettings},
         JitOptimization,
     },
     JitRuntime,

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

@@ -3,10 +3,7 @@ use burn_tensor::repr::{FloatOperationDescription, OperationDescription};
 
 use crate::{
     fusion::{
-        on_write::{
-            builder::{FuseOnWriteBuilder, FuseSettings},
-            ir::ElemwisePrecision,
-        },
+        on_write::{builder::FuseOnWriteBuilder, ir::ElemwisePrecision, settings::FuseSettings},
         JitOptimization,
     },
     JitRuntime,

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

@@ -1,7 +1,7 @@
 use super::{
     ir::{Arg, BinaryElemwiseArgs, ElemwiseOp, ElemwisePrecision, UnaryElemwiseArgs},
-    trace::FuseOnWriteTrace,
-    trace_builder::FuseOnWriteTraceBuilder,
+    settings::FuseSettings,
+    trace::{FuseOnWriteTrace, FuseOnWriteTraceBuilder},
 };
 use burn_fusion::{OptimizationBuilder, OptimizationProperties, OptimizationStatus};
 use burn_tensor::{
@@ -13,7 +13,6 @@ use burn_tensor::{
     Element,
 };
 use cubecl::ir::Elem;
-use serde::{Deserialize, Serialize};
 
 /// Fused element wise operations that are normally memory bound.
 pub(crate) struct FuseOnWriteBuilder {
@@ -26,25 +25,6 @@ pub(crate) struct FuseOnWriteBuilder {
     max_bindings: u32,
 }
 
-/// Controls which operations can be fused.
-#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
-pub struct FuseSettings {
-    /// Enables broadcasting of shapes.
-    pub broadcast: bool,
-    /// Enables output shape updates.
-    ///
-    /// When broadcast is enabled, the output shape can become bigger after a fusion,
-    /// therefore an update is needed.
-    pub output_shape_updates: bool,
-    /// Enables mix vectorization factor.
-    ///
-    /// Useful when the last dimension is broadcasted for one of the tensors, which would limit the
-    /// vectorization factor to be 1 without this setting enabled.
-    pub mix_vectorization: bool,
-    /// Enables the reuse of input buffers.
-    pub inplace: bool,
-}
-
 struct TryFuseBuilder {
     builder: FuseOnWriteTraceBuilder,
     max_bindings: u32,

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

@@ -2,7 +2,6 @@ pub(crate) mod builder;
 pub(crate) mod io;
 pub(crate) mod ir;
 pub(crate) mod kernel;
-pub(super) mod position;
+pub(crate) mod settings;
 
 pub mod trace;
-pub(crate) mod trace_builder;

crates/burn-jit/src/fusion/on_write/position.rs

@@ -0,0 +1,20 @@
+use serde::{Deserialize, Serialize};
+
+/// Controls which operations can be fused.
+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+pub struct FuseSettings {
+    /// Enables broadcasting of shapes.
+    pub broadcast: bool,
+    /// Enables output shape updates.
+    ///
+    /// When broadcast is enabled, the output shape can become bigger after a fusion,
+    /// therefore an update is needed.
+    pub output_shape_updates: bool,
+    /// Enables mix vectorization factor.
+    ///
+    /// Useful when the last dimension is broadcasted for one of the tensors, which would limit the
+    /// vectorization factor to be 1 without this setting enabled.
+    pub mix_vectorization: bool,
+    /// Enables the reuse of input buffers.
+    pub inplace: bool,
+}