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
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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251

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

//! Pad index optimization.
//!
//! This module provides index padding optimizations for kernels,
//! which pad global indices to hardware-friendly sizes.
//!
//! Padding indices is useful for:
//!
//! - Enabling efficient tiling when tensor dimensions aren't multiples
//!   of the tile size
//! - Guarding out-of-range loads
//! - Skipping out-of-range stores with conditionals

use crate::{
    Set,
    backend::DeviceInfo,
    dtype::Constant,
    kernel::{BOp, IDX_T, Kernel, MemLayout, Op, OpId, RangeKind},
    shape::Dim,
};

use super::autotune::Optimization;

/// Pad indices to hardware-friendly sizes (e.g., 32 for CUDA warps).
#[derive(Debug)]
pub struct PadIndex {
    /// Pairs of (index_op_id, target_size) for each padding.
    pub factors: Vec<(OpId, Dim)>,
}

impl Optimization for PadIndex {
    fn nconfigs(&self) -> u64 {
        self.factors.len() as u64
    }

    fn apply(&self, kernel: &mut Kernel, config: u64) {
        if self.factors.is_empty() {
            return;
        }
        let (idx_id, pad_to) = self.factors[config as usize];
        let Op::Range { kind, .. } = kernel.ops[idx_id].op else {
            unreachable!()
        };
        let current_len = match kind {
            RangeKind::Group(len) => kernel.resolve_const(len).and_then(crate::dtype::Constant::as_dim).unwrap(),
            RangeKind::Local(len) => i64::from(len),
            // opt_pad_index never collects warp factors (warp views can't be padded).
            RangeKind::Warp(_) => unreachable!("PadIndex never collects warp factors"),
        };
        let pad_len = (pad_to - current_len % pad_to) % pad_to;
        if pad_len > 0 {
            kernel.pad_range(idx_id, pad_len);
        }
    }
}

impl Kernel {
    /// Pads a global index to the next multiple of `tile_size`, guarding out-of-range loads
    /// and skipping out-of-range stores.
    ///
    /// This extends `Op::Index { len: current_len, .. }` to `len: current_len + pad_len`
    /// so the grid covers full tiles.  OOB reads are redirected to element 0 (safe).
    /// OOB stores are wrapped in `Op::If { .. }` / `Op::EndIf` and skipped entirely.
    ///
    /// Useful for tiling: when a tensor dimension isn't a multiple of the tile size,
    /// pad the index so the grid covers full tiles, and OOB threads compute garbage
    /// but never write it to memory.
    ///
    /// # Panics
    /// - If `gidx_id` is not an `Op::Index` node.
    pub(crate) fn pad_range(&mut self, gidx_id: OpId, pad_len: Dim) {
        if pad_len == 0 {
            return;
        }

        // 1. Extend the index length
        let Op::Range { axis, kind } = self.ops[gidx_id].op else {
            panic!("pad_index: op is not an Index");
        };
        let (current_len, new_kind) = match kind {
            RangeKind::Group(len) => match self.resolve_const(len).and_then(crate::dtype::Constant::as_dim) {
                Some(current_len) => {
                    let new_len = self.insert_before(gidx_id, Op::Const(Constant::idx(current_len + pad_len)));
                    (current_len, RangeKind::Group(new_len))
                }
                None => return,
            },
            RangeKind::Local(len) => {
                let current_len = Dim::from(len);
                let new_len = len + u32::try_from(pad_len).expect("pad_len too large for local index");
                (current_len, RangeKind::Local(new_len))
            }
            RangeKind::Warp(_) => {
                panic!("pad_range: cannot pad a warp view; pad its underlying local range instead")
            }
        };
        self.ops[gidx_id].op = Op::Range { axis, kind: new_kind };

        // 2. Create limit constant for comparison
        let limit = self.insert_before(gidx_id, Op::Const(Constant::idx(current_len)));

        // 3. Walk all ops to guard loads and stores depending on this index
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);

            // Redirect OOB stores to trash element at index `limit`
            if let Op::Store { dst, src: x, index: store_idx, layout } = self.ops[op_id].op.clone()
                && self.depends_on(store_idx, gidx_id, &mut Set::default())
            {
                let buf_len: Option<Dim> = match &self.ops[dst].op {
                    Op::Param { .. } => Some(self.shape(dst).iter().product()),
                    _ => None,
                };
                if let Some(buf_len) = buf_len {
                    let clen = self.insert_before(op_id, Op::Const(Constant::idx(buf_len)));
                    let cond = self.insert_before(op_id, Op::Binary { x: gidx_id, y: limit, bop: BOp::Cmplt });
                    let cast_cond = self.insert_before(op_id, Op::Cast { x: cond, dtype: IDX_T });
                    let one = self.insert_before(op_id, Op::Const(Constant::idx(1)));
                    let not_cond = self.insert_before(op_id, Op::Binary { x: one, y: cast_cond, bop: BOp::Sub });
                    let idx_term = self.insert_before(op_id, Op::Binary { x: store_idx, y: cast_cond, bop: BOp::Mul });
                    let lim_term = self.insert_before(op_id, Op::Binary { x: clen, y: not_cond, bop: BOp::Mul });
                    let safe_idx = self.insert_before(op_id, Op::Binary { x: idx_term, y: lim_term, bop: BOp::Add });
                    self.ops[op_id].op = Op::Store { dst, src: x, index: safe_idx, layout };
                }
            }

            // Guard loads: redirect OOB reads to element 0 (safe)
            if let Op::Load { src, index: load_idx, layout } = self.ops[op_id].op.clone() {
                if layout != MemLayout::Scalar {
                    op_id = next;
                    continue;
                }
                if self.depends_on(load_idx, gidx_id, &mut Set::default()) {
                    let cond = self.insert_before(op_id, Op::Binary { x: gidx_id, y: limit, bop: BOp::Cmplt });
                    let cast_idx = self.insert_before(op_id, Op::Cast { x: cond, dtype: IDX_T });
                    let safe_idx = self.insert_before(op_id, Op::Binary { x: load_idx, y: cast_idx, bop: BOp::Mul });
                    let safe_load = self.insert_before(op_id, Op::Load { src, index: safe_idx, layout });
                    self.remap(op_id, safe_load);
                    self.remove_op(op_id);
                }
            }

            op_id = next;
        }
    }

    #[allow(unused)]
    pub(crate) fn pad_loop(&mut self, loop_id: OpId, pad_len: Dim) {
        if pad_len == 0 {
            return;
        }

        // 1. Extend the loop length
        let Op::Loop { len } = &self.ops[loop_id].op else {
            panic!("pad_loop: op is not a Loop");
        };
        let current_len = self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap();
        let new_len = self.insert_before(loop_id, Op::Const(Constant::idx(current_len + pad_len)));
        self.ops[loop_id].op = Op::Loop { len: new_len };

        // 2. Create limit constant for comparison
        let limit = self.insert_before(loop_id, Op::Const(Constant::idx(current_len)));

        // 3. Walk all ops to guard loads and stores depending on this loop
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);

            // Redirect OOB stores to trash element at index `buf_len`
            if let Op::Store { dst, src: x, index: store_idx, layout } = self.ops[op_id].op.clone()
                && self.depends_on(store_idx, loop_id, &mut Set::default())
            {
                let buf_len: Option<Dim> = match self.ops[dst].op {
                    Op::Param { .. } => todo!(),
                    Op::Storage { len, .. } => Some(len),
                    _ => None,
                };
                if let Some(buf_len) = buf_len {
                    let clen = self.insert_before(op_id, Op::Const(Constant::idx(buf_len)));
                    let cond = self.insert_before(op_id, Op::Binary { x: loop_id, y: limit, bop: BOp::Cmplt });
                    let cast_cond = self.insert_before(op_id, Op::Cast { x: cond, dtype: IDX_T });
                    let one = self.insert_before(op_id, Op::Const(Constant::idx(1)));
                    let not_cond = self.insert_before(op_id, Op::Binary { x: one, y: cast_cond, bop: BOp::Sub });
                    let idx_term = self.insert_before(op_id, Op::Binary { x: store_idx, y: cast_cond, bop: BOp::Mul });
                    let lim_term = self.insert_before(op_id, Op::Binary { x: clen, y: not_cond, bop: BOp::Mul });
                    let safe_idx = self.insert_before(op_id, Op::Binary { x: idx_term, y: lim_term, bop: BOp::Add });
                    self.ops[op_id].op = Op::Store { dst, src: x, index: safe_idx, layout };
                }
            }

            // Guard loads: redirect OOB reads to element 0 (safe)
            if let Op::Load { src, index: load_idx, layout } = self.ops[op_id].op.clone() {
                if layout != MemLayout::Scalar {
                    op_id = next;
                    continue;
                }
                if self.depends_on(load_idx, loop_id, &mut Set::default()) {
                    let cond = self.insert_before(op_id, Op::Binary { x: loop_id, y: limit, bop: BOp::Cmplt });
                    let cast_idx = self.insert_before(op_id, Op::Cast { x: cond, dtype: IDX_T });
                    let safe_idx = self.insert_before(op_id, Op::Binary { x: load_idx, y: cast_idx, bop: BOp::Mul });
                    let safe_load = self.insert_before(op_id, Op::Load { src, index: safe_idx, layout });
                    self.remap(op_id, safe_load);
                    self.remove_op(op_id);
                }
            }

            op_id = next;
        }
    }

    /// Make the `PadIndex` optimization: scan the kernel for indices whose
    /// lengths are not multiples of hardware-friendly sizes.
    pub fn opt_pad_index(&self, _dev_info: &DeviceInfo) -> Box<dyn Optimization> {
        let mut factors = Vec::new();
        let mut op_id = self.head;
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let Op::Range { kind, .. } = self.ops[op_id].op {
                let len = match kind {
                    RangeKind::Group(len) => match self.resolve_const(len).and_then(crate::dtype::Constant::as_dim) {
                        Some(len) => len,
                        None => continue,
                    },
                    RangeKind::Local(len) => Dim::from(len),
                    // Warp views can't be padded — skip them entirely.
                    RangeKind::Warp(_) => continue,
                };
                for pad_to in [8, 16, 32] {
                    if len % pad_to as Dim != 0 {
                        factors.push((op_id, pad_to));
                    }
                }
            }
            op_id = next;
        }
        Box::new(PadIndex { factors })
    }

    pub(crate) fn depends_on(&self, expr: OpId, target: OpId, visited: &mut Set<OpId>) -> bool {
        if expr == target || !visited.insert(expr) {
            return expr == target;
        }
        match self.at(expr) {
            Op::Const(_) | Op::Range { .. } | Op::Storage { .. } | Op::Loop { .. } | Op::EndLoop => false,
            op => op.parameters().any(|p| self.depends_on(p, target, visited)),
        }
    }
}