Add f16 support in the wgpu backend

crates/burn-jit/Cargo.toml

@@ -33,6 +33,7 @@ burn-fusion = { path = "../burn-fusion", version = "0.13.0", optional = true }
 
 bytemuck = { workspace = true }
 derive-new = { workspace = true }
+half = { workspace = true, features = ["bytemuck"] }
 log = { workspace = true }
 num-traits = { workspace = true }
 rand = { workspace = true }

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

@@ -21,6 +21,7 @@ pub enum Visibility {
 #[allow(missing_docs)]
 pub enum Elem {
     Float,
+    Half,
     Int,
     UInt,
     Bool,
@@ -36,6 +37,7 @@ impl Display for Elem {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Float => f.write_str("float"),
+            Self::Half => f.write_str("half"),
             Self::Int => f.write_str("int"),
             Self::UInt => f.write_str("uint"),
             Self::Bool => f.write_str("bool"),

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

@@ -316,6 +316,7 @@ fn create_scalar_handles<R: Runtime, E1: JitElement, E2: JitElement, E3: JitElem
         Elem::Int => 1,
         Elem::UInt => 2,
         Elem::Bool => panic!("Bool scalars are not supported"),
+        Elem::Half => panic!("Half scalars are not supported"),
     };
     let scalar_priorities: [usize; 3] = [
         element_priority(E1::gpu_elem()),

crates/burn-jit/src/element.rs

@@ -1,5 +1,6 @@
 use crate::codegen::dialect::gpu;
 use burn_tensor::Element;
+use half::f16;
 
 /// The base element trait for the jit backend.
 pub trait JitElement:
@@ -90,5 +91,27 @@ impl JitElement for f32 {
     }
 }
 
+impl JitElement for f16 {
+    fn type_name() -> &'static str {
+        "f16"
+    }
+    fn as_bytes(slice: &[Self]) -> &[u8] {
+        bytemuck::cast_slice(slice)
+    }
+    fn from_bytes(bytes: &[u8]) -> &[Self] {
+        bytemuck::cast_slice(bytes)
+    }
+    fn gpu_elem() -> gpu::Elem {
+        gpu::Elem::Half
+    }
+    fn maximum_value() -> Self {
+        f16::MAX
+    }
+    fn minimum_value() -> Self {
+        f16::MIN
+    }
+}
+
 impl FloatElement for f32 {}
 impl IntElement for i32 {}
+impl FloatElement for f16 {}

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

@@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
 #[derive(Default, Clone, Serialize, Deserialize)]
 pub struct Scalars {
     pub(crate) num_float: usize,
+    pub(crate) num_half: usize,
     pub(crate) num_int: usize,
     pub(crate) num_uint: usize,
     pub(crate) num_bool: usize,

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

@@ -95,6 +95,13 @@ impl TraceBuilder {
                 self.scalars.num_float += 1;
                 var
             }
+            gpu::Elem::Half => {
+                let var = self
+                    .scope
+                    .read_scalar(self.scalars.num_half as u16, elem_type);
+                self.scalars.num_half += 1;
+                var
+            }
             gpu::Elem::Int => {
                 let var = self
                     .scope

crates/burn-wgpu/Cargo.toml

@@ -26,6 +26,7 @@ burn-common = { path = "../burn-common", version = "0.13.0" }
 burn-fusion = { path = "../burn-fusion", version = "0.13.0", optional = true }
 
 bytemuck = { workspace = true }
+half = { workspace = true, features = ["bytemuck"] }
 wgpu = { workspace = true, features = ["fragile-send-sync-non-atomic-wasm"] }
 pollster = { workspace = true }
 

crates/burn-wgpu/src/compiler/wgsl/base.rs

@@ -42,6 +42,7 @@ pub enum Variable {
 #[derive(Debug, Clone, PartialEq, Eq, Copy)]
 pub enum Elem {
     F32,
+    F16,
     I32,
     U32,
     Bool,
@@ -165,6 +166,7 @@ impl Elem {
     pub fn size(&self) -> usize {
         match self {
             Self::F32 => core::mem::size_of::<f32>(),
+            Self::F16 => core::mem::size_of::<half::f16>(),
             Self::I32 => core::mem::size_of::<i32>(),
             Self::U32 => core::mem::size_of::<u32>(),
             Self::Bool => core::mem::size_of::<bool>(),
@@ -176,6 +178,7 @@ impl Display for Elem {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::F32 => f.write_str("f32"),
+            Self::F16 => f.write_str("f16"),
             Self::I32 => f.write_str("i32"),
             Self::U32 => f.write_str("u32"),
             Self::Bool => f.write_str("bool"),
@@ -218,6 +221,7 @@ impl Display for Variable {
             }
             Variable::ConstantScalar(number, elem) => match elem {
                 Elem::F32 => f.write_fmt(format_args!("{number}f")),
+                Elem::F16 => f.write_fmt(format_args!("{number}f")),
                 Elem::I32 => f.write_fmt(format_args!("{number}i")),
                 Elem::U32 => f.write_fmt(format_args!("{number}u")),
                 Elem::Bool => f.write_fmt(format_args!("bool({number})")),

crates/burn-wgpu/src/compiler/wgsl/compiler.rs

@@ -78,6 +78,11 @@ impl<F: FloatElement, I: IntElement> WgslCompiler<F, I> {
         self.num_inputs = value.inputs.len();
         self.num_outputs = value.outputs.len();
 
+        let features = match F::gpu_elem() {
+            gpu::Elem::Half => vec![wgsl::Feature::ShaderF16],
+            _ => vec![],
+        };
+
         let instructions = self.compile_scope(&mut value.body);
         let extensions = register_extensions(&instructions);
         let body = wgsl::Body {
@@ -114,6 +119,7 @@ impl<F: FloatElement, I: IntElement> WgslCompiler<F, I> {
             workgroup_id: self.workgroup_id,
             body,
             extensions,
+            features,
         }
     }
 
@@ -129,6 +135,7 @@ impl<F: FloatElement, I: IntElement> WgslCompiler<F, I> {
     fn compile_elem(value: gpu::Elem) -> wgsl::Elem {
         match value {
             gpu::Elem::Float => F::wgpu_elem(),
+            gpu::Elem::Half => F::wgpu_elem(),
             gpu::Elem::Int => I::wgpu_elem(),
             gpu::Elem::UInt => wgsl::Elem::U32,
             gpu::Elem::Bool => wgsl::Elem::Bool,

crates/burn-wgpu/src/compiler/wgsl/shader.rs

@@ -60,6 +60,11 @@ impl LocalArray {
     }
 }
 
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum Feature {
+    ShaderF16,
+}
+
 #[derive(Debug, Clone)]
 pub struct ComputeShader {
     pub inputs: Vec<Binding>,
@@ -75,10 +80,15 @@ pub struct ComputeShader {
     pub workgroup_id: bool,
     pub body: Body,
     pub extensions: Vec<Extension>,
+    pub features: Vec<Feature>,
 }
 
 impl Display for ComputeShader {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for feature in self.features.iter() {
+            f.write_fmt(format_args!("{feature};\n"))?;
+        }
+
         Self::format_bindings(f, "input", &self.inputs, 0)?;
         Self::format_bindings(f, "output", &self.outputs, self.inputs.len())?;
 
@@ -218,3 +228,11 @@ impl Display for Visibility {
         }
     }
 }
+
+impl Display for Feature {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Feature::ShaderF16 => f.write_str("enable f16"),
+        }
+    }
+}

crates/burn-wgpu/src/element.rs

@@ -1,6 +1,7 @@
 use burn_jit::JitElement;
 
 use crate::compiler::wgsl;
+use half::f16;
 
 /// The base element trait for the wgpu backend.
 pub trait WgpuElement: JitElement {
@@ -31,5 +32,12 @@ impl WgpuElement for f32 {
     }
 }
 
+impl WgpuElement for f16 {
+    fn wgpu_elem() -> wgsl::Elem {
+        wgsl::Elem::F16
+    }
+}
+
 impl FloatElement for f32 {}
 impl IntElement for i32 {}
+impl FloatElement for f16 {}