rlx_ir/ops/
axial_rope2d.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! SAM2-style axial 2-D RoPE on `[batch, seq, num_heads * head_dim]`.
6
7use crate::{Graph, NodeId, Op};
8
9/// Apply axial 2-D RoPE on flattened `[nh, n_tokens, head_dim]` layout.
10pub fn apply_axial_rope2d(
11    x: &[f32],
12    num_heads: usize,
13    n_tokens: usize,
14    head_dim: usize,
15    end_x: usize,
16    end_y: usize,
17    theta: f32,
18    repeat_factor: usize,
19) -> Vec<f32> {
20    debug_assert!(head_dim.is_multiple_of(4));
21    let half = head_dim / 2;
22    let q4 = head_dim / 4;
23    let spatial = end_x * end_y;
24    let repeat = repeat_factor.max(1);
25    debug_assert_eq!(n_tokens, spatial * repeat);
26
27    let mut freqs = vec![0f32; q4];
28    for i in 0..q4 {
29        freqs[i] = 1.0 / theta.powf((4 * i) as f32 / head_dim as f32);
30    }
31    let mut cs_x = vec![0f32; spatial * q4];
32    let mut sn_x = vec![0f32; spatial * q4];
33    let mut cs_y = vec![0f32; spatial * q4];
34    let mut sn_y = vec![0f32; spatial * q4];
35    for pos in 0..spatial {
36        let tx = (pos % end_x) as f32;
37        let ty = (pos / end_x) as f32;
38        for c in 0..q4 {
39            let ax = tx * freqs[c];
40            let ay = ty * freqs[c];
41            cs_x[pos * q4 + c] = ax.cos();
42            sn_x[pos * q4 + c] = ax.sin();
43            cs_y[pos * q4 + c] = ay.cos();
44            sn_y[pos * q4 + c] = ay.sin();
45        }
46    }
47    let mut cos_x = vec![0f32; n_tokens * q4];
48    let mut sin_x = vec![0f32; n_tokens * q4];
49    let mut cos_y = vec![0f32; n_tokens * q4];
50    let mut sin_y = vec![0f32; n_tokens * q4];
51    for tok in 0..n_tokens {
52        let pos = tok / repeat;
53        for c in 0..q4 {
54            cos_x[tok * q4 + c] = cs_x[pos * q4 + c];
55            sin_x[tok * q4 + c] = sn_x[pos * q4 + c];
56            cos_y[tok * q4 + c] = cs_y[pos * q4 + c];
57            sin_y[tok * q4 + c] = sn_y[pos * q4 + c];
58        }
59    }
60
61    // `[batch, seq, num_heads * head_dim]` layout (token-major, heads interleaved).
62    let hs = num_heads * head_dim;
63    let mut out = x.to_vec();
64    for tok in 0..n_tokens {
65        for h in 0..num_heads {
66            let base = tok * hs + h * head_dim;
67            for c in 0..q4 {
68                let ix0 = base + 2 * c;
69                let ix1 = base + 2 * c + 1;
70                let x0 = out[ix0];
71                let x1 = out[ix1];
72                let co = cos_x[tok * q4 + c];
73                let si = sin_x[tok * q4 + c];
74                out[ix0] = x0 * co - x1 * si;
75                out[ix1] = x0 * si + x1 * co;
76            }
77            for c in 0..q4 {
78                let ix0 = base + half + 2 * c;
79                let ix1 = base + half + 2 * c + 1;
80                let x0 = out[ix0];
81                let x1 = out[ix1];
82                let co = cos_y[tok * q4 + c];
83                let si = sin_y[tok * q4 + c];
84                out[ix0] = x0 * co - x1 * si;
85                out[ix1] = x0 * si + x1 * co;
86            }
87        }
88    }
89    out
90}
91
92impl Graph {
93    /// `x`: `[1, seq, num_heads * head_dim]` → same shape.
94    pub fn axial_rope2d(
95        &mut self,
96        x: NodeId,
97        end_x: usize,
98        end_y: usize,
99        head_dim: usize,
100        num_heads: usize,
101        theta: f32,
102        repeat_factor: usize,
103    ) -> NodeId {
104        let s = crate::shape::unary_shape(self.shape(x));
105        self.push(
106            Op::AxialRope2d {
107                end_x,
108                end_y,
109                head_dim,
110                num_heads,
111                theta,
112                repeat_factor,
113            },
114            vec![x],
115            s,
116            None,
117        )
118    }
119}
rlx_ir/ops/axial_rope2d.rs

rlx_ir/ops/
axial_rope2d.rs
