1use std::collections::HashMap;
20
21use thiserror::Error;
22use vyre::ir::{BufferDecl, DataType, Expr as IrExpr, Node, Program};
23
24use super::lex::tokens::{ANDAND, EQ, GE, GT, LE, LT, MINUS, NE, OROR, PERCENT, PLUS, SLASH, STAR};
25use super::parse::{Expr, Module, Stmt, Type};
26use super::sema::{BindingId, Resolution};
27
28#[derive(Debug, Clone, Error)]
30pub enum RustLowerError {
31 #[error("Rust lowering needs at least one function to use as the entry kernel")]
33 NoEntryFunction,
34 #[error(
36 "Rust to Vyre IR lowering does not support {0} yet; not emitting a miscompiled Program"
37 )]
38 Unsupported(String),
39}
40
41pub fn lower(module: &Module, resolution: &Resolution) -> Result<Program, RustLowerError> {
47 lower_entry(module, resolution, LowerMode::Scalar)
48}
49
50pub fn lower_batched(
62 module: &Module,
63 resolution: &Resolution,
64 lane_count: u32,
65) -> Result<Program, RustLowerError> {
66 if lane_count == 0 {
67 return Err(RustLowerError::Unsupported(
68 "batched Rust lowering with zero lanes".to_string(),
69 ));
70 }
71 lower_entry(module, resolution, LowerMode::Batched { lane_count })
72}
73
74#[derive(Clone, Copy)]
75enum LowerMode {
76 Scalar,
77 Batched { lane_count: u32 },
78}
79
80impl LowerMode {
81 fn buffer_count(self) -> u32 {
82 match self {
83 Self::Scalar => 1,
84 Self::Batched { lane_count } => lane_count,
85 }
86 }
87
88 fn workgroup_size(self) -> [u32; 3] {
89 match self {
90 Self::Scalar => [1, 1, 1],
91 Self::Batched { .. } => [256, 1, 1],
92 }
93 }
94
95 fn lane_index(self) -> IrExpr {
96 match self {
97 Self::Scalar => IrExpr::u32(0),
98 Self::Batched { .. } => IrExpr::var(BATCH_LANE_VAR),
99 }
100 }
101}
102
103const BATCH_LANE_VAR: &str = "__rust_lane";
104
105fn lower_entry(
106 module: &Module,
107 resolution: &Resolution,
108 mode: LowerMode,
109) -> Result<Program, RustLowerError> {
110 let entry_index = module
111 .functions
112 .len()
113 .checked_sub(1)
114 .ok_or(RustLowerError::NoEntryFunction)?;
115 let func = &module.functions[entry_index];
116
117 let def_to_id: HashMap<u32, BindingId> = resolution
118 .bindings
119 .iter()
120 .enumerate()
121 .map(|(id, b)| (b.def_offset, id))
122 .collect();
123
124 let mut buffers = Vec::with_capacity(func.params.len() + 1);
125 let mut entry_nodes = Vec::new();
126 for (i, (offset, ty)) in func.params.iter().enumerate() {
127 let dtype = scalar_dtype(ty)?;
128 let buf = format!("p{i}");
129 buffers.push(BufferDecl::read(&buf, i as u32, dtype).with_count(mode.buffer_count()));
130 let binding = def_to_id
131 .get(offset)
132 .copied()
133 .ok_or_else(|| RustLowerError::Unsupported("unresolved parameter".to_string()))?;
134 entry_nodes.push(Node::let_bind(
135 format!("v{binding}"),
136 IrExpr::load(buf, mode.lane_index()),
137 ));
138 }
139 let out_dtype = scalar_dtype(&func.ret)?;
140 buffers.push(
141 BufferDecl::output("out", func.params.len() as u32, out_dtype)
142 .with_count(mode.buffer_count()),
143 );
144
145 let ctx = LowerCtx {
146 module,
147 resolution,
148 def_to_id: &def_to_id,
149 output_index: mode.lane_index(),
150 };
151 entry_nodes.extend(ctx.lower_stmts(&func.body, Subst::Local(None))?);
152 let entry_nodes = match mode {
153 LowerMode::Scalar => entry_nodes,
154 LowerMode::Batched { lane_count } => vec![
155 Node::let_bind(BATCH_LANE_VAR, IrExpr::gid_x()),
156 Node::if_then(
157 IrExpr::lt(IrExpr::var(BATCH_LANE_VAR), IrExpr::u32(lane_count)),
158 entry_nodes,
159 ),
160 ],
161 };
162 Ok(Program::wrapped(
163 buffers,
164 mode.workgroup_size(),
165 entry_nodes,
166 ))
167}
168
169fn scalar_dtype(ty: &Type) -> Result<DataType, RustLowerError> {
171 match ty {
172 Type::I32 => Ok(DataType::I32),
173 Type::Bool => Ok(DataType::Bool),
174 Type::Unit => Err(RustLowerError::Unsupported(
175 "unit-typed parameter or return".to_string(),
176 )),
177 Type::Ref { inner, .. } => scalar_dtype(inner),
181 }
182}
183
184struct LowerCtx<'a> {
185 module: &'a Module,
186 resolution: &'a Resolution,
187 def_to_id: &'a HashMap<u32, BindingId>,
188 output_index: IrExpr,
189}
190
191#[derive(Clone, Copy)]
206enum Subst<'a> {
207 Local(Option<&'a HashMap<BindingId, IrExpr>>),
209 Inline(&'a HashMap<BindingId, IrExpr>),
211}
212
213impl LowerCtx<'_> {
214 fn lower_stmts(&self, stmts: &[Stmt], subst: Subst<'_>) -> Result<Vec<Node>, RustLowerError> {
215 let mut nodes = Vec::new();
216 for stmt in stmts {
217 match stmt {
218 Stmt::Let { name, init, .. } => {
219 let binding = self.def_to_id.get(name).copied().ok_or_else(|| {
220 RustLowerError::Unsupported("unresolved let binding".to_string())
221 })?;
222 nodes.push(Node::let_bind(
223 format!("v{binding}"),
224 self.lower_value(init, subst)?,
225 ));
226 }
227 Stmt::Return(Some(expr)) => {
228 nodes.push(Node::store(
229 "out",
230 self.output_index.clone(),
231 self.lower_value(expr, subst)?,
232 ));
233 return Ok(nodes);
234 }
235 Stmt::Return(None) => return Ok(nodes),
236 Stmt::Assign { name, value } => {
237 let binding = self.resolution.uses.get(name).copied().ok_or_else(|| {
238 RustLowerError::Unsupported("unresolved assignment target".to_string())
239 })?;
240 nodes.push(Node::assign(
241 format!("v{binding}"),
242 self.lower_value(value, subst)?,
243 ));
244 }
245 Stmt::Expr(Expr::If {
246 cond,
247 then_block,
248 else_block,
249 }) => {
250 let then_nodes = self.lower_stmts(block_stmts(then_block), subst)?;
251 let else_nodes = match else_block {
252 Some(block) => self.lower_stmts(block_stmts(block), subst)?,
253 None => Vec::new(),
254 };
255 nodes.push(Node::if_then_else(
256 self.lower_value(cond, subst)?,
257 then_nodes,
258 else_nodes,
259 ));
260 let then_div = stmts_diverge(block_stmts(then_block));
261 let else_div = else_block
262 .as_ref()
263 .is_some_and(|b| stmts_diverge(block_stmts(b)));
264 if then_div && else_div {
265 return Ok(nodes);
266 }
267 }
268 Stmt::While { cond, body } => {
269 nodes.extend(self.lower_while(cond, body, subst)?);
270 }
271 Stmt::For {
272 name,
273 start,
274 end,
275 body,
276 } => {
277 nodes.extend(self.lower_for_range(*name, start, end, body, subst)?);
278 }
279 Stmt::Expr(_) => {}
281 }
282 }
283 Ok(nodes)
284 }
285
286 fn counted_loop_trip(lo: &IrExpr, hi: &IrExpr) -> IrExpr {
295 let span_u32 = IrExpr::sub(
296 IrExpr::cast(DataType::U32, hi.clone()),
297 IrExpr::cast(DataType::U32, lo.clone()),
298 );
299 IrExpr::select(IrExpr::gt(hi.clone(), lo.clone()), span_u32, IrExpr::u32(0))
300 }
301
302 fn counted_loop_induction(lo: &IrExpr, loop_var: &str) -> IrExpr {
306 IrExpr::add(
307 lo.clone(),
308 IrExpr::cast(DataType::I32, IrExpr::var(loop_var.to_string())),
309 )
310 }
311
312 fn lower_while(
317 &self,
318 cond: &Expr,
319 body: &[Stmt],
320 subst: Subst<'_>,
321 ) -> Result<Vec<Node>, RustLowerError> {
322 let bad = || {
323 RustLowerError::Unsupported(
324 "while loop that is not a canonical `while i < BOUND { ...; i = i + 1; }` counting loop"
325 .to_string(),
326 )
327 };
328 let (i_off, bound) = match cond {
330 Expr::Binary { op, lhs, rhs } if *op == LT => match lhs.as_ref() {
331 Expr::Var(off) => (*off, rhs.as_ref()),
332 _ => return Err(bad()),
333 },
334 _ => return Err(bad()),
335 };
336 let b_i = self.resolution.uses.get(&i_off).copied().ok_or_else(bad)?;
337 let Some((last, init_stmts)) = body.split_last() else {
339 return Err(bad());
340 };
341 let inc_ok = matches!(last, Stmt::Assign { name, value }
342 if self.resolution.uses.get(name).copied() == Some(b_i)
343 && matches!(value, Expr::Binary { op, lhs, rhs }
344 if *op == PLUS
345 && matches!(lhs.as_ref(), Expr::Var(o) if self.resolution.uses.get(o).copied() == Some(b_i))
346 && matches!(rhs.as_ref(), Expr::LiteralInt(_, 1))));
347 if !inc_ok {
348 return Err(bad());
349 }
350 if stmts_assign_binding(init_stmts, b_i, self.resolution) {
352 return Err(bad());
353 }
354 for v in expr_var_bindings(bound, self.resolution) {
357 if stmts_assign_binding(body, v, self.resolution) {
358 return Err(bad());
359 }
360 }
361 let loop_var = format!("v{b_i}__w");
362 let mut inner: HashMap<BindingId, IrExpr> = match subst {
367 Subst::Local(Some(m)) => m.clone(),
368 Subst::Local(None) => HashMap::new(),
369 Subst::Inline(m) => m.clone(),
370 };
371 inner.insert(
383 b_i,
384 Self::counted_loop_induction(&IrExpr::var(format!("v{b_i}")), &loop_var),
385 );
386 let inner_subst = match subst {
387 Subst::Inline(_) => Subst::Inline(&inner),
388 Subst::Local(_) => Subst::Local(Some(&inner)),
389 };
390 let from_i32 = IrExpr::var(format!("v{b_i}"));
396 let to_i32 = self.lower_value(bound, subst)?;
397 let trip = Self::counted_loop_trip(&from_i32, &to_i32);
398 let from = IrExpr::u32(0);
399 let to = trip;
400 let loop_body = self.lower_stmts(init_stmts, inner_subst)?;
401 let post = IrExpr::select(
405 IrExpr::gt(to_i32.clone(), from_i32.clone()),
406 to_i32,
407 from_i32,
408 );
409 Ok(vec![
410 Node::loop_for(loop_var, from, to, loop_body),
411 Node::assign(format!("v{b_i}"), post),
412 ])
413 }
414
415 fn lower_for_range(
420 &self,
421 name: u32,
422 start: &Expr,
423 end: &Expr,
424 body: &[Stmt],
425 subst: Subst<'_>,
426 ) -> Result<Vec<Node>, RustLowerError> {
427 let b_i = self.def_to_id.get(&name).copied().ok_or_else(|| {
428 RustLowerError::Unsupported("unresolved for-loop binding".to_string())
429 })?;
430 let start_name = format!("v{b_i}__for_start");
431 let end_name = format!("v{b_i}__for_end");
432 let loop_var = format!("v{b_i}__for");
433
434 let start_i32 = IrExpr::var(start_name.clone());
435 let end_i32 = IrExpr::var(end_name.clone());
436 let trip = Self::counted_loop_trip(&start_i32, &end_i32);
437
438 let mut inner: HashMap<BindingId, IrExpr> = match subst {
439 Subst::Local(Some(m)) => m.clone(),
440 Subst::Local(None) => HashMap::new(),
441 Subst::Inline(m) => m.clone(),
442 };
443 inner.insert(b_i, Self::counted_loop_induction(&start_i32, &loop_var));
444 let inner_subst = match subst {
445 Subst::Inline(_) => Subst::Inline(&inner),
446 Subst::Local(_) => Subst::Local(Some(&inner)),
447 };
448 let loop_body = self.lower_stmts(body, inner_subst)?;
449
450 Ok(vec![
451 Node::let_bind(start_name, self.lower_value(start, subst)?),
452 Node::let_bind(end_name, self.lower_value(end, subst)?),
453 Node::loop_for(loop_var, IrExpr::u32(0), trip, loop_body),
454 ])
455 }
456
457 fn lower_value(&self, expr: &Expr, subst: Subst<'_>) -> Result<IrExpr, RustLowerError> {
463 match expr {
464 Expr::LiteralInt(_, value) => Ok(IrExpr::i32(*value as i32)),
465 Expr::LiteralBool(_, value) => Ok(IrExpr::bool(*value)),
466 Expr::Var(offset) => {
467 let binding = self.resolution.uses.get(offset).copied().ok_or_else(|| {
468 RustLowerError::Unsupported("unresolved variable use".to_string())
469 })?;
470 match subst {
471 Subst::Local(Some(map)) => Ok(map
474 .get(&binding)
475 .cloned()
476 .unwrap_or_else(|| IrExpr::var(format!("v{binding}")))),
477 Subst::Local(None) => Ok(IrExpr::var(format!("v{binding}"))),
478 Subst::Inline(map) => map.get(&binding).cloned().ok_or_else(|| {
480 RustLowerError::Unsupported("callee variable not substituted".to_string())
481 }),
482 }
483 }
484 Expr::Binary { op, lhs, rhs } => {
485 let l = self.lower_value(lhs, subst)?;
486 let r = self.lower_value(rhs, subst)?;
487 Ok(match *op {
488 PLUS => IrExpr::add(l, r),
489 MINUS => IrExpr::sub(l, r),
490 STAR => IrExpr::mul(l, r),
491 SLASH => IrExpr::div(l, r),
492 PERCENT => IrExpr::cast(DataType::I32, IrExpr::rem(l, r)),
496 EQ => IrExpr::eq(l, r),
497 NE => IrExpr::ne(l, r),
498 LT => IrExpr::lt(l, r),
499 GT => IrExpr::gt(l, r),
500 LE => IrExpr::le(l, r),
501 GE => IrExpr::ge(l, r),
502 ANDAND => IrExpr::and(l, r),
503 OROR => IrExpr::or(l, r),
504 other => {
505 return Err(RustLowerError::Unsupported(format!(
506 "binary operator {other}"
507 )))
508 }
509 })
510 }
511 Expr::Call { name, args } => self.lower_call(name, args, subst),
512 Expr::Borrow { expr, .. } => self.lower_value(expr, subst),
516 Expr::Deref(inner) => self.lower_value(inner, subst),
517 Expr::Not(inner) => Ok(IrExpr::not(self.lower_value(inner, subst)?)),
518 Expr::Neg(inner) => Ok(IrExpr::sub(IrExpr::i32(0), self.lower_value(inner, subst)?)),
523 Expr::Block(_) | Expr::If { .. } => Err(RustLowerError::Unsupported(
524 "block/if used as a value".to_string(),
525 )),
526 }
527 }
528
529 fn lower_call(
534 &self,
535 name: &u32,
536 args: &[Expr],
537 caller_subst: Subst<'_>,
538 ) -> Result<IrExpr, RustLowerError> {
539 let callee_index = self
540 .resolution
541 .calls
542 .get(name)
543 .copied()
544 .ok_or_else(|| RustLowerError::Unsupported("unresolved call".to_string()))?;
545 let callee = &self.module.functions[callee_index];
546 if args.len() != callee.params.len() {
547 return Err(RustLowerError::Unsupported(
548 "call arity mismatch".to_string(),
549 ));
550 }
551 let mut subst: HashMap<BindingId, IrExpr> = HashMap::new();
552 for (i, (offset, _)) in callee.params.iter().enumerate() {
553 let binding = self.def_to_id.get(offset).copied().ok_or_else(|| {
554 RustLowerError::Unsupported("unresolved callee parameter".to_string())
555 })?;
556 subst.insert(binding, self.lower_value(&args[i], caller_subst)?);
557 }
558 for stmt in &callee.body {
559 match stmt {
560 Stmt::Let {
561 name: offset, init, ..
562 } => {
563 let value = self.lower_value(init, Subst::Inline(&subst))?;
564 let binding = self.def_to_id.get(offset).copied().ok_or_else(|| {
565 RustLowerError::Unsupported("unresolved callee binding".to_string())
566 })?;
567 subst.insert(binding, value);
568 }
569 Stmt::Return(Some(expr)) => return self.lower_value(expr, Subst::Inline(&subst)),
570 _ => {
571 return Err(RustLowerError::Unsupported(
572 "call to a callee with control flow or no terminal return".to_string(),
573 ))
574 }
575 }
576 }
577 Err(RustLowerError::Unsupported(
578 "call to a callee with no return".to_string(),
579 ))
580 }
581}
582
583fn stmts_assign_binding(stmts: &[Stmt], b: BindingId, res: &Resolution) -> bool {
585 stmts.iter().any(|s| match s {
586 Stmt::Assign { name, .. } => res.uses.get(name).copied() == Some(b),
587 Stmt::Expr(Expr::If {
588 then_block,
589 else_block,
590 ..
591 }) => {
592 stmts_assign_binding(block_stmts(then_block), b, res)
593 || else_block
594 .as_ref()
595 .is_some_and(|e| stmts_assign_binding(block_stmts(e), b, res))
596 }
597 Stmt::While { body, .. } => stmts_assign_binding(body, b, res),
598 Stmt::For { body, .. } => stmts_assign_binding(body, b, res),
599 _ => false,
600 })
601}
602
603fn expr_var_bindings(expr: &Expr, res: &Resolution) -> Vec<BindingId> {
605 let mut out = Vec::new();
606 collect_var_bindings(expr, res, &mut out);
607 out
608}
609
610fn collect_var_bindings(expr: &Expr, res: &Resolution, out: &mut Vec<BindingId>) {
611 match expr {
612 Expr::Var(off) => {
613 if let Some(&id) = res.uses.get(off) {
614 out.push(id);
615 }
616 }
617 Expr::Binary { lhs, rhs, .. } => {
618 collect_var_bindings(lhs, res, out);
619 collect_var_bindings(rhs, res, out);
620 }
621 Expr::Borrow { expr, .. } => collect_var_bindings(expr, res, out),
622 Expr::Deref(inner) => collect_var_bindings(inner, res, out),
623 Expr::Not(inner) => collect_var_bindings(inner, res, out),
624 Expr::Neg(inner) => collect_var_bindings(inner, res, out),
625 Expr::Call { args, .. } => {
626 for a in args {
627 collect_var_bindings(a, res, out);
628 }
629 }
630 _ => {}
631 }
632}
633
634fn block_stmts(expr: &Expr) -> &[Stmt] {
636 match expr {
637 Expr::Block(stmts) => stmts,
638 _ => &[],
639 }
640}
641
642fn stmts_diverge(stmts: &[Stmt]) -> bool {
645 stmts.iter().any(|stmt| match stmt {
646 Stmt::Return(_) => true,
647 Stmt::Expr(Expr::If {
648 then_block,
649 else_block: Some(else_block),
650 ..
651 }) => stmts_diverge(block_stmts(then_block)) && stmts_diverge(block_stmts(else_block)),
652 _ => false,
653 })
654}