#![cfg(feature = "cuda-oxide-backend")]
use std::collections::HashMap;
use cranelift_codegen::ir::{self, InstructionData, Opcode};
use crate::lowered_ir::{LoweredOp, LoweredType, LoweredValueId};
fn lane_to_lowered(ty: ir::Type) -> Option<LoweredType> {
Some(match ty {
ir::types::I8 => LoweredType::I8,
ir::types::I16 => LoweredType::I16,
ir::types::I32 => LoweredType::I32,
ir::types::I64 => LoweredType::I64,
ir::types::F32 => LoweredType::F32,
ir::types::F64 => LoweredType::F64,
_ => return None,
})
}
fn vector_lane_type(ty: ir::Type) -> Option<LoweredType> {
if !ty.is_vector() {
return None;
}
lane_to_lowered(ty.lane_type())
}
fn fresh_id(next: &mut LoweredValueId) -> Option<LoweredValueId> {
let id = *next;
*next = next.checked_add(1)?;
Some(id)
}
fn map_value(
value_map: &HashMap<ir::Value, LoweredValueId>,
v: ir::Value,
) -> Option<LoweredValueId> {
value_map.get(&v).copied()
}
pub fn lower_vector_inst(
inst: ir::Inst,
func: &ir::Function,
value_map: &HashMap<ir::Value, LoweredValueId>,
next_value_id: &mut LoweredValueId,
) -> Option<LoweredOp> {
let data = &func.dfg.insts[inst];
match data {
InstructionData::Binary { opcode, args } => {
let result_ty = func.dfg.value_type(func.dfg.first_result(inst));
let lane_ty = vector_lane_type(result_ty)?;
let lhs = map_value(value_map, args[0])?;
let rhs = map_value(value_map, args[1])?;
let result = fresh_id(next_value_id)?;
match opcode {
Opcode::Fmin => Some(LoweredOp::VMin {
lane_ty,
signed: true,
lhs,
rhs,
result,
}),
Opcode::Smin => Some(LoweredOp::VMin {
lane_ty,
signed: true,
lhs,
rhs,
result,
}),
Opcode::Umin => Some(LoweredOp::VMin {
lane_ty,
signed: false,
lhs,
rhs,
result,
}),
Opcode::Fmax => Some(LoweredOp::VMax {
lane_ty,
signed: true,
lhs,
rhs,
result,
}),
Opcode::Smax => Some(LoweredOp::VMax {
lane_ty,
signed: true,
lhs,
rhs,
result,
}),
Opcode::Umax => Some(LoweredOp::VMax {
lane_ty,
signed: false,
lhs,
rhs,
result,
}),
_ => None,
}
}
InstructionData::Unary { opcode, arg } => match opcode {
Opcode::Splat => {
let result_value = func.dfg.first_result(inst);
let result_ty = func.dfg.value_type(result_value);
let lane_ty = vector_lane_type(result_ty)?;
let src = map_value(value_map, *arg)?;
let result = fresh_id(next_value_id)?;
Some(LoweredOp::VSplat {
lane_ty,
src,
result,
})
}
Opcode::VallTrue => {
let input_ty = func.dfg.value_type(*arg);
if !input_ty.is_vector() {
return None;
}
let src = map_value(value_map, *arg)?;
let result = fresh_id(next_value_id)?;
Some(LoweredOp::VAllTrue { src, result })
}
Opcode::VanyTrue => {
let input_ty = func.dfg.value_type(*arg);
if !input_ty.is_vector() {
return None;
}
let src = map_value(value_map, *arg)?;
let result = fresh_id(next_value_id)?;
Some(LoweredOp::VAnyTrue { src, result })
}
_ => None,
},
InstructionData::Ternary { opcode, args } => {
if *opcode != Opcode::Bitselect {
return None;
}
let result_ty = func.dfg.value_type(func.dfg.first_result(inst));
let lane_ty = vector_lane_type(result_ty)?;
let cond = map_value(value_map, args[0])?;
let then_v = map_value(value_map, args[1])?;
let else_v = map_value(value_map, args[2])?;
let result = fresh_id(next_value_id)?;
Some(LoweredOp::VSelect {
lane_ty,
cond,
then_v,
else_v,
result,
})
}
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use cranelift_codegen::ir::types;
use cranelift_codegen::ir::{
Function, InstructionData, Opcode, Signature, UserFuncName, Value,
};
use cranelift_codegen::isa::CallConv;
fn skeleton(param_tys: &[ir::Type]) -> (Function, Vec<Value>) {
let mut func =
Function::with_name_signature(UserFuncName::user(0, 0), Signature::new(CallConv::Fast));
let block = func.dfg.make_block();
let mut values = Vec::with_capacity(param_tys.len());
for &ty in param_tys {
values.push(func.dfg.append_block_param(block, ty));
}
(func, values)
}
fn push_binary(
func: &mut Function,
opcode: Opcode,
lhs: Value,
rhs: Value,
ctrl_ty: ir::Type,
) -> (ir::Inst, Value) {
let inst = func.dfg.make_inst(InstructionData::Binary {
opcode,
args: [lhs, rhs],
});
func.dfg.make_inst_results(inst, ctrl_ty);
let result = func.dfg.first_result(inst);
(inst, result)
}
fn push_unary(
func: &mut Function,
opcode: Opcode,
arg: Value,
ctrl_ty: ir::Type,
) -> (ir::Inst, Value) {
let inst = func.dfg.make_inst(InstructionData::Unary { opcode, arg });
func.dfg.make_inst_results(inst, ctrl_ty);
let result = func.dfg.first_result(inst);
(inst, result)
}
fn push_ternary(
func: &mut Function,
opcode: Opcode,
args: [Value; 3],
ctrl_ty: ir::Type,
) -> (ir::Inst, Value) {
let inst = func
.dfg
.make_inst(InstructionData::Ternary { opcode, args });
func.dfg.make_inst_results(inst, ctrl_ty);
let result = func.dfg.first_result(inst);
(inst, result)
}
fn seed_map(params: &[Value]) -> (HashMap<ir::Value, LoweredValueId>, LoweredValueId) {
let mut map = HashMap::new();
for (i, &v) in params.iter().enumerate() {
map.insert(v, i as LoweredValueId);
}
(map, params.len() as LoweredValueId)
}
#[test]
fn fmin_vector_lowers_to_vmin() {
let (mut func, params) = skeleton(&[types::F32X4, types::F32X4]);
let (inst, _res) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
match op {
LoweredOp::VMin {
lane_ty,
signed,
lhs,
rhs,
result,
} => {
assert_eq!(lane_ty, LoweredType::F32);
assert!(signed);
assert_eq!(lhs, 0);
assert_eq!(rhs, 1);
assert_eq!(result, 2);
}
other => panic!("expected VMin, got {other:?}"),
}
assert_eq!(next, 3);
}
#[test]
fn fmax_vector_lowers_to_vmax() {
let (mut func, params) = skeleton(&[types::F64X2, types::F64X2]);
let (inst, _) = push_binary(&mut func, Opcode::Fmax, params[0], params[1], types::F64X2);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMax {
lane_ty: LoweredType::F64,
..
}
));
}
#[test]
fn smin_vector_lowers_to_vmin_with_signed_lane() {
let (mut func, params) = skeleton(&[types::I32X4, types::I32X4]);
let (inst, _) = push_binary(&mut func, Opcode::Smin, params[0], params[1], types::I32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMin {
lane_ty: LoweredType::I32,
signed: true,
..
}
));
}
#[test]
fn umin_vector_lowers_to_unsigned_vmin() {
let (mut func, params) = skeleton(&[types::I32X4, types::I32X4]);
let (inst, _) = push_binary(&mut func, Opcode::Umin, params[0], params[1], types::I32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMin {
lane_ty: LoweredType::I32,
signed: false,
..
}
));
}
#[test]
fn smin_and_umin_are_distinguishable() {
let (mut sf, sp) = skeleton(&[types::I32X4, types::I32X4]);
let (si, _) = push_binary(&mut sf, Opcode::Smin, sp[0], sp[1], types::I32X4);
let (smap, mut sn) = seed_map(&sp);
let s = lower_vector_inst(si, &sf, &smap, &mut sn).expect("smin");
let (mut uf, up) = skeleton(&[types::I32X4, types::I32X4]);
let (ui, _) = push_binary(&mut uf, Opcode::Umin, up[0], up[1], types::I32X4);
let (umap, mut un) = seed_map(&up);
let u = lower_vector_inst(ui, &uf, &umap, &mut un).expect("umin");
let s_signed = matches!(s, LoweredOp::VMin { signed: true, .. });
let u_unsigned = matches!(u, LoweredOp::VMin { signed: false, .. });
assert!(s_signed && u_unsigned, "smin must be signed, umin unsigned");
}
#[test]
fn umax_vector_lowers_to_vmax_with_unsigned_lane() {
let (mut func, params) = skeleton(&[types::I16X8, types::I16X8]);
let (inst, _) = push_binary(&mut func, Opcode::Umax, params[0], params[1], types::I16X8);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMax {
lane_ty: LoweredType::I16,
signed: false,
..
}
));
}
#[test]
fn smax_i8x16_lowers_to_vmax() {
let (mut func, params) = skeleton(&[types::I8X16, types::I8X16]);
let (inst, _) = push_binary(&mut func, Opcode::Smax, params[0], params[1], types::I8X16);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMax {
lane_ty: LoweredType::I8,
signed: true,
..
}
));
}
#[test]
fn umin_i64x2_lowers_to_vmin() {
let (mut func, params) = skeleton(&[types::I64X2, types::I64X2]);
let (inst, _) = push_binary(&mut func, Opcode::Umin, params[0], params[1], types::I64X2);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VMin {
lane_ty: LoweredType::I64,
signed: false,
..
}
));
}
#[test]
fn scalar_fmin_returns_none() {
let (mut func, params) = skeleton(&[types::F32, types::F32]);
let (inst, _) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32);
let (map, mut next) = seed_map(¶ms);
assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
}
#[test]
fn scalar_smin_returns_none() {
let (mut func, params) = skeleton(&[types::I32, types::I32]);
let (inst, _) = push_binary(&mut func, Opcode::Smin, params[0], params[1], types::I32);
let (map, mut next) = seed_map(¶ms);
assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
}
#[test]
fn splat_f32_lowers_to_vsplat() {
let (mut func, params) = skeleton(&[types::F32]);
let (inst, _) = push_unary(&mut func, Opcode::Splat, params[0], types::F32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
match op {
LoweredOp::VSplat {
lane_ty,
src,
result,
} => {
assert_eq!(lane_ty, LoweredType::F32);
assert_eq!(src, 0);
assert_eq!(result, 1);
}
other => panic!("expected VSplat, got {other:?}"),
}
}
#[test]
fn splat_i32_lowers_to_vsplat() {
let (mut func, params) = skeleton(&[types::I32]);
let (inst, _) = push_unary(&mut func, Opcode::Splat, params[0], types::I32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(
op,
LoweredOp::VSplat {
lane_ty: LoweredType::I32,
..
}
));
}
#[test]
fn vall_true_lowers() {
let (mut func, params) = skeleton(&[types::I32X4]);
let (inst, _) = push_unary(&mut func, Opcode::VallTrue, params[0], types::I32X4);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
match op {
LoweredOp::VAllTrue { src, result } => {
assert_eq!(src, 0);
assert_eq!(result, 1);
}
other => panic!("expected VAllTrue, got {other:?}"),
}
}
#[test]
fn vany_true_lowers() {
let (mut func, params) = skeleton(&[types::I8X16]);
let (inst, _) = push_unary(&mut func, Opcode::VanyTrue, params[0], types::I8X16);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
assert!(matches!(op, LoweredOp::VAnyTrue { src: 0, result: 1 }));
}
#[test]
fn bitselect_vector_lowers_to_vselect() {
let (mut func, params) = skeleton(&[types::I32X4, types::I32X4, types::I32X4]);
let (inst, _) = push_ternary(
&mut func,
Opcode::Bitselect,
[params[0], params[1], params[2]],
types::I32X4,
);
let (map, mut next) = seed_map(¶ms);
let op = lower_vector_inst(inst, &func, &map, &mut next).expect("must lower");
match op {
LoweredOp::VSelect {
lane_ty,
cond,
then_v,
else_v,
result,
} => {
assert_eq!(lane_ty, LoweredType::I32);
assert_eq!(cond, 0);
assert_eq!(then_v, 1);
assert_eq!(else_v, 2);
assert_eq!(result, 3);
}
other => panic!("expected VSelect, got {other:?}"),
}
}
#[test]
fn scalar_bitselect_returns_none() {
let (mut func, params) = skeleton(&[types::I32, types::I32, types::I32]);
let (inst, _) = push_ternary(
&mut func,
Opcode::Bitselect,
[params[0], params[1], params[2]],
types::I32,
);
let (map, mut next) = seed_map(¶ms);
assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
}
#[test]
fn unrelated_op_returns_none() {
let (mut func, params) = skeleton(&[types::I32, types::I32]);
let (inst, _) = push_binary(&mut func, Opcode::Iadd, params[0], params[1], types::I32);
let (map, mut next) = seed_map(¶ms);
assert!(lower_vector_inst(inst, &func, &map, &mut next).is_none());
}
#[test]
fn fresh_id_allocation_is_monotone() {
let (mut func, params) = skeleton(&[types::F32X4, types::F32X4]);
let (inst, _) = push_binary(&mut func, Opcode::Fmin, params[0], params[1], types::F32X4);
let (map, mut next) = seed_map(¶ms);
let before = next;
let _ = lower_vector_inst(inst, &func, &map, &mut next).unwrap();
assert_eq!(next, before + 1);
}
}