ramtej/nac3
1
0
Fork 0
forked from M-Labs/nac3
Code Pull Requests Activity
nac3/nac3core/src/codegen/builtin_fns.rs

1956 lines
68 KiB
Rust
Raw Blame History

use inkwell::{
types::BasicTypeEnum,
values::{BasicValue, BasicValueEnum, IntValue},
FloatPredicate, IntPredicate, OptimizationLevel,
};
use itertools::Itertools;
use super::{
classes::RangeValue,
expr::destructure_range,
extern_fns, irrt,
irrt::calculate_len_for_slice_range,
llvm_intrinsics,
macros::codegen_unreachable,
model::*,
object::{
any::AnyObject,
list::ListObject,
ndarray::{NDArrayObject, NDArrayOut, ScalarOrNDArray},
tuple::TupleObject,
},
CodeGenContext, CodeGenerator,
};
use crate::{
toplevel::helper::PrimDef,
typecheck::typedef::{Type, TypeEnum},
};
/// Shorthand for [`unreachable!()`] when a type of argument is not supported.
///
/// The generated message will contain the function name and the name of the unsupported type.
fn unsupported_type(ctx: &CodeGenContext<'_, '_>, fn_name: &str, tys: &[Type]) -> ! {
codegen_unreachable!(
ctx,
"{fn_name}() not supported for '{}'",
tys.iter().map(|ty| format!("'{}'", ctx.unifier.stringify(*ty))).join(", "),
)
}
/// Invokes the `len` builtin function.
pub fn call_len<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<IntValue<'ctx>, String> {
let (arg_ty, arg) = n;
Ok(if ctx.unifier.unioned(arg_ty, ctx.primitives.range) {
let arg = RangeValue::from_ptr_val(arg.into_pointer_value(), Some("range"));
let (start, end, step) = destructure_range(ctx, arg);
calculate_len_for_slice_range(generator, ctx, start, end, step)
} else {
let arg = AnyObject { ty: arg_ty, value: arg };
let len: Instance<'ctx, Int<Int32>> = match &*ctx.unifier.get_ty(arg_ty) {
TypeEnum::TTuple { .. } => {
let tuple = TupleObject::from_object(ctx, arg);
tuple.len(generator, ctx).truncate_or_bit_cast(generator, ctx, Int32)
}
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() =>
{
let ndarray = NDArrayObject::from_object(generator, ctx, arg);
ndarray.len(generator, ctx).truncate_or_bit_cast(generator, ctx, Int32)
}
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
{
let list = ListObject::from_object(generator, ctx, arg);
list.len(generator, ctx).truncate_or_bit_cast(generator, ctx, Int32)
}
_ => unsupported_type(ctx, "len", &[arg_ty]),
};
len.value
})
}
/// Invokes the `int32` builtin function.
pub fn call_int32<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let llvm_i32 = ctx.ctx.i32_type();
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.bool));
ctx.builder.build_int_z_extend(n, llvm_i32, "zext").map(Into::into).unwrap()
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 32 => {
debug_assert!([ctx.primitives.int32, ctx.primitives.uint32,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
n.into()
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 64 => {
debug_assert!([ctx.primitives.int64, ctx.primitives.uint64,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
ctx.builder.build_int_truncate(n, llvm_i32, "trunc").map(Into::into).unwrap()
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let to_int64 =
ctx.builder.build_float_to_signed_int(n, ctx.ctx.i64_type(), "").unwrap();
ctx.builder.build_int_truncate(to_int64, llvm_i32, "conv").map(Into::into).unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.int32 },
|generator, ctx, scalar| call_int32(generator, ctx, (ndarray.dtype, scalar)),
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, "int32", &[n_ty]),
})
}
/// Invokes the `int64` builtin function.
pub fn call_int64<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let llvm_i64 = ctx.ctx.i64_type();
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8 | 32) => {
debug_assert!([ctx.primitives.bool, ctx.primitives.int32, ctx.primitives.uint32,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
if ctx.unifier.unioned(n_ty, ctx.primitives.int32) {
ctx.builder.build_int_s_extend(n, llvm_i64, "sext").map(Into::into).unwrap()
} else {
ctx.builder.build_int_z_extend(n, llvm_i64, "zext").map(Into::into).unwrap()
}
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 64 => {
debug_assert!([ctx.primitives.int64, ctx.primitives.uint64,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
n.into()
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
ctx.builder
.build_float_to_signed_int(n, ctx.ctx.i64_type(), "fptosi")
.map(Into::into)
.unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.int64 },
|generator, ctx, scalar| call_int64(generator, ctx, (ndarray.dtype, scalar)),
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, "int64", &[n_ty]),
})
}
/// Invokes the `uint32` builtin function.
pub fn call_uint32<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let llvm_i32 = ctx.ctx.i32_type();
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.bool));
ctx.builder.build_int_z_extend(n, llvm_i32, "zext").map(Into::into).unwrap()
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 32 => {
debug_assert!([ctx.primitives.int32, ctx.primitives.uint32,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
n.into()
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 64 => {
debug_assert!(
ctx.unifier.unioned(n_ty, ctx.primitives.int64)
|| ctx.unifier.unioned(n_ty, ctx.primitives.uint64)
);
ctx.builder.build_int_truncate(n, llvm_i32, "trunc").map(Into::into).unwrap()
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let n_gez = ctx
.builder
.build_float_compare(FloatPredicate::OGE, n, n.get_type().const_zero(), "")
.unwrap();
let to_int32 = ctx.builder.build_float_to_signed_int(n, llvm_i32, "").unwrap();
let to_uint64 =
ctx.builder.build_float_to_unsigned_int(n, ctx.ctx.i64_type(), "").unwrap();
ctx.builder
.build_select(
n_gez,
ctx.builder.build_int_truncate(to_uint64, llvm_i32, "").unwrap(),
to_int32,
"conv",
)
.unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.uint32 },
|generator, ctx, scalar| call_uint32(generator, ctx, (ndarray.dtype, scalar)),
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, "uint32", &[n_ty]),
})
}
/// Invokes the `uint64` builtin function.
pub fn call_uint64<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let llvm_i64 = ctx.ctx.i64_type();
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8 | 32) => {
debug_assert!([ctx.primitives.bool, ctx.primitives.int32, ctx.primitives.uint32,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
if ctx.unifier.unioned(n_ty, ctx.primitives.int32) {
ctx.builder.build_int_s_extend(n, llvm_i64, "sext").map(Into::into).unwrap()
} else {
ctx.builder.build_int_z_extend(n, llvm_i64, "zext").map(Into::into).unwrap()
}
}
BasicValueEnum::IntValue(n) if n.get_type().get_bit_width() == 64 => {
debug_assert!([ctx.primitives.int64, ctx.primitives.uint64,]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
n.into()
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let val_gez = ctx
.builder
.build_float_compare(FloatPredicate::OGE, n, n.get_type().const_zero(), "")
.unwrap();
let to_int64 = ctx.builder.build_float_to_signed_int(n, llvm_i64, "").unwrap();
let to_uint64 = ctx.builder.build_float_to_unsigned_int(n, llvm_i64, "").unwrap();
ctx.builder.build_select(val_gez, to_uint64, to_int64, "conv").unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.uint64 },
|generator, ctx, scalar| call_uint64(generator, ctx, (ndarray.dtype, scalar)),
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, "uint64", &[n_ty]),
})
}
/// Invokes the `float` builtin function.
pub fn call_float<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let llvm_f64 = ctx.ctx.f64_type();
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8 | 32 | 64) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
if [ctx.primitives.bool, ctx.primitives.int32, ctx.primitives.int64]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty))
{
ctx.builder
.build_signed_int_to_float(n, llvm_f64, "sitofp")
.map(Into::into)
.unwrap()
} else {
ctx.builder
.build_unsigned_int_to_float(n, llvm_f64, "uitofp")
.map(Into::into)
.unwrap()
}
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
n.into()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.float },
|generator, ctx, scalar| call_float(generator, ctx, (ndarray.dtype, scalar)),
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, "float", &[n_ty]),
})
}
/// Invokes the `round` builtin function.
pub fn call_round<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
ret_elem_ty: Type,
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "round";
let (n_ty, n) = n;
let llvm_ret_elem_ty = ctx.get_llvm_abi_type(generator, ret_elem_ty).into_int_type();
Ok(match n {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let val = llvm_intrinsics::call_float_round(ctx, n, None);
ctx.builder
.build_float_to_signed_int(val, llvm_ret_elem_ty, FN_NAME)
.map(Into::into)
.unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ret_elem_ty },
|generator, ctx, scalar| {
call_round(generator, ctx, (ndarray.dtype, scalar), ret_elem_ty)
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[n_ty]),
})
}
/// Invokes the `np_round` builtin function.
pub fn call_numpy_round<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_round";
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
llvm_intrinsics::call_float_rint(ctx, n, None).into()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.float },
|generator, ctx, scalar| {
call_numpy_round(generator, ctx, (ndarray.dtype, scalar))
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[n_ty]),
})
}
/// Invokes the `bool` builtin function.
pub fn call_bool<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "bool";
let (n_ty, n) = n;
Ok(match n {
BasicValueEnum::IntValue(n) if matches!(n.get_type().get_bit_width(), 1 | 8) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.bool));
n.into()
}
BasicValueEnum::IntValue(n) => {
debug_assert!([
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
]
.iter()
.any(|ty| ctx.unifier.unioned(n_ty, *ty)));
ctx.builder
.build_int_compare(IntPredicate::NE, n, n.get_type().const_zero(), FN_NAME)
.map(Into::into)
.unwrap()
}
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
ctx.builder
.build_float_compare(FloatPredicate::UNE, n, n.get_type().const_zero(), FN_NAME)
.map(Into::into)
.unwrap()
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ctx.primitives.bool },
|generator, ctx, scalar| {
let elem = call_bool(generator, ctx, (ndarray.dtype, scalar))?;
Ok(generator.bool_to_i8(ctx, elem.into_int_value()).into())
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[n_ty]),
})
}
/// Invokes the `floor` builtin function.
pub fn call_floor<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
ret_elem_ty: Type,
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "floor";
let (n_ty, n) = n;
let llvm_ret_elem_ty = ctx.get_llvm_abi_type(generator, ret_elem_ty);
Ok(match n {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let val = llvm_intrinsics::call_float_floor(ctx, n, None);
if let BasicTypeEnum::IntType(llvm_ret_elem_ty) = llvm_ret_elem_ty {
ctx.builder
.build_float_to_signed_int(val, llvm_ret_elem_ty, FN_NAME)
.map(Into::into)
.unwrap()
} else {
val.into()
}
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ret_elem_ty },
|generator, ctx, scalar| {
call_floor(generator, ctx, (ndarray.dtype, scalar), ret_elem_ty)
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[n_ty]),
})
}
/// Invokes the `ceil` builtin function.
pub fn call_ceil<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
ret_elem_ty: Type,
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "ceil";
let (n_ty, n) = n;
let llvm_ret_elem_ty = ctx.get_llvm_abi_type(generator, ret_elem_ty);
Ok(match n {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let val = llvm_intrinsics::call_float_ceil(ctx, n, None);
if let BasicTypeEnum::IntType(llvm_ret_elem_ty) = llvm_ret_elem_ty {
ctx.builder
.build_float_to_signed_int(val, llvm_ret_elem_ty, FN_NAME)
.map(Into::into)
.unwrap()
} else {
val.into()
}
}
_ if n_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: n_ty, value: n };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let result = ndarray
.map(
generator,
ctx,
NDArrayOut::NewNDArray { dtype: ret_elem_ty },
|generator, ctx, scalar| {
call_ceil(generator, ctx, (ndarray.dtype, scalar), ret_elem_ty)
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[n_ty]),
})
}
/// Invokes the `min` builtin function.
pub fn call_min<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
m: (Type, BasicValueEnum<'ctx>),
n: (Type, BasicValueEnum<'ctx>),
) -> BasicValueEnum<'ctx> {
const FN_NAME: &str = "min";
let (m_ty, m) = m;
let (n_ty, n) = n;
let common_ty = if ctx.unifier.unioned(m_ty, n_ty) {
m_ty
} else {
unsupported_type(ctx, FN_NAME, &[m_ty, n_ty])
};
match (m, n) {
(BasicValueEnum::IntValue(m), BasicValueEnum::IntValue(n)) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty, *ty)));
if [ctx.primitives.int32, ctx.primitives.int64]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty, *ty))
{
llvm_intrinsics::call_int_smin(ctx, m, n, Some(FN_NAME)).into()
} else {
llvm_intrinsics::call_int_umin(ctx, m, n, Some(FN_NAME)).into()
}
}
(BasicValueEnum::FloatValue(m), BasicValueEnum::FloatValue(n)) => {
debug_assert!(ctx.unifier.unioned(common_ty, ctx.primitives.float));
llvm_intrinsics::call_float_minnum(ctx, m, n, Some(FN_NAME)).into()
}
_ => unsupported_type(ctx, FN_NAME, &[m_ty, n_ty]),
}
}
/// Invokes the `np_minimum` builtin function.
pub fn call_numpy_minimum<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>),
x2: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_minimum";
let (x1_ty, x1) = x1;
let (x2_ty, x2) = x2;
let common_ty = if ctx.unifier.unioned(x1_ty, x2_ty) { Some(x1_ty) } else { None };
Ok(match (x1, x2) {
(BasicValueEnum::IntValue(x1), BasicValueEnum::IntValue(x2)) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
ctx.primitives.float,
]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty.unwrap(), *ty)));
call_min(ctx, (x1_ty, x1.into()), (x2_ty, x2.into()))
}
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
debug_assert!(ctx.unifier.unioned(common_ty.unwrap(), ctx.primitives.float));
call_min(ctx, (x1_ty, x1.into()), (x2_ty, x2.into()))
}
_ if [&x1_ty, &x2_ty]
.into_iter()
.any(|ty| ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())) =>
{
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1).to_ndarray(generator, ctx);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2).to_ndarray(generator, ctx);
debug_assert!(ctx.unifier.unioned(x1.dtype, x2.dtype));
let common_dtype = x1.dtype;
let result = NDArrayObject::broadcast_starmap(
generator,
ctx,
&[x1, x2],
NDArrayOut::NewNDArray { dtype: common_dtype },
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
Ok(call_min(ctx, (x1.dtype, x1_scalar), (x2.dtype, x2_scalar)))
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
})
}
/// Invokes the `max` builtin function.
pub fn call_max<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
m: (Type, BasicValueEnum<'ctx>),
n: (Type, BasicValueEnum<'ctx>),
) -> BasicValueEnum<'ctx> {
const FN_NAME: &str = "max";
let (m_ty, m) = m;
let (n_ty, n) = n;
let common_ty = if ctx.unifier.unioned(m_ty, n_ty) {
m_ty
} else {
unsupported_type(ctx, FN_NAME, &[m_ty, n_ty])
};
match (m, n) {
(BasicValueEnum::IntValue(m), BasicValueEnum::IntValue(n)) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty, *ty)));
if [ctx.primitives.int32, ctx.primitives.int64]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty, *ty))
{
llvm_intrinsics::call_int_smax(ctx, m, n, Some(FN_NAME)).into()
} else {
llvm_intrinsics::call_int_umax(ctx, m, n, Some(FN_NAME)).into()
}
}
(BasicValueEnum::FloatValue(m), BasicValueEnum::FloatValue(n)) => {
debug_assert!(ctx.unifier.unioned(common_ty, ctx.primitives.float));
llvm_intrinsics::call_float_maxnum(ctx, m, n, Some(FN_NAME)).into()
}
_ => unsupported_type(ctx, FN_NAME, &[m_ty, n_ty]),
}
}
/// Invokes the `np_max`, `np_min`, `np_argmax`, `np_argmin` functions
/// * `fn_name`: Can be one of `"np_argmin"`, `"np_argmax"`, `"np_max"`, `"np_min"`
pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
a: (Type, BasicValueEnum<'ctx>),
fn_name: &str,
) -> Result<BasicValueEnum<'ctx>, String> {
debug_assert!(["np_argmin", "np_argmax", "np_max", "np_min"].iter().any(|f| *f == fn_name));
let llvm_int64 = ctx.ctx.i64_type();
let (a_ty, a) = a;
Ok(match a {
BasicValueEnum::IntValue(_) | BasicValueEnum::FloatValue(_) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
ctx.primitives.float,
]
.iter()
.any(|ty| ctx.unifier.unioned(a_ty, *ty)));
match fn_name {
"np_argmin" | "np_argmax" => llvm_int64.const_zero().into(),
"np_max" | "np_min" => a,
_ => codegen_unreachable!(ctx),
}
}
_ if a_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) => {
let ndarray = AnyObject { ty: a_ty, value: a };
let ndarray = NDArrayObject::from_object(generator, ctx, ndarray);
let dtype_llvm = ctx.get_llvm_type(generator, ndarray.dtype);
let zero = Int(SizeT).const_0(generator, ctx.ctx);
if ctx.registry.llvm_options.opt_level == OptimizationLevel::None {
let size_isnt_zero =
ndarray.size(generator, ctx).compare(ctx, IntPredicate::NE, zero);
ctx.make_assert(
generator,
size_isnt_zero.value,
"0:ValueError",
format!("zero-size array to reduction operation {fn_name}").as_str(),
[None, None, None],
ctx.current_loc,
);
}
let extremum = generator.gen_var_alloc(ctx, dtype_llvm, None)?;
let extremum_idx = Int(SizeT).var_alloca(generator, ctx, None)?;
let first_value = ndarray.get_nth_scalar(generator, ctx, zero).value;
ctx.builder.build_store(extremum, first_value).unwrap();
extremum_idx.store(ctx, zero);
// The first element is iterated, but this doesn't matter.
ndarray
.foreach(generator, ctx, |generator, ctx, _hooks, nditer| {
let old_extremum = ctx.builder.build_load(extremum, "").unwrap();
let old_extremum_idx = extremum_idx.load(generator, ctx);
let curr_value = nditer.get_scalar(generator, ctx).value;
let curr_idx = nditer.get_index(generator, ctx);
let new_extremum = match fn_name {
"np_argmin" | "np_min" => call_min(
ctx,
(ndarray.dtype, old_extremum),
(ndarray.dtype, curr_value),
),
"np_argmax" | "np_max" => call_max(
ctx,
(ndarray.dtype, old_extremum),
(ndarray.dtype, curr_value),
),
_ => codegen_unreachable!(ctx),
};
let new_extremum_idx = match (old_extremum, new_extremum) {
(BasicValueEnum::IntValue(m), BasicValueEnum::IntValue(n)) => ctx
.builder
.build_select(
ctx.builder.build_int_compare(IntPredicate::NE, m, n, "").unwrap(),
curr_idx.value,
old_extremum_idx.value,
"",
)
.unwrap(),
(BasicValueEnum::FloatValue(m), BasicValueEnum::FloatValue(n)) => ctx
.builder
.build_select(
ctx.builder
.build_float_compare(FloatPredicate::ONE, m, n, "")
.unwrap(),
curr_idx.value,
old_extremum_idx.value,
"",
)
.unwrap(),
_ => unsupported_type(ctx, fn_name, &[ndarray.dtype, ndarray.dtype]),
};
ctx.builder.build_store(extremum, new_extremum).unwrap();
let new_extremum_idx =
unsafe { Int(SizeT).believe_value(new_extremum_idx.into_int_value()) };
extremum_idx.store(ctx, new_extremum_idx);
Ok(())
})
.unwrap();
match fn_name {
"np_argmin" | "np_argmax" => extremum_idx
.load(generator, ctx)
.s_extend_or_bit_cast(generator, ctx, Int64)
.value
.as_basic_value_enum(),
"np_max" | "np_min" => ctx.builder.build_load(extremum, "").unwrap(),
_ => codegen_unreachable!(ctx),
}
}
_ => unsupported_type(ctx, fn_name, &[a_ty]),
})
}
/// Invokes the `np_maximum` builtin function.
pub fn call_numpy_maximum<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>),
x2: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_maximum";
let (x1_ty, x1) = x1;
let (x2_ty, x2) = x2;
let common_ty = if ctx.unifier.unioned(x1_ty, x2_ty) { Some(x1_ty) } else { None };
Ok(match (x1, x2) {
(BasicValueEnum::IntValue(x1), BasicValueEnum::IntValue(x2)) => {
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
ctx.primitives.float,
]
.iter()
.any(|ty| ctx.unifier.unioned(common_ty.unwrap(), *ty)));
call_max(ctx, (x1_ty, x1.into()), (x2_ty, x2.into()))
}
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
debug_assert!(ctx.unifier.unioned(common_ty.unwrap(), ctx.primitives.float));
call_max(ctx, (x1_ty, x1.into()), (x2_ty, x2.into()))
}
_ if [&x1_ty, &x2_ty]
.into_iter()
.any(|ty| ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())) =>
{
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1).to_ndarray(generator, ctx);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2).to_ndarray(generator, ctx);
debug_assert!(ctx.unifier.unioned(x1.dtype, x2.dtype));
let common_dtype = x1.dtype;
let result = NDArrayObject::broadcast_starmap(
generator,
ctx,
&[x1, x2],
NDArrayOut::NewNDArray { dtype: common_dtype },
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
Ok(call_max(ctx, (x1.dtype, x1_scalar), (x2.dtype, x2_scalar)))
},
)
.unwrap();
result.instance.value.as_basic_value_enum()
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
})
}
/// Helper function to create a built-in elementwise unary numpy function that takes in either an ndarray or a scalar.
///
/// * `(arg_ty, arg_val)`: The [`Type`] and llvm value of the input argument.
/// * `fn_name`: The name of the function, only used when throwing an error with [`unsupported_type`]
/// * `get_ret_elem_type`: A function that takes in the input scalar [`Type`], and returns the function's return scalar [`Type`].
/// Return a constant [`Type`] here if the return type does not depend on the input type.
/// * `on_scalar`: The function that acts on the scalars of the input. Returns [`Option::None`]
/// if the scalar type & value are faulty and should panic with [`unsupported_type`].
fn helper_call_numpy_unary_elementwise<'ctx, OnScalarFn, RetElemFn, G>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(arg_ty, arg_val): (Type, BasicValueEnum<'ctx>),
fn_name: &str,
get_ret_elem_type: &RetElemFn,
on_scalar: &OnScalarFn,
) -> Result<BasicValueEnum<'ctx>, String>
where
G: CodeGenerator + ?Sized,
OnScalarFn: Fn(
&mut G,
&mut CodeGenContext<'ctx, '_>,
Type,
BasicValueEnum<'ctx>,
) -> Option<BasicValueEnum<'ctx>>,
RetElemFn: Fn(&mut CodeGenContext<'ctx, '_>, Type) -> Type,
{
let arg = AnyObject { ty: arg_ty, value: arg_val };
let arg = ScalarOrNDArray::split_object(generator, ctx, arg);
let dtype = arg.get_dtype();
let ret_ty = get_ret_elem_type(ctx, dtype);
let result = arg.map(generator, ctx, ret_ty, |generator, ctx, scalar| {
let Some(result) = on_scalar(generator, ctx, dtype, scalar) else {
unsupported_type(ctx, fn_name, &[arg_ty])
};
Ok(result)
})?;
Ok(result.to_basic_value_enum())
}
pub fn call_abs<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
n: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "abs";
helper_call_numpy_unary_elementwise(
generator,
ctx,
n,
FN_NAME,
&|_ctx, elem_ty| elem_ty,
&|_generator, ctx, val_ty, val| match val {
BasicValueEnum::IntValue(n) => Some({
debug_assert!([
ctx.primitives.bool,
ctx.primitives.int32,
ctx.primitives.uint32,
ctx.primitives.int64,
ctx.primitives.uint64,
]
.iter()
.any(|ty| ctx.unifier.unioned(val_ty, *ty)));
if [ctx.primitives.int32, ctx.primitives.int64]
.iter()
.any(|ty| ctx.unifier.unioned(val_ty, *ty))
{
llvm_intrinsics::call_int_abs(
ctx,
n,
ctx.ctx.bool_type().const_zero(),
Some(FN_NAME),
)
.into()
} else {
n.into()
}
}),
BasicValueEnum::FloatValue(n) => Some({
debug_assert!(ctx.unifier.unioned(val_ty, ctx.primitives.float));
llvm_intrinsics::call_float_fabs(ctx, n, Some(FN_NAME)).into()
}),
_ => None,
},
)
}
/// Macro to conveniently generate numpy functions with [`helper_call_numpy_unary_elementwise`].
///
/// Arguments:
/// * `$name:ident`: The identifier of the rust function to be generated.
/// * `$fn_name:literal`: To be passed to the `fn_name` parameter of [`helper_call_numpy_unary_elementwise`]
/// * `$get_ret_elem_type:expr`: To be passed to the `get_ret_elem_type` parameter of [`helper_call_numpy_unary_elementwise`].
/// But there is no need to make it a reference.
/// * `$on_scalar:expr`: To be passed to the `on_scalar` parameter of [`helper_call_numpy_unary_elementwise`].
/// But there is no need to make it a reference.
macro_rules! create_helper_call_numpy_unary_elementwise {
($name:ident, $fn_name:literal, $get_ret_elem_type:expr, $on_scalar:expr) => {
#[allow(clippy::redundant_closure_call)]
pub fn $name<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
arg: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
helper_call_numpy_unary_elementwise(
generator,
ctx,
arg,
$fn_name,
&$get_ret_elem_type,
&$on_scalar,
)
}
};
}
/// A specialized version of [`create_helper_call_numpy_unary_elementwise`] to generate functions that takes in float and returns boolean (as an `i8`) elementwise.
///
/// Arguments:
/// * `$name:ident`: The identifier of the rust function to be generated.
/// * `$fn_name:literal`: To be passed to the `fn_name` parameter of [`helper_call_numpy_unary_elementwise`].
/// * `$on_scalar:expr`: The closure (see below for its type) that acts on float scalar values and returns
/// the boolean results of LLVM type `i1`. The returned `i1` value will be converted into an `i8`.
///
/// ```ignore
/// // Type of `$on_scalar:expr`
/// fn on_scalar<'ctx, G: CodeGenerator + ?Sized>(
/// generator: &mut G,
/// ctx: &mut CodeGenContext<'ctx, '_>,
/// arg: FloatValue<'ctx>
/// ) -> IntValue<'ctx> // of LLVM type `i1`
/// ```
macro_rules! create_helper_call_numpy_unary_elementwise_float_to_bool {
($name:ident, $fn_name:literal, $on_scalar:expr) => {
create_helper_call_numpy_unary_elementwise!(
$name,
$fn_name,
|ctx, _| ctx.primitives.bool,
|generator, ctx, n_ty, val| {
match val {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(n_ty, ctx.primitives.float));
let ret = $on_scalar(generator, ctx, n);
Some(generator.bool_to_i8(ctx, ret).into())
}
_ => None,
}
}
);
};
}
/// A specialized version of [`create_helper_call_numpy_unary_elementwise`] to generate functions that takes in float and returns float elementwise.
///
/// Arguments:
/// * `$name:ident`: The identifier of the rust function to be generated.
/// * `$fn_name:literal`: To be passed to the `fn_name` parameter of [`helper_call_numpy_unary_elementwise`].
/// * `$on_scalar:expr`: The closure (see below for its type) that acts on float scalar values and returns float results.
///
/// ```ignore
/// // Type of `$on_scalar:expr`
/// fn on_scalar<'ctx, G: CodeGenerator + ?Sized>(
/// generator: &mut G,
/// ctx: &mut CodeGenContext<'ctx, '_>,
/// arg: FloatValue<'ctx>
/// ) -> FloatValue<'ctx>
/// ```
macro_rules! create_helper_call_numpy_unary_elementwise_float_to_float {
($name:ident, $fn_name:literal, $elem_call:expr) => {
create_helper_call_numpy_unary_elementwise!(
$name,
$fn_name,
|ctx, _| ctx.primitives.float,
|_generator, ctx, val_ty, val| {
match val {
BasicValueEnum::FloatValue(n) => {
debug_assert!(ctx.unifier.unioned(val_ty, ctx.primitives.float));
Some($elem_call(ctx, n, Option::<&str>::None).into())
}
_ => None,
}
}
);
};
}
create_helper_call_numpy_unary_elementwise_float_to_bool!(
call_numpy_isnan,
"np_isnan",
irrt::call_isnan
);
create_helper_call_numpy_unary_elementwise_float_to_bool!(
call_numpy_isinf,
"np_isinf",
irrt::call_isinf
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_sin,
"np_sin",
llvm_intrinsics::call_float_sin
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_cos,
"np_cos",
llvm_intrinsics::call_float_cos
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_tan,
"np_tan",
extern_fns::call_tan
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arcsin,
"np_arcsin",
extern_fns::call_asin
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arccos,
"np_arccos",
extern_fns::call_acos
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arctan,
"np_arctan",
extern_fns::call_atan
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_sinh,
"np_sinh",
extern_fns::call_sinh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_cosh,
"np_cosh",
extern_fns::call_cosh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_tanh,
"np_tanh",
extern_fns::call_tanh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arcsinh,
"np_arcsinh",
extern_fns::call_asinh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arccosh,
"np_arccosh",
extern_fns::call_acosh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_arctanh,
"np_arctanh",
extern_fns::call_atanh
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_exp,
"np_exp",
llvm_intrinsics::call_float_exp
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_exp2,
"np_exp2",
llvm_intrinsics::call_float_exp2
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_expm1,
"np_expm1",
extern_fns::call_expm1
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_log,
"np_log",
llvm_intrinsics::call_float_log
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_log2,
"np_log2",
llvm_intrinsics::call_float_log2
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_log10,
"np_log10",
llvm_intrinsics::call_float_log10
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_sqrt,
"np_sqrt",
llvm_intrinsics::call_float_sqrt
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_cbrt,
"np_cbrt",
extern_fns::call_cbrt
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_fabs,
"np_fabs",
llvm_intrinsics::call_float_fabs
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_numpy_rint,
"np_rint",
llvm_intrinsics::call_float_rint
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_erf,
"sp_spec_erf",
extern_fns::call_erf
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_erfc,
"sp_spec_erfc",
extern_fns::call_erfc
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_gamma,
"sp_spec_gamma",
|ctx, val, _| irrt::call_gamma(ctx, val)
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_gammaln,
"sp_spec_gammaln",
|ctx, val, _| irrt::call_gammaln(ctx, val)
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_j0,
"sp_spec_j0",
|ctx, val, _| irrt::call_j0(ctx, val)
);
create_helper_call_numpy_unary_elementwise_float_to_float!(
call_scipy_special_j1,
"sp_spec_j1",
extern_fns::call_j1
);
/// Invokes the `np_arctan2` builtin function.
pub fn call_numpy_arctan2<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_arctan2";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(extern_fns::call_atan2(ctx, x1, x2, None).as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_copysign` builtin function.
pub fn call_numpy_copysign<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_copysign";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(llvm_intrinsics::call_float_copysign(ctx, x1, x2, None)
.as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_fmax` builtin function.
pub fn call_numpy_fmax<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_fmax";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(llvm_intrinsics::call_float_maxnum(ctx, x1, x2, None).as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_fmin` builtin function.
pub fn call_numpy_fmin<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_fmin";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(llvm_intrinsics::call_float_minnum(ctx, x1, x2, None).as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_ldexp` builtin function.
pub fn call_numpy_ldexp<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_ldexp";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1_scalar), BasicValueEnum::IntValue(x2_scalar)) => {
debug_assert!(ctx.unifier.unioned(x1.get_dtype(), ctx.primitives.float));
debug_assert!(ctx.unifier.unioned(x2.get_dtype(), ctx.primitives.int32));
Ok(extern_fns::call_ldexp(ctx, x1_scalar, x2_scalar, None)
.as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_hypot` builtin function.
pub fn call_numpy_hypot<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_hypot";
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(extern_fns::call_hypot(ctx, x1, x2, None).as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_nextafter` builtin function.
pub fn call_numpy_nextafter<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>),
x2: (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_nextafter";
let (x1_ty, x1) = x1;
let (x2_ty, x2) = x2;
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = ScalarOrNDArray::split_object(generator, ctx, x1);
let x2 = AnyObject { ty: x2_ty, value: x2 };
let x2 = ScalarOrNDArray::split_object(generator, ctx, x2);
let result = ScalarOrNDArray::broadcasting_starmap(
generator,
ctx,
&[x1, x2],
ctx.primitives.float,
|_generator, ctx, scalars| {
let x1_scalar = scalars[0];
let x2_scalar = scalars[1];
match (x1_scalar, x2_scalar) {
(BasicValueEnum::FloatValue(x1), BasicValueEnum::FloatValue(x2)) => {
Ok(extern_fns::call_nextafter(ctx, x1, x2, None).as_basic_value_enum())
}
_ => unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]),
}
},
)
.unwrap();
Ok(result.to_basic_value_enum())
}
/// Invokes the `np_linalg_cholesky` linalg function
pub fn call_np_linalg_cholesky<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let out = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
out.copy_shape_from_ndarray(generator, ctx, x1);
out.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let out_c = out.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_cholesky(
ctx,
x1_c.value.as_basic_value_enum(),
out_c.value.as_basic_value_enum(),
None,
);
Ok(out.instance.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_qr` linalg function
pub fn call_np_linalg_qr<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let d0 = x1_shape.get_index_const(generator, ctx, 0);
let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk = unsafe {
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None))
};
let q = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, dk]);
q.create_data(generator, ctx);
let r = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk, d1]);
r.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let q_c = q.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let r_c = r.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_qr(
ctx,
x1_c.value.as_basic_value_enum(),
q_c.value.as_basic_value_enum(),
r_c.value.as_basic_value_enum(),
None,
);
let q = q.to_any(ctx);
let r = r.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [q, r]);
Ok(tuple.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_svd` linalg function
pub fn call_np_linalg_svd<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let d0 = x1_shape.get_index_const(generator, ctx, 0);
let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk = unsafe {
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None))
};
let u = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, d0]);
u.create_data(generator, ctx);
let s = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk]);
s.create_data(generator, ctx);
let vh = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d1, d1]);
vh.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let u_c = u.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let s_c = s.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let vh_c = vh.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_svd(
ctx,
x1_c.value.as_basic_value_enum(),
u_c.value.as_basic_value_enum(),
s_c.value.as_basic_value_enum(),
vh_c.value.as_basic_value_enum(),
None,
);
let u = u.to_any(ctx);
let s = s.to_any(ctx);
let vh = vh.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [u, s, vh]);
Ok(tuple.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_inv` linalg function
pub fn call_np_linalg_inv<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let out = NDArrayObject::alloca(generator, ctx, x1.dtype, 2);
out.copy_shape_from_ndarray(generator, ctx, x1);
out.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let out_c = out.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_inv(
ctx,
x1_c.value.as_basic_value_enum(),
out_c.value.as_basic_value_enum(),
None,
);
Ok(out.instance.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_pinv` linalg function
pub fn call_np_linalg_pinv<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let d0 = x1_shape.get_index_const(generator, ctx, 0);
let d1 = x1_shape.get_index_const(generator, ctx, 1);
let out = NDArrayObject::alloca_dynamic_shape(generator, ctx, x1.dtype, &[d1, d0]);
out.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let out_c = out.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_pinv(
ctx,
x1_c.value.as_basic_value_enum(),
out_c.value.as_basic_value_enum(),
None,
);
Ok(out.instance.value.as_basic_value_enum())
}
/// Invokes the `sp_linalg_lu` linalg function
pub fn call_sp_linalg_lu<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let d0 = x1_shape.get_index_const(generator, ctx, 0);
let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk = unsafe {
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None))
};
let l = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, dk]);
l.create_data(generator, ctx);
let u = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk, d1]);
u.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let l_c = l.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let u_c = u.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_sp_linalg_lu(
ctx,
x1_c.value.as_basic_value_enum(),
l_c.value.as_basic_value_enum(),
u_c.value.as_basic_value_enum(),
None,
);
let l = l.to_any(ctx);
let u = u.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [l, u]);
Ok(tuple.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_matrix_power` linalg function
pub fn call_np_linalg_matrix_power<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
(x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
// x2 is a float, but we are promoting this to a 1D ndarray (.shape == [1]) for uniformity in function call.
let x2 = call_float(generator, ctx, (x2_ty, x2)).unwrap();
let x2 = AnyObject { ty: ctx.primitives.float, value: x2 };
let x2 = NDArrayObject::make_unsized(generator, ctx, x2); // x2.shape == []
let x2 = x2.atleast_nd(generator, ctx, 1); // x2.shape == [1]
let out = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
out.copy_shape_from_ndarray(generator, ctx, x1);
out.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let x2_c = x2.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let out_c = out.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_matrix_power(
ctx,
x1_c.value.as_basic_value_enum(),
x2_c.value.as_basic_value_enum(),
out_c.value.as_basic_value_enum(),
None,
);
Ok(out.instance.value.as_basic_value_enum())
}
/// Invokes the `np_linalg_det` linalg function
pub fn call_np_linalg_det<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
// The output is a float64, but we are using an ndarray (shape == [1]) for uniformity in function call.
let det = NDArrayObject::alloca_constant_shape(generator, ctx, ctx.primitives.float, &[1]);
det.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let out_c = det.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_np_linalg_det(
ctx,
x1_c.value.as_basic_value_enum(),
out_c.value.as_basic_value_enum(),
None,
);
// Get the determinant out of `out`
let zero = Int(SizeT).const_0(generator, ctx.ctx);
let det = det.get_nth_scalar(generator, ctx, zero);
Ok(det.value)
}
/// Invokes the `sp_linalg_schur` linalg function
pub fn call_sp_linalg_schur<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
assert_eq!(x1.ndims, 2);
let t = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
t.copy_shape_from_ndarray(generator, ctx, x1);
t.create_data(generator, ctx);
let z = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
z.copy_shape_from_ndarray(generator, ctx, x1);
z.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let t_c = t.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let z_c = z.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_sp_linalg_schur(
ctx,
x1_c.value.as_basic_value_enum(),
t_c.value.as_basic_value_enum(),
z_c.value.as_basic_value_enum(),
None,
);
let t = t.to_any(ctx);
let z = z.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [t, z]);
Ok(tuple.value.as_basic_value_enum())
}
/// Invokes the `sp_linalg_hessenberg` linalg function
pub fn call_sp_linalg_hessenberg<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> {
let x1 = AnyObject { ty: x1_ty, value: x1 };
let x1 = NDArrayObject::from_object(generator, ctx, x1);
assert_eq!(x1.ndims, 2);
let h = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
h.copy_shape_from_ndarray(generator, ctx, x1);
h.create_data(generator, ctx);
let q = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
q.copy_shape_from_ndarray(generator, ctx, x1);
q.create_data(generator, ctx);
let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
let h_c = h.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let q_c = q.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
extern_fns::call_sp_linalg_hessenberg(
ctx,
x1_c.value.as_basic_value_enum(),
h_c.value.as_basic_value_enum(),
q_c.value.as_basic_value_enum(),
None,
);
let h = h.to_any(ctx);
let q = q.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [h, q]);
Ok(tuple.value.as_basic_value_enum())
}
View Git Blame Copy Permalink