rlx_ir/hir/
mod.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! **HIR** — high-level IR.
17//!
18//! Block-oriented IR for model authors and external graph builders.
19//! HIR captures fusion-friendly patterns (SwiGLU FFN, linear layers,
20//! residual RMSNorm) as first-class ops and lowers to MIR via
21//! [`HirModule::lower_to_mir`].
22
23mod blocks;
24mod conv;
25mod fusion;
26mod graph_ext;
27mod lower;
28mod window;
29
30pub use blocks::lower_llama_decoder_block;
31pub use blocks::lower_qwen35_mtp_head;
32pub use fusion::FusionPolicy;
33pub use graph_ext::{GridMode, GridPad, HirGraphExt, HirMut};
34pub use window::{window_token_gather_bsn, window_token_scatter_bsn};
35
36use crate::mir::MirModule;
37use crate::op::Activation;
38use crate::op::MaskKind;
39use crate::quant::QuantScheme;
40use crate::{Op, Shape};
41
42pub use lower::LowerError;
43
44/// Stable node identifier within a HIR module.
45#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
46#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
47pub struct HirNodeId(pub u32);
48
49impl std::fmt::Display for HirNodeId {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        write!(f, "h{}", self.0)
52    }
53}
54
55/// High-level operation — blocks and escape hatches.
56#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
57#[derive(Debug, Clone, PartialEq)]
58pub enum HirOp {
59    Input {
60        name: String,
61    },
62    Param {
63        name: String,
64    },
65    Constant {
66        data: Vec<u8>,
67    },
68
69    /// `matmul → add(bias)? → activation?`
70    /// Inputs: `[x, weight]` or `[x, weight, bias]`.
71    Linear {
72        activation: Option<Activation>,
73        has_bias: bool,
74    },
75
76    /// Emit [`Op::FusedMatMulBiasAct`] directly.
77    /// Inputs: `[x, weight, bias]`.
78    LinearFused {
79        activation: Option<Activation>,
80    },
81
82    /// Two matmuls sharing the same input (QKV / SwiGLU gate+up).
83    /// Inputs: `[x, w_first, w_second]`. `slot` selects which output.
84    SharedLinearPair {
85        slot: u8,
86    },
87
88    /// Full SwiGLU FFN.
89    /// Inputs: `[x, up_w, gate_w, down_w]`.
90    SwiGLU,
91
92    /// `add(x, residual)` then RMSNorm.
93    /// Inputs: `[x, residual, gamma, beta]`.
94    ResidualRmsNorm {
95        eps: f32,
96    },
97
98    /// Scaled dot-product attention.
99    /// Inputs: `[q, k, v, mask?]` — mask omitted when `mask == None`.
100    Attention {
101        num_heads: usize,
102        head_dim: usize,
103        mask: MaskKind,
104    },
105
106    /// Causal depthwise Conv1d on `[batch, seq, channels]` tensors.
107    /// Inputs: `[input, weight, left_pad]` — see [`conv::lower_depthwise_conv1d_causal`].
108    DepthwiseConv1dCausal {
109        kernel_size: usize,
110    },
111
112    /// Fused dequant + matmul. GGUF schemes take `[x, packed_w]`; legacy
113    /// Int8/NVFP4 schemes take `[x, w_q, scale, zp]`.
114    DequantMatMul {
115        scheme: QuantScheme,
116    },
117
118    /// Gated DeltaNet linear-attention scan (Qwen3.5 trunk).
119    /// Inputs: `[q, k, v, g, beta]` or with carry `[…, state]`.
120    GatedDeltaNet {
121        state_size: usize,
122        carry_state: bool,
123    },
124
125    /// Multi-layer (optionally bidirectional) LSTM. Packed weights;
126    /// `+ [h0, c0]` when `carry`. See [`crate::Op::Lstm`].
127    Lstm {
128        hidden_size: usize,
129        num_layers: usize,
130        bidirectional: bool,
131        carry: bool,
132    },
133
134    /// Multi-layer (optionally bidirectional) GRU. Inputs
135    /// `[x, w_ih, w_hh, b_ih, b_hh]` (`+ [h0]` when `carry`). See
136    /// [`crate::Op::Gru`].
137    Gru {
138        hidden_size: usize,
139        num_layers: usize,
140        bidirectional: bool,
141        carry: bool,
142    },
143
144    /// Rotary position embedding. Inputs: `[x, cos, sin]`.
145    RoPE {
146        head_dim: usize,
147        n_rot: usize,
148    },
149
150    /// RMS normalization without residual. Inputs: `[x, gamma, beta]`.
151    RmsNorm {
152        eps: f32,
153    },
154
155    /// LLaMA-style pre-norm decoder block: attn (GQA) + SwiGLU FFN.
156    /// Inputs (causal): `[x, ln1_g, ln1_b, q_w, k_w, v_w, o_w, ln2_g, ln2_b,
157    /// gate_w, up_w, down_w, cos, sin]`. With `MaskKind::Custom` or `Bias`
158    /// append `mask`.
159    LlamaDecoderBlock {
160        num_heads: usize,
161        head_dim: usize,
162        num_kv_heads: usize,
163        eps: f32,
164        mask: MaskKind,
165        /// RoPE pairing flavor (GGUF Llama → [`crate::op::RopeStyle::GptJ`]).
166        rope_style: crate::op::RopeStyle,
167    },
168
169    /// Qwen3.5 MTP draft head: hnorm∥enorm → eh_proj → full-attn → LM.
170    /// See [`blocks::lower_qwen35_mtp_head`] for the input layout.
171    Qwen35MtpHead {
172        num_heads: usize,
173        num_kv_heads: usize,
174        head_dim: usize,
175        n_rot: usize,
176        n_embd: usize,
177        n_ff: usize,
178        mtp_vocab: usize,
179        eps: f32,
180    },
181
182    /// Escape hatch — embed a single MIR op verbatim.
183    Mir(Op),
184}
185
186/// One node in a HIR module.
187#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
188#[derive(Debug, Clone)]
189pub struct HirNode {
190    pub id: HirNodeId,
191    pub op: HirOp,
192    pub inputs: Vec<HirNodeId>,
193    pub shape: Shape,
194    pub name: Option<String>,
195}
196
197/// High-level module — model builder output.
198#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
199#[derive(Debug, Clone)]
200pub struct HirModule {
201    pub name: String,
202    nodes: Vec<HirNode>,
203    pub outputs: Vec<HirNodeId>,
204    /// How block ops lower to MIR. Default: [`FusionPolicy::Direct`]
205    /// for new model code (fusion as a first-class citizen).
206    pub fusion_policy: FusionPolicy,
207}
208
209impl HirModule {
210    pub fn new(name: impl Into<String>) -> Self {
211        Self {
212            name: name.into(),
213            nodes: Vec::new(),
214            outputs: Vec::new(),
215            fusion_policy: FusionPolicy::Direct,
216        }
217    }
218
219    pub fn with_fusion_policy(mut self, policy: FusionPolicy) -> Self {
220        self.fusion_policy = policy;
221        self
222    }
223
224    pub fn len(&self) -> usize {
225        self.nodes.len()
226    }
227
228    pub fn is_empty(&self) -> bool {
229        self.nodes.is_empty()
230    }
231
232    pub fn nodes(&self) -> &[HirNode] {
233        &self.nodes
234    }
235
236    pub fn node(&self, id: HirNodeId) -> &HirNode {
237        &self.nodes[id.0 as usize]
238    }
239
240    pub fn node_mut(&mut self, id: HirNodeId) -> &mut HirNode {
241        &mut self.nodes[id.0 as usize]
242    }
243
244    /// Build a named block — sets `HirNode::name` on the returned node.
245    pub fn named(
246        &mut self,
247        name: impl Into<String>,
248        build: impl FnOnce(&mut Self) -> HirNodeId,
249    ) -> HirNodeId {
250        let id = build(self);
251        self.node_mut(id).name = Some(name.into());
252        id
253    }
254
255    /// Stamp a source-op provenance `name` onto every node created at or
256    /// after index `from` whose label is only the generic op default (or
257    /// unset). Importers call this after lowering one source op (an ONNX
258    /// node, a torch aten call) so all the HIR nodes it expanded into share
259    /// that op's identity — the label then flows to every derived MIR node
260    /// via `tag_hir_subgraph`, so a NaN in any of them localizes back to the
261    /// user's original op instead of a generic `"mir"`.
262    ///
263    /// Explicitly-meaningful names are preserved: `Input`/`Param` nodes have
264    /// no default label, so their weight names (which differ from `None`) are
265    /// left untouched even if they fall inside the range.
266    pub fn label_nodes_since(&mut self, from: usize, name: &str) {
267        if name.is_empty() {
268            return;
269        }
270        for i in from..self.nodes.len() {
271            let is_generic = {
272                let n = &self.nodes[i];
273                n.name.is_none() || n.name.as_deref() == default_hir_block_label(&n.op).as_deref()
274            };
275            if is_generic {
276                self.nodes[i].name = Some(name.to_string());
277            }
278        }
279    }
280
281    fn push_block(
282        &mut self,
283        op: HirOp,
284        inputs: Vec<HirNodeId>,
285        shape: Shape,
286        name: Option<String>,
287    ) -> HirNodeId {
288        let name = name.or_else(|| default_hir_block_label(&op));
289        self.push(op, inputs, shape, name)
290    }
291
292    fn push(
293        &mut self,
294        op: HirOp,
295        inputs: Vec<HirNodeId>,
296        shape: Shape,
297        name: Option<String>,
298    ) -> HirNodeId {
299        let id = HirNodeId(self.nodes.len() as u32);
300        self.nodes.push(HirNode {
301            id,
302            op,
303            inputs,
304            shape,
305            name,
306        });
307        id
308    }
309
310    pub fn input(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId {
311        self.push(HirOp::Input { name: name.into() }, vec![], shape, None)
312    }
313
314    /// `[batch, seq, hidden]` input with symbolic leading axes.
315    pub fn input_batch_seq(
316        &mut self,
317        name: impl Into<String>,
318        batch: u32,
319        seq: u32,
320        hidden: usize,
321        dtype: crate::DType,
322    ) -> HirNodeId {
323        self.input(name, Shape::batch_seq(batch, seq, hidden, dtype))
324    }
325
326    pub fn param(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId {
327        self.push(HirOp::Param { name: name.into() }, vec![], shape, None)
328    }
329
330    pub fn linear(
331        &mut self,
332        x: HirNodeId,
333        weight: HirNodeId,
334        bias: Option<HirNodeId>,
335        activation: Option<Activation>,
336        out_shape: Shape,
337    ) -> HirNodeId {
338        let mut inputs = vec![x, weight];
339        if let Some(b) = bias {
340            inputs.push(b);
341        }
342        self.push_block(
343            HirOp::Linear {
344                activation,
345                has_bias: bias.is_some(),
346            },
347            inputs,
348            out_shape,
349            None,
350        )
351    }
352
353    /// Emit [`HirOp::LinearFused`] — fused matmul+bias+act at MIR level.
354    pub fn linear_fused(
355        &mut self,
356        x: HirNodeId,
357        weight: HirNodeId,
358        bias: HirNodeId,
359        activation: Option<Activation>,
360        out_shape: Shape,
361    ) -> HirNodeId {
362        self.push_block(
363            HirOp::LinearFused { activation },
364            vec![x, weight, bias],
365            out_shape,
366            None,
367        )
368    }
369
370    /// Two matmuls sharing `x`. Returns `(first, second)` in weight order.
371    pub fn shared_linear_pair(
372        &mut self,
373        x: HirNodeId,
374        w_first: HirNodeId,
375        w_second: HirNodeId,
376        out_shape: Shape,
377    ) -> (HirNodeId, HirNodeId) {
378        let inputs = vec![x, w_first, w_second];
379        let first = self.push_block(
380            HirOp::SharedLinearPair { slot: 0 },
381            inputs.clone(),
382            out_shape.clone(),
383            None,
384        );
385        let second = self.push_block(HirOp::SharedLinearPair { slot: 1 }, inputs, out_shape, None);
386        (first, second)
387    }
388
389    pub fn swiglu_ffn(
390        &mut self,
391        x: HirNodeId,
392        up_w: HirNodeId,
393        gate_w: HirNodeId,
394        down_w: HirNodeId,
395        out_shape: Shape,
396    ) -> HirNodeId {
397        self.push_block(
398            HirOp::SwiGLU,
399            vec![x, up_w, gate_w, down_w],
400            out_shape,
401            None,
402        )
403    }
404
405    pub fn residual_rms_norm(
406        &mut self,
407        x: HirNodeId,
408        residual: HirNodeId,
409        gamma: HirNodeId,
410        beta: HirNodeId,
411        eps: f32,
412        out_shape: Shape,
413    ) -> HirNodeId {
414        self.push_block(
415            HirOp::ResidualRmsNorm { eps },
416            vec![x, residual, gamma, beta],
417            out_shape,
418            None,
419        )
420    }
421
422    /// Scaled dot-product attention — see [`HirOp::Attention`].
423    pub fn attention(
424        &mut self,
425        q: HirNodeId,
426        k: HirNodeId,
427        v: HirNodeId,
428        mask: Option<HirNodeId>,
429        num_heads: usize,
430        head_dim: usize,
431        mask_kind: MaskKind,
432        out_shape: Shape,
433    ) -> HirNodeId {
434        let mut inputs = vec![q, k, v];
435        if let Some(m) = mask {
436            inputs.push(m);
437        }
438        self.push_block(
439            HirOp::Attention {
440                num_heads,
441                head_dim,
442                mask: mask_kind,
443            },
444            inputs,
445            out_shape,
446            None,
447        )
448    }
449
450    /// Causal depthwise Conv1d — Conformer / Wav2Vec2-BERT conv module.
451    ///
452    /// `input` and `left_pad` are `[B, S, C]` / `[B, K-1, C]`; `weight` is
453    /// `[C, 1, 1, K]` in grouped Conv2d layout.
454    pub fn depthwise_conv1d_causal(
455        &mut self,
456        input: HirNodeId,
457        weight: HirNodeId,
458        left_pad: HirNodeId,
459        kernel_size: usize,
460        out_shape: Shape,
461    ) -> HirNodeId {
462        self.push_block(
463            HirOp::DepthwiseConv1dCausal { kernel_size },
464            vec![input, weight, left_pad],
465            out_shape,
466            None,
467        )
468    }
469
470    /// Fused dequant + matmul — see [`HirOp::DequantMatMul`].
471    pub fn dequant_matmul(
472        &mut self,
473        x: HirNodeId,
474        w: HirNodeId,
475        scale: Option<HirNodeId>,
476        zp: Option<HirNodeId>,
477        scheme: QuantScheme,
478        out_shape: Shape,
479    ) -> HirNodeId {
480        let mut inputs = vec![x, w];
481        if !scheme.is_gguf() {
482            inputs.push(scale.expect("DequantMatMul: scale required for non-GGUF schemes"));
483            inputs.push(zp.expect("DequantMatMul: zp required for non-GGUF schemes"));
484        }
485        self.push_block(HirOp::DequantMatMul { scheme }, inputs, out_shape, None)
486    }
487
488    /// Gated DeltaNet without carry state (prefill / reset per batch).
489    pub fn gated_delta_net(
490        &mut self,
491        q: HirNodeId,
492        k: HirNodeId,
493        v: HirNodeId,
494        g: HirNodeId,
495        beta: HirNodeId,
496        state_size: usize,
497        out_shape: Shape,
498    ) -> HirNodeId {
499        self.push_block(
500            HirOp::GatedDeltaNet {
501                state_size,
502                carry_state: false,
503            },
504            vec![q, k, v, g, beta],
505            out_shape,
506            None,
507        )
508    }
509
510    /// Gated DeltaNet with decode carry — threads `state` in/out.
511    pub fn gated_delta_net_carry(
512        &mut self,
513        q: HirNodeId,
514        k: HirNodeId,
515        v: HirNodeId,
516        g: HirNodeId,
517        beta: HirNodeId,
518        state: HirNodeId,
519        state_size: usize,
520        out_shape: Shape,
521    ) -> HirNodeId {
522        self.push_block(
523            HirOp::GatedDeltaNet {
524                state_size,
525                carry_state: true,
526            },
527            vec![q, k, v, g, beta, state],
528            out_shape,
529            None,
530        )
531    }
532
533    /// Multi-layer (optionally bidirectional) LSTM with packed weights.
534    /// `out_shape` = `[batch, seq, D*hidden]`. See [`crate::Op::Lstm`].
535    #[allow(clippy::too_many_arguments)]
536    pub fn lstm(
537        &mut self,
538        x: HirNodeId,
539        w_ih: HirNodeId,
540        w_hh: HirNodeId,
541        bias: HirNodeId,
542        hidden_size: usize,
543        num_layers: usize,
544        bidirectional: bool,
545        out_shape: Shape,
546    ) -> HirNodeId {
547        self.push_block(
548            HirOp::Lstm {
549                hidden_size,
550                num_layers,
551                bidirectional,
552                carry: false,
553            },
554            vec![x, w_ih, w_hh, bias],
555            out_shape,
556            None,
557        )
558    }
559
560    /// Multi-layer (optionally bidirectional) GRU. Inputs
561    /// `[x, w_ih, w_hh, b_ih, b_hh]`. `out_shape` = `[batch, seq, D*hidden]`.
562    /// See [`crate::Op::Gru`].
563    #[allow(clippy::too_many_arguments)]
564    pub fn gru(
565        &mut self,
566        x: HirNodeId,
567        w_ih: HirNodeId,
568        w_hh: HirNodeId,
569        b_ih: HirNodeId,
570        b_hh: HirNodeId,
571        hidden_size: usize,
572        num_layers: usize,
573        bidirectional: bool,
574        out_shape: Shape,
575    ) -> HirNodeId {
576        self.push_block(
577            HirOp::Gru {
578                hidden_size,
579                num_layers,
580                bidirectional,
581                carry: false,
582            },
583            vec![x, w_ih, w_hh, b_ih, b_hh],
584            out_shape,
585            None,
586        )
587    }
588
589    /// LSTM threading decode state: `h0`/`c0` `[L*D, batch, hidden]` in,
590    /// `hn`/`cn` written back in place.
591    #[allow(clippy::too_many_arguments)]
592    pub fn lstm_carry(
593        &mut self,
594        x: HirNodeId,
595        w_ih: HirNodeId,
596        w_hh: HirNodeId,
597        bias: HirNodeId,
598        h0: HirNodeId,
599        c0: HirNodeId,
600        hidden_size: usize,
601        num_layers: usize,
602        bidirectional: bool,
603        out_shape: Shape,
604    ) -> HirNodeId {
605        self.push_block(
606            HirOp::Lstm {
607                hidden_size,
608                num_layers,
609                bidirectional,
610                carry: true,
611            },
612            vec![x, w_ih, w_hh, bias, h0, c0],
613            out_shape,
614            None,
615        )
616    }
617
618    /// Rotary position embedding.
619    pub fn rope(
620        &mut self,
621        x: HirNodeId,
622        cos: HirNodeId,
623        sin: HirNodeId,
624        head_dim: usize,
625        n_rot: usize,
626        out_shape: Shape,
627    ) -> HirNodeId {
628        self.push_block(
629            HirOp::RoPE { head_dim, n_rot },
630            vec![x, cos, sin],
631            out_shape,
632            None,
633        )
634    }
635
636    /// RMS normalization (no residual add).
637    pub fn rms_norm(
638        &mut self,
639        x: HirNodeId,
640        gamma: HirNodeId,
641        beta: HirNodeId,
642        eps: f32,
643        out_shape: Shape,
644    ) -> HirNodeId {
645        self.push_block(
646            HirOp::RmsNorm { eps },
647            vec![x, gamma, beta],
648            out_shape,
649            None,
650        )
651    }
652
653    /// LLaMA / LLaMA-3.2 decoder layer (pre-norm GQA + SwiGLU).
654    pub fn llama_decoder_block(
655        &mut self,
656        x: HirNodeId,
657        ln1_g: HirNodeId,
658        ln1_b: HirNodeId,
659        q_w: HirNodeId,
660        k_w: HirNodeId,
661        v_w: HirNodeId,
662        o_w: HirNodeId,
663        ln2_g: HirNodeId,
664        ln2_b: HirNodeId,
665        gate_w: HirNodeId,
666        up_w: HirNodeId,
667        down_w: HirNodeId,
668        cos: HirNodeId,
669        sin: HirNodeId,
670        mask: Option<HirNodeId>,
671        num_heads: usize,
672        head_dim: usize,
673        num_kv_heads: usize,
674        eps: f32,
675        mask_kind: MaskKind,
676        rope_style: crate::op::RopeStyle,
677        out_shape: Shape,
678    ) -> HirNodeId {
679        let mut ins = vec![
680            x, ln1_g, ln1_b, q_w, k_w, v_w, o_w, ln2_g, ln2_b, gate_w, up_w, down_w, cos, sin,
681        ];
682        if let Some(m) = mask {
683            ins.push(m);
684        }
685        self.push_block(
686            HirOp::LlamaDecoderBlock {
687                num_heads,
688                head_dim,
689                num_kv_heads,
690                eps,
691                mask: mask_kind,
692                rope_style,
693            },
694            ins,
695            out_shape,
696            Some("llama_decoder_block".into()),
697        )
698    }
699
700    /// Standard pre-norm transformer decoder block — alias for
701    /// [`Self::llama_decoder_block`] (LLaMA / GPT-style layers).
702    pub fn transformer_block(
703        &mut self,
704        x: HirNodeId,
705        ln1_g: HirNodeId,
706        ln1_b: HirNodeId,
707        q_w: HirNodeId,
708        k_w: HirNodeId,
709        v_w: HirNodeId,
710        o_w: HirNodeId,
711        ln2_g: HirNodeId,
712        ln2_b: HirNodeId,
713        gate_w: HirNodeId,
714        up_w: HirNodeId,
715        down_w: HirNodeId,
716        cos: HirNodeId,
717        sin: HirNodeId,
718        mask: Option<HirNodeId>,
719        num_heads: usize,
720        head_dim: usize,
721        num_kv_heads: usize,
722        eps: f32,
723        mask_kind: MaskKind,
724        out_shape: Shape,
725    ) -> HirNodeId {
726        let id = self.llama_decoder_block(
727            x,
728            ln1_g,
729            ln1_b,
730            q_w,
731            k_w,
732            v_w,
733            o_w,
734            ln2_g,
735            ln2_b,
736            gate_w,
737            up_w,
738            down_w,
739            cos,
740            sin,
741            mask,
742            num_heads,
743            head_dim,
744            num_kv_heads,
745            eps,
746            mask_kind,
747            // Generic transformer alias keeps the HF/NeoX rotate-half flavor.
748            crate::op::RopeStyle::NeoX,
749            out_shape,
750        );
751        self.node_mut(id).name = Some("transformer_block".into());
752        id
753    }
754
755    /// Qwen3.5 MTP draft head — see [`blocks::lower_qwen35_mtp_head`].
756    #[allow(clippy::too_many_arguments)]
757    pub fn qwen35_mtp_head(
758        &mut self,
759        h_pre_norm: HirNodeId,
760        input_ids: HirNodeId,
761        cos: HirNodeId,
762        sin: HirNodeId,
763        last_token_idx: HirNodeId,
764        embed_w: HirNodeId,
765        hnorm_w: HirNodeId,
766        hnorm_b: HirNodeId,
767        enorm_w: HirNodeId,
768        enorm_b: HirNodeId,
769        eh_w: HirNodeId,
770        fa_attn_norm_w: HirNodeId,
771        fa_attn_norm_b: HirNodeId,
772        fa_q_gate_w: HirNodeId,
773        fa_k_w: HirNodeId,
774        fa_v_w: HirNodeId,
775        fa_q_norm_w: HirNodeId,
776        fa_q_norm_b: HirNodeId,
777        fa_k_norm_w: HirNodeId,
778        fa_k_norm_b: HirNodeId,
779        fa_o_w: HirNodeId,
780        fa_post_norm_w: HirNodeId,
781        fa_post_norm_b: HirNodeId,
782        fa_gate_w: HirNodeId,
783        fa_up_w: HirNodeId,
784        fa_down_w: HirNodeId,
785        head_norm_w: HirNodeId,
786        head_norm_b: HirNodeId,
787        lm_head_w: HirNodeId,
788        num_heads: usize,
789        num_kv_heads: usize,
790        head_dim: usize,
791        n_rot: usize,
792        n_embd: usize,
793        n_ff: usize,
794        mtp_vocab: usize,
795        eps: f32,
796        out_shape: Shape,
797    ) -> HirNodeId {
798        self.push_block(
799            HirOp::Qwen35MtpHead {
800                num_heads,
801                num_kv_heads,
802                head_dim,
803                n_rot,
804                n_embd,
805                n_ff,
806                mtp_vocab,
807                eps,
808            },
809            vec![
810                h_pre_norm,
811                input_ids,
812                cos,
813                sin,
814                last_token_idx,
815                embed_w,
816                hnorm_w,
817                hnorm_b,
818                enorm_w,
819                enorm_b,
820                eh_w,
821                fa_attn_norm_w,
822                fa_attn_norm_b,
823                fa_q_gate_w,
824                fa_k_w,
825                fa_v_w,
826                fa_q_norm_w,
827                fa_q_norm_b,
828                fa_k_norm_w,
829                fa_k_norm_b,
830                fa_o_w,
831                fa_post_norm_w,
832                fa_post_norm_b,
833                fa_gate_w,
834                fa_up_w,
835                fa_down_w,
836                head_norm_w,
837                head_norm_b,
838                lm_head_w,
839            ],
840            out_shape,
841            Some("qwen35_mtp_head".into()),
842        )
843    }
844
845    /// Escape hatch — embed a single MIR [`Op`] verbatim.
846    pub fn mir(&mut self, op: Op, inputs: Vec<HirNodeId>, shape: Shape) -> HirNodeId {
847        self.push(HirOp::Mir(op), inputs, shape, None)
848    }
849
850    pub fn set_outputs(&mut self, outputs: Vec<HirNodeId>) {
851        self.outputs = outputs;
852    }
853
854    /// Lower this module to MIR.
855    pub fn lower_to_mir(self) -> Result<MirModule, LowerError> {
856        lower::lower_module(self)
857    }
858
859    /// Lower with [`FusionPolicy::for_autodiff`] — primitive MIR chains
860    /// that need less unfuse work before `rlx_opt::prepare_graph_for_ad`.
861    pub fn lower_for_autodiff(self) -> Result<MirModule, LowerError> {
862        self.with_fusion_policy(FusionPolicy::for_autodiff())
863            .lower_to_mir()
864    }
865
866    /// Wrap an existing MIR [`Graph`] as a HIR module (`HirOp::Mir` per node).
867    /// Enables `Session::compile_hir` for legacy graph builders during migration.
868    pub fn wrap_mir_graph(graph: crate::Graph) -> Self {
869        use std::collections::HashMap;
870        let mut hir = Self::new(graph.name.clone()).with_fusion_policy(FusionPolicy::Direct);
871        let mut map: HashMap<crate::NodeId, HirNodeId> = HashMap::new();
872        for node in graph.nodes() {
873            let inputs: Vec<HirNodeId> = node.inputs.iter().map(|&id| map[&id]).collect();
874            let id = hir.mir(node.op.clone(), inputs, node.shape.clone());
875            map.insert(node.id, id);
876        }
877        let outputs: Vec<HirNodeId> = graph.outputs.iter().map(|&id| map[&id]).collect();
878        hir.set_outputs(outputs);
879        hir
880    }
881}
882
883pub(crate) fn default_hir_block_label(op: &HirOp) -> Option<String> {
884    Some(match op {
885        HirOp::Linear { .. } => "linear".into(),
886        HirOp::LinearFused { .. } => "linear_fused".into(),
887        HirOp::SharedLinearPair { slot } => return Some(format!("shared_linear_pair[{slot}]")),
888        HirOp::SwiGLU => "swiglu_ffn".into(),
889        HirOp::ResidualRmsNorm { .. } => "residual_rms_norm".into(),
890        HirOp::Attention { .. } => "attention".into(),
891        HirOp::DepthwiseConv1dCausal { .. } => "depthwise_conv1d_causal".into(),
892        HirOp::DequantMatMul { scheme } => format!("dequant_matmul({scheme})"),
893        HirOp::GatedDeltaNet {
894            carry_state: true, ..
895        } => "gated_delta_net_carry".into(),
896        HirOp::GatedDeltaNet { .. } => "gated_delta_net".into(),
897        HirOp::Lstm { .. } => "lstm".into(),
898        HirOp::Gru { .. } => "gru".into(),
899        HirOp::RoPE { .. } => "rope".into(),
900        HirOp::RmsNorm { .. } => "rms_norm".into(),
901        HirOp::Mir(_) => "mir".into(),
902        HirOp::LlamaDecoderBlock { .. } => "llama_decoder_block".into(),
903        HirOp::Qwen35MtpHead { .. } => "qwen35_mtp_head".into(),
904        HirOp::Input { .. } | HirOp::Param { .. } | HirOp::Constant { .. } => return None,
905    })
906}
907
908impl std::fmt::Display for HirModule {
909    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
910        writeln!(f, "hir @{} {{", self.name)?;
911        for node in &self.nodes {
912            write!(f, "  {} = {:?}", node.id, node.op)?;
913            if !node.inputs.is_empty() {
914                write!(f, "(")?;
915                for (i, inp) in node.inputs.iter().enumerate() {
916                    if i > 0 {
917                        write!(f, ", ")?;
918                    }
919                    write!(f, "{inp}")?;
920                }
921                write!(f, ")")?;
922            }
923            writeln!(f, " : {}", node.shape)?;
924        }
925        if !self.outputs.is_empty() {
926            write!(f, "  return ")?;
927            for (i, o) in self.outputs.iter().enumerate() {
928                if i > 0 {
929                    write!(f, ", ")?;
930                }
931                write!(f, "{o}")?;
932            }
933            writeln!(f)?;
934        }
935        write!(f, "}}")
936    }
937}
938
939#[cfg(test)]
940mod tests {
941    use super::*;
942    use crate::DType;
943
944    fn f32_shape(d: &[usize]) -> Shape {
945        Shape::new(d, DType::F32)
946    }
947
948    #[test]
949    fn hir_depthwise_conv1d_causal_lowers_to_grouped_conv() {
950        use crate::Op;
951
952        let mut hir = HirModule::new("dw");
953        let x = hir.input("x", f32_shape(&[2, 8, 16]));
954        let w = hir.param("w", f32_shape(&[16, 1, 1, 3]));
955        let pad = hir.param("pad", f32_shape(&[2, 2, 16]));
956        let out = hir.depthwise_conv1d_causal(x, w, pad, 3, f32_shape(&[2, 8, 16]));
957        hir.outputs = vec![out];
958
959        let g = hir.lower_to_mir().expect("lower").into_graph();
960        assert!(g.nodes().iter().any(|n| matches!(n.op, Op::Conv { .. })));
961        assert!(g.nodes().iter().any(|n| matches!(n.op, Op::Concat { .. })));
962    }
963
964    #[test]
965    fn hir_swiglu_lowers_to_fusable_mir() {
966        use crate::Op;
967        use crate::hir::FusionPolicy;
968
969        let mut hir = HirModule::new("ffn").with_fusion_policy(FusionPolicy::Fusable);
970        let x = hir.input("x", f32_shape(&[4, 768]));
971        let up_w = hir.param("up", f32_shape(&[768, 2048]));
972        let gate_w = hir.param("gate", f32_shape(&[768, 2048]));
973        let down_w = hir.param("down", f32_shape(&[2048, 768]));
974        let out = hir.swiglu_ffn(x, up_w, gate_w, down_w, f32_shape(&[4, 768]));
975        hir.set_outputs(vec![out]);
976
977        let mir = hir.lower_to_mir().expect("lower");
978        let g = mir.into_graph();
979        assert!(g.nodes().iter().any(|n| matches!(n.op, Op::MatMul)));
980        assert_eq!(g.len(), 9);
981    }
982
983    #[test]
984    fn label_nodes_since_flows_source_name_to_all_mir_nodes() {
985        // Mirrors the importer path: build a few raw `mir` ops (which default
986        // to the generic "mir" label), stamp them with the source-op name,
987        // then lower and confirm EVERY derived MIR node localizes back to that
988        // name instead of "mir" — this is exactly what the NaN localizer reads.
989        use crate::hir::HirMut;
990        use crate::op::Activation;
991        use crate::{Op, node_label};
992
993        let mut hir = HirModule::new("m");
994        {
995            let mut b = HirMut::new(&mut hir);
996            let x = b.input("x", f32_shape(&[4]));
997            let from = b.0.len(); // stamp everything created after the input
998            let e = b.activation(Activation::Exp, x, f32_shape(&[4]));
999            let two = b.add(e, e);
1000            let out = b.div(two, e);
1001            b.0.label_nodes_since(from, "Softmax_7");
1002            b.set_outputs(vec![out]);
1003        }
1004        let g = hir.lower_to_mir().expect("lower").into_graph();
1005
1006        let mut compute = 0;
1007        for node in g.nodes() {
1008            match node.op {
1009                Op::Input { .. } | Op::Param { .. } | Op::Constant { .. } => {}
1010                _ => {
1011                    compute += 1;
1012                    assert_eq!(
1013                        node_label(&g, node.id),
1014                        "Softmax_7",
1015                        "node {} {:?} lost source provenance",
1016                        node.id,
1017                        node.op
1018                    );
1019                }
1020            }
1021        }
1022        assert!(
1023            compute >= 3,
1024            "expected exp/add/div compute nodes, got {compute}"
1025        );
1026
1027        // The input keeps its own identity, not the source-op label.
1028        let x_mir = g.input_id("x").expect("input x present");
1029        assert_eq!(node_label(&g, x_mir), "x");
1030    }
1031
1032    #[test]
1033    fn hir_gdn_dequant_rope_rms_lowers() {
1034        use crate::Op;
1035        use crate::quant::QuantScheme;
1036
1037        let mut hir = HirModule::new("qwen_block");
1038        let q = hir.input("q", f32_shape(&[1, 4, 2, 8]));
1039        let k = hir.param("k", f32_shape(&[1, 4, 2, 8]));
1040        let v = hir.param("v", f32_shape(&[1, 4, 2, 8]));
1041        let g_in = hir.param("g", f32_shape(&[1, 4, 2]));
1042        let beta = hir.param("beta", f32_shape(&[1, 4, 2]));
1043        let scan = hir.gated_delta_net(q, k, v, g_in, beta, 8, f32_shape(&[1, 4, 2, 8]));
1044
1045        let cos = hir.param("cos", f32_shape(&[1, 4, 8]));
1046        let sin = hir.param("sin", f32_shape(&[1, 4, 8]));
1047        let x = hir.input("x", f32_shape(&[1, 4, 8]));
1048        let rotated = hir.rope(x, cos, sin, 8, 8, f32_shape(&[1, 4, 8]));
1049
1050        let gamma = hir.param("gamma", f32_shape(&[8]));
1051        let beta_n = hir.param("beta_n", f32_shape(&[8]));
1052        let normed = hir.rms_norm(rotated, gamma, beta_n, 1e-6, f32_shape(&[1, 4, 8]));
1053
1054        let x_in = hir.input("hidden", f32_shape(&[4, 128]));
1055        let w = hir.param("w_q", f32_shape(&[1024]));
1056        let proj = hir.dequant_matmul(
1057            x_in,
1058            w,
1059            None,
1060            None,
1061            QuantScheme::GgufQ4K,
1062            f32_shape(&[4, 128]),
1063        );
1064        hir.set_outputs(vec![scan, normed, proj]);
1065
1066        let g = hir.lower_to_mir().expect("lower").into_graph();
1067        assert!(g.nodes().iter().any(|n| matches!(
1068            n.op,
1069            Op::GatedDeltaNet {
1070                carry_state: false,
1071                ..
1072            }
1073        )));
1074        assert!(g.nodes().iter().any(|n| matches!(n.op, Op::Rope { .. })));
1075        assert!(g.nodes().iter().any(|n| matches!(n.op, Op::RmsNorm { .. })));
1076        assert!(g.nodes().iter().any(|n| matches!(
1077            n.op,
1078            Op::DequantMatMul {
1079                scheme: QuantScheme::GgufQ4K
1080            }
1081        )));
1082    }
1083}
rlx_ir/hir/mod.rs

rlx_ir/hir/
mod.rs
