1#![cfg(feature = "cuda-oxide-backend")]
49
50use std::collections::HashMap;
51
52use cranelift_codegen::ir::{self, InstructionData, Opcode};
53
54use crate::lowered_ir::{LoweredOp, LoweredType, LoweredValueId};
55
56fn lane_to_lowered(ty: ir::Type) -> Option<LoweredType> {
63 Some(match ty {
64 ir::types::I8 => LoweredType::I8,
65 ir::types::I16 => LoweredType::I16,
66 ir::types::I32 => LoweredType::I32,
67 ir::types::I64 => LoweredType::I64,
68 ir::types::F32 => LoweredType::F32,
69 ir::types::F64 => LoweredType::F64,
70 _ => return None,
71 })
72}
73
74fn vector_lane_type(ty: ir::Type) -> Option<LoweredType> {
79 if !ty.is_vector() {
80 return None;
81 }
82 lane_to_lowered(ty.lane_type())
83}
84
85fn fresh_id(next: &mut LoweredValueId) -> Option<LoweredValueId> {
98 let id = *next;
99 *next = next.checked_add(1)?;
100 Some(id)
101}
102
103fn map_value(
110 value_map: &HashMap<ir::Value, LoweredValueId>,
111 v: ir::Value,
112) -> Option<LoweredValueId> {
113 value_map.get(&v).copied()
114}
115
116pub fn lower_vector_inst(
138 inst: ir::Inst,
139 func: &ir::Function,
140 value_map: &HashMap<ir::Value, LoweredValueId>,
141 next_value_id: &mut LoweredValueId,
142) -> Option<LoweredOp> {
143 let data = &func.dfg.insts[inst];
144 match data {
145 InstructionData::Binary { opcode, args } => {
147 let result_ty = func.dfg.value_type(func.dfg.first_result(inst));
151 let lane_ty = vector_lane_type(result_ty)?;
152 let lhs = map_value(value_map, args[0])?;
153 let rhs = map_value(value_map, args[1])?;
154 let result = fresh_id(next_value_id)?;
155 match opcode {
160 Opcode::Fmin => Some(LoweredOp::VMin {
161 lane_ty,
162 signed: true,
163 lhs,
164 rhs,
165 result,
166 }),
167 Opcode::Smin => Some(LoweredOp::VMin {
168 lane_ty,
169 signed: true,
170 lhs,
171 rhs,
172 result,
173 }),
174 Opcode::Umin => Some(LoweredOp::VMin {
175 lane_ty,
176 signed: false,
177 lhs,
178 rhs,
179 result,
180 }),
181 Opcode::Fmax => Some(LoweredOp::VMax {
182 lane_ty,
183 signed: true,
184 lhs,
185 rhs,
186 result,
187 }),
188 Opcode::Smax => Some(LoweredOp::VMax {
189 lane_ty,
190 signed: true,
191 lhs,
192 rhs,
193 result,
194 }),
195 Opcode::Umax => Some(LoweredOp::VMax {
196 lane_ty,
197 signed: false,
198 lhs,
199 rhs,
200 result,
201 }),
202 _ => None,
203 }
204 }
205
206 InstructionData::Unary { opcode, arg } => match opcode {
208 Opcode::Splat => {
209 let result_value = func.dfg.first_result(inst);
213 let result_ty = func.dfg.value_type(result_value);
214 let lane_ty = vector_lane_type(result_ty)?;
215 let src = map_value(value_map, *arg)?;
216 let result = fresh_id(next_value_id)?;
217 Some(LoweredOp::VSplat {
218 lane_ty,
219 src,
220 result,
221 })
222 }
223 Opcode::VallTrue => {
224 let input_ty = func.dfg.value_type(*arg);
230 if !input_ty.is_vector() {
233 return None;
234 }
235 let src = map_value(value_map, *arg)?;
236 let result = fresh_id(next_value_id)?;
237 Some(LoweredOp::VAllTrue { src, result })
238 }
239 Opcode::VanyTrue => {
240 let input_ty = func.dfg.value_type(*arg);
241 if !input_ty.is_vector() {
242 return None;
243 }
244 let src = map_value(value_map, *arg)?;
245 let result = fresh_id(next_value_id)?;
246 Some(LoweredOp::VAnyTrue { src, result })
247 }
248 _ => None,
249 },
250
251 InstructionData::Ternary { opcode, args } => {
253 if *opcode != Opcode::Bitselect {
254 return None;
255 }
256 let result_ty = func.dfg.value_type(func.dfg.first_result(inst));
257 let lane_ty = vector_lane_type(result_ty)?;
262 let cond = map_value(value_map, args[0])?;
266 let then_v = map_value(value_map, args[1])?;
267 let else_v = map_value(value_map, args[2])?;
268 let result = fresh_id(next_value_id)?;
269 Some(LoweredOp::VSelect {
270 lane_ty,
271 cond,
272 then_v,
273 else_v,
274 result,
275 })
276 }
277
278 _ => None,
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use cranelift_codegen::ir::types;
287 use cranelift_codegen::ir::{
288 Function, InstructionData, Opcode, Signature, UserFuncName, Value,
289 };
290 use cranelift_codegen::isa::CallConv;
291
292 fn skeleton(param_tys: &[ir::Type]) -> (Function, Vec<Value>) {
300 let mut func =
301 Function::with_name_signature(UserFuncName::user(0, 0), Signature::new(CallConv::Fast));
302 let block = func.dfg.make_block();
303 let mut values = Vec::with_capacity(param_tys.len());
304 for &ty in param_tys {
305 values.push(func.dfg.append_block_param(block, ty));
306 }
307 (func, values)
308 }
309
310 fn push_binary(
314 func: &mut Function,
315 opcode: Opcode,
316 lhs: Value,
317 rhs: Value,
318 ctrl_ty: ir::Type,
319 ) -> (ir::Inst, Value) {
320 let inst = func.dfg.make_inst(InstructionData::Binary {
321 opcode,
322 args: [lhs, rhs],
323 });
324 func.dfg.make_inst_results(inst, ctrl_ty);
325 let result = func.dfg.first_result(inst);
326 (inst, result)
327 }
328
329 fn push_unary(
331 func: &mut Function,
332 opcode: Opcode,
333 arg: Value,
334 ctrl_ty: ir::Type,
335 ) -> (ir::Inst, Value) {
336 let inst = func.dfg.make_inst(InstructionData::Unary { opcode, arg });
337 func.dfg.make_inst_results(inst, ctrl_ty);
338 let result = func.dfg.first_result(inst);
339 (inst, result)
340 }
341
342 fn push_ternary(
344 func: &mut Function,
345 opcode: Opcode,
346 args: [Value; 3],
347 ctrl_ty: ir::Type,
348 ) -> (ir::Inst, Value) {
349 let inst = func
350 .dfg
351 .make_inst(InstructionData::Ternary { opcode, args });
352 func.dfg.make_inst_results(inst, ctrl_ty);
353 let result = func.dfg.first_result(inst);
354 (inst, result)
355 }
356
357 fn seed_map(params: &[Value]) -> (HashMap<ir::Value, LoweredValueId>, LoweredValueId) {
360 let mut map = HashMap::new();
361 for (i, &v) in params.iter().enumerate() {
362 map.insert(v, i as LoweredValueId);
363 }
364 (map, params.len() as LoweredValueId)
365 }
366
367 #[test]
370 fn fmin_vector_lowers_to_vmin() {
371 let (mut func, params) = skeleton(&[types::F32X4, types::F32X4]);
372 let (inst, _res) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32X4);
373 let (map, mut next) = seed_map(¶ms);
374 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
375 match op {
376 LoweredOp::VMin {
377 lane_ty,
378 signed,
379 lhs,
380 rhs,
381 result,
382 } => {
383 assert_eq!(lane_ty, LoweredType::F32);
384 assert!(signed);
386 assert_eq!(lhs, 0);
387 assert_eq!(rhs, 1);
388 assert_eq!(result, 2);
389 }
390 other => panic!("expected VMin, got {other:?}"),
391 }
392 assert_eq!(next, 3);
393 }
394
395 #[test]
396 fn fmax_vector_lowers_to_vmax() {
397 let (mut func, params) = skeleton(&[types::F64X2, types::F64X2]);
398 let (inst, _) = push_binary(&mut func, Opcode::Fmax, params[0], params[1], types::F64X2);
399 let (map, mut next) = seed_map(¶ms);
400 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
401 assert!(matches!(
402 op,
403 LoweredOp::VMax {
404 lane_ty: LoweredType::F64,
405 ..
406 }
407 ));
408 }
409
410 #[test]
412 fn smin_vector_lowers_to_vmin_with_signed_lane() {
413 let (mut func, params) = skeleton(&[types::I32X4, types::I32X4]);
414 let (inst, _) = push_binary(&mut func, Opcode::Smin, params[0], params[1], types::I32X4);
415 let (map, mut next) = seed_map(¶ms);
416 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
417 assert!(matches!(
418 op,
419 LoweredOp::VMin {
420 lane_ty: LoweredType::I32,
421 signed: true,
422 ..
423 }
424 ));
425 }
426
427 #[test]
431 fn umin_vector_lowers_to_unsigned_vmin() {
432 let (mut func, params) = skeleton(&[types::I32X4, types::I32X4]);
433 let (inst, _) = push_binary(&mut func, Opcode::Umin, params[0], params[1], types::I32X4);
434 let (map, mut next) = seed_map(¶ms);
435 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
436 assert!(matches!(
437 op,
438 LoweredOp::VMin {
439 lane_ty: LoweredType::I32,
440 signed: false,
441 ..
442 }
443 ));
444 }
445
446 #[test]
449 fn smin_and_umin_are_distinguishable() {
450 let (mut sf, sp) = skeleton(&[types::I32X4, types::I32X4]);
451 let (si, _) = push_binary(&mut sf, Opcode::Smin, sp[0], sp[1], types::I32X4);
452 let (smap, mut sn) = seed_map(&sp);
453 let s = lower_vector_inst(si, &sf, &smap, &mut sn).expect("smin");
454
455 let (mut uf, up) = skeleton(&[types::I32X4, types::I32X4]);
456 let (ui, _) = push_binary(&mut uf, Opcode::Umin, up[0], up[1], types::I32X4);
457 let (umap, mut un) = seed_map(&up);
458 let u = lower_vector_inst(ui, &uf, &umap, &mut un).expect("umin");
459
460 let s_signed = matches!(s, LoweredOp::VMin { signed: true, .. });
461 let u_unsigned = matches!(u, LoweredOp::VMin { signed: false, .. });
462 assert!(s_signed && u_unsigned, "smin must be signed, umin unsigned");
463 }
464
465 #[test]
466 fn umax_vector_lowers_to_vmax_with_unsigned_lane() {
467 let (mut func, params) = skeleton(&[types::I16X8, types::I16X8]);
468 let (inst, _) = push_binary(&mut func, Opcode::Umax, params[0], params[1], types::I16X8);
469 let (map, mut next) = seed_map(¶ms);
470 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
471 assert!(matches!(
472 op,
473 LoweredOp::VMax {
474 lane_ty: LoweredType::I16,
475 signed: false,
476 ..
477 }
478 ));
479 }
480
481 #[test]
482 fn smax_i8x16_lowers_to_vmax() {
483 let (mut func, params) = skeleton(&[types::I8X16, types::I8X16]);
484 let (inst, _) = push_binary(&mut func, Opcode::Smax, params[0], params[1], types::I8X16);
485 let (map, mut next) = seed_map(¶ms);
486 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
487 assert!(matches!(
488 op,
489 LoweredOp::VMax {
490 lane_ty: LoweredType::I8,
491 signed: true,
492 ..
493 }
494 ));
495 }
496
497 #[test]
498 fn umin_i64x2_lowers_to_vmin() {
499 let (mut func, params) = skeleton(&[types::I64X2, types::I64X2]);
500 let (inst, _) = push_binary(&mut func, Opcode::Umin, params[0], params[1], types::I64X2);
501 let (map, mut next) = seed_map(¶ms);
502 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
503 assert!(matches!(
504 op,
505 LoweredOp::VMin {
506 lane_ty: LoweredType::I64,
507 signed: false,
508 ..
509 }
510 ));
511 }
512
513 #[test]
516 fn scalar_fmin_returns_none() {
517 let (mut func, params) = skeleton(&[types::F32, types::F32]);
518 let (inst, _) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32);
519 let (map, mut next) = seed_map(¶ms);
520 assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
521 }
522
523 #[test]
525 fn scalar_smin_returns_none() {
526 let (mut func, params) = skeleton(&[types::I32, types::I32]);
527 let (inst, _) = push_binary(&mut func, Opcode::Smin, params[0], params[1], types::I32);
528 let (map, mut next) = seed_map(¶ms);
529 assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
530 }
531
532 #[test]
535 fn splat_f32_lowers_to_vsplat() {
536 let (mut func, params) = skeleton(&[types::F32]);
538 let (inst, _) = push_unary(&mut func, Opcode::Splat, params[0], types::F32X4);
539 let (map, mut next) = seed_map(¶ms);
540 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
541 match op {
542 LoweredOp::VSplat {
543 lane_ty,
544 src,
545 result,
546 } => {
547 assert_eq!(lane_ty, LoweredType::F32);
548 assert_eq!(src, 0);
549 assert_eq!(result, 1);
550 }
551 other => panic!("expected VSplat, got {other:?}"),
552 }
553 }
554
555 #[test]
556 fn splat_i32_lowers_to_vsplat() {
557 let (mut func, params) = skeleton(&[types::I32]);
558 let (inst, _) = push_unary(&mut func, Opcode::Splat, params[0], types::I32X4);
559 let (map, mut next) = seed_map(¶ms);
560 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
561 assert!(matches!(
562 op,
563 LoweredOp::VSplat {
564 lane_ty: LoweredType::I32,
565 ..
566 }
567 ));
568 }
569
570 #[test]
573 fn vall_true_lowers() {
574 let (mut func, params) = skeleton(&[types::I32X4]);
575 let (inst, _) = push_unary(&mut func, Opcode::VallTrue, params[0], types::I32X4);
576 let (map, mut next) = seed_map(¶ms);
577 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
578 match op {
579 LoweredOp::VAllTrue { src, result } => {
580 assert_eq!(src, 0);
581 assert_eq!(result, 1);
582 }
583 other => panic!("expected VAllTrue, got {other:?}"),
584 }
585 }
586
587 #[test]
588 fn vany_true_lowers() {
589 let (mut func, params) = skeleton(&[types::I8X16]);
590 let (inst, _) = push_unary(&mut func, Opcode::VanyTrue, params[0], types::I8X16);
591 let (map, mut next) = seed_map(¶ms);
592 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
593 assert!(matches!(op, LoweredOp::VAnyTrue { src: 0, result: 1 }));
594 }
595
596 #[test]
599 fn bitselect_vector_lowers_to_vselect() {
600 let (mut func, params) = skeleton(&[types::I32X4, types::I32X4, types::I32X4]);
601 let (inst, _) = push_ternary(
602 &mut func,
603 Opcode::Bitselect,
604 [params[0], params[1], params[2]],
605 types::I32X4,
606 );
607 let (map, mut next) = seed_map(¶ms);
608 let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
609 match op {
610 LoweredOp::VSelect {
611 lane_ty,
612 cond,
613 then_v,
614 else_v,
615 result,
616 } => {
617 assert_eq!(lane_ty, LoweredType::I32);
618 assert_eq!(cond, 0);
619 assert_eq!(then_v, 1);
620 assert_eq!(else_v, 2);
621 assert_eq!(result, 3);
622 }
623 other => panic!("expected VSelect, got {other:?}"),
624 }
625 }
626
627 #[test]
629 fn scalar_bitselect_returns_none() {
630 let (mut func, params) = skeleton(&[types::I32, types::I32, types::I32]);
631 let (inst, _) = push_ternary(
632 &mut func,
633 Opcode::Bitselect,
634 [params[0], params[1], params[2]],
635 types::I32,
636 );
637 let (map, mut next) = seed_map(¶ms);
638 assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
639 }
640
641 #[test]
646 fn unrelated_op_returns_none() {
647 let (mut func, params) = skeleton(&[types::I32, types::I32]);
648 let (inst, _) = push_binary(&mut func, Opcode::Iadd, params[0], params[1], types::I32);
649 let (map, mut next) = seed_map(¶ms);
650 assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
651 }
652
653 #[test]
656 fn fresh_id_allocation_is_monotone() {
657 let (mut func, params) = skeleton(&[types::F32X4, types::F32X4]);
658 let (inst, _) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32X4);
659 let (map, mut next) = seed_map(¶ms);
660 let before = next;
661 let _ = lower_vector_inst(inst, &func, &map, &mut next).unwrap();
662 assert_eq!(next, before + 1);
663 }
664}