vyre_lower/
pre_emit.rs

1//! Canonical pre-emit lowering pipeline.
2//!
3//! This is the single production boundary from high-level `Program` IR to
4//! emitter-ready `KernelDescriptor`: inline calls, run semantic Program
5//! optimization, lower to descriptor form, verify, run descriptor cleanup, and
6//! verify again. Backends should not assemble their own partial version of
7//! this sequence.
8
9use crate::descriptor::KernelDescriptor;
10use crate::lower::lower;
11use crate::rewrites::OptimizationStats;
12use crate::{verify_then_optimize, VerifyFailure};
13use std::fmt;
14use vyre_foundation::ir::Program;
15
16/// Program + descriptor pair produced by the canonical pre-emit pipeline.
17#[derive(Debug, Clone)]
18pub struct LoweredKernel {
19    /// Program after call inlining and IR-semantic optimization.
20    pub program: Program,
21    /// Verified descriptor after descriptor-level cleanup rewrites.
22    pub descriptor: KernelDescriptor,
23    /// Descriptor rewrite statistics collected from the cleanup phase.
24    pub descriptor_stats: OptimizationStats,
25}
26
27/// Error raised by the canonical pre-emit pipeline.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct PreEmitError {
30    message: String,
31}
32
33impl PreEmitError {
34    fn new(message: impl Into<String>) -> Self {
35        let message = message.into();
36        debug_assert!(message.contains("Fix:"));
37        Self { message }
38    }
39
40    /// Return the actionable diagnostic.
41    #[must_use]
42    pub fn message(&self) -> &str {
43        &self.message
44    }
45}
46
47impl fmt::Display for PreEmitError {
48    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
49        formatter.write_str(&self.message)
50    }
51}
52
53impl std::error::Error for PreEmitError {}
54
55/// Inline calls and run the semantic Program-level optimizer.
56///
57/// This prepares high-level IR for descriptor lowering while preserving the
58/// distinction between Layer-1 semantic rewrites and lowered descriptor
59/// cleanup.
60///
61/// # Errors
62///
63/// Returns [`PreEmitError`] when call inlining fails.
64pub fn prepare_program_for_emit(program: &Program) -> Result<Program, PreEmitError> {
65    let pruned = vyre_foundation::optimizer::pre_lowering::optimize(program.clone());
66    let pruned = lower_single_rank_collectives_for_emit(pruned)?;
67    let inlined = vyre_foundation::ir::inline_calls(&pruned).map_err(|error| {
68        PreEmitError::new(format!(
69            "call inlining failed before descriptor lowering: {error}. Fix: register every Expr::Call target with the active dialect resolver or eliminate the call before backend emission."
70        ))
71    })?;
72    lower_single_rank_collectives_for_emit(vyre_foundation::optimizer::pre_lowering::optimize(
73        inlined,
74    ))
75}
76
77fn lower_single_rank_collectives_for_emit(program: Program) -> Result<Program, PreEmitError> {
78    match vyre_foundation::transform::collectives::lower_single_rank_collectives(&program) {
79        Ok(Some(lowered)) => Ok(lowered),
80        Ok(None) => Ok(program),
81        Err(error) => Err(PreEmitError::new(format!(
82            "single-rank collective lowering failed before descriptor lowering: {error}. Fix: route true multi-rank collectives through a backend transport path or lower them before pre-emit."
83        ))),
84    }
85}
86
87/// Run the complete canonical pre-emit pipeline.
88///
89/// # Errors
90///
91/// Returns [`PreEmitError`] when inlining, descriptor lowering, input
92/// verification, descriptor cleanup, or output verification fails.
93pub fn lower_for_emit(program: &Program) -> Result<LoweredKernel, PreEmitError> {
94    let program = prepare_program_for_emit(program)?;
95    let descriptor = lower(&program).map_err(|error| {
96        PreEmitError::new(format!(
97            "KernelDescriptor lowering failed after semantic Program optimization: {error}. Fix: add the missing neutral descriptor mapping before any concrete backend emits this Program."
98        ))
99    })?;
100    let (descriptor, descriptor_stats) = verify_then_optimize(&descriptor).map_err(|error| {
101        PreEmitError::new(format!(
102            "KernelDescriptor verification/cleanup failed in the canonical pre-emit pipeline: {}. Fix: repair vyre-lower so descriptor validation succeeds before concrete emission.",
103            format_verify_failure(&error)
104        ))
105    })?;
106    Ok(LoweredKernel {
107        program,
108        descriptor,
109        descriptor_stats,
110    })
111}
112
113fn format_verify_failure(error: &VerifyFailure) -> String {
114    let stage = match error {
115        VerifyFailure::Input(_) => "input",
116        VerifyFailure::Output(_) => "output",
117    };
118    let mut out = format!("{stage} descriptor invalid");
119    for (index, err) in error.errors().iter().take(4).enumerate() {
120        if index == 0 {
121            out.push_str(": ");
122        } else {
123            out.push_str("; ");
124        }
125        out.push_str(&format!("{err:?}"));
126    }
127    if error.errors().len() > 4 {
128        out.push_str("; ...");
129    }
130    out
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use crate::{KernelBody, KernelOpKind};
137    use vyre_foundation::ir::{
138        BufferAccess, BufferDecl, CollectiveOp, CommGroup, DataType, Expr, Ident, Node,
139    };
140
141    #[test]
142    fn lower_for_emit_runs_program_and_descriptor_pipeline() {
143        let buffer =
144            BufferDecl::storage("out", 0, BufferAccess::ReadWrite, DataType::U32).with_count(16);
145        let program = Program::wrapped(
146            vec![buffer],
147            [64, 1, 1],
148            vec![Node::Store {
149                buffer: Ident::from("out"),
150                index: Expr::InvocationId { axis: 0 },
151                value: Expr::LitU32(7),
152            }],
153        );
154
155        let lowered = lower_for_emit(&program).expect("Fix: pre-emit lowering must pass");
156
157        assert_eq!(lowered.program.workgroup_size(), [64, 1, 1]);
158        assert_eq!(lowered.descriptor.dispatch.workgroup_size, [64, 1, 1]);
159        assert_eq!(lowered.descriptor.bindings.slots.len(), 1);
160        assert!(crate::verify::verify(&lowered.descriptor).is_ok());
161        assert!(lowered.descriptor_stats.iterations >= 1);
162    }
163
164    #[test]
165    fn lower_for_emit_rejects_invalid_descriptor_before_backend_emit() {
166        let program = Program::wrapped(Vec::new(), [0, 1, 1], Vec::new());
167
168        let error = lower_for_emit(&program).expect_err("zero dispatch must fail");
169
170        assert!(error.message().contains("KernelDescriptor"));
171        assert!(error.message().contains("Fix:"));
172    }
173
174    #[test]
175    fn lower_for_emit_lowers_world_allgather_before_descriptor_lowering() {
176        let program = Program::wrapped(
177            vec![
178                BufferDecl::read("input", 0, DataType::U32).with_count(4),
179                BufferDecl::output("out", 1, DataType::U32).with_count(4),
180            ],
181            [64, 1, 1],
182            vec![Node::AllGather {
183                input: "input".into(),
184                output: "out".into(),
185                group: CommGroup::WORLD,
186            }],
187        );
188
189        let lowered = lower_for_emit(&program).expect(
190            "Fix: canonical pre-emit must lower WORLD AllGather before descriptor lowering.",
191        );
192
193        assert!(!lowered.program.stats().distributed_collectives());
194        assert!(crate::verify::verify(&lowered.descriptor).is_ok());
195    }
196
197    #[test]
198    fn lower_for_emit_rejects_transport_collectives_before_descriptor_lowering() {
199        let program = Program::wrapped(
200            vec![
201                BufferDecl::read("input", 0, DataType::U32).with_count(4),
202                BufferDecl::output("out", 1, DataType::U32).with_count(4),
203            ],
204            [64, 1, 1],
205            vec![Node::ReduceScatter {
206                input: "input".into(),
207                output: "out".into(),
208                op: CollectiveOp::Sum,
209                group: CommGroup(7),
210            }],
211        );
212
213        let error = lower_for_emit(&program)
214            .expect_err("Fix: canonical pre-emit must reject collectives that need transport.");
215
216        assert!(error.message().contains("Multi-rank collective transport"));
217    }
218
219    #[test]
220    fn concrete_backend_program_descriptor_lowering_is_single_sourced() {
221        use std::fmt::Write as _;
222
223        let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
224            .parent()
225            .expect("Fix: vyre-lower must live under the vyre workspace root.");
226        let mut files = Vec::default();
227        for rel in [
228            "vyre-driver/src",
229            "vyre-driver-cuda/src",
230            "vyre-driver-metal/src",
231            "vyre-driver-wgpu/src",
232            "vyre-emit-naga/src",
233            "vyre-emit-ptx/src",
234        ] {
235            collect_rust_files(&root.join(rel), &mut files)
236                .expect("Fix: boundary scan must read concrete backend source trees.");
237        }
238
239        let mut violations = Vec::default();
240        for path in files {
241            let text = std::fs::read_to_string(&path)
242                .expect("Fix: boundary scan must read Rust source files.");
243            for (line_no, line) in text.lines().enumerate() {
244                let trimmed = line.trim_start();
245                if trimmed.starts_with("//") {
246                    continue;
247                }
248                if trimmed.starts_with("#[cfg(test)]") {
249                    break;
250                }
251                if line.contains("vyre_lower::lower(")
252                    || line.contains("vyre_lower::prepare_program_for_emit(")
253                    || line.contains("vyre_foundation::optimizer::pre_lowering::optimize(")
254                    || line.contains("vyre_foundation::ir::inline_calls(")
255                {
256                    let mut violation = String::default();
257                    let _ = write!(
258                        &mut violation,
259                        "{}:{}:{}",
260                        path.strip_prefix(root).unwrap_or(&path).display(),
261                        line_no + 1,
262                        trimmed
263                    );
264                    violations.push(violation);
265                }
266            }
267        }
268
269        assert!(
270            violations.is_empty(),
271            "Fix: concrete backend Program input must enter through vyre_lower::lower_for_emit or consume a typed KernelDescriptor artifact, not raw lowering steps: {violations:?}"
272        );
273    }
274
275    fn collect_rust_files(
276        dir: &std::path::Path,
277        out: &mut Vec<std::path::PathBuf>,
278    ) -> std::io::Result<()> {
279        for entry in std::fs::read_dir(dir)? {
280            let entry = entry?;
281            let path = entry.path();
282            if path.is_dir() {
283                collect_rust_files(&path, out)?;
284            } else if path.extension().is_some_and(|extension| extension == "rs") {
285                out.push(path);
286            }
287        }
288        Ok(())
289    }
290
291    #[test]
292    fn lower_for_emit_preserves_loop_carrier_swap_snapshot() {
293        let program = Program::wrapped(
294            vec![
295                BufferDecl::storage("instrs", 0, BufferAccess::ReadOnly, DataType::U32)
296                    .with_count(7),
297                BufferDecl::output("results", 1, DataType::U32).with_count(1),
298            ],
299            [256, 1, 1],
300            vec![
301                Node::let_bind("tid", Expr::gid_x()),
302                Node::if_then(
303                    Expr::lt(Expr::var("tid"), Expr::u32(1)),
304                    vec![
305                        Node::let_bind("base", Expr::mul(Expr::var("tid"), Expr::u32(7))),
306                        Node::let_bind("s0", Expr::u32(0)),
307                        Node::let_bind("s1", Expr::u32(0)),
308                        Node::let_bind("s2", Expr::u32(0)),
309                        Node::let_bind("s3", Expr::u32(0)),
310                        Node::Loop {
311                            var: "pc".into(),
312                            from: Expr::u32(0),
313                            to: Expr::u32(7),
314                            body: vec![
315                                Node::let_bind(
316                                    "instr",
317                                    Expr::load(
318                                        "instrs",
319                                        Expr::add(Expr::var("base"), Expr::var("pc")),
320                                    ),
321                                ),
322                                Node::let_bind(
323                                    "op",
324                                    Expr::bitand(Expr::var("instr"), Expr::u32(0xFF)),
325                                ),
326                                Node::let_bind("imm", Expr::shr(Expr::var("instr"), Expr::u32(8))),
327                                Node::if_then(
328                                    Expr::eq(Expr::var("op"), Expr::u32(0)),
329                                    vec![
330                                        Node::assign("s3", Expr::var("s2")),
331                                        Node::assign("s2", Expr::var("s1")),
332                                        Node::assign("s1", Expr::var("s0")),
333                                        Node::assign("s0", Expr::var("imm")),
334                                    ],
335                                ),
336                                Node::if_then(
337                                    Expr::eq(Expr::var("op"), Expr::u32(1)),
338                                    vec![
339                                        Node::assign(
340                                            "s0",
341                                            Expr::add(Expr::var("s0"), Expr::var("s1")),
342                                        ),
343                                        Node::assign("s1", Expr::var("s2")),
344                                        Node::assign("s2", Expr::var("s3")),
345                                        Node::assign("s3", Expr::u32(0)),
346                                    ],
347                                ),
348                                Node::if_then(
349                                    Expr::eq(Expr::var("op"), Expr::u32(2)),
350                                    vec![
351                                        Node::assign(
352                                            "s0",
353                                            Expr::mul(Expr::var("s0"), Expr::var("s1")),
354                                        ),
355                                        Node::assign("s1", Expr::var("s2")),
356                                        Node::assign("s2", Expr::var("s3")),
357                                        Node::assign("s3", Expr::u32(0)),
358                                    ],
359                                ),
360                                Node::if_then(
361                                    Expr::eq(Expr::var("op"), Expr::u32(3)),
362                                    vec![
363                                        Node::assign("s3", Expr::var("s2")),
364                                        Node::assign("s2", Expr::var("s1")),
365                                        Node::assign("s1", Expr::var("s0")),
366                                    ],
367                                ),
368                                Node::if_then(
369                                    Expr::eq(Expr::var("op"), Expr::u32(4)),
370                                    vec![
371                                        Node::let_bind("tmp", Expr::var("s0")),
372                                        Node::assign("s0", Expr::var("s1")),
373                                        Node::assign("s1", Expr::var("tmp")),
374                                    ],
375                                ),
376                            ],
377                        },
378                        Node::store("results", Expr::var("tid"), Expr::var("s0")),
379                    ],
380                ),
381            ],
382        );
383
384        let lowered = lower_for_emit(&program).expect("Fix: pre-emit lowering must pass");
385
386        assert!(
387            body_has_s1_end_from_copy(&lowered.descriptor.body),
388            "Fix: lowering must preserve `let tmp = s0` as a Copy snapshot so SWAP writes s1 from old s0 instead of the post-assign s0 carrier"
389        );
390    }
391
392    fn body_has_s1_end_from_copy(body: &KernelBody) -> bool {
393        body.ops.iter().any(|op| {
394            let KernelOpKind::LoopCarrierEnd { name } = &op.kind else {
395                return false;
396            };
397            name.as_ref() == "s1"
398                && op.operands.first().copied().is_some_and(|operand| {
399                    body.ops.iter().any(|producer| {
400                        producer.result == Some(operand)
401                            && matches!(producer.kind, KernelOpKind::Copy)
402                    })
403                })
404        }) || body.child_bodies.iter().any(body_has_s1_end_from_copy)
405    }
406}
vyre_lower/pre_emit.rs

vyre_lower/
pre_emit.rs
