1#![cfg(feature = "cuda-oxide-backend")]
40
41use std::collections::HashMap;
42
43use cranelift_codegen::ir::{self, InstructionData, Opcode};
44
45use crate::lowered_ir::{LoweredOp, LoweredType, LoweredValueId};
46
47#[derive(Debug, thiserror::Error, PartialEq, Eq)]
49pub enum MemLowerError {
50 #[error("memory lowering: Cranelift value {0:?} not in value_map")]
56 UnmappedValue(ir::Value),
57
58 #[error("memory lowering: unsupported Cranelift type `{0}`")]
63 UnsupportedType(String),
64
65 #[error("memory lowering: SSA value-id counter overflowed u32::MAX")]
74 SsaOverflow,
75}
76
77#[derive(Debug)]
87pub struct MemLowerContext<'a> {
88 pub linear_memory_base: LoweredValueId,
95
96 pub value_map: &'a HashMap<ir::Value, LoweredValueId>,
100
101 pub stack_slot_map: &'a mut HashMap<ir::StackSlot, LoweredValueId>,
105
106 pub next_value_id: &'a mut LoweredValueId,
109}
110
111impl<'a> MemLowerContext<'a> {
112 pub fn fresh_id(&mut self) -> Result<LoweredValueId, MemLowerError> {
121 let id = *self.next_value_id;
122 *self.next_value_id = id.checked_add(1).ok_or(MemLowerError::SsaOverflow)?;
123 Ok(id)
124 }
125
126 pub fn lookup_value(&self, v: ir::Value) -> Result<LoweredValueId, MemLowerError> {
130 self.value_map
131 .get(&v)
132 .copied()
133 .ok_or(MemLowerError::UnmappedValue(v))
134 }
135}
136
137pub fn lower_memory_inst(
159 inst: ir::Inst,
160 func: &ir::Function,
161 ctx: &mut MemLowerContext<'_>,
162) -> Result<Option<Vec<LoweredOp>>, MemLowerError> {
163 let data = &func.dfg.insts[inst];
164 match data {
165 InstructionData::Load {
166 opcode: Opcode::Load,
167 arg,
168 offset,
169 ..
170 } => {
171 let base = ctx.lookup_value(*arg)?;
172 let result_val = func.dfg.first_result(inst);
173 let ty = cranelift_type_to_lowered(func.dfg.value_type(result_val))?;
174 let result_id = ctx.fresh_id()?;
175 Ok(Some(vec![LoweredOp::Load {
179 ty,
180 base,
181 offset: (*offset).into(),
182 result: result_id,
183 }]))
184 }
185 InstructionData::Store {
186 opcode: Opcode::Store,
187 args,
188 offset,
189 ..
190 } => {
191 let value_v = args[0];
194 let base_v = args[1];
195 let value = ctx.lookup_value(value_v)?;
196 let base = ctx.lookup_value(base_v)?;
197 let ty = cranelift_type_to_lowered(func.dfg.value_type(value_v))?;
198 Ok(Some(vec![LoweredOp::Store {
199 ty,
200 value,
201 base,
202 offset: (*offset).into(),
203 }]))
204 }
205 InstructionData::StackLoad {
206 opcode: Opcode::StackLoad,
207 stack_slot,
208 offset,
209 } => {
210 let mut emitted = Vec::with_capacity(2);
211 let result_val = func.dfg.first_result(inst);
212 let ty = cranelift_type_to_lowered(func.dfg.value_type(result_val))?;
213 let alloca_ptr = ensure_stack_alloca(*stack_slot, func, ty.clone(), ctx, &mut emitted)?;
214 let result_id = ctx.fresh_id()?;
215 emitted.push(LoweredOp::Load {
216 ty,
217 base: alloca_ptr,
218 offset: (*offset).into(),
219 result: result_id,
220 });
221 Ok(Some(emitted))
222 }
223 InstructionData::StackStore {
224 opcode: Opcode::StackStore,
225 arg,
226 stack_slot,
227 offset,
228 } => {
229 let value = ctx.lookup_value(*arg)?;
230 let ty = cranelift_type_to_lowered(func.dfg.value_type(*arg))?;
231 let mut emitted = Vec::with_capacity(2);
232 let alloca_ptr = ensure_stack_alloca(*stack_slot, func, ty.clone(), ctx, &mut emitted)?;
233 emitted.push(LoweredOp::Store {
234 ty,
235 value,
236 base: alloca_ptr,
237 offset: (*offset).into(),
238 });
239 Ok(Some(emitted))
240 }
241 _ => Ok(None),
242 }
243}
244
245fn ensure_stack_alloca(
254 stack_slot: ir::StackSlot,
255 func: &ir::Function,
256 ty: LoweredType,
257 ctx: &mut MemLowerContext<'_>,
258 out: &mut Vec<LoweredOp>,
259) -> Result<LoweredValueId, MemLowerError> {
260 if let Some(existing) = ctx.stack_slot_map.get(&stack_slot) {
261 return Ok(*existing);
262 }
263 let bytes = func.sized_stack_slots[stack_slot].size;
264 let ptr_id = ctx.fresh_id()?;
265 out.push(LoweredOp::StackAlloc {
266 ty,
267 bytes,
268 result: ptr_id,
269 });
270 ctx.stack_slot_map.insert(stack_slot, ptr_id);
271 Ok(ptr_id)
272}
273
274fn cranelift_type_to_lowered(ty: ir::Type) -> Result<LoweredType, MemLowerError> {
282 use cranelift_codegen::ir::types;
283 Ok(match ty {
284 types::I8 => LoweredType::I8,
285 types::I16 => LoweredType::I16,
286 types::I32 => LoweredType::I32,
287 types::I64 => LoweredType::I64,
288 types::F32 => LoweredType::F32,
289 types::F64 => LoweredType::F64,
290 other => return Err(MemLowerError::UnsupportedType(other.to_string())),
291 })
292}
293
294#[cfg(test)]
295mod tests {
296 use super::*;
297 use cranelift_codegen::ir::immediates::Offset32;
298 use cranelift_codegen::ir::{
299 types, AbiParam, Function, MemFlags, Signature, StackSlotData, StackSlotKind, UserFuncName,
300 };
301 use cranelift_codegen::isa::CallConv;
302
303 fn skeleton() -> (
308 Function,
309 ir::Value,
310 ir::Value,
311 HashMap<ir::Value, LoweredValueId>,
312 ) {
313 let mut sig = Signature::new(CallConv::SystemV);
314 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I32)); let mut func = Function::with_name_signature(UserFuncName::default(), sig);
317 let block = func.dfg.make_block();
318 let base_v = func.dfg.append_block_param(block, types::I64);
319 let val_v = func.dfg.append_block_param(block, types::I32);
320 let mut map = HashMap::new();
321 map.insert(base_v, 100);
322 map.insert(val_v, 101);
323 (func, base_v, val_v, map)
324 }
325
326 #[test]
327 fn lower_load_produces_single_load_op() {
328 let (mut func, base_v, _val_v, value_map) = skeleton();
329 let inst = func.dfg.make_inst(InstructionData::Load {
331 opcode: Opcode::Load,
332 arg: base_v,
333 flags: MemFlags::new(),
334 offset: Offset32::new(8),
335 });
336 func.dfg.make_inst_results(inst, types::I32);
337
338 let mut stack_map = HashMap::new();
339 let mut next_id: LoweredValueId = 200;
340 let mut ctx = MemLowerContext {
341 linear_memory_base: 100,
342 value_map: &value_map,
343 stack_slot_map: &mut stack_map,
344 next_value_id: &mut next_id,
345 };
346 let out = lower_memory_inst(inst, &func, &mut ctx)
347 .expect("lowering must succeed")
348 .expect("inst is a memory op");
349 assert_eq!(out.len(), 1);
350 match &out[0] {
351 LoweredOp::Load {
352 ty,
353 base,
354 offset,
355 result,
356 } => {
357 assert_eq!(*ty, LoweredType::I32);
358 assert_eq!(*base, 100);
359 assert_eq!(*offset, 8);
360 assert_eq!(*result, 200);
361 }
362 other => panic!("expected Load, got {other:?}"),
363 }
364 assert_eq!(next_id, 201);
365 }
366
367 #[test]
372 fn fresh_id_overflow_is_structured_error_not_wrap() {
373 let (mut func, base_v, _val_v, value_map) = skeleton();
374 let inst = func.dfg.make_inst(InstructionData::Load {
375 opcode: Opcode::Load,
376 arg: base_v,
377 flags: MemFlags::new(),
378 offset: Offset32::new(0),
379 });
380 func.dfg.make_inst_results(inst, types::I32);
381
382 let mut stack_map = HashMap::new();
383 let mut next_id: LoweredValueId = u32::MAX;
384 let mut ctx = MemLowerContext {
385 linear_memory_base: 100,
386 value_map: &value_map,
387 stack_slot_map: &mut stack_map,
388 next_value_id: &mut next_id,
389 };
390 let err = lower_memory_inst(inst, &func, &mut ctx)
391 .expect_err("counter at u32::MAX must overflow, not wrap");
392 assert_eq!(err, MemLowerError::SsaOverflow);
393 }
394
395 #[test]
396 fn lower_store_produces_single_store_op() {
397 let (mut func, base_v, val_v, value_map) = skeleton();
398 let inst = func.dfg.make_inst(InstructionData::Store {
400 opcode: Opcode::Store,
401 args: [val_v, base_v],
402 flags: MemFlags::new(),
403 offset: Offset32::new(-4),
404 });
405 func.dfg.make_inst_results(inst, types::I32);
407
408 let mut stack_map = HashMap::new();
409 let mut next_id: LoweredValueId = 300;
410 let mut ctx = MemLowerContext {
411 linear_memory_base: 100,
412 value_map: &value_map,
413 stack_slot_map: &mut stack_map,
414 next_value_id: &mut next_id,
415 };
416 let out = lower_memory_inst(inst, &func, &mut ctx)
417 .expect("lowering must succeed")
418 .expect("inst is a memory op");
419 assert_eq!(out.len(), 1);
420 match &out[0] {
421 LoweredOp::Store {
422 ty,
423 value,
424 base,
425 offset,
426 } => {
427 assert_eq!(*ty, LoweredType::I32);
428 assert_eq!(*value, 101);
429 assert_eq!(*base, 100);
430 assert_eq!(*offset, -4);
431 }
432 other => panic!("expected Store, got {other:?}"),
433 }
434 assert_eq!(next_id, 300);
436 }
437
438 #[test]
439 fn lower_stack_load_first_touch_emits_alloca_then_load() {
440 let (mut func, _base_v, _val_v, value_map) = skeleton();
441 let ss =
442 func.create_sized_stack_slot(StackSlotData::new(StackSlotKind::ExplicitSlot, 16, 0));
443 let inst = func.dfg.make_inst(InstructionData::StackLoad {
444 opcode: Opcode::StackLoad,
445 stack_slot: ss,
446 offset: Offset32::new(0),
447 });
448 func.dfg.make_inst_results(inst, types::I32);
449
450 let mut stack_map = HashMap::new();
451 let mut next_id: LoweredValueId = 400;
452 let mut ctx = MemLowerContext {
453 linear_memory_base: 100,
454 value_map: &value_map,
455 stack_slot_map: &mut stack_map,
456 next_value_id: &mut next_id,
457 };
458 let out = lower_memory_inst(inst, &func, &mut ctx)
459 .expect("lowering must succeed")
460 .expect("inst is a memory op");
461 assert_eq!(out.len(), 2, "first touch emits alloca + load");
462 match &out[0] {
463 LoweredOp::StackAlloc { ty, bytes, result } => {
464 assert_eq!(*ty, LoweredType::I32);
465 assert_eq!(*bytes, 16);
466 assert_eq!(*result, 400);
467 }
468 other => panic!("expected StackAlloc, got {other:?}"),
469 }
470 match &out[1] {
471 LoweredOp::Load {
472 ty,
473 base,
474 offset,
475 result,
476 } => {
477 assert_eq!(*ty, LoweredType::I32);
478 assert_eq!(*base, 400, "Load reuses the alloca's pointer id");
479 assert_eq!(*offset, 0);
480 assert_eq!(*result, 401);
481 }
482 other => panic!("expected Load, got {other:?}"),
483 }
484 assert_eq!(stack_map.get(&ss).copied(), Some(400));
485 assert_eq!(next_id, 402);
486 }
487
488 #[test]
489 fn lower_stack_load_second_touch_reuses_alloca() {
490 let (mut func, _base_v, _val_v, value_map) = skeleton();
491 let ss =
492 func.create_sized_stack_slot(StackSlotData::new(StackSlotKind::ExplicitSlot, 8, 0));
493 let inst = func.dfg.make_inst(InstructionData::StackLoad {
494 opcode: Opcode::StackLoad,
495 stack_slot: ss,
496 offset: Offset32::new(0),
497 });
498 func.dfg.make_inst_results(inst, types::I32);
499
500 let mut stack_map = HashMap::new();
501 stack_map.insert(ss, 999);
503 let mut next_id: LoweredValueId = 500;
504 let mut ctx = MemLowerContext {
505 linear_memory_base: 100,
506 value_map: &value_map,
507 stack_slot_map: &mut stack_map,
508 next_value_id: &mut next_id,
509 };
510 let out = lower_memory_inst(inst, &func, &mut ctx)
511 .expect("lowering must succeed")
512 .expect("inst is a memory op");
513 assert_eq!(out.len(), 1, "subsequent touches skip the alloca");
514 match &out[0] {
515 LoweredOp::Load { base, result, .. } => {
516 assert_eq!(*base, 999);
517 assert_eq!(*result, 500);
518 }
519 other => panic!("expected Load, got {other:?}"),
520 }
521 assert_eq!(next_id, 501);
522 }
523
524 #[test]
525 fn lower_stack_store_first_touch_emits_alloca_then_store() {
526 let (mut func, _base_v, val_v, value_map) = skeleton();
527 let ss =
528 func.create_sized_stack_slot(StackSlotData::new(StackSlotKind::ExplicitSlot, 4, 0));
529 let inst = func.dfg.make_inst(InstructionData::StackStore {
530 opcode: Opcode::StackStore,
531 arg: val_v,
532 stack_slot: ss,
533 offset: Offset32::new(0),
534 });
535 func.dfg.make_inst_results(inst, types::I32);
536
537 let mut stack_map = HashMap::new();
538 let mut next_id: LoweredValueId = 600;
539 let mut ctx = MemLowerContext {
540 linear_memory_base: 100,
541 value_map: &value_map,
542 stack_slot_map: &mut stack_map,
543 next_value_id: &mut next_id,
544 };
545 let out = lower_memory_inst(inst, &func, &mut ctx)
546 .expect("lowering must succeed")
547 .expect("inst is a memory op");
548 assert_eq!(out.len(), 2, "first touch emits alloca + store");
549 match &out[0] {
550 LoweredOp::StackAlloc { ty, bytes, result } => {
551 assert_eq!(*ty, LoweredType::I32);
552 assert_eq!(*bytes, 4);
553 assert_eq!(*result, 600);
554 }
555 other => panic!("expected StackAlloc, got {other:?}"),
556 }
557 match &out[1] {
558 LoweredOp::Store {
559 ty,
560 value,
561 base,
562 offset,
563 } => {
564 assert_eq!(*ty, LoweredType::I32);
565 assert_eq!(*value, 101);
566 assert_eq!(*base, 600);
567 assert_eq!(*offset, 0);
568 }
569 other => panic!("expected Store, got {other:?}"),
570 }
571 assert_eq!(next_id, 601);
573 }
574
575 #[test]
576 fn lower_returns_none_for_non_memory_op() {
577 let mut sig = Signature::new(CallConv::SystemV);
579 sig.params.push(AbiParam::new(types::I32));
580 let mut func = Function::with_name_signature(UserFuncName::default(), sig);
581 let block = func.dfg.make_block();
582 let _arg = func.dfg.append_block_param(block, types::I32);
583 let inst = func.dfg.make_inst(InstructionData::UnaryImm {
584 opcode: Opcode::Iconst,
585 imm: 42i64.into(),
586 });
587 func.dfg.make_inst_results(inst, types::I32);
588
589 let value_map: HashMap<ir::Value, LoweredValueId> = HashMap::new();
590 let mut stack_map = HashMap::new();
591 let mut next_id: LoweredValueId = 700;
592 let mut ctx = MemLowerContext {
593 linear_memory_base: 0,
594 value_map: &value_map,
595 stack_slot_map: &mut stack_map,
596 next_value_id: &mut next_id,
597 };
598 let out = lower_memory_inst(inst, &func, &mut ctx).expect("lowering must succeed");
599 assert!(out.is_none(), "iconst is not a memory op");
600 assert_eq!(next_id, 700);
601 }
602
603 #[test]
604 fn unmapped_base_returns_error() {
605 let mut sig = Signature::new(CallConv::SystemV);
607 sig.params.push(AbiParam::new(types::I64));
608 let mut func = Function::with_name_signature(UserFuncName::default(), sig);
609 let block = func.dfg.make_block();
610 let base_v = func.dfg.append_block_param(block, types::I64);
611 let inst = func.dfg.make_inst(InstructionData::Load {
612 opcode: Opcode::Load,
613 arg: base_v,
614 flags: MemFlags::new(),
615 offset: Offset32::new(0),
616 });
617 func.dfg.make_inst_results(inst, types::I32);
618
619 let value_map: HashMap<ir::Value, LoweredValueId> = HashMap::new();
620 let mut stack_map = HashMap::new();
621 let mut next_id: LoweredValueId = 800;
622 let mut ctx = MemLowerContext {
623 linear_memory_base: 0,
624 value_map: &value_map,
625 stack_slot_map: &mut stack_map,
626 next_value_id: &mut next_id,
627 };
628 let err = lower_memory_inst(inst, &func, &mut ctx).expect_err("unmapped base must error");
629 assert_eq!(err, MemLowerError::UnmappedValue(base_v));
630 }
631
632 #[test]
633 fn cranelift_type_to_lowered_supports_scalars() {
634 assert_eq!(
635 cranelift_type_to_lowered(types::I8).unwrap(),
636 LoweredType::I8
637 );
638 assert_eq!(
639 cranelift_type_to_lowered(types::I32).unwrap(),
640 LoweredType::I32
641 );
642 assert_eq!(
643 cranelift_type_to_lowered(types::F64).unwrap(),
644 LoweredType::F64
645 );
646 }
647
648 #[test]
649 fn cranelift_type_to_lowered_rejects_vector() {
650 let err = cranelift_type_to_lowered(types::I32X4).unwrap_err();
651 assert!(
652 matches!(err, MemLowerError::UnsupportedType(_)),
653 "vector types are deferred to wave-2"
654 );
655 }
656}