From 7c3fe6674751fdbc8eb1344956e9af9523b0bf83 Mon Sep 17 00:00:00 2001 From: sgt0 <140186177+sgt0@users.noreply.github.com> Date: Wed, 1 Jul 2026 02:18:09 +0000 Subject: [PATCH] codegen: better support for half precision floats I didn't see any native f16<->f32 conversion functions in Cranelift, so the vectorized paths decode/encode the bit patterns inline based on an implementation described here: . --- crates/cranexpr-cli/src/main.rs | 1 + crates/cranexpr-codegen/src/compiler.rs | 134 ++++++++++++++++ crates/cranexpr-codegen/src/component_type.rs | 6 +- crates/cranexpr-codegen/src/translate_simd.rs | 151 +++++++++++++++++- src/component_type.rs | 6 +- src/expr.rs | 37 +++-- 6 files changed, 315 insertions(+), 20 deletions(-) diff --git a/crates/cranexpr-cli/src/main.rs b/crates/cranexpr-cli/src/main.rs index 0a372c5..c1da24e 100644 --- a/crates/cranexpr-cli/src/main.rs +++ b/crates/cranexpr-cli/src/main.rs @@ -29,6 +29,7 @@ struct Args { fn parse_component_type(s: &str) -> Option { match s { "f32" => Some(ComponentType::F32), + "f16" => Some(ComponentType::F16), "u8" => Some(ComponentType::U8), "u16" => Some(ComponentType::U16), _ => None, diff --git a/crates/cranexpr-codegen/src/compiler.rs b/crates/cranexpr-codegen/src/compiler.rs index 57d7849..38608ad 100644 --- a/crates/cranexpr-codegen/src/compiler.rs +++ b/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, + ) -> Vec { + 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::(), 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::(), 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::(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 { + 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::(), 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::(), 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::(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::(), 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::(), 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::(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)] diff --git a/crates/cranexpr-codegen/src/component_type.rs b/crates/cranexpr-codegen/src/component_type.rs index da009b7..73fa733 100644 --- a/crates/cranexpr-codegen/src/component_type.rs +++ b/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 for types::Type { match value { ComponentType::U8 => types::I8, ComponentType::U16 => types::I16, + ComponentType::F16 => types::F16, ComponentType::F32 => types::F32, } } diff --git a/crates/cranexpr-codegen/src/translate_simd.rs b/crates/cranexpr-codegen/src/translate_simd.rs index 1f802ee..5b9f6d6 100644 --- a/crates/cranexpr-codegen/src/translate_simd.rs +++ b/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 +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 +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), } } diff --git a/src/component_type.rs b/src/component_type.rs index 254583c..8f5f38c 100644 --- a/src/component_type.rs +++ b/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!(), + }, } } } diff --git a/src/expr.rs b/src/expr.rs index 876d407..ac0037f 100644 --- a/src/expr.rs +++ b/src/expr.rs @@ -318,17 +318,32 @@ impl Filter for CranexprFilter { }, _ => unreachable!(), }, - SampleType::Float => unsafe { - expr.invoke( - dst.as_mut_slice::(plane_idx as i32), - dst_stride, - &src_ptrs, - &src_strides, - width, - height, - n, - &frame_props, - ); + SampleType::Float => match self.vi.format.bytes_per_sample { + 2 => unsafe { + expr.invoke( + dst.as_mut_slice::(plane_idx as i32), + dst_stride, + &src_ptrs, + &src_strides, + width, + height, + n, + &frame_props, + ); + }, + 4 => unsafe { + expr.invoke( + dst.as_mut_slice::(plane_idx as i32), + dst_stride, + &src_ptrs, + &src_strides, + width, + height, + n, + &frame_props, + ); + }, + _ => unreachable!(), }, } }