Splitted the JIT stuff from the Wgpu stuff (#1417)

Author: nathanielsimard

Cargo.lock

@@ -0,0 +1,57 @@
+[package]
+authors = ["nathanielsimard <[email protected]>"]
+categories = ["science"]
+description = "Generic backend that can be compiled just-in-time to any shader language target"
+edition.workspace = true
+keywords = ["deep-learning", "machine-learning", "gpu"]
+license.workspace = true
+name = "burn-jit"
+readme.workspace = true
+repository = "https://github.com/tracel-ai/burn/tree/main/burn-jit"
+version.workspace = true
+
+[features]
+default = ["autotune", "std", "burn-compute/default", "fusion"]
+std = []
+doc = ["default"]
+autotune = []
+fusion = ["burn-fusion"]
+export_tests = [
+  "burn-tensor-testgen",
+  "serial_test",
+  "burn-autodiff/export_tests",
+  "burn-tensor/export_tests",
+  "burn-ndarray",
+  "fusion",
+]
+
+[dependencies]
+burn-common = { path = "../burn-common", version = "0.13.0" }
+burn-tensor = { path = "../burn-tensor", version = "0.13.0" }
+burn-fusion = { path = "../burn-fusion", version = "0.13.0", optional = true }
+
+bytemuck = { workspace = true }
+derive-new = { workspace = true }
+log = { workspace = true }
+num-traits = { workspace = true }
+rand = { workspace = true }
+spin = { workspace = true }
+
+# Template
+serde = { workspace = true }
+text_placeholder = { workspace = true, features = ["struct_context"] }
+
+hashbrown = { workspace = true }
+burn-compute = { path = "../burn-compute", version = "0.13.0", default-features = false, features = [
+  "channel-mutex",
+  "std",
+] }
+burn-tensor-testgen = { path = "../burn-tensor-testgen", version = "0.13.0", optional = true }
+
+# When exporting tests
+serial_test = { workspace = true, optional = true }
+burn-autodiff = { path = "../burn-autodiff", version = "0.13.0", default-features = false, optional = true }
+burn-ndarray = { path = "../burn-ndarray", version = "0.13.0", optional = true }
+
+[package.metadata.docs.rs]
+features = ["doc"]

crates/burn-jit/Cargo.toml

@@ -0,0 +1,4 @@
+# Burn JIT Backend
+
+Generic backend that can be compiled just-in-time (JIT) to any shader language target
+In progress: At the moment, only WGSL compilation is supported, and some kernels still rely on pure WGSL

crates/burn-wgpu/LICENSE-APACHE crates/burn-jit/LICENSE-APACHE

@@ -5,7 +5,7 @@ use std::{marker::PhantomData, sync::Mutex};
 
 pub(crate) static SEED: Mutex<Option<StdRng>> = Mutex::new(None);
 
-/// Tensor backend that uses the [wgpu] crate for executing GPU compute shaders.
+/// Generic tensor backend that can be compiled just-in-time to any shader runtime
 #[derive(new)]
 pub struct JitBackend<R: Runtime> {
     _runtime: PhantomData<R>,

crates/burn-wgpu/LICENSE-MIT crates/burn-jit/LICENSE-MIT

@@ -2,16 +2,24 @@ use super::dialect::gpu;
 use crate::{FloatElement, IntElement};
 use std::fmt::Display;
 
+/// Compiles the [gpu representation](gpu::ComputeShader) into its own representation that can be
+/// formatted into tokens.
 pub trait Compiler: Sync + Send + 'static + Clone + Default + core::fmt::Debug {
+    /// The representation for the compiled code.
     type Representation: Display;
+    /// The float element type used for compilation.
     type Float: FloatElement;
+    /// The int element type used for compilation.
     type Int: IntElement;
+    /// The compiler that can be used to generate full precision shaders.
     type FullPrecisionCompiler: Compiler<
         Representation = Self::Representation,
         Float = f32,
         Int = i32,
     >;
 
+    /// Compiles the [gpu shader](gpu::ComputeShader) into the compiler's representation.
     fn compile(shader: gpu::ComputeShader) -> Self::Representation;
+    /// The size of the given element in bytes.
     fn elem_size(elem: gpu::Elem) -> usize;
 }

crates/burn-jit/README.md

@@ -4,34 +4,37 @@ use serde::{Deserialize, Serialize};
 /// All branching types.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub enum Branch {
-    // An if statement.
+    /// An if statement.
     If(If),
-    // An if else statement.
+    /// An if else statement.
     IfElse(IfElse),
-    // A range loop.
+    /// A range loop.
     RangeLoop(RangeLoop),
-    // A loop.
+    /// A loop.
     Loop(Loop),
-    // A return statement.
+    /// A return statement.
     Return,
-    // A break statement.
+    /// A break statement.
     Break,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct If {
     pub cond: Variable,
     pub scope: Scope,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct IfElse {
     pub cond: Variable,
     pub scope_if: Scope,
     pub scope_else: Scope,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct RangeLoop {
     pub i: Variable,
     pub start: Variable,
@@ -40,6 +43,7 @@ pub struct RangeLoop {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct Loop {
     pub scope: Scope,
 }

crates/burn-wgpu/src/backend.rs crates/burn-jit/src/backend.rs

@@ -0,0 +1,21 @@
+mod branch;
+mod macros;
+mod operation;
+mod procedure;
+mod processing;
+mod scope;
+mod shader;
+mod synchronization;
+mod variable;
+mod vectorization;
+
+pub use branch::*;
+pub use operation::*;
+pub use procedure::*;
+pub use scope::*;
+pub use shader::*;
+pub use synchronization::*;
+pub use variable::*;
+pub use vectorization::*;
+
+pub(crate) use macros::gpu;

crates/burn-wgpu/src/codegen/compilation.rs crates/burn-jit/src/codegen/compilation.rs

@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
 ///
 /// [Procedure] expansions can safely use all operation variants.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-#[allow(dead_code)] // Some variants might not be used with different flags
+#[allow(dead_code, missing_docs)] // Some variants might not be used with different flags
 pub enum Operation {
     Operator(Operator),
     Procedure(Procedure),
@@ -21,7 +21,7 @@ pub enum Operation {
 
 /// All operators that can be used in a GPU compute shader.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-#[allow(dead_code)] // Some variants might not be used with different flags
+#[allow(dead_code, missing_docs)] // Some variants might not be used with different flags
 pub enum Operator {
     Add(BinaryOperator),
     Sub(BinaryOperator),
@@ -57,6 +57,7 @@ pub enum Operator {
 
 /// All metadata that can be access in a shader.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub enum Metadata {
     /// The stride of an array at the given dimension.
     Stride {
@@ -77,19 +78,22 @@ pub enum Metadata {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct BinaryOperator {
     pub lhs: Variable,
     pub rhs: Variable,
     pub out: Variable,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct UnaryOperator {
     pub input: Variable,
     pub out: Variable,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct ClampOperator {
     pub input: Variable,
     pub min_value: Variable,
@@ -98,11 +102,13 @@ pub struct ClampOperator {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct ReadGlobalOperator {
     pub variable: Variable,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct ReadGlobalWithLayoutOperator {
     pub variable: Variable,
     pub tensor_read_pos: usize,

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

@@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize};
 
 /// Assign value to a variable based on a given condition.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct ConditionalAssign {
     pub cond: Variable,
     pub lhs: Variable,
@@ -11,6 +12,7 @@ pub struct ConditionalAssign {
 }
 
 impl ConditionalAssign {
+    #[allow(missing_docs)]
     pub fn expand(self, scope: &mut Scope) {
         let cond = self.cond;
         let lhs = self.lhs;

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

@@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
 /// Tensor operations that can't be executed with a simple [operator](super::super::Operator) should use a
 /// procedure.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub enum Procedure {
     ReadGlobalWithLayout(ReadGlobalWithLayout),
     IndexOffsetGlobalWithLayout(IndexOffsetGlobalWithLayout),
@@ -16,7 +17,7 @@ pub enum Procedure {
 }
 
 impl Procedure {
-    pub fn vectorize(&self, vectorization: Vectorization) -> Self {
+    pub(crate) fn vectorize(&self, vectorization: Vectorization) -> Self {
         match self {
             Procedure::ReadGlobalWithLayout(op) => {
                 Procedure::ReadGlobalWithLayout(op.vectorize(vectorization))

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

@@ -0,0 +1,9 @@
+mod assign;
+mod base;
+mod read;
+mod write;
+
+pub use assign::*;
+pub use base::*;
+pub use read::*;
+pub use write::*;

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

@@ -23,6 +23,7 @@ pub struct ReadGlobalWithLayout {
 }
 
 impl ReadGlobal {
+    #[allow(missing_docs)]
     pub fn expand(self, scope: &mut Scope) {
         scope.register(Operator::Index(BinaryOperator {
             lhs: self.global,
@@ -61,6 +62,7 @@ impl ReadGlobalWithLayout {
         })
     }
 
+    #[allow(missing_docs)]
     pub fn expand(self, scope: &mut Scope) {
         let outputs = self.outs;
         let tensors = self.globals;
@@ -107,6 +109,7 @@ impl ReadGlobalWithLayout {
 
 /// Calculate the index offset for all tensor variables provided compatible with the given layout.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct IndexOffsetGlobalWithLayout {
     /// Tensor [variables](Variable), same length as [indexes](Self::indexes).
     pub tensors: Vec<Variable>,
@@ -123,6 +126,7 @@ pub struct IndexOffsetGlobalWithLayout {
 }
 
 impl IndexOffsetGlobalWithLayout {
+    #[allow(missing_docs)]
     pub fn expand(self, scope: &mut Scope) {
         let layout = self.layout;
         let index_item_ty = Item::Scalar(Elem::UInt);

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

@@ -3,12 +3,14 @@ use serde::{Deserialize, Serialize};
 
 /// Write to a global array.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[allow(missing_docs)]
 pub struct WriteGlobal {
     pub input: Variable,
     pub global: Variable,
 }
 
 impl WriteGlobal {
+    #[allow(missing_docs)]
     pub fn expand(self, scope: &mut Scope) {
         let output = self.global;
         let input = self.input;