Migrate/jit/matmul tiling 2d

crates/burn-jit/src/codegen/compiler.rs

@@ -22,4 +22,6 @@ pub trait Compiler: Sync + Send + 'static + Clone + Default + core::fmt::Debug {
     fn compile(shader: gpu::ComputeShader) -> Self::Representation;
     /// The size of the given element in bytes.
     fn elem_size(elem: gpu::Elem) -> usize;
+    /// The maximal size of a shared memory
+    fn max_shared_memory_size() -> usize;
 }

crates/burn-jit/src/codegen/dialect/gpu/branch.rs

@@ -119,3 +119,15 @@ impl Loop {
         parent_scope.register(Branch::Loop(op));
     }
 }
+
+#[allow(missing_docs)]
+pub struct UnrolledRangeLoop;
+
+impl UnrolledRangeLoop {
+    /// Registers an unrolled range loop to the given scope.
+    pub fn register<F: Fn(Variable, &mut Scope)>(scope: &mut Scope, start: u32, end: u32, func: F) {
+        for i in start..end {
+            func(i.into(), scope);
+        }
+    }
+}

crates/burn-jit/src/codegen/dialect/gpu/macros.rs

@@ -293,6 +293,12 @@ macro_rules! gpu {
             gpu!(unary $input, $out)
         ));
     };
+    // out = vec4(a, b, c, d)
+    ($scope:expr, $out:ident = vec4($a:ident,$b:ident,$c:ident,$d:ident)) => {
+        $scope.register($crate::codegen::dialect::gpu::Operator::AssignVec4(
+            $crate::codegen::dialect::gpu::AssignVec4Operator{a:$a,b:$b,c:$c,d:$d,out:$out}
+        ));
+    };
     // out = input
     ($scope:expr, $out:ident = $input:ident) => {
         gpu!($scope, $out = cast($input))
@@ -326,10 +332,18 @@ macro_rules! gpu {
             out: $out.into(),
         });
     };
-    // range(start, end).for_each(|scope| { ... })
+    // range(start, end).for_each(|i, scope| { ... })
     ($scope:expr, range($start:expr, $end:expr).for_each($arg:expr)) => {
         $crate::codegen::dialect::gpu::RangeLoop::register($scope, $start.into(), $end.into(), $arg);
     };
+    // range(start, end, unroll).for_each(|i, scope| { ... })
+    ($scope:expr, range($start:expr, $end:expr, $unroll:expr).for_each($arg:expr)) => {
+        if $unroll {
+            $crate::codegen::dialect::gpu::UnrolledRangeLoop::register($scope, $start.into(), $end.into(), $arg);
+        } else {
+            $crate::codegen::dialect::gpu::RangeLoop::register($scope, $start.into(), $end.into(), $arg);
+        }
+    };
     // loop(|scope| { ... })
     ($scope:expr, loop($arg:expr)) => {
         $crate::codegen::dialect::gpu::Loop::register($scope, $arg);

crates/burn-jit/src/codegen/dialect/gpu/operation.rs

@@ -36,6 +36,7 @@ pub enum Operator {
     Tanh(UnaryOperator),
     Powf(BinaryOperator),
     Sqrt(UnaryOperator),
+    Ceil(UnaryOperator),
     Erf(UnaryOperator),
     Recip(UnaryOperator),
     Equal(BinaryOperator),
@@ -58,6 +59,7 @@ pub enum Operator {
     BitwiseXor(BinaryOperator),
     ShiftLeft(BinaryOperator),
     ShiftRight(BinaryOperator),
+    AssignVec4(AssignVec4Operator),
 }
 
 /// All metadata that can be access in a shader.
@@ -106,6 +108,16 @@ pub struct ClampOperator {
     pub out: Variable,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
+pub struct AssignVec4Operator {
+    pub a: Variable,
+    pub b: Variable,
+    pub c: Variable,
+    pub d: Variable,
+    pub out: Variable,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[allow(missing_docs)]
 pub struct ReadGlobalOperator {

crates/burn-jit/src/codegen/dialect/gpu/scope.rs

@@ -20,7 +20,8 @@ pub struct Scope {
     pub depth: u8,
     pub operations: Vec<Operation>,
     locals: Vec<Variable>,
-    shared: Vec<Variable>,
+    shared_memories: Vec<Variable>,
+    local_arrays: Vec<Variable>,
     reads_global: Vec<(Variable, ReadingStrategy, Variable)>,
     index_offset_with_output_layout_position: Vec<usize>,
     writes_global: Vec<(Variable, Variable)>,
@@ -48,7 +49,8 @@ impl Scope {
             depth: 0,
             operations: Vec::new(),
             locals: Vec::new(),
-            shared: Vec::new(),
+            local_arrays: Vec::new(),
+            shared_memories: Vec::new(),
             reads_global: Vec::new(),
             index_offset_with_output_layout_position: Vec::new(),
             writes_global: Vec::new(),
@@ -213,7 +215,8 @@ impl Scope {
             depth: self.depth + 1,
             operations: Vec::new(),
             locals: Vec::new(),
-            shared: Vec::new(),
+            shared_memories: Vec::new(),
+            local_arrays: Vec::new(),
             reads_global: Vec::new(),
             index_offset_with_output_layout_position: Vec::new(),
             writes_global: Vec::new(),
@@ -308,7 +311,11 @@ impl Scope {
     }
 
     fn new_shared_index(&self) -> u16 {
-        self.shared.len() as u16
+        self.shared_memories.len() as u16
+    }
+
+    fn new_local_array_index(&self) -> u16 {
+        self.local_arrays.len() as u16
     }
 
     fn read_input_strategy(
@@ -339,7 +346,16 @@ impl Scope {
         let item = item.into();
         let index = self.new_shared_index();
         let shared_memory = Variable::SharedMemory(index, item, shared_memory_size);
-        self.shared.push(shared_memory);
+        self.shared_memories.push(shared_memory);
         shared_memory
     }
+
+    /// Create a local array of the given [item type](Item).
+    pub fn create_local_array<I: Into<Item>>(&mut self, item: I, array_size: u32) -> Variable {
+        let item = item.into();
+        let index = self.new_local_array_index();
+        let local_array = Variable::LocalArray(index, item, self.depth, array_size);
+        self.local_arrays.push(local_array);
+        local_array
+    }
 }

crates/burn-jit/src/codegen/dialect/gpu/variable.rs

@@ -11,6 +11,7 @@ pub enum Variable {
     LocalScalar(u16, Elem, u8),
     ConstantScalar(f64, Elem),
     SharedMemory(u16, Item, u32),
+    LocalArray(u16, Item, u8, u32),
     Id,
     LocalInvocationIndex,
     LocalInvocationIdX,
@@ -41,6 +42,7 @@ impl Variable {
             Variable::GlobalOutputArray(idx, _) => Some(*idx),
             Variable::ConstantScalar(_, _) => None,
             Variable::SharedMemory(idx, _, _) => Some(*idx),
+            Variable::LocalArray(idx, _, _, _) => Some(*idx),
             Variable::Id => None,
             Variable::LocalInvocationIndex => None,
             Variable::LocalInvocationIdX => None,
@@ -70,6 +72,7 @@ impl Variable {
             Variable::LocalScalar(_, elem, _) => Item::Scalar(*elem),
             Variable::ConstantScalar(_, elem) => Item::Scalar(*elem),
             Variable::SharedMemory(_, item, _) => *item,
+            Variable::LocalArray(_, item, _, _) => *item,
             Variable::Id => Item::Scalar(Elem::UInt),
             Variable::Rank => Item::Scalar(Elem::UInt),
             Variable::LocalInvocationIndex => Item::Scalar(Elem::UInt),

crates/burn-jit/src/codegen/dialect/gpu/vectorization.rs

@@ -1,4 +1,7 @@
-use super::{BinaryOperator, ClampOperator, Item, Operation, Operator, UnaryOperator, Variable};
+use super::{
+    AssignVec4Operator, BinaryOperator, ClampOperator, Item, Operation, Operator, UnaryOperator,
+    Variable,
+};
 
 /// Define a vectorization scheme.
 #[allow(dead_code)]
@@ -52,6 +55,7 @@ impl Operator {
             Operator::Tanh(op) => Operator::Tanh(op.vectorize(vectorization)),
             Operator::Powf(op) => Operator::Powf(op.vectorize(vectorization)),
             Operator::Sqrt(op) => Operator::Sqrt(op.vectorize(vectorization)),
+            Operator::Ceil(op) => Operator::Ceil(op.vectorize(vectorization)),
             Operator::Erf(op) => Operator::Erf(op.vectorize(vectorization)),
             Operator::Recip(op) => Operator::Recip(op.vectorize(vectorization)),
             Operator::Equal(op) => Operator::Equal(op.vectorize(vectorization)),
@@ -79,6 +83,7 @@ impl Operator {
             Operator::BitwiseXor(op) => Operator::BitwiseXor(op.vectorize(vectorization)),
             Operator::ShiftLeft(op) => Operator::ShiftLeft(op.vectorize(vectorization)),
             Operator::ShiftRight(op) => Operator::ShiftRight(op.vectorize(vectorization)),
+            Operator::AssignVec4(op) => Operator::AssignVec4(op.vectorize(vectorization)),
         }
     }
 }
@@ -113,6 +118,18 @@ impl ClampOperator {
     }
 }
 
+impl AssignVec4Operator {
+    pub(crate) fn vectorize(&self, vectorization: Vectorization) -> Self {
+        Self {
+            a: self.a,
+            b: self.b,
+            c: self.c,
+            d: self.d,
+            out: self.out.vectorize(vectorization),
+        }
+    }
+}
+
 impl Variable {
     pub(crate) fn vectorize(&self, vectorize: Vectorization) -> Self {
         match self {
@@ -130,6 +147,12 @@ impl Variable {
                 item.vectorize(vectorize),
                 item.vectorized_size(vectorize, *size),
             ),
+            Variable::LocalArray(index, item, name, size) => Variable::LocalArray(
+                *index,
+                item.vectorize(vectorize),
+                *name,
+                item.vectorized_size(vectorize, *size),
+            ),
             Variable::ConstantScalar(_, _) => *self,
             Variable::GlobalScalar(_, _) => *self,
             Variable::Id => *self,

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

@@ -247,6 +247,11 @@ impl TraceBuilder {
                         &mut local_tensor_ids_input,
                         &mut local_tensor_ids_output,
                     ),
+                    gpu::Operator::Ceil(op) => mark_unary(
+                        op,
+                        &mut local_tensor_ids_input,
+                        &mut local_tensor_ids_output,
+                    ),
                     gpu::Operator::Log(op) => mark_unary(
                         op,
                         &mut local_tensor_ids_input,
@@ -351,6 +356,13 @@ impl TraceBuilder {
                         &mut local_tensor_ids_input,
                         &mut local_tensor_ids_output,
                     ),
+                    gpu::Operator::AssignVec4(op) => {
+                        mark(&op.a, &mut local_tensor_ids_input);
+                        mark(&op.b, &mut local_tensor_ids_input);
+                        mark(&op.c, &mut local_tensor_ids_input);
+                        mark(&op.d, &mut local_tensor_ids_input);
+                        mark(&op.out, &mut local_tensor_ids_output);
+                    }
                 },
                 Operation::Procedure(proc) => {
                     match proc {