Struct HirModule 

pub struct HirModule {
    pub name: String,
    pub outputs: Vec<HirNodeId>,
    pub fusion_policy: FusionPolicy,
    /* private fields */
}

High-level module — model builder output.

Fields

name: String

outputs: Vec<HirNodeId>

fusion_policy: FusionPolicy

How block ops lower to MIR. Default: FusionPolicy::Direct for new model code (fusion as a first-class citizen).

Implementations

impl HirModule

pub fn new(name: impl Into<String>) -> Self

pub fn with_fusion_policy(self, policy: FusionPolicy) -> Self

pub fn len(&self) -> usize

pub fn is_empty(&self) -> bool

pub fn nodes(&self) -> &[HirNode]

pub fn node(&self, id: HirNodeId) -> &HirNode

pub fn node_mut(&mut self, id: HirNodeId) -> &mut HirNode

pub fn named( &mut self, name: impl Into<String>, build: impl FnOnce(&mut Self) -> HirNodeId, ) -> HirNodeId

Build a named block — sets HirNode::name on the returned node.

pub fn label_nodes_since(&mut self, from: usize, name: &str)

Stamp a source-op provenance name onto every node created at or after index from whose label is only the generic op default (or unset). Importers call this after lowering one source op (an ONNX node, a torch aten call) so all the HIR nodes it expanded into share that op’s identity — the label then flows to every derived MIR node via tag_hir_subgraph, so a NaN in any of them localizes back to the user’s original op instead of a generic "mir".

Explicitly-meaningful names are preserved: Input/Param nodes have no default label, so their weight names (which differ from None) are left untouched even if they fall inside the range.

pub fn input(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId

pub fn input_batch_seq( &mut self, name: impl Into<String>, batch: u32, seq: u32, hidden: usize, dtype: DType, ) -> HirNodeId

[batch, seq, hidden] input with symbolic leading axes.

pub fn param(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId

pub fn linear( &mut self, x: HirNodeId, weight: HirNodeId, bias: Option<HirNodeId>, activation: Option<Activation>, out_shape: Shape, ) -> HirNodeId

pub fn linear_fused( &mut self, x: HirNodeId, weight: HirNodeId, bias: HirNodeId, activation: Option<Activation>, out_shape: Shape, ) -> HirNodeId

Emit HirOp::LinearFused — fused matmul+bias+act at MIR level.

pub fn shared_linear_pair( &mut self, x: HirNodeId, w_first: HirNodeId, w_second: HirNodeId, out_shape: Shape, ) -> (HirNodeId, HirNodeId)

Two matmuls sharing x. Returns (first, second) in weight order.

pub fn swiglu_ffn( &mut self, x: HirNodeId, up_w: HirNodeId, gate_w: HirNodeId, down_w: HirNodeId, out_shape: Shape, ) -> HirNodeId

pub fn residual_rms_norm( &mut self, x: HirNodeId, residual: HirNodeId, gamma: HirNodeId, beta: HirNodeId, eps: f32, out_shape: Shape, ) -> HirNodeId

pub fn attention( &mut self, q: HirNodeId, k: HirNodeId, v: HirNodeId, mask: Option<HirNodeId>, num_heads: usize, head_dim: usize, mask_kind: MaskKind, out_shape: Shape, ) -> HirNodeId

Scaled dot-product attention — see HirOp::Attention.

pub fn depthwise_conv1d_causal( &mut self, input: HirNodeId, weight: HirNodeId, left_pad: HirNodeId, kernel_size: usize, out_shape: Shape, ) -> HirNodeId

Causal depthwise Conv1d — Conformer / Wav2Vec2-BERT conv module.

input and left_pad are [B, S, C] / [B, K-1, C]; weight is [C, 1, 1, K] in grouped Conv2d layout.

pub fn dequant_matmul( &mut self, x: HirNodeId, w: HirNodeId, scale: Option<HirNodeId>, zp: Option<HirNodeId>, scheme: QuantScheme, out_shape: Shape, ) -> HirNodeId

Fused dequant + matmul — see HirOp::DequantMatMul.

pub fn gated_delta_net( &mut self, q: HirNodeId, k: HirNodeId, v: HirNodeId, g: HirNodeId, beta: HirNodeId, state_size: usize, out_shape: Shape, ) -> HirNodeId

Gated DeltaNet without carry state (prefill / reset per batch).

pub fn gated_delta_net_carry( &mut self, q: HirNodeId, k: HirNodeId, v: HirNodeId, g: HirNodeId, beta: HirNodeId, state: HirNodeId, state_size: usize, out_shape: Shape, ) -> HirNodeId

Gated DeltaNet with decode carry — threads state in/out.

pub fn lstm( &mut self, x: HirNodeId, w_ih: HirNodeId, w_hh: HirNodeId, bias: HirNodeId, hidden_size: usize, num_layers: usize, bidirectional: bool, out_shape: Shape, ) -> HirNodeId

Multi-layer (optionally bidirectional) LSTM with packed weights. out_shape = [batch, seq, D*hidden]. See crate::Op::Lstm.

pub fn gru( &mut self, x: HirNodeId, w_ih: HirNodeId, w_hh: HirNodeId, b_ih: HirNodeId, b_hh: HirNodeId, hidden_size: usize, num_layers: usize, bidirectional: bool, out_shape: Shape, ) -> HirNodeId

Multi-layer (optionally bidirectional) GRU. Inputs [x, w_ih, w_hh, b_ih, b_hh]. out_shape = [batch, seq, D*hidden]. See crate::Op::Gru.

pub fn lstm_carry( &mut self, x: HirNodeId, w_ih: HirNodeId, w_hh: HirNodeId, bias: HirNodeId, h0: HirNodeId, c0: HirNodeId, hidden_size: usize, num_layers: usize, bidirectional: bool, out_shape: Shape, ) -> HirNodeId

LSTM threading decode state: h0/c0 [L*D, batch, hidden] in, hn/cn written back in place.

pub fn rope( &mut self, x: HirNodeId, cos: HirNodeId, sin: HirNodeId, head_dim: usize, n_rot: usize, out_shape: Shape, ) -> HirNodeId

Rotary position embedding.

pub fn rms_norm( &mut self, x: HirNodeId, gamma: HirNodeId, beta: HirNodeId, eps: f32, out_shape: Shape, ) -> HirNodeId

RMS normalization (no residual add).

pub fn llama_decoder_block( &mut self, x: HirNodeId, ln1_g: HirNodeId, ln1_b: HirNodeId, q_w: HirNodeId, k_w: HirNodeId, v_w: HirNodeId, o_w: HirNodeId, ln2_g: HirNodeId, ln2_b: HirNodeId, gate_w: HirNodeId, up_w: HirNodeId, down_w: HirNodeId, cos: HirNodeId, sin: HirNodeId, mask: Option<HirNodeId>, num_heads: usize, head_dim: usize, num_kv_heads: usize, eps: f32, mask_kind: MaskKind, rope_style: RopeStyle, out_shape: Shape, ) -> HirNodeId

LLaMA / LLaMA-3.2 decoder layer (pre-norm GQA + SwiGLU).

pub fn transformer_block( &mut self, x: HirNodeId, ln1_g: HirNodeId, ln1_b: HirNodeId, q_w: HirNodeId, k_w: HirNodeId, v_w: HirNodeId, o_w: HirNodeId, ln2_g: HirNodeId, ln2_b: HirNodeId, gate_w: HirNodeId, up_w: HirNodeId, down_w: HirNodeId, cos: HirNodeId, sin: HirNodeId, mask: Option<HirNodeId>, num_heads: usize, head_dim: usize, num_kv_heads: usize, eps: f32, mask_kind: MaskKind, out_shape: Shape, ) -> HirNodeId

Standard pre-norm transformer decoder block — alias for Self::llama_decoder_block (LLaMA / GPT-style layers).

pub fn qwen35_mtp_head( &mut self, h_pre_norm: HirNodeId, input_ids: HirNodeId, cos: HirNodeId, sin: HirNodeId, last_token_idx: HirNodeId, embed_w: HirNodeId, hnorm_w: HirNodeId, hnorm_b: HirNodeId, enorm_w: HirNodeId, enorm_b: HirNodeId, eh_w: HirNodeId, fa_attn_norm_w: HirNodeId, fa_attn_norm_b: HirNodeId, fa_q_gate_w: HirNodeId, fa_k_w: HirNodeId, fa_v_w: HirNodeId, fa_q_norm_w: HirNodeId, fa_q_norm_b: HirNodeId, fa_k_norm_w: HirNodeId, fa_k_norm_b: HirNodeId, fa_o_w: HirNodeId, fa_post_norm_w: HirNodeId, fa_post_norm_b: HirNodeId, fa_gate_w: HirNodeId, fa_up_w: HirNodeId, fa_down_w: HirNodeId, head_norm_w: HirNodeId, head_norm_b: HirNodeId, lm_head_w: HirNodeId, num_heads: usize, num_kv_heads: usize, head_dim: usize, n_rot: usize, n_embd: usize, n_ff: usize, mtp_vocab: usize, eps: f32, out_shape: Shape, ) -> HirNodeId

Qwen3.5 MTP draft head — see blocks::lower_qwen35_mtp_head.

pub fn mir(&mut self, op: Op, inputs: Vec<HirNodeId>, shape: Shape) -> HirNodeId

Escape hatch — embed a single MIR Op verbatim.

pub fn set_outputs(&mut self, outputs: Vec<HirNodeId>)

pub fn lower_to_mir(self) -> Result<MirModule, LowerError>

Lower this module to MIR.

pub fn lower_for_autodiff(self) -> Result<MirModule, LowerError>

Lower with FusionPolicy::for_autodiff — primitive MIR chains that need less unfuse work before rlx_opt::prepare_graph_for_ad.

pub fn wrap_mir_graph(graph: Graph) -> Self

Wrap an existing MIR [Graph] as a HIR module (HirOp::Mir per node). Enables Session::compile_hir for legacy graph builders during migration.

impl HirModule

pub fn inspect(&self) -> String

Text dump for inspection. Alias for inspect_hir.

Trait Implementations

impl Clone for HirModule

fn clone(&self) -> HirModule

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for HirModule

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Display for HirModule

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl From<HirModule> for GraphModule

fn from(hir: HirModule) -> Self

Converts to this type from the input type.