lyken
/
nac3
forked from M-Labs/nac3
1
0
Fork 0
Code Pull Requests Activity

core/ndstrides: implement nalgebra functions

This commit is contained in:
lyken 2024-08-21 09:53:56 +08:00
parent 1562a938a1
commit 54f7e1edfd
No known key found for this signature in database
GPG Key ID: 3BD5FC6AC8325DD8
1 changed files with 222 additions and 393 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

615
nac3core/src/codegen/builtin_fns.rs
View File

@ -1,17 +1,14 @@
use inkwell::types::BasicTypeEnum; use inkwell::types::BasicTypeEnum;
use inkwell::values::{BasicValue, BasicValueEnum, IntValue, PointerValue}; use inkwell::values::{BasicValue, BasicValueEnum, IntValue};
use inkwell::{FloatPredicate, IntPredicate, OptimizationLevel}; use inkwell::{FloatPredicate, IntPredicate, OptimizationLevel};
use itertools::Itertools; use itertools::Itertools;
use crate::codegen::classes::{ use crate::codegen::classes::RangeValue;
NDArrayValue, ProxyValue, RangeValue, UntypedArrayLikeAccessor, UntypedArrayLikeMutator,
};
use crate::codegen::expr::destructure_range; use crate::codegen::expr::destructure_range;
use crate::codegen::irrt::calculate_len_for_slice_range; use crate::codegen::irrt::calculate_len_for_slice_range;
use crate::codegen::object::ndarray::{NDArrayOut, ScalarOrNDArray}; use crate::codegen::object::ndarray::{NDArrayOut, ScalarOrNDArray};
use crate::codegen::{extern_fns, irrt, llvm_intrinsics, numpy, CodeGenContext, CodeGenerator}; use crate::codegen::{extern_fns, irrt, llvm_intrinsics, CodeGenContext, CodeGenerator};
use crate::toplevel::helper::PrimDef; use crate::toplevel::helper::PrimDef;
use crate::toplevel::numpy::unpack_ndarray_var_tys;
use crate::typecheck::typedef::Type; use crate::typecheck::typedef::Type;
use super::model::*; use super::model::*;
@ -1607,500 +1604,332 @@ pub fn call_numpy_nextafter<'ctx, G: CodeGenerator + ?Sized>(
Ok(result.to_basic_value_enum()) Ok(result.to_basic_value_enum())
} }
/// Allocates a struct with the fields specified by `out_matrices` and returns a pointer to it
fn build_output_struct<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
out_matrices: Vec<BasicValueEnum<'ctx>>,
) -> PointerValue<'ctx> {
let field_ty =
out_matrices.iter().map(BasicValueEnum::get_type).collect::<Vec<BasicTypeEnum>>();
let out_ty = ctx.ctx.struct_type(&field_ty, false);
let out_ptr = ctx.builder.build_alloca(out_ty, "").unwrap();
for (i, v) in out_matrices.into_iter().enumerate() {
unsafe {
let ptr = ctx
.builder
.build_in_bounds_gep(
out_ptr,
&[
ctx.ctx.i32_type().const_zero(),
ctx.ctx.i32_type().const_int(i as u64, false),
],
"",
)
.unwrap();
ctx.builder.build_store(ptr, v).unwrap();
}
}
out_ptr
}
/// Invokes the `np_linalg_cholesky` linalg function /// Invokes the `np_linalg_cholesky` linalg function
pub fn call_np_linalg_cholesky<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_cholesky<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_cholesky"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let out = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); out.copy_shape_from_ndarray(generator, ctx, x1);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); out.create_data(generator, ctx);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
unsupported_type(ctx, FN_NAME, &[x1_ty]); 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,
);
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); Ok(out.instance.value.as_basic_value_enum())
let dim0 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let dim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let out = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim1])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_np_linalg_cholesky(ctx, x1, out, None);
Ok(out)
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `np_linalg_qr` linalg function /// Invokes the `np_linalg_qr` linalg function
pub fn call_np_linalg_qr<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_qr<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_qr"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); let d0 = x1_shape.get_index_const(generator, ctx, 0);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk =
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None));
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let q = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, dk]);
unimplemented!("{FN_NAME} operates on float type NdArrays only"); q.create_data(generator, ctx);
};
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); let r = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk, d1]);
let dim0 = unsafe { r.create_data(generator, ctx);
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let dim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let k = llvm_intrinsics::call_int_smin(ctx, dim0, dim1, None);
let out_q = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, k]) let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
.unwrap() let q_c = q.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.as_base_value() let r_c = r.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.as_basic_value_enum(); extern_fns::call_np_linalg_qr(
let out_r = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[k, dim1]) ctx,
.unwrap() x1_c.value.as_basic_value_enum(),
.as_base_value() q_c.value.as_basic_value_enum(),
.as_basic_value_enum(); r_c.value.as_basic_value_enum(),
None,
);
extern_fns::call_np_linalg_qr(ctx, x1, out_q, out_r, None); let q = q.to_any(ctx);
let r = r.to_any(ctx);
let out_ptr = build_output_struct(ctx, vec![out_q, out_r]); let tuple = TupleObject::from_objects(generator, ctx, [q, r]);
Ok(tuple.value.as_basic_value_enum())
Ok(ctx.builder.build_load(out_ptr, "QR_Factorization_result").map(Into::into).unwrap())
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `np_linalg_svd` linalg function /// Invokes the `np_linalg_svd` linalg function
pub fn call_np_linalg_svd<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_svd<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_svd"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); let d0 = x1_shape.get_index_const(generator, ctx, 0);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk =
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None));
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let u = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, d0]);
unsupported_type(ctx, FN_NAME, &[x1_ty]); u.create_data(generator, ctx);
};
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); let s = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk]);
s.create_data(generator, ctx);
let dim0 = unsafe { let vh = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d1, d1]);
n1.dim_sizes() vh.create_data(generator, ctx);
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let dim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let k = llvm_intrinsics::call_int_smin(ctx, dim0, dim1, None);
let out_u = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim0]) let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
.unwrap() let u_c = u.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.as_base_value() let s_c = s.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.as_basic_value_enum(); let vh_c = vh.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
let out_s = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[k]) extern_fns::call_np_linalg_svd(
.unwrap() ctx,
.as_base_value() x1_c.value.as_basic_value_enum(),
.as_basic_value_enum(); u_c.value.as_basic_value_enum(),
let out_vh = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim1, dim1]) s_c.value.as_basic_value_enum(),
.unwrap() vh_c.value.as_basic_value_enum(),
.as_base_value() None,
.as_basic_value_enum(); );
extern_fns::call_np_linalg_svd(ctx, x1, out_u, out_s, out_vh, None); let u = u.to_any(ctx);
let s = s.to_any(ctx);
let out_ptr = build_output_struct(ctx, vec![out_u, out_s, out_vh]); let vh = vh.to_any(ctx);
let tuple = TupleObject::from_objects(generator, ctx, [u, s, vh]);
Ok(ctx.builder.build_load(out_ptr, "SVD_Factorization_result").map(Into::into).unwrap()) Ok(tuple.value.as_basic_value_enum())
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `np_linalg_inv` linalg function /// Invokes the `np_linalg_inv` linalg function
pub fn call_np_linalg_inv<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_inv<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_inv"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let out = NDArrayObject::alloca(generator, ctx, x1.dtype, 2);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); out.copy_shape_from_ndarray(generator, ctx, x1);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); out.create_data(generator, ctx);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
unsupported_type(ctx, FN_NAME, &[x1_ty]); 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,
);
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); Ok(out.instance.value.as_basic_value_enum())
let dim0 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let dim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let out = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim1])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_np_linalg_inv(ctx, x1, out, None);
Ok(out)
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `np_linalg_pinv` linalg function /// Invokes the `np_linalg_pinv` linalg function
pub fn call_np_linalg_pinv<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_pinv<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_pinv"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); let d0 = x1_shape.get_index_const(generator, ctx, 0);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); let d1 = x1_shape.get_index_const(generator, ctx, 1);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let out = NDArrayObject::alloca_dynamic_shape(generator, ctx, x1.dtype, &[d1, d0]);
unsupported_type(ctx, FN_NAME, &[x1_ty]); out.create_data(generator, ctx);
};
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); 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,
);
let dim0 = unsafe { Ok(out.instance.value.as_basic_value_enum())
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let dim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let out = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim1, dim0])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_np_linalg_pinv(ctx, x1, out, None);
Ok(out)
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `sp_linalg_lu` linalg function /// Invokes the `sp_linalg_lu` linalg function
pub fn call_sp_linalg_lu<'ctx, G: CodeGenerator + ?Sized>( pub fn call_sp_linalg_lu<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "sp_linalg_lu"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx);
if let BasicValueEnum::PointerValue(n1) = x1 { let x1_shape = x1.instance.get(generator, ctx, |f| f.shape);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); let d0 = x1_shape.get_index_const(generator, ctx, 0);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); let d1 = x1_shape.get_index_const(generator, ctx, 1);
let dk =
Int(SizeT).believe_value(llvm_intrinsics::call_int_smin(ctx, d0.value, d1.value, None));
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let l = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[d0, dk]);
unsupported_type(ctx, FN_NAME, &[x1_ty]); l.create_data(generator, ctx);
};
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); let u = NDArrayObject::alloca_dynamic_shape(generator, ctx, ctx.primitives.float, &[dk, d1]);
u.create_data(generator, ctx);
let dim0 = unsafe { let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
n1.dim_sizes() let l_c = l.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None) let u_c = u.make_contiguous_ndarray(generator, ctx, Float(Float64)); // Shares `data`.
.into_int_value() extern_fns::call_sp_linalg_lu(
}; ctx,
let dim1 = unsafe { x1_c.value.as_basic_value_enum(),
n1.dim_sizes() l_c.value.as_basic_value_enum(),
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None) u_c.value.as_basic_value_enum(),
.into_int_value() None,
}; );
let k = llvm_intrinsics::call_int_smin(ctx, dim0, dim1, None);
let out_l = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, k]) let l = l.to_any(ctx);
.unwrap() let u = u.to_any(ctx);
.as_base_value() let tuple = TupleObject::from_objects(generator, ctx, [l, u]);
.as_basic_value_enum(); Ok(tuple.value.as_basic_value_enum())
let out_u = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[k, dim1])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_sp_linalg_lu(ctx, x1, out_l, out_u, None);
let out_ptr = build_output_struct(ctx, vec![out_l, out_u]);
Ok(ctx.builder.build_load(out_ptr, "LU_Factorization_result").map(Into::into).unwrap())
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `np_linalg_matrix_power` linalg function /// Invokes the `np_linalg_matrix_power` linalg function
pub fn call_np_linalg_matrix_power<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_matrix_power<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
x2: (Type, BasicValueEnum<'ctx>), (x2_ty, x2): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_matrix_power"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let (x2_ty, x2) = x2;
// 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 = 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 llvm_usize = generator.get_size_type(ctx.ctx); let out = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
if let (BasicValueEnum::PointerValue(n1), BasicValueEnum::FloatValue(n2)) = (x1, x2) { out.copy_shape_from_ndarray(generator, ctx, x1);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); out.create_data(generator, ctx);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]); 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,
);
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); Ok(out.instance.value.as_basic_value_enum())
// Changing second parameter to a `NDArray` for uniformity in function call
let n2_array = numpy::create_ndarray_const_shape(
generator,
ctx,
elem_ty,
&[llvm_usize.const_int(1, false)],
)
.unwrap();
unsafe {
n2_array.data().set_unchecked(
ctx,
generator,
&llvm_usize.const_zero(),
n2.as_basic_value_enum(),
);
};
let n2_array = n2_array.as_base_value().as_basic_value_enum();
let outdim0 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
.into_int_value()
};
let outdim1 = unsafe {
n1.dim_sizes()
.get_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
.into_int_value()
};
let out = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[outdim0, outdim1])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_np_linalg_matrix_power(ctx, x1, n2_array, out, None);
Ok(out)
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty])
}
} }
/// Invokes the `np_linalg_det` linalg function /// Invokes the `np_linalg_det` linalg function
pub fn call_np_linalg_det<'ctx, G: CodeGenerator + ?Sized>( pub fn call_np_linalg_det<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "np_linalg_matrix_power"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx); // The output is a float64, but we are using an ndarray (shape == [1]) for uniformity in function call.
if let BasicValueEnum::PointerValue(_) = x1 { let det = NDArrayObject::alloca_constant_shape(generator, ctx, ctx.primitives.float, &[1]);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); det.create_data(generator, ctx);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let x1_c = x1.make_contiguous_ndarray(generator, ctx, Float(Float64));
unsupported_type(ctx, FN_NAME, &[x1_ty]); 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,
);
// Changing second parameter to a `NDArray` for uniformity in function call // Get the determinant out of `out`
let out = numpy::create_ndarray_const_shape( let zero = Int(SizeT).const_0(generator, ctx.ctx);
generator, let det = det.get_nth_scalar(generator, ctx, zero);
ctx, Ok(det.value)
elem_ty,
&[llvm_usize.const_int(1, false)],
)
.unwrap();
extern_fns::call_np_linalg_det(ctx, x1, out.as_base_value().as_basic_value_enum(), None);
let res =
unsafe { out.data().get_unchecked(ctx, generator, &llvm_usize.const_zero(), None) };
Ok(res)
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `sp_linalg_schur` linalg function /// Invokes the `sp_linalg_schur` linalg function
pub fn call_sp_linalg_schur<'ctx, G: CodeGenerator + ?Sized>( pub fn call_sp_linalg_schur<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "sp_linalg_schur"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx); assert_eq!(x1.ndims, 2);
if let BasicValueEnum::PointerValue(n1) = x1 { let t = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); t.copy_shape_from_ndarray(generator, ctx, x1);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); t.create_data(generator, ctx);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let z = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
unsupported_type(ctx, FN_NAME, &[x1_ty]); z.copy_shape_from_ndarray(generator, ctx, x1);
}; z.create_data(generator, ctx);
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); 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 dim0 = unsafe { let t = t.to_any(ctx);
n1.dim_sizes() let z = z.to_any(ctx);
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None) let tuple = TupleObject::from_objects(generator, ctx, [t, z]);
.into_int_value() Ok(tuple.value.as_basic_value_enum())
};
let out_t = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim0])
.unwrap()
.as_base_value()
.as_basic_value_enum();
let out_z = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim0])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_sp_linalg_schur(ctx, x1, out_t, out_z, None);
let out_ptr = build_output_struct(ctx, vec![out_t, out_z]);
Ok(ctx.builder.build_load(out_ptr, "Schur_Factorization_result").map(Into::into).unwrap())
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }
/// Invokes the `sp_linalg_hessenberg` linalg function /// Invokes the `sp_linalg_hessenberg` linalg function
pub fn call_sp_linalg_hessenberg<'ctx, G: CodeGenerator + ?Sized>( pub fn call_sp_linalg_hessenberg<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
x1: (Type, BasicValueEnum<'ctx>), (x1_ty, x1): (Type, BasicValueEnum<'ctx>),
) -> Result<BasicValueEnum<'ctx>, String> { ) -> Result<BasicValueEnum<'ctx>, String> {
const FN_NAME: &str = "sp_linalg_hessenberg"; let x1 = AnyObject { ty: x1_ty, value: x1 };
let (x1_ty, x1) = x1; let x1 = NDArrayObject::from_object(generator, ctx, x1);
let llvm_usize = generator.get_size_type(ctx.ctx); assert_eq!(x1.ndims, 2);
if let BasicValueEnum::PointerValue(n1) = x1 { let h = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, x1_ty); h.copy_shape_from_ndarray(generator, ctx, x1);
let n1_elem_ty = ctx.get_llvm_type(generator, elem_ty); h.create_data(generator, ctx);
let BasicTypeEnum::FloatType(_) = n1_elem_ty else { let q = NDArrayObject::alloca(generator, ctx, ctx.primitives.float, 2);
unsupported_type(ctx, FN_NAME, &[x1_ty]); q.copy_shape_from_ndarray(generator, ctx, x1);
}; q.create_data(generator, ctx);
let n1 = NDArrayValue::from_ptr_val(n1, llvm_usize, None); 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 dim0 = unsafe { let h = h.to_any(ctx);
n1.dim_sizes() let q = q.to_any(ctx);
.get_unchecked(ctx, generator, &llvm_usize.const_zero(), None) let tuple = TupleObject::from_objects(generator, ctx, [h, q]);
.into_int_value() Ok(tuple.value.as_basic_value_enum())
};
let out_h = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim0])
.unwrap()
.as_base_value()
.as_basic_value_enum();
let out_q = numpy::create_ndarray_const_shape(generator, ctx, elem_ty, &[dim0, dim0])
.unwrap()
.as_base_value()
.as_basic_value_enum();
extern_fns::call_sp_linalg_hessenberg(ctx, x1, out_h, out_q, None);
let out_ptr = build_output_struct(ctx, vec![out_h, out_q]);
Ok(ctx
.builder
.build_load(out_ptr, "Hessenberg_decomposition_result")
.map(Into::into)
.unwrap())
} else {
unsupported_type(ctx, FN_NAME, &[x1_ty])
}
} }