zyx 0.17.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

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0

//! Pattern matching of subgraph structures onto specialized kernels.
//!
//! The graph op set is intentionally tiny (Const, Leaf, Expand, Permute,
//! Reshape, Reduce, Cast, Unary, Binary, ToDevice, Kernel). Every tensor
//! expression collapses into this fixed set, which means any computation —
//! however it was written by the user — decomposes into a handful of canonical
//! subgraphs. This module matches those subgraphs so they can be replaced by
//! specialized AOT kernels (cblas, cublas, ...) instead of being fused into
//! generic zyx kernels. The graph measures both alternatives and picks the
//! fastest path through extraction.

#![allow(unused)]

use crate::{
    graph::{Graph, Node},
    kernel::{BOp, OpId},
    shape::Dim,
};

mod matmul;

impl Graph {
    /// Resolves a class's symbolic shape to numeric dims, or `None` if any
    /// dim is not a constant.
    fn const_shape(&self, cid: OpId) -> Option<Vec<Dim>> {
        self.shape(cid)
            .into_iter()
            .map(|dim| match &self.nodes[dim].node {
                Node::Const { value: c, .. } => c.as_dim(),
                _ => None,
            })
            .collect()
    }

    /// Finds a `Reduce(Add)` over the single trailing axis of a 3D product.
    /// Returns the product class and the contraction dim `k`.
    fn reduce_add_last(&self, cid: OpId) -> Option<(OpId, Dim)> {
        self.class_nodes(cid).find_map(|nid| match &self.nodes[nid].node {
            Node::Reduce { x, rop: BOp::Add, axes } => {
                let prod_shape = self.const_shape(*x)?;
                if prod_shape.len() == 3 && axes.len() == 1 && axes[0] == prod_shape.len() - 1 {
                    Some((*x, prod_shape[2]))
                } else {
                    None
                }
            }
            _ => None,
        })
    }

    /// Finds the elementwise `Mul` beneath the product class, through an optional
    /// accumulator `Cast` (`dot` casts the product before reducing).
    fn mul_of(&self, cid: OpId) -> Option<(OpId, OpId)> {
        if let Some((x, y)) = self.class_nodes(cid).find_map(|nid| match &self.nodes[nid].node {
            Node::Binary { x, y, bop: BOp::Mul } => Some((*x, *y)),
            _ => None,
        }) {
            return Some((x, y));
        }
        let x = self.class_nodes(cid).find_map(|nid| match &self.nodes[nid].node {
            Node::Cast { x, .. } => Some(*x),
            _ => None,
        })?;
        self.class_nodes(x).find_map(|nid| match &self.nodes[nid].node {
            Node::Binary { x, y, bop: BOp::Mul } => Some((*x, *y)),
            _ => None,
        })
    }

    /// Finds the source of an `Expand` and its class shape.
    ///
    /// The shape is read from the source class itself, so matching does not
    /// depend on how the shape was produced (e.g. a `Reshape` in the canonical
    /// matmul form, or an eager tensor already at the broadcast shape).
    fn expand_src(&self, cid: OpId) -> Option<(OpId, Vec<Dim>)> {
        let x = self.class_nodes(cid).find_map(|nid| match &self.nodes[nid].node {
            Node::Expand { x, .. } => Some(*x),
            _ => None,
        })?;
        Some((x, self.const_shape(x)?))
    }

    /// Finds the source of a 2D `Permute [1, 0]` (a `[n, k]` transposed from
    /// `[k, n]`), looking through shape-only wrappers such as the `Reshape` to
    /// `[1, n, k]` in the broadcast matmul form.
    fn transpose_src(&self, cid: OpId) -> Option<OpId> {
        self.class_nodes(cid).find_map(|nid| match &self.nodes[nid].node {
            Node::Reshape { x, .. } => self.transpose_src(*x),
            Node::Permute { x, axes } if axes.len() == 2 && axes[0] == 1 && axes[1] == 0 => Some(*x),
            _ => None,
        })
    }
}