zyx 0.16.0

Zyx machine learning library
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only

//! Convert exp2(x * log2(e)) to exp(x).
//!
//! This module provides optimization for converting `exp2(x * log2(e))`
//! to `exp(x)`, which allows backends like Tenstorrent to use their
//! native `exp_tile` instead of the unsupported `exp2_tile`.

use std::f64::consts::{LN_2, LOG2_E};

use crate::scalar::{bf16, f16};
use crate::{
    dtype::Constant,
    kernel::{BOp, Kernel, Op, OpId, UOp},
};

fn constant_is_ln_2(c: &Constant) -> bool {
    let val = match *c {
        Constant::BF16(x) => bf16::from_le_bytes(x).to_f32() as f64,
        Constant::F16(x) => f16::from_le_bytes(x).to_f32() as f64,
        Constant::F32(x) => f32::from_le_bytes(x) as f64,
        Constant::F64(x) => f64::from_le_bytes(x),
        _ => return false,
    };
    (val - LN_2).abs() < 1e-6
}

fn constant_is_log2_e(c: &Constant) -> bool {
    let val = match *c {
        Constant::BF16(x) => bf16::from_le_bytes(x).to_f32() as f64,
        Constant::F16(x) => f16::from_le_bytes(x).to_f32() as f64,
        Constant::F32(x) => f32::from_le_bytes(x) as f64,
        Constant::F64(x) => f64::from_le_bytes(x),
        _ => return false,
    };
    (val - LOG2_E).abs() < 1e-6
}

impl Kernel {
    /// Finds `exp2(x * log2(e))` and replaces it with `exp(x)`.
    ///
    /// This recognizes the pattern produced by `tensor.exp()` which is
    /// implemented as `(x * log2(e)).exp2()`. Converting back to `exp`
    /// allows Tenstorrent to use its native `exp_tile` instead of the
    /// unsupported `exp2_tile`.
    pub fn exp2_to_exp(&mut self) {
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let &Op::Unary { x, uop: UOp::Exp2 } = self.at(op_id) {
                if let &Op::Binary { x: left, y: right, bop: BOp::Mul } = self.at(x) {
                    let input = match (self.at(left), self.at(right)) {
                        (&Op::Const(c), _) if constant_is_log2_e(&c) => right,
                        (_, &Op::Const(c)) if constant_is_log2_e(&c) => left,
                        _ => OpId::NULL,
                    };
                    if input != OpId::NULL {
                        self.ops[op_id].op = Op::Unary { x: input, uop: UOp::Exp };
                    }
                }
            }
            op_id = next;
        }
    }

    /// Finds `log2(x) * ln(2)` and replaces it with `ln(x)`.
    ///
    /// This recognizes the pattern produced by `tensor.ln()` which is
    /// implemented as `log2(x) * (1/log2(e))` = `log2(x) * ln(2)`.
    /// Converting to `ln` allows backends that lack native `log2` to
    /// use a native `ln` function instead.
    pub fn log2_to_ln(&mut self) {
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let &Op::Binary { x: left, y: right, bop: BOp::Mul } = self.at(op_id) {
                let ((&Op::Unary { x: log2_op, uop: UOp::Log2 }, const_op)
                | (const_op, &Op::Unary { x: log2_op, uop: UOp::Log2 })) = (self.at(left), self.at(right))
                else {
                    op_id = next;
                    continue;
                };
                if let &Op::Const(c) = const_op {
                    if constant_is_ln_2(&c) {
                        self.ops[op_id].op = Op::Unary { x: log2_op, uop: UOp::Ln };
                    }
                }
            }
            op_id = next;
        }
    }

    /// Exp to exp2
    /// Converts `exp(x * ln(e))` to `exp2(x)`
    /// This allows backends with native `exp2` but not `exp` to use exp2 instead
    pub fn exp_to_exp2(&mut self) {
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let &Op::Unary { x, uop: UOp::Exp } = self.at(op_id) {
                let dtype = self.dtype(x);
                let y = self.insert_before(op_id, Op::Const(Constant::F64(LOG2_E.to_le_bytes()).cast(dtype)));
                let z = self.insert_before(op_id, Op::Binary { x, y, bop: BOp::Mul });
                self.ops[op_id].op = Op::Unary { x: z, uop: UOp::Exp2 };
            }
            op_id = next;
        }
    }

    /// Ln to log2
    /// Converts `ln(x)` to `log2(x) * ln(2)`
    /// This allows backends with native `log2` but not `ln` to use log2 instead
    pub fn ln_to_log2(&mut self) {
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let &Op::Unary { x, uop: UOp::Ln } = self.at(op_id) {
                let dtype = self.dtype(x);
                let y = self.insert_before(op_id, Op::Const(Constant::F64(LN_2.to_le_bytes()).cast(dtype)));
                let x = self.insert_before(op_id, Op::Unary { x, uop: UOp::Log2 });
                self.ops[op_id].op = Op::Binary { x, y, bop: BOp::Mul };
            }
            op_id = next;
        }
    }
}