codegen: better support for half precision floats

crates/cranexpr-cli/src/main.rs

@@ -29,6 +29,7 @@ struct Args {
 fn parse_component_type(s: &str) -> Option<ComponentType> {
   match s {
     "f32" => Some(ComponentType::F32),
+    "f16" => Some(ComponentType::F16),
     "u8" => Some(ComponentType::U8),
     "u16" => Some(ComponentType::U16),
     _ => None,

crates/cranexpr-codegen/src/compiler.rs

@@ -1421,6 +1421,140 @@ mod tests {
     assert_eq!(actual[0], 2);
   }
 
+  /// Runs `expr` with an `F16` source plane (given as raw bit patterns) into
+  /// an `F32` destination plane.
+  fn run_f16_src_to_f32(
+    expr: &str,
+    halfbits: &[u16],
+    width: usize,
+    height: usize,
+    boundary: Option<BoundaryMode>,
+  ) -> Vec<f32> {
+    let ast = cranexpr_parser::parse_expr(expr).unwrap();
+    let compiled = compile_jit(
+      &ast,
+      ComponentType::F32,
+      &[ComponentType::F16],
+      boundary,
+      &[],
+    )
+    .expect("should compile expr");
+
+    let (mut src_buf, src_stride) = alloc_aligned_plane(width * size_of::<u16>(), height);
+    fill_plane(src_buf.as_mut_slice(), src_stride, halfbits, width, height);
+    let (mut dst_buf, dst_stride) = alloc_aligned_plane(width * size_of::<f32>(), height);
+
+    let src_ptrs = [src_buf.as_ptr()];
+    let src_strides = [src_stride as i64];
+    unsafe {
+      compiled.invoke(
+        dst_buf.as_mut_slice(),
+        dst_stride as i64,
+        &src_ptrs,
+        &src_strides,
+        width as i32,
+        height as i32,
+        0,
+        &[],
+      );
+    }
+    unpack_plane::<f32>(dst_buf.as_slice(), dst_stride, width, height)
+  }
+
+  /// Runs `expr` with an `F32` source plane into an `F16` destination plane,
+  /// returning the raw `f16` bit patterns.
+  fn run_f32_src_to_f16(expr: &str, inputs: &[f32], width: usize, height: usize) -> Vec<u16> {
+    let ast = cranexpr_parser::parse_expr(expr).unwrap();
+    let compiled = compile_jit(&ast, ComponentType::F16, &[ComponentType::F32], None, &[])
+      .expect("should compile expr");
+
+    let (mut src_buf, src_stride) = alloc_aligned_plane(width * size_of::<f32>(), height);
+    fill_plane(src_buf.as_mut_slice(), src_stride, inputs, width, height);
+    let (mut dst_buf, dst_stride) = alloc_aligned_plane(width * size_of::<u16>(), height);
+
+    let src_ptrs = [src_buf.as_ptr()];
+    let src_strides = [src_stride as i64];
+    unsafe {
+      compiled.invoke(
+        dst_buf.as_mut_slice(),
+        dst_stride as i64,
+        &src_ptrs,
+        &src_strides,
+        width as i32,
+        height as i32,
+        0,
+        &[],
+      );
+    }
+    unpack_plane::<u16>(dst_buf.as_slice(), dst_stride, width, height)
+  }
+
+  #[rstest]
+  fn test_f16_load_to_f32() {
+    // Zero, +1, +0.5, +2, -1, +10, ~1/3, smallest subnormal.
+    let halfbits = [
+      0x0000, 0x3c00, 0x3800, 0x4000, 0xbc00, 0x4900, 0x3555, 0x0001,
+    ];
+    let expected: [f32; 8] = [0.0, 1.0, 0.5, 2.0, -1.0, 10.0, 0.333_251_95, 5.960_464_5e-8];
+    let out = run_f16_src_to_f32("x", &halfbits, halfbits.len(), 1, None);
+    for (i, (&a, &e)) in out.iter().zip(expected.iter()).enumerate() {
+      assert_eq!(a.to_bits(), e.to_bits(), "lane {i}: got {a}, want {e}");
+    }
+  }
+
+  #[rstest]
+  fn test_f16_store_from_f32() {
+    let inputs = [0.0, 1.0, 0.5, 2.0, -1.0, 10.0, 0.333_333_3, 1.337];
+    // Round-to-nearest-even f16 bit patterns of the inputs.
+    let expected = [
+      0x0000, 0x3c00, 0x3800, 0x4000, 0xbc00, 0x4900, 0x3555, 0x3d59,
+    ];
+    let out = run_f32_src_to_f16("x", &inputs, inputs.len(), 1);
+    assert_eq!(out, expected);
+  }
+
+  #[rstest]
+  fn test_f16_roundtrip_arith() {
+    // load -> multiply -> store, all through f16.
+    let ast = cranexpr_parser::parse_expr("x 2 *").unwrap();
+    let compiled = compile_jit(&ast, ComponentType::F16, &[ComponentType::F16], None, &[])
+      .expect("should compile expr");
+
+    // 1, 0.5, -3, 7 as f16 bit patterns.
+    let halfbits = [0x3c00u16, 0x3800, 0xc200, 0x4700];
+    let width = halfbits.len();
+    let (mut src_buf, src_stride) = alloc_aligned_plane(width * size_of::<u16>(), 1);
+    fill_plane(src_buf.as_mut_slice(), src_stride, &halfbits, width, 1);
+    let (mut dst_buf, dst_stride) = alloc_aligned_plane(width * size_of::<u16>(), 1);
+
+    let src_ptrs = [src_buf.as_ptr()];
+    let src_strides = [src_stride as i64];
+    unsafe {
+      compiled.invoke(
+        dst_buf.as_mut_slice(),
+        dst_stride as i64,
+        &src_ptrs,
+        &src_strides,
+        width as i32,
+        1,
+        0,
+        &[],
+      );
+    }
+    let out = unpack_plane::<u16>(dst_buf.as_slice(), dst_stride, width, 1);
+    // 2, 1, -6, 14 as f16 bit patterns.
+    assert_eq!(out, [0x4000, 0x3c00, 0xc600, 0x4b00]);
+  }
+
+  #[rstest]
+  fn test_f16_scalar_load_via_mirror() {
+    let halfbits = [0x4900u16; 8]; // all 10.0
+    let out = run_f16_src_to_f32("x[-2,0]", &halfbits, 8, 1, Some(BoundaryMode::Mirror));
+    for (i, &v) in out.iter().enumerate() {
+      assert_eq!(v.to_bits(), 10.0_f32.to_bits(), "lane {i}");
+    }
+  }
+
   #[rstest]
   #[case("1 1 and", 1.0)]
   #[case("1 0 and", 0.0)]

crates/cranexpr-codegen/src/component_type.rs

@@ -4,6 +4,7 @@ use cranelift::prelude::*;
 pub enum ComponentType {
   U8,
   U16,
+  F16,
   F32,
 }
 
@@ -12,7 +13,7 @@ impl ComponentType {
   pub const fn bytes(self) -> usize {
     match self {
       Self::U8 => 1,
-      Self::U16 => 2,
+      Self::U16 | Self::F16 => 2,
       Self::F32 => 4,
     }
   }
@@ -22,7 +23,7 @@ impl ComponentType {
     match self {
       Self::U8 => 255.0,
       Self::U16 => 65535.0,
-      Self::F32 => 1.0,
+      Self::F16 | Self::F32 => 1.0,
     }
   }
 }
@@ -32,6 +33,7 @@ impl From<ComponentType> for types::Type {
     match value {
       ComponentType::U8 => types::I8,
       ComponentType::U16 => types::I16,
+      ComponentType::F16 => types::F16,
       ComponentType::F32 => types::F32,
     }
   }

crates/cranexpr-codegen/src/translate_simd.rs

@@ -27,6 +27,11 @@ pub(crate) fn load_pixel_vec_f32x4(
 ) -> Value {
   let ptr = src_ptr.offset_value(fx, offset);
   match src_type {
+    ComponentType::F16 => {
+      let addr = ptr.get_addr(fx);
+      let halfbits = fx.bcx.ins().uload16x4(SRC_MEMFLAGS, addr, 0);
+      half_to_float_simd(fx, halfbits)
+    }
     ComponentType::F32 => {
       let flags = if aligned {
         SRC_MEMFLAGS.with_aligned()
@@ -60,6 +65,15 @@ pub(crate) fn store_pixel_vec_f32x4(
 ) {
   let ptr = dest_ptr.offset_value(fx, offset);
   match dst_type {
+    ComponentType::F16 => {
+      let halfbits = float_to_half_simd(fx, value);
+      // Narrow the 16-bit half patterns down to I16X8.
+      let narrowed = fx.bcx.ins().unarrow(halfbits, halfbits);
+      let byte_flags = MemFlagsData::new().with_endianness(Endianness::Little);
+      let as_i64x2 = fx.bcx.ins().bitcast(types::I64X2, byte_flags, narrowed);
+      let low = fx.bcx.ins().extractlane(as_i64x2, 0);
+      ptr.store(fx, low, MemFlagsData::trusted());
+    }
     ComponentType::F32 => {
       // VapourSynth guarantees row strides are aligned to at least 32 bytes.
       ptr.store(fx, value, MemFlagsData::trusted().with_aligned());
@@ -92,7 +106,7 @@ pub(crate) fn store_pixel_vec_f32x4(
           let low = fx.bcx.ins().extractlane(as_i32x4, 0);
           ptr.store(fx, low, MemFlagsData::trusted());
         }
-        ComponentType::F32 => unreachable!(),
+        ComponentType::F16 | ComponentType::F32 => unreachable!(),
       }
     }
   }
@@ -304,6 +318,124 @@ fn splat_f32(fx: &mut FunctionCx<'_, '_>, val: f32) -> Value {
   fx.bcx.ins().splat(VEC_TYPE, scalar)
 }
 
+/// Splats an `i32` constant across an `I32X4` vector.
+fn splat_i32(fx: &mut FunctionCx<'_, '_>, val: i64) -> Value {
+  let scalar = fx.bcx.ins().iconst(types::I32, val);
+  fx.bcx.ins().splat(types::I32X4, scalar)
+}
+
+/// Converts an `I32X4` of zero-extended IEEE-754 binary16 bit patterns into an
+/// `F32X4`.
+///
+/// Based on <https://gist.github.com/rygorous/2156668>
+fn half_to_float_simd(fx: &mut FunctionCx<'_, '_>, h: Value) -> Value {
+  let no_flags = MemFlagsData::new();
+
+  let mant_exp_mask = splat_i32(fx, 0x7fff);
+  let mant_exp = fx.bcx.ins().band(h, mant_exp_mask);
+  let shifted = fx.bcx.ins().ishl_imm(mant_exp, 13);
+  let as_f = fx.bcx.ins().bitcast(VEC_TYPE, no_flags, shifted);
+
+  let magic = splat_f32(fx, f32::from_bits(0x7780_0000));
+  let scaled = fx.bcx.ins().fmul(as_f, magic);
+
+  let was_infnan = splat_f32(fx, f32::from_bits(0x4780_0000));
+  let infnan_mask = fx
+    .bcx
+    .ins()
+    .fcmp(FloatCC::GreaterThanOrEqual, scaled, was_infnan);
+  let scaled_bits = fx.bcx.ins().bitcast(types::I32X4, no_flags, scaled);
+  let inf_exp = splat_i32(fx, 255 << 23);
+  let with_inf = fx.bcx.ins().bor(scaled_bits, inf_exp);
+  let merged = fx.bcx.ins().bitselect(infnan_mask, with_inf, scaled_bits);
+
+  let sign_mask = splat_i32(fx, 0x8000);
+  let sign = fx.bcx.ins().band(h, sign_mask);
+  let sign = fx.bcx.ins().ishl_imm(sign, 16);
+  let result = fx.bcx.ins().bor(merged, sign);
+  fx.bcx.ins().bitcast(VEC_TYPE, no_flags, result)
+}
+
+/// Converts an `F32X4` into an `I32X4` of IEEE-754 binary16 bit patterns, held
+/// in the low 16 bits of each lane.
+///
+/// Based on <https://gist.github.com/rygorous/2156668>
+fn float_to_half_simd(fx: &mut FunctionCx<'_, '_>, f: Value) -> Value {
+  let no_flags = MemFlagsData::new();
+  let f_bits = fx.bcx.ins().bitcast(types::I32X4, no_flags, f);
+
+  let sign_mask = splat_i32(fx, i64::from(0x8000_0000_u32 as i32));
+  let sign = fx.bcx.ins().band(f_bits, sign_mask);
+  let abs_mask = splat_i32(fx, 0x7fff_ffff);
+  let abs = fx.bcx.ins().band(f_bits, abs_mask);
+
+  // Inf/NaN: NaN -> 0x7e00, Inf -> 0x7c00.
+  let f16max = splat_i32(fx, 0x4780_0000);
+  let is_infnan = fx
+    .bcx
+    .ins()
+    .icmp(IntCC::UnsignedGreaterThanOrEqual, abs, f16max);
+  let f32infty = splat_i32(fx, 0x7f80_0000);
+  let is_nan = fx.bcx.ins().icmp(IntCC::UnsignedGreaterThan, abs, f32infty);
+  let qnan = splat_i32(fx, 0x7e00);
+  let inf = splat_i32(fx, 0x7c00);
+  let infnan_o = fx.bcx.ins().bitselect(is_nan, qnan, inf);
+
+  // Subnormal/zero: align the mantissa with a magic add, then subtract bias.
+  let sub_thresh = splat_i32(fx, 113 << 23);
+  let is_subnormal = fx.bcx.ins().icmp(IntCC::UnsignedLessThan, abs, sub_thresh);
+  let abs_f = fx.bcx.ins().bitcast(VEC_TYPE, no_flags, abs);
+  let denorm_magic = splat_f32(fx, f32::from_bits(0x3f00_0000));
+  let sub_f = fx.bcx.ins().fadd(abs_f, denorm_magic);
+  let sub_bits = fx.bcx.ins().bitcast(types::I32X4, no_flags, sub_f);
+  let denorm_magic_bits = splat_i32(fx, 0x3f00_0000);
+  let sub_o = fx.bcx.ins().isub(sub_bits, denorm_magic_bits);
+
+  // Normalized: bias the exponent and round to nearest even.
+  let shifted13 = fx.bcx.ins().ushr_imm(abs, 13);
+  let one = splat_i32(fx, 1);
+  let mant_odd = fx.bcx.ins().band(shifted13, one);
+  let bias = splat_i32(fx, i64::from(0xc800_0fff_u32 as i32));
+  let norm = fx.bcx.ins().iadd(abs, bias);
+  let norm = fx.bcx.ins().iadd(norm, mant_odd);
+  let norm_o = fx.bcx.ins().ushr_imm(norm, 13);
+
+  let o = fx.bcx.ins().bitselect(is_subnormal, sub_o, norm_o);
+  let o = fx.bcx.ins().bitselect(is_infnan, infnan_o, o);
+  let sign_h = fx.bcx.ins().ushr_imm(sign, 16);
+  fx.bcx.ins().bor(o, sign_h)
+}
+
+/// Scalar counterpart of [`half_to_float_simd`]: converts a zero-extended
+/// `I32` binary16 bit pattern into an `F32`.
+fn half_to_float_scalar(fx: &mut FunctionCx<'_, '_>, h: Value) -> Value {
+  let no_flags = MemFlagsData::new();
+
+  let mant_exp_mask = fx.bcx.ins().iconst(types::I32, 0x7fff);
+  let mant_exp = fx.bcx.ins().band(h, mant_exp_mask);
+  let shifted = fx.bcx.ins().ishl_imm(mant_exp, 13);
+  let as_f = fx.bcx.ins().bitcast(types::F32, no_flags, shifted);
+
+  let magic = fx.bcx.ins().f32const(f32::from_bits(0x7780_0000));
+  let scaled = fx.bcx.ins().fmul(as_f, magic);
+
+  let was_infnan = fx.bcx.ins().f32const(f32::from_bits(0x4780_0000));
+  let infnan = fx
+    .bcx
+    .ins()
+    .fcmp(FloatCC::GreaterThanOrEqual, scaled, was_infnan);
+  let scaled_bits = fx.bcx.ins().bitcast(types::I32, no_flags, scaled);
+  let inf_exp = fx.bcx.ins().iconst(types::I32, 255 << 23);
+  let with_inf = fx.bcx.ins().bor(scaled_bits, inf_exp);
+  let merged = fx.bcx.ins().select(infnan, with_inf, scaled_bits);
+
+  let sign_mask = fx.bcx.ins().iconst(types::I32, 0x8000);
+  let sign = fx.bcx.ins().band(h, sign_mask);
+  let sign = fx.bcx.ins().ishl_imm(sign, 16);
+  let result = fx.bcx.ins().bor(merged, sign);
+  fx.bcx.ins().bitcast(types::F32, no_flags, result)
+}
+
 /// Per-lane x offsets used to synthesize the SIMD X coordinate vector.
 ///
 /// Recall that in the scalar implementation, each pixel has one X coordinate.
@@ -1123,12 +1255,19 @@ fn load_scalar_as_f32(
   offset: Value,
   src_type: ComponentType,
 ) -> Value {
-  let val = src_ptr
-    .offset_value(fx, offset)
-    .load(fx, src_type.into(), SRC_MEMFLAGS);
+  let ptr = src_ptr.offset_value(fx, offset);
   match src_type {
-    ComponentType::U8 | ComponentType::U16 => fx.bcx.ins().fcvt_from_uint(types::F32, val),
-    ComponentType::F32 => val,
+    ComponentType::F16 => {
+      // Load the raw f16 bits as an integer and convert via bit twiddling.
+      let halfbits = ptr.load(fx, types::I16, SRC_MEMFLAGS);
+      let halfbits = fx.bcx.ins().uextend(types::I32, halfbits);
+      half_to_float_scalar(fx, halfbits)
+    }
+    ComponentType::U8 | ComponentType::U16 => {
+      let val = ptr.load(fx, src_type.into(), SRC_MEMFLAGS);
+      fx.bcx.ins().fcvt_from_uint(types::F32, val)
+    }
+    ComponentType::F32 => ptr.load(fx, types::F32, SRC_MEMFLAGS),
   }
 }
 

src/component_type.rs

@@ -16,7 +16,11 @@ impl FromVideoFormat for ComponentType {
         2 => Self::U16,
         _ => unreachable!(),
       },
-      SampleType::Float => Self::F32,
+      SampleType::Float => match format.video_format().bytes_per_sample {
+        2 => Self::F16,
+        4 => Self::F32,
+        _ => unreachable!(),
+      },
     }
   }
 }