boostr 0.1.0

ML framework built on numr - attention, quantization, model architectures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

/// Common WGSL helper functions for quantized operations.
///
/// WGSL has no u8 array type — must use array<u32> and extract bytes via bitwise ops.
/// WGSL has no native f16 — need software f16_to_f32 via IEEE 754 bit manipulation.
///
/// NOTE: WGSL does not allow passing storage pointers to functions, so read_u8/read_i8/read_f16
/// are provided as Rust functions that generate inline WGSL expressions. Use them in format! strings.
pub fn common_helpers() -> &'static str {
    r#"
fn f16_to_f32(bits: u32) -> f32 {
    let sign = (bits >> 15u) & 1u;
    let exp = (bits >> 10u) & 0x1Fu;
    let mant = bits & 0x3FFu;

    if (exp == 0u) {
        if (mant == 0u) {
            // Zero
            if (sign == 1u) { return -0.0; }
            return 0.0;
        }
        // Subnormal: value = (-1)^sign * 2^-14 * (mant / 1024)
        let val = f32(mant) / 1024.0 * 0.00006103515625; // 2^-14
        if (sign == 1u) { return -val; }
        return val;
    }
    if (exp == 31u) {
        // Inf or NaN
        if (mant == 0u) {
            if (sign == 1u) { return bitcast<f32>(0xFF800000u); } // -inf
            return bitcast<f32>(0x7F800000u); // inf
        }
        return bitcast<f32>(0x7FC00000u); // NaN
    }

    // Normal: value = (-1)^sign * 2^(exp-15) * (1 + mant/1024)
    let f_exp = f32(i32(exp) - 15);
    let f_mant = 1.0 + f32(mant) / 1024.0;
    let val = f_mant * exp2(f_exp);
    if (sign == 1u) { return -val; }
    return val;
}
"#
}

/// Generate inline WGSL expression to read a u8 from a u32 storage buffer.
/// `buf_name` is the WGSL variable name for the storage buffer (array<u32>).
/// `byte_idx_expr` is a WGSL expression for the byte index.
/// Returns a WGSL expression of type u32.
pub fn read_u8_inline(buf_name: &str, byte_idx_expr: &str) -> String {
    format!(
        "(({buf}[{idx} / 4u] >> (({idx} % 4u) * 8u)) & 0xFFu)",
        buf = buf_name,
        idx = byte_idx_expr
    )
}

/// Generate inline WGSL expression to read an i8 from a u32 storage buffer.
pub fn read_i8_inline(buf_name: &str, byte_idx_expr: &str) -> String {
    let u8_expr = read_u8_inline(buf_name, byte_idx_expr);
    format!("select(i32({u8_expr}), i32({u8_expr}) - 256, {u8_expr} >= 128u)")
}

/// Generate inline WGSL expression to read an f16 from a u32 storage buffer.
pub fn read_f16_inline(buf_name: &str, byte_idx_expr: &str) -> String {
    format!(
        "f16_to_f32(({buf}[{idx} / 4u] >> (({idx} % 4u) * 8u)) & 0xFFFFu)",
        buf = buf_name,
        idx = byte_idx_expr
    )
}