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

Merge remote-tracking branch 'upstream/master'

Browse Source
This commit is contained in:
ram
2025-01-10 04:00:40 +00:00
parent 0d5be81ef4 8baf111734
commit c3bcf101d4
89 changed files with 5561 additions and 5196 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
Cargo.lock generated
View File

@@ -126,9 +126,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]]
name = "cc"
version = "1.2.4"
version = "1.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9157bbaa6b165880c27a4293a474c91cdcf265cc68cc829bf10be0964a391caf"
checksum = "a012a0df96dd6d06ba9a1b29d6402d1a5d77c6befd2566afdc26e10603dc93d7"
dependencies = [
"shlex",
]
@@ -170,7 +170,7 @@ dependencies = [
"heck 0.5.0",
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -187,14 +187,14 @@ checksum = "5b63caa9aa9397e2d9480a9b13673856c78d8ac123288526c37d7839f2a86990"
[[package]]
name = "console"
version = "0.15.8"
version = "0.15.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0e1f83fc076bd6dd27517eacdf25fef6c4dfe5f1d7448bafaaf3a26f13b5e4eb"
checksum = "ea3c6ecd8059b57859df5c69830340ed3c41d30e3da0c1cbed90a96ac853041b"
dependencies = [
"encode_unicode",
"lazy_static",
"libc",
"windows-sys 0.52.0",
"once_cell",
"windows-sys 0.59.0",
]
[[package]]
@@ -221,18 +221,18 @@ dependencies = [
[[package]]
name = "crossbeam-channel"
version = "0.5.13"
version = "0.5.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "33480d6946193aa8033910124896ca395333cae7e2d1113d1fef6c3272217df2"
checksum = "06ba6d68e24814cb8de6bb986db8222d3a027d15872cabc0d18817bc3c0e4471"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.5"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d"
checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
@@ -249,18 +249,18 @@ dependencies = [
[[package]]
name = "crossbeam-queue"
version = "0.3.11"
version = "0.3.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df0346b5d5e76ac2fe4e327c5fd1118d6be7c51dfb18f9b7922923f287471e35"
checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.20"
version = "0.8.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80"
checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
[[package]]
name = "crypto-common"
@@ -305,9 +305,9 @@ dependencies = [
[[package]]
name = "encode_unicode"
version = "0.3.6"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f"
checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0"
[[package]]
name = "equivalent"
@@ -378,9 +378,9 @@ dependencies = [
[[package]]
name = "glob"
version = "0.3.1"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2"
[[package]]
name = "hashbrown"
@@ -417,11 +417,11 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "home"
version = "0.5.9"
version = "0.5.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3d1354bf6b7235cb4a0576c2619fd4ed18183f689b12b006a0ee7329eeff9a5"
checksum = "589533453244b0995c858700322199b2becb13b627df2851f64a2775d024abcf"
dependencies = [
"windows-sys 0.52.0",
"windows-sys 0.59.0",
]
[[package]]
@@ -472,7 +472,7 @@ checksum = "9dd28cfd4cfba665d47d31c08a6ba637eed16770abca2eccbbc3ca831fef1e44"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -559,9 +559,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.168"
version = "0.2.169"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5aaeb2981e0606ca11d79718f8bb01164f1d6ed75080182d3abf017e6d244b6d"
checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a"
[[package]]
name = "libloading"
@@ -678,7 +678,7 @@ dependencies = [
"proc-macro-error",
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
"trybuild",
]
@@ -799,7 +799,7 @@ dependencies = [
"phf_shared 0.11.2",
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -927,7 +927,7 @@ dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -940,14 +940,14 @@ dependencies = [
"proc-macro2",
"pyo3-build-config",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
name = "quote"
version = "1.0.37"
version = "1.0.38"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc"
dependencies = [
"proc-macro2",
]
@@ -1062,9 +1062,9 @@ dependencies = [
[[package]]
name = "rustversion"
version = "1.0.18"
version = "1.0.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0e819f2bc632f285be6d7cd36e25940d45b2391dd6d9b939e79de557f7014248"
checksum = "f7c45b9784283f1b2e7fb61b42047c2fd678ef0960d4f6f1eba131594cc369d4"
[[package]]
name = "ryu"
@@ -1095,29 +1095,29 @@ checksum = "3cb6eb87a131f756572d7fb904f6e7b68633f09cca868c5df1c4b8d1a694bbba"
[[package]]
name = "serde"
version = "1.0.216"
version = "1.0.217"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e"
checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.216"
version = "1.0.217"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e"
checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
name = "serde_json"
version = "1.0.133"
version = "1.0.134"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7fceb2473b9166b2294ef05efcb65a3db80803f0b03ef86a5fc88a2b85ee377"
checksum = "d00f4175c42ee48b15416f6193a959ba3a0d67fc699a0db9ad12df9f83991c7d"
dependencies = [
"itoa",
"memchr",
@@ -1226,7 +1226,7 @@ dependencies = [
"proc-macro2",
"quote",
"rustversion",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -1242,9 +1242,9 @@ dependencies = [
[[package]]
name = "syn"
version = "2.0.90"
version = "2.0.94"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31"
checksum = "987bc0be1cdea8b10216bd06e2ca407d40b9543468fafd3ddfb02f36e77f71f3"
dependencies = [
"proc-macro2",
"quote",
@@ -1265,12 +1265,13 @@ checksum = "42a4d50cdb458045afc8131fd91b64904da29548bcb63c7236e0844936c13078"
[[package]]
name = "tempfile"
version = "3.14.0"
version = "3.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "28cce251fcbc87fac86a866eeb0d6c2d536fc16d06f184bb61aeae11aa4cee0c"
checksum = "9a8a559c81686f576e8cd0290cd2a24a2a9ad80c98b3478856500fcbd7acd704"
dependencies = [
"cfg-if",
"fastrand",
"getrandom",
"once_cell",
"rustix",
"windows-sys 0.59.0",
@@ -1278,9 +1279,9 @@ dependencies = [
[[package]]
name = "term"
version = "1.0.0"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4df4175de05129f31b80458c6df371a15e7fc3fd367272e6bf938e5c351c7ea0"
checksum = "a3bb6001afcea98122260987f8b7b5da969ecad46dbf0b5453702f776b491a41"
dependencies = [
"home",
"windows-sys 0.52.0",
@@ -1325,7 +1326,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]
[[package]]
@@ -1603,9 +1604,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
[[package]]
name = "winnow"
version = "0.6.20"
version = "0.6.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "36c1fec1a2bb5866f07c25f68c26e565c4c200aebb96d7e55710c19d3e8ac49b"
checksum = "39281189af81c07ec09db316b302a3e67bf9bd7cbf6c820b50e35fee9c2fa980"
dependencies = [
"memchr",
]
@@ -1637,5 +1638,5 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.90",
"syn 2.0.94",
]

6
flake.lock generated
View File

@@ -2,11 +2,11 @@
"nodes": {
"nixpkgs": {
"locked": {
"lastModified": 1733940404,
"narHash": "sha256-Pj39hSoUA86ZePPF/UXiYHHM7hMIkios8TYG29kQT4g=",
"lastModified": 1735834308,
"narHash": "sha256-dklw3AXr3OGO4/XT1Tu3Xz9n/we8GctZZ75ZWVqAVhk=",
"owner": "NixOS",
"repo": "nixpkgs",
"rev": "5d67ea6b4b63378b9c13be21e2ec9d1afc921713",
"rev": "6df24922a1400241dae323af55f30e4318a6ca65",
"type": "github"
},
"original": {

29
nac3artiq/demo/numpy_primitives_decay.py Normal file
View File

@@ -0,0 +1,29 @@
from min_artiq import *
import numpy
from numpy import int32
@nac3
class NumpyBoolDecay:
core: KernelInvariant[Core]
np_true: KernelInvariant[bool]
np_false: KernelInvariant[bool]
np_int: KernelInvariant[int32]
np_float: KernelInvariant[float]
np_str: KernelInvariant[str]
def __init__(self):
self.core = Core()
self.np_true = numpy.True_
self.np_false = numpy.False_
self.np_int = numpy.int32(0)
self.np_float = numpy.float64(0.0)
self.np_str = numpy.str_("")
@kernel
def run(self):
pass
if __name__ == "__main__":
NumpyBoolDecay().run()

10
nac3artiq/src/codegen.rs
View File

@@ -162,7 +162,7 @@ impl<'a> ArtiqCodeGenerator<'a> {
}
}
impl<'b> CodeGenerator for ArtiqCodeGenerator<'b> {
impl CodeGenerator for ArtiqCodeGenerator<'_> {
fn get_name(&self) -> &str {
&self.name
}
@@ -464,7 +464,7 @@ fn format_rpc_arg<'ctx>(
let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, arg_ty);
let ndims = extract_ndims(&ctx.unifier, ndims);
let dtype = ctx.get_llvm_type(generator, elem_ty);
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype, Some(ndims))
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype, ndims)
.map_value(arg.into_pointer_value(), None);
let ndims = llvm_usize.const_int(ndims, false);
@@ -597,7 +597,7 @@ fn format_rpc_ret<'ctx>(
let (dtype, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty);
let dtype_llvm = ctx.get_llvm_type(generator, dtype);
let ndims = extract_ndims(&ctx.unifier, ndims);
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype_llvm, Some(ndims))
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype_llvm, ndims)
.construct_uninitialized(generator, ctx, None);
// NOTE: Current content of `ndarray`:
@@ -1505,7 +1505,7 @@ pub fn call_rtio_log_impl<'ctx>(
/// Generates a call to `core_log`.
pub fn gen_core_log<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
obj: &Option<(Type, ValueEnum<'ctx>)>,
obj: Option<&(Type, ValueEnum<'ctx>)>,
fun: (&FunSignature, DefinitionId),
args: &[(Option<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,
@@ -1522,7 +1522,7 @@ pub fn gen_core_log<'ctx>(
/// Generates a call to `rtio_log`.
pub fn gen_rtio_log<'ctx>(
ctx: &mut CodeGenContext<'ctx, '_>,
obj: &Option<(Type, ValueEnum<'ctx>)>,
obj: Option<&(Type, ValueEnum<'ctx>)>,
fun: (&FunSignature, DefinitionId),
args: &[(Option<StrRef>, ValueEnum<'ctx>)],
generator: &mut dyn CodeGenerator,

4
nac3artiq/src/lib.rs
View File

@@ -330,7 +330,7 @@ impl Nac3 {
vars: into_var_map([arg_ty]),
},
Arc::new(GenCall::new(Box::new(move |ctx, obj, fun, args, generator| {
gen_core_log(ctx, &obj, fun, &args, generator)?;
gen_core_log(ctx, obj.as_ref(), fun, &args, generator)?;
Ok(None)
}))),
@@ -360,7 +360,7 @@ impl Nac3 {
vars: into_var_map([arg_ty]),
},
Arc::new(GenCall::new(Box::new(move |ctx, obj, fun, args, generator| {
gen_rtio_log(ctx, &obj, fun, &args, generator)?;
gen_rtio_log(ctx, obj.as_ref(), fun, &args, generator)?;
Ok(None)
}))),

64
nac3artiq/src/symbol_resolver.rs
View File

@@ -931,10 +931,13 @@ impl InnerResolver {
|_| Ok(Ok(extracted_ty)),
)
} else if unifier.unioned(extracted_ty, primitives.bool) {
obj.extract::<bool>().map_or_else(
|_| Ok(Err(format!("{obj} is not in the range of bool"))),
|_| Ok(Ok(extracted_ty)),
)
if obj.extract::<bool>().is_ok()
|| obj.call_method("__bool__", (), None)?.extract::<bool>().is_ok()
{
Ok(Ok(extracted_ty))
} else {
Ok(Err(format!("{obj} is not in the range of bool")))
}
} else if unifier.unioned(extracted_ty, primitives.float) {
obj.extract::<f64>().map_or_else(
|_| Ok(Err(format!("{obj} is not in the range of float64"))),
@@ -974,10 +977,14 @@ impl InnerResolver {
let val: u64 = obj.extract().unwrap();
self.id_to_primitive.write().insert(id, PrimitiveValue::U64(val));
Ok(Some(ctx.ctx.i64_type().const_int(val, false).into()))
} else if ty_id == self.primitive_ids.bool || ty_id == self.primitive_ids.np_bool_ {
} else if ty_id == self.primitive_ids.bool {
let val: bool = obj.extract().unwrap();
self.id_to_primitive.write().insert(id, PrimitiveValue::Bool(val));
Ok(Some(ctx.ctx.i8_type().const_int(u64::from(val), false).into()))
} else if ty_id == self.primitive_ids.np_bool_ {
let val: bool = obj.call_method("__bool__", (), None)?.extract().unwrap();
self.id_to_primitive.write().insert(id, PrimitiveValue::Bool(val));
Ok(Some(ctx.ctx.i8_type().const_int(u64::from(val), false).into()))
} else if ty_id == self.primitive_ids.string || ty_id == self.primitive_ids.np_str_ {
let val: String = obj.extract().unwrap();
self.id_to_primitive.write().insert(id, PrimitiveValue::Str(val.clone()));
@@ -1107,7 +1114,7 @@ impl InnerResolver {
self.global_value_ids.write().insert(id, obj.into());
}
let ndims = llvm_ndarray.ndims().unwrap();
let ndims = llvm_ndarray.ndims();
// Obtain the shape of the ndarray
let shape_tuple: &PyTuple = obj.getattr("shape")?.downcast()?;
@@ -1124,7 +1131,10 @@ impl InnerResolver {
super::CompileError::new_err(format!("Error getting element {i}: {e}"))
})?
.unwrap();
let value = value.into_int_value();
let value = ctx
.builder
.build_int_z_extend(value.into_int_value(), llvm_usize, "")
.unwrap();
Ok(value)
})
.collect::<Result<Vec<_>, PyErr>>()?;
@@ -1203,8 +1213,16 @@ impl InnerResolver {
data_global.set_initializer(&data);
// Get the constant itemsize.
let itemsize = dtype.size_of().unwrap();
let itemsize = itemsize.get_zero_extended_constant().unwrap();
//
// NOTE: dtype.size_of() may return a non-constant, where `TargetData::get_store_size`
// will always return a constant size.
let itemsize = ctx
.registry
.llvm_options
.create_target_machine()
.map(|tm| tm.get_target_data().get_store_size(&dtype))
.unwrap();
assert_ne!(itemsize, 0);
// Create the strides needed for ndarray.strides
let strides = make_contiguous_strides(itemsize, ndims, &shape_u64s);
@@ -1214,7 +1232,7 @@ impl InnerResolver {
// create a global for ndarray.strides and initialize it
let strides_global = ctx.module.add_global(
llvm_i8.array_type(ndims as u32),
llvm_usize.array_type(ndims as u32),
Some(AddressSpace::default()),
&format!("${id_str}.strides"),
);
@@ -1230,9 +1248,30 @@ impl InnerResolver {
let ndarray_ndims = llvm_usize.const_int(ndims, false);
// calling as_pointer_value on shape and strides returns [i64 x ndims]*
// convert into i64* to conform with expected layout of ndarray
let ndarray_shape = shape_global.as_pointer_value();
let ndarray_shape = unsafe {
ctx.builder
.build_in_bounds_gep(
ndarray_shape,
&[llvm_usize.const_zero(), llvm_usize.const_zero()],
"",
)
.unwrap()
};
let ndarray_strides = strides_global.as_pointer_value();
let ndarray_strides = unsafe {
ctx.builder
.build_in_bounds_gep(
ndarray_strides,
&[llvm_usize.const_zero(), llvm_usize.const_zero()],
"",
)
.unwrap()
};
let ndarray = llvm_ndarray
.as_base_type()
@@ -1413,9 +1452,12 @@ impl InnerResolver {
} else if ty_id == self.primitive_ids.uint64 {
let val: u64 = obj.extract()?;
Ok(SymbolValue::U64(val))
} else if ty_id == self.primitive_ids.bool || ty_id == self.primitive_ids.np_bool_ {
} else if ty_id == self.primitive_ids.bool {
let val: bool = obj.extract()?;
Ok(SymbolValue::Bool(val))
} else if ty_id == self.primitive_ids.np_bool_ {
let val: bool = obj.call_method("__bool__", (), None)?.extract()?;
Ok(SymbolValue::Bool(val))
} else if ty_id == self.primitive_ids.string || ty_id == self.primitive_ids.np_str_ {
let val: String = obj.extract()?;
Ok(SymbolValue::Str(val))

7
nac3core/irrt/irrt.cpp
View File

@@ -1,7 +1,6 @@
#include "irrt/exception.hpp"
#include "irrt/list.hpp"
#include "irrt/math.hpp"
#include "irrt/ndarray.hpp"
#include "irrt/range.hpp"
#include "irrt/slice.hpp"
#include "irrt/string.hpp"
@@ -9,4 +8,8 @@
#include "irrt/ndarray/def.hpp"
#include "irrt/ndarray/iter.hpp"
#include "irrt/ndarray/indexing.hpp"
#include "irrt/string.hpp"
#include "irrt/ndarray/array.hpp"
#include "irrt/ndarray/reshape.hpp"
#include "irrt/ndarray/broadcast.hpp"
#include "irrt/ndarray/transpose.hpp"
#include "irrt/ndarray/matmul.hpp"

2
nac3core/irrt/irrt/int_types.hpp
View File

@@ -21,7 +21,5 @@ using uint64_t = unsigned _ExtInt(64);
#endif
// NDArray indices are always `uint32_t`.
using NDIndexInt = uint32_t;
// The type of an index or a value describing the length of a range/slice is always `int32_t`.
using SliceIndex = int32_t;

15
nac3core/irrt/irrt/list.hpp
View File

@@ -2,6 +2,21 @@
#include "irrt/int_types.hpp"
#include "irrt/math_util.hpp"
#include "irrt/slice.hpp"
namespace {
/**
* @brief A list in NAC3.
*
* The `items` field is opaque. You must rely on external contexts to
* know how to interpret it.
*/
template<typename SizeT>
struct List {
uint8_t* items;
SizeT len;
};
} // namespace
extern "C" {
// Handle list assignment and dropping part of the list when

2
nac3core/irrt/irrt/math.hpp
View File

@@ -1,5 +1,7 @@
#pragma once
#include "irrt/int_types.hpp"
namespace {
// adapted from GNU Scientific Library: https://git.savannah.gnu.org/cgit/gsl.git/tree/sys/pow_int.c
// need to make sure `exp >= 0` before calling this function

151
nac3core/irrt/irrt/ndarray.hpp
View File

@@ -1,151 +0,0 @@
#pragma once
#include "irrt/int_types.hpp"
// TODO: To be deleted since NDArray with strides is done.
namespace {
template<typename SizeT>
SizeT __nac3_ndarray_calc_size_impl(const SizeT* list_data, SizeT list_len, SizeT begin_idx, SizeT end_idx) {
__builtin_assume(end_idx <= list_len);
SizeT num_elems = 1;
for (SizeT i = begin_idx; i < end_idx; ++i) {
SizeT val = list_data[i];
__builtin_assume(val > 0);
num_elems *= val;
}
return num_elems;
}
template<typename SizeT>
void __nac3_ndarray_calc_nd_indices_impl(SizeT index, const SizeT* dims, SizeT num_dims, NDIndexInt* idxs) {
SizeT stride = 1;
for (SizeT dim = 0; dim < num_dims; dim++) {
SizeT i = num_dims - dim - 1;
__builtin_assume(dims[i] > 0);
idxs[i] = (index / stride) % dims[i];
stride *= dims[i];
}
}
template<typename SizeT>
SizeT __nac3_ndarray_flatten_index_impl(const SizeT* dims,
SizeT num_dims,
const NDIndexInt* indices,
SizeT num_indices) {
SizeT idx = 0;
SizeT stride = 1;
for (SizeT i = 0; i < num_dims; ++i) {
SizeT ri = num_dims - i - 1;
if (ri < num_indices) {
idx += stride * indices[ri];
}
__builtin_assume(dims[i] > 0);
stride *= dims[ri];
}
return idx;
}
template<typename SizeT>
void __nac3_ndarray_calc_broadcast_impl(const SizeT* lhs_dims,
SizeT lhs_ndims,
const SizeT* rhs_dims,
SizeT rhs_ndims,
SizeT* out_dims) {
SizeT max_ndims = lhs_ndims > rhs_ndims ? lhs_ndims : rhs_ndims;
for (SizeT i = 0; i < max_ndims; ++i) {
const SizeT* lhs_dim_sz = i < lhs_ndims ? &lhs_dims[lhs_ndims - i - 1] : nullptr;
const SizeT* rhs_dim_sz = i < rhs_ndims ? &rhs_dims[rhs_ndims - i - 1] : nullptr;
SizeT* out_dim = &out_dims[max_ndims - i - 1];
if (lhs_dim_sz == nullptr) {
*out_dim = *rhs_dim_sz;
} else if (rhs_dim_sz == nullptr) {
*out_dim = *lhs_dim_sz;
} else if (*lhs_dim_sz == 1) {
*out_dim = *rhs_dim_sz;
} else if (*rhs_dim_sz == 1) {
*out_dim = *lhs_dim_sz;
} else if (*lhs_dim_sz == *rhs_dim_sz) {
*out_dim = *lhs_dim_sz;
} else {
__builtin_unreachable();
}
}
}
template<typename SizeT>
void __nac3_ndarray_calc_broadcast_idx_impl(const SizeT* src_dims,
SizeT src_ndims,
const NDIndexInt* in_idx,
NDIndexInt* out_idx) {
for (SizeT i = 0; i < src_ndims; ++i) {
SizeT src_i = src_ndims - i - 1;
out_idx[src_i] = src_dims[src_i] == 1 ? 0 : in_idx[src_i];
}
}
} // namespace
extern "C" {
uint32_t __nac3_ndarray_calc_size(const uint32_t* list_data, uint32_t list_len, uint32_t begin_idx, uint32_t end_idx) {
return __nac3_ndarray_calc_size_impl(list_data, list_len, begin_idx, end_idx);
}
uint64_t
__nac3_ndarray_calc_size64(const uint64_t* list_data, uint64_t list_len, uint64_t begin_idx, uint64_t end_idx) {
return __nac3_ndarray_calc_size_impl(list_data, list_len, begin_idx, end_idx);
}
void __nac3_ndarray_calc_nd_indices(uint32_t index, const uint32_t* dims, uint32_t num_dims, NDIndexInt* idxs) {
__nac3_ndarray_calc_nd_indices_impl(index, dims, num_dims, idxs);
}
void __nac3_ndarray_calc_nd_indices64(uint64_t index, const uint64_t* dims, uint64_t num_dims, NDIndexInt* idxs) {
__nac3_ndarray_calc_nd_indices_impl(index, dims, num_dims, idxs);
}
uint32_t
__nac3_ndarray_flatten_index(const uint32_t* dims, uint32_t num_dims, const NDIndexInt* indices, uint32_t num_indices) {
return __nac3_ndarray_flatten_index_impl(dims, num_dims, indices, num_indices);
}
uint64_t __nac3_ndarray_flatten_index64(const uint64_t* dims,
uint64_t num_dims,
const NDIndexInt* indices,
uint64_t num_indices) {
return __nac3_ndarray_flatten_index_impl(dims, num_dims, indices, num_indices);
}
void __nac3_ndarray_calc_broadcast(const uint32_t* lhs_dims,
uint32_t lhs_ndims,
const uint32_t* rhs_dims,
uint32_t rhs_ndims,
uint32_t* out_dims) {
return __nac3_ndarray_calc_broadcast_impl(lhs_dims, lhs_ndims, rhs_dims, rhs_ndims, out_dims);
}
void __nac3_ndarray_calc_broadcast64(const uint64_t* lhs_dims,
uint64_t lhs_ndims,
const uint64_t* rhs_dims,
uint64_t rhs_ndims,
uint64_t* out_dims) {
return __nac3_ndarray_calc_broadcast_impl(lhs_dims, lhs_ndims, rhs_dims, rhs_ndims, out_dims);
}
void __nac3_ndarray_calc_broadcast_idx(const uint32_t* src_dims,
uint32_t src_ndims,
const NDIndexInt* in_idx,
NDIndexInt* out_idx) {
__nac3_ndarray_calc_broadcast_idx_impl(src_dims, src_ndims, in_idx, out_idx);
}
void __nac3_ndarray_calc_broadcast_idx64(const uint64_t* src_dims,
uint64_t src_ndims,
const NDIndexInt* in_idx,
NDIndexInt* out_idx) {
__nac3_ndarray_calc_broadcast_idx_impl(src_dims, src_ndims, in_idx, out_idx);
}
} // namespace

132
nac3core/irrt/irrt/ndarray/array.hpp Normal file
View File

@@ -0,0 +1,132 @@
#pragma once
#include "irrt/debug.hpp"
#include "irrt/exception.hpp"
#include "irrt/int_types.hpp"
#include "irrt/list.hpp"
#include "irrt/ndarray/basic.hpp"
#include "irrt/ndarray/def.hpp"
namespace {
namespace ndarray::array {
/**
* @brief In the context of `np.array(<list>)`, deduce the ndarray's shape produced by `<list>` and raise
* an exception if there is anything wrong with `<shape>` (e.g., inconsistent dimensions `np.array([[1.0, 2.0],
* [3.0]])`)
*
* If this function finds no issues with `<list>`, the deduced shape is written to `shape`. The caller has the
* responsibility to allocate `[SizeT; ndims]` for `shape`. The caller must also initialize `shape` with `-1`s because
* of implementation details.
*/
template<typename SizeT>
void set_and_validate_list_shape_helper(SizeT axis, List<SizeT>* list, SizeT ndims, SizeT* shape) {
if (shape[axis] == -1) {
// Dimension is unspecified. Set it.
shape[axis] = list->len;
} else {
// Dimension is specified. Check.
if (shape[axis] != list->len) {
// Mismatch, throw an error.
// NOTE: NumPy's error message is more complex and needs more PARAMS to display.
raise_exception(SizeT, EXN_VALUE_ERROR,
"The requested array has an inhomogenous shape "
"after {0} dimension(s).",
axis, shape[axis], list->len);
}
}
if (axis + 1 == ndims) {
// `list` has type `list[ItemType]`
// Do nothing
} else {
// `list` has type `list[list[...]]`
List<SizeT>** lists = (List<SizeT>**)(list->items);
for (SizeT i = 0; i < list->len; i++) {
set_and_validate_list_shape_helper<SizeT>(axis + 1, lists[i], ndims, shape);
}
}
}
/**
* @brief See `set_and_validate_list_shape_helper`.
*/
template<typename SizeT>
void set_and_validate_list_shape(List<SizeT>* list, SizeT ndims, SizeT* shape) {
for (SizeT axis = 0; axis < ndims; axis++) {
shape[axis] = -1; // Sentinel to say this dimension is unspecified.
}
set_and_validate_list_shape_helper<SizeT>(0, list, ndims, shape);
}
/**
* @brief In the context of `np.array(<list>)`, copied the contents stored in `list` to `ndarray`.
*
* `list` is assumed to be "legal". (i.e., no inconsistent dimensions)
*
* # Notes on `ndarray`
* The caller is responsible for allocating space for `ndarray`.
* Here is what this function expects from `ndarray` when called:
* - `ndarray->data` has to be allocated, contiguous, and may contain uninitialized values.
* - `ndarray->itemsize` has to be initialized.
* - `ndarray->ndims` has to be initialized.
* - `ndarray->shape` has to be initialized.
* - `ndarray->strides` is ignored, but note that `ndarray->data` is contiguous.
* When this function call ends:
* - `ndarray->data` is written with contents from `<list>`.
*/
template<typename SizeT>
void write_list_to_array_helper(SizeT axis, SizeT* index, List<SizeT>* list, NDArray<SizeT>* ndarray) {
debug_assert_eq(SizeT, list->len, ndarray->shape[axis]);
if (IRRT_DEBUG_ASSERT_BOOL) {
if (!ndarray::basic::is_c_contiguous(ndarray)) {
raise_debug_assert(SizeT, "ndarray is not C-contiguous", ndarray->strides[0], ndarray->strides[1],
NO_PARAM);
}
}
if (axis + 1 == ndarray->ndims) {
// `list` has type `list[scalar]`
// `ndarray` is contiguous, so we can do this, and this is fast.
uint8_t* dst = static_cast<uint8_t*>(ndarray->data) + (ndarray->itemsize * (*index));
__builtin_memcpy(dst, list->items, ndarray->itemsize * list->len);
*index += list->len;
} else {
// `list` has type `list[list[...]]`
List<SizeT>** lists = (List<SizeT>**)(list->items);
for (SizeT i = 0; i < list->len; i++) {
write_list_to_array_helper<SizeT>(axis + 1, index, lists[i], ndarray);
}
}
}
/**
* @brief See `write_list_to_array_helper`.
*/
template<typename SizeT>
void write_list_to_array(List<SizeT>* list, NDArray<SizeT>* ndarray) {
SizeT index = 0;
write_list_to_array_helper<SizeT>((SizeT)0, &index, list, ndarray);
}
} // namespace ndarray::array
} // namespace
extern "C" {
using namespace ndarray::array;
void __nac3_ndarray_array_set_and_validate_list_shape(List<int32_t>* list, int32_t ndims, int32_t* shape) {
set_and_validate_list_shape(list, ndims, shape);
}
void __nac3_ndarray_array_set_and_validate_list_shape64(List<int64_t>* list, int64_t ndims, int64_t* shape) {
set_and_validate_list_shape(list, ndims, shape);
}
void __nac3_ndarray_array_write_list_to_array(List<int32_t>* list, NDArray<int32_t>* ndarray) {
write_list_to_array(list, ndarray);
}
void __nac3_ndarray_array_write_list_to_array64(List<int64_t>* list, NDArray<int64_t>* ndarray) {
write_list_to_array(list, ndarray);
}
}

6
nac3core/irrt/irrt/ndarray/basic.hpp
View File

@@ -6,8 +6,7 @@
#include "irrt/ndarray/def.hpp"
namespace {
namespace ndarray {
namespace basic {
namespace ndarray::basic {
/**
* @brief Assert that `shape` does not contain negative dimensions.
*
@@ -247,8 +246,7 @@ void copy_data(const NDArray<SizeT>* src_ndarray, NDArray<SizeT>* dst_ndarray) {
ndarray::basic::set_pelement_value(dst_ndarray, dst_element, src_element);
}
}
} // namespace basic
} // namespace ndarray
} // namespace ndarray::basic
} // namespace
extern "C" {

165
nac3core/irrt/irrt/ndarray/broadcast.hpp Normal file
View File

@@ -0,0 +1,165 @@
#pragma once
#include "irrt/int_types.hpp"
#include "irrt/ndarray/def.hpp"
#include "irrt/slice.hpp"
namespace {
template<typename SizeT>
struct ShapeEntry {
SizeT ndims;
SizeT* shape;
};
} // namespace
namespace {
namespace ndarray::broadcast {
/**
* @brief Return true if `src_shape` can broadcast to `dst_shape`.
*
* See https://numpy.org/doc/stable/user/basics.broadcasting.html
*/
template<typename SizeT>
bool can_broadcast_shape_to(SizeT target_ndims, const SizeT* target_shape, SizeT src_ndims, const SizeT* src_shape) {
if (src_ndims > target_ndims) {
return false;
}
for (SizeT i = 0; i < src_ndims; i++) {
SizeT target_dim = target_shape[target_ndims - i - 1];
SizeT src_dim = src_shape[src_ndims - i - 1];
if (!(src_dim == 1 || target_dim == src_dim)) {
return false;
}
}
return true;
}
/**
* @brief Performs `np.broadcast_shapes(<shapes>)`
*
* @param num_shapes Number of entries in `shapes`
* @param shapes The list of shape to do `np.broadcast_shapes` on.
* @param dst_ndims The length of `dst_shape`.
* `dst_ndims` must be `max([shape.ndims for shape in shapes])`, but the caller has to calculate it/provide it.
* for this function since they should already know in order to allocate `dst_shape` in the first place.
* @param dst_shape The resulting shape. Must be pre-allocated by the caller. This function calculate the result
* of `np.broadcast_shapes` and write it here.
*/
template<typename SizeT>
void broadcast_shapes(SizeT num_shapes, const ShapeEntry<SizeT>* shapes, SizeT dst_ndims, SizeT* dst_shape) {
for (SizeT dst_axis = 0; dst_axis < dst_ndims; dst_axis++) {
dst_shape[dst_axis] = 1;
}
#ifdef IRRT_DEBUG_ASSERT
SizeT max_ndims_found = 0;
#endif
for (SizeT i = 0; i < num_shapes; i++) {
ShapeEntry<SizeT> entry = shapes[i];
// Check pre-condition: `dst_ndims` must be `max([shape.ndims for shape in shapes])`
debug_assert(SizeT, entry.ndims <= dst_ndims);
#ifdef IRRT_DEBUG_ASSERT
max_ndims_found = max(max_ndims_found, entry.ndims);
#endif
for (SizeT j = 0; j < entry.ndims; j++) {
SizeT entry_axis = entry.ndims - j - 1;
SizeT dst_axis = dst_ndims - j - 1;
SizeT entry_dim = entry.shape[entry_axis];
SizeT dst_dim = dst_shape[dst_axis];
if (dst_dim == 1) {
dst_shape[dst_axis] = entry_dim;
} else if (entry_dim == 1 || entry_dim == dst_dim) {
// Do nothing
} else {
raise_exception(SizeT, EXN_VALUE_ERROR,
"shape mismatch: objects cannot be broadcast "
"to a single shape.",
NO_PARAM, NO_PARAM, NO_PARAM);
}
}
}
#ifdef IRRT_DEBUG_ASSERT
// Check pre-condition: `dst_ndims` must be `max([shape.ndims for shape in shapes])`
debug_assert_eq(SizeT, max_ndims_found, dst_ndims);
#endif
}
/**
* @brief Perform `np.broadcast_to(<ndarray>, <target_shape>)` and appropriate assertions.
*
* This function attempts to broadcast `src_ndarray` to a new shape defined by `dst_ndarray.shape`,
* and return the result by modifying `dst_ndarray`.
*
* # Notes on `dst_ndarray`
* The caller is responsible for allocating space for the resulting ndarray.
* Here is what this function expects from `dst_ndarray` when called:
* - `dst_ndarray->data` does not have to be initialized.
* - `dst_ndarray->itemsize` does not have to be initialized.
* - `dst_ndarray->ndims` must be initialized, determining the length of `dst_ndarray->shape`
* - `dst_ndarray->shape` must be allocated, and must contain the desired target broadcast shape.
* - `dst_ndarray->strides` must be allocated, through it can contain uninitialized values.
* When this function call ends:
* - `dst_ndarray->data` is set to `src_ndarray->data` (`dst_ndarray` is just a view to `src_ndarray`)
* - `dst_ndarray->itemsize` is set to `src_ndarray->itemsize`
* - `dst_ndarray->ndims` is unchanged.
* - `dst_ndarray->shape` is unchanged.
* - `dst_ndarray->strides` is updated accordingly by how ndarray broadcast_to works.
*/
template<typename SizeT>
void broadcast_to(const NDArray<SizeT>* src_ndarray, NDArray<SizeT>* dst_ndarray) {
if (!ndarray::broadcast::can_broadcast_shape_to(dst_ndarray->ndims, dst_ndarray->shape, src_ndarray->ndims,
src_ndarray->shape)) {
raise_exception(SizeT, EXN_VALUE_ERROR, "operands could not be broadcast together", NO_PARAM, NO_PARAM,
NO_PARAM);
}
dst_ndarray->data = src_ndarray->data;
dst_ndarray->itemsize = src_ndarray->itemsize;
for (SizeT i = 0; i < dst_ndarray->ndims; i++) {
SizeT src_axis = src_ndarray->ndims - i - 1;
SizeT dst_axis = dst_ndarray->ndims - i - 1;
if (src_axis < 0 || (src_ndarray->shape[src_axis] == 1 && dst_ndarray->shape[dst_axis] != 1)) {
// Freeze the steps in-place
dst_ndarray->strides[dst_axis] = 0;
} else {
dst_ndarray->strides[dst_axis] = src_ndarray->strides[src_axis];
}
}
}
} // namespace ndarray::broadcast
} // namespace
extern "C" {
using namespace ndarray::broadcast;
void __nac3_ndarray_broadcast_to(NDArray<int32_t>* src_ndarray, NDArray<int32_t>* dst_ndarray) {
broadcast_to(src_ndarray, dst_ndarray);
}
void __nac3_ndarray_broadcast_to64(NDArray<int64_t>* src_ndarray, NDArray<int64_t>* dst_ndarray) {
broadcast_to(src_ndarray, dst_ndarray);
}
void __nac3_ndarray_broadcast_shapes(int32_t num_shapes,
const ShapeEntry<int32_t>* shapes,
int32_t dst_ndims,
int32_t* dst_shape) {
broadcast_shapes(num_shapes, shapes, dst_ndims, dst_shape);
}
void __nac3_ndarray_broadcast_shapes64(int64_t num_shapes,
const ShapeEntry<int64_t>* shapes,
int64_t dst_ndims,
int64_t* dst_shape) {
broadcast_shapes(num_shapes, shapes, dst_ndims, dst_shape);
}
}

9
nac3core/irrt/irrt/ndarray/indexing.hpp
View File

@@ -65,8 +65,7 @@ struct NDIndex {
} // namespace
namespace {
namespace ndarray {
namespace indexing {
namespace ndarray::indexing {
/**
* @brief Perform ndarray "basic indexing" (https://numpy.org/doc/stable/user/basics.indexing.html#basic-indexing)
*
@@ -162,7 +161,8 @@ void index(SizeT num_indices, const NDIndex* indices, const NDArray<SizeT>* src_
Range<int32_t> range = slice->indices_checked<SizeT>(src_ndarray->shape[src_axis]);
dst_ndarray->data = static_cast<uint8_t*>(dst_ndarray->data) + (SizeT)range.start * src_ndarray->strides[src_axis];
dst_ndarray->data =
static_cast<uint8_t*>(dst_ndarray->data) + (SizeT)range.start * src_ndarray->strides[src_axis];
dst_ndarray->strides[dst_axis] = ((SizeT)range.step) * src_ndarray->strides[src_axis];
dst_ndarray->shape[dst_axis] = (SizeT)range.len<SizeT>();
@@ -197,8 +197,7 @@ void index(SizeT num_indices, const NDIndex* indices, const NDArray<SizeT>* src_
debug_assert_eq(SizeT, src_ndarray->ndims, src_axis);
debug_assert_eq(SizeT, dst_ndarray->ndims, dst_axis);
}
} // namespace indexing
} // namespace ndarray
} // namespace ndarray::indexing
} // namespace
extern "C" {

98
nac3core/irrt/irrt/ndarray/matmul.hpp Normal file
View File

@@ -0,0 +1,98 @@
#pragma once
#include "irrt/debug.hpp"
#include "irrt/exception.hpp"
#include "irrt/int_types.hpp"
#include "irrt/ndarray/basic.hpp"
#include "irrt/ndarray/broadcast.hpp"
#include "irrt/ndarray/iter.hpp"
// NOTE: Everything would be much easier and elegant if einsum is implemented.
namespace {
namespace ndarray::matmul {
/**
* @brief Perform the broadcast in `np.einsum("...ij,...jk->...ik", a, b)`.
*
* Example:
* Suppose `a_shape == [1, 97, 4, 2]`
* and `b_shape == [99, 98, 1, 2, 5]`,
*
* ...then `new_a_shape == [99, 98, 97, 4, 2]`,
* `new_b_shape == [99, 98, 97, 2, 5]`,
* and `dst_shape == [99, 98, 97, 4, 5]`.
* ^^^^^^^^^^ ^^^^
* (broadcasted) (4x2 @ 2x5 => 4x5)
*
* @param a_ndims Length of `a_shape`.
* @param a_shape Shape of `a`.
* @param b_ndims Length of `b_shape`.
* @param b_shape Shape of `b`.
* @param final_ndims Should be equal to `max(a_ndims, b_ndims)`. This is the length of `new_a_shape`,
* `new_b_shape`, and `dst_shape` - the number of dimensions after broadcasting.
*/
template<typename SizeT>
void calculate_shapes(SizeT a_ndims,
SizeT* a_shape,
SizeT b_ndims,
SizeT* b_shape,
SizeT final_ndims,
SizeT* new_a_shape,
SizeT* new_b_shape,
SizeT* dst_shape) {
debug_assert(SizeT, a_ndims >= 2);
debug_assert(SizeT, b_ndims >= 2);
debug_assert_eq(SizeT, max(a_ndims, b_ndims), final_ndims);
// Check that a and b are compatible for matmul
if (a_shape[a_ndims - 1] != b_shape[b_ndims - 2]) {
// This is a custom error message. Different from NumPy.
raise_exception(SizeT, EXN_VALUE_ERROR, "Cannot multiply LHS (shape ?x{0}) with RHS (shape {1}x?})",
a_shape[a_ndims - 1], b_shape[b_ndims - 2], NO_PARAM);
}
const SizeT num_entries = 2;
ShapeEntry<SizeT> entries[num_entries] = {{.ndims = a_ndims - 2, .shape = a_shape},
{.ndims = b_ndims - 2, .shape = b_shape}};
// TODO: Optimize this
ndarray::broadcast::broadcast_shapes<SizeT>(num_entries, entries, final_ndims - 2, new_a_shape);
ndarray::broadcast::broadcast_shapes<SizeT>(num_entries, entries, final_ndims - 2, new_b_shape);
ndarray::broadcast::broadcast_shapes<SizeT>(num_entries, entries, final_ndims - 2, dst_shape);
new_a_shape[final_ndims - 2] = a_shape[a_ndims - 2];
new_a_shape[final_ndims - 1] = a_shape[a_ndims - 1];
new_b_shape[final_ndims - 2] = b_shape[b_ndims - 2];
new_b_shape[final_ndims - 1] = b_shape[b_ndims - 1];
dst_shape[final_ndims - 2] = a_shape[a_ndims - 2];
dst_shape[final_ndims - 1] = b_shape[b_ndims - 1];
}
} // namespace ndarray::matmul
} // namespace
extern "C" {
using namespace ndarray::matmul;
void __nac3_ndarray_matmul_calculate_shapes(int32_t a_ndims,
int32_t* a_shape,
int32_t b_ndims,
int32_t* b_shape,
int32_t final_ndims,
int32_t* new_a_shape,
int32_t* new_b_shape,
int32_t* dst_shape) {
calculate_shapes(a_ndims, a_shape, b_ndims, b_shape, final_ndims, new_a_shape, new_b_shape, dst_shape);
}
void __nac3_ndarray_matmul_calculate_shapes64(int64_t a_ndims,
int64_t* a_shape,
int64_t b_ndims,
int64_t* b_shape,
int64_t final_ndims,
int64_t* new_a_shape,
int64_t* new_b_shape,
int64_t* dst_shape) {
calculate_shapes(a_ndims, a_shape, b_ndims, b_shape, final_ndims, new_a_shape, new_b_shape, dst_shape);
}
}

97
nac3core/irrt/irrt/ndarray/reshape.hpp Normal file
View File

@@ -0,0 +1,97 @@
#pragma once
#include "irrt/exception.hpp"
#include "irrt/int_types.hpp"
#include "irrt/ndarray/def.hpp"
namespace {
namespace ndarray::reshape {
/**
* @brief Perform assertions on and resolve unknown dimensions in `new_shape` in `np.reshape(<ndarray>, new_shape)`
*
* If `new_shape` indeed contains unknown dimensions (specified with `-1`, just like numpy), `new_shape` will be
* modified to contain the resolved dimension.
*
* To perform assertions on and resolve unknown dimensions in `new_shape`, we don't need the actual
* `<ndarray>` object itself, but only the `.size` of the `<ndarray>`.
*
* @param size The `.size` of `<ndarray>`
* @param new_ndims Number of elements in `new_shape`
* @param new_shape Target shape to reshape to
*/
template<typename SizeT>
void resolve_and_check_new_shape(SizeT size, SizeT new_ndims, SizeT* new_shape) {
// Is there a -1 in `new_shape`?
bool neg1_exists = false;
// Location of -1, only initialized if `neg1_exists` is true
SizeT neg1_axis_i;
// The computed ndarray size of `new_shape`
SizeT new_size = 1;
for (SizeT axis_i = 0; axis_i < new_ndims; axis_i++) {
SizeT dim = new_shape[axis_i];
if (dim < 0) {
if (dim == -1) {
if (neg1_exists) {
// Multiple `-1` found. Throw an error.
raise_exception(SizeT, EXN_VALUE_ERROR, "can only specify one unknown dimension", NO_PARAM,
NO_PARAM, NO_PARAM);
} else {
neg1_exists = true;
neg1_axis_i = axis_i;
}
} else {
// TODO: What? In `np.reshape` any negative dimensions is
// treated like its `-1`.
//
// Try running `np.zeros((3, 4)).reshape((-999, 2))`
//
// It is not documented by numpy.
// Throw an error for now...
raise_exception(SizeT, EXN_VALUE_ERROR, "Found non -1 negative dimension {0} on axis {1}", dim, axis_i,
NO_PARAM);
}
} else {
new_size *= dim;
}
}
bool can_reshape;
if (neg1_exists) {
// Let `x` be the unknown dimension
// Solve `x * <new_size> = <size>`
if (new_size == 0 && size == 0) {
// `x` has infinitely many solutions
can_reshape = false;
} else if (new_size == 0 && size != 0) {
// `x` has no solutions
can_reshape = false;
} else if (size % new_size != 0) {
// `x` has no integer solutions
can_reshape = false;
} else {
can_reshape = true;
new_shape[neg1_axis_i] = size / new_size; // Resolve dimension
}
} else {
can_reshape = (new_size == size);
}
if (!can_reshape) {
raise_exception(SizeT, EXN_VALUE_ERROR, "cannot reshape array of size {0} into given shape", size, NO_PARAM,
NO_PARAM);
}
}
} // namespace ndarray::reshape
} // namespace
extern "C" {
void __nac3_ndarray_reshape_resolve_and_check_new_shape(int32_t size, int32_t new_ndims, int32_t* new_shape) {
ndarray::reshape::resolve_and_check_new_shape(size, new_ndims, new_shape);
}
void __nac3_ndarray_reshape_resolve_and_check_new_shape64(int64_t size, int64_t new_ndims, int64_t* new_shape) {
ndarray::reshape::resolve_and_check_new_shape(size, new_ndims, new_shape);
}
}

143
nac3core/irrt/irrt/ndarray/transpose.hpp Normal file
View File

@@ -0,0 +1,143 @@
#pragma once
#include "irrt/debug.hpp"
#include "irrt/exception.hpp"
#include "irrt/int_types.hpp"
#include "irrt/ndarray/def.hpp"
#include "irrt/slice.hpp"
/*
* Notes on `np.transpose(<array>, <axes>)`
*
* TODO: `axes`, if specified, can actually contain negative indices,
* but it is not documented in numpy.
*
* Supporting it for now.
*/
namespace {
namespace ndarray::transpose {
/**
* @brief Do assertions on `<axes>` in `np.transpose(<array>, <axes>)`.
*
* Note that `np.transpose`'s `<axe>` argument is optional. If the argument
* is specified but the user, use this function to do assertions on it.
*
* @param ndims The number of dimensions of `<array>`
* @param num_axes Number of elements in `<axes>` as specified by the user.
* This should be equal to `ndims`. If not, a "ValueError: axes don't match array" is thrown.
* @param axes The user specified `<axes>`.
*/
template<typename SizeT>
void assert_transpose_axes(SizeT ndims, SizeT num_axes, const SizeT* axes) {
if (ndims != num_axes) {
raise_exception(SizeT, EXN_VALUE_ERROR, "axes don't match array", NO_PARAM, NO_PARAM, NO_PARAM);
}
// TODO: Optimize this
bool* axe_specified = (bool*)__builtin_alloca(sizeof(bool) * ndims);
for (SizeT i = 0; i < ndims; i++)
axe_specified[i] = false;
for (SizeT i = 0; i < ndims; i++) {
SizeT axis = slice::resolve_index_in_length(ndims, axes[i]);
if (axis == -1) {
// TODO: numpy actually throws a `numpy.exceptions.AxisError`
raise_exception(SizeT, EXN_VALUE_ERROR, "axis {0} is out of bounds for array of dimension {1}", axis, ndims,
NO_PARAM);
}
if (axe_specified[axis]) {
raise_exception(SizeT, EXN_VALUE_ERROR, "repeated axis in transpose", NO_PARAM, NO_PARAM, NO_PARAM);
}
axe_specified[axis] = true;
}
}
/**
* @brief Create a transpose view of `src_ndarray` and perform proper assertions.
*
* This function is very similar to doing `dst_ndarray = np.transpose(src_ndarray, <axes>)`.
* If `<axes>` is supposed to be `None`, caller can pass in a `nullptr` to `<axes>`.
*
* The transpose view created is returned by modifying `dst_ndarray`.
*
* The caller is responsible for setting up `dst_ndarray` before calling this function.
* Here is what this function expects from `dst_ndarray` when called:
* - `dst_ndarray->data` does not have to be initialized.
* - `dst_ndarray->itemsize` does not have to be initialized.
* - `dst_ndarray->ndims` must be initialized, must be equal to `src_ndarray->ndims`.
* - `dst_ndarray->shape` must be allocated, through it can contain uninitialized values.
* - `dst_ndarray->strides` must be allocated, through it can contain uninitialized values.
* When this function call ends:
* - `dst_ndarray->data` is set to `src_ndarray->data` (`dst_ndarray` is just a view to `src_ndarray`)
* - `dst_ndarray->itemsize` is set to `src_ndarray->itemsize`
* - `dst_ndarray->ndims` is unchanged
* - `dst_ndarray->shape` is updated according to how `np.transpose` works
* - `dst_ndarray->strides` is updated according to how `np.transpose` works
*
* @param src_ndarray The NDArray to build a transpose view on
* @param dst_ndarray The resulting NDArray after transpose. Further details in the comments above,
* @param num_axes Number of elements in axes. Unused if `axes` is nullptr.
* @param axes Axes permutation. Set it to `nullptr` if `<axes>` is `None`.
*/
template<typename SizeT>
void transpose(const NDArray<SizeT>* src_ndarray, NDArray<SizeT>* dst_ndarray, SizeT num_axes, const SizeT* axes) {
debug_assert_eq(SizeT, src_ndarray->ndims, dst_ndarray->ndims);
const auto ndims = src_ndarray->ndims;
if (axes != nullptr)
assert_transpose_axes(ndims, num_axes, axes);
dst_ndarray->data = src_ndarray->data;
dst_ndarray->itemsize = src_ndarray->itemsize;
// Check out https://ajcr.net/stride-guide-part-2/ to see how `np.transpose` works behind the scenes.
if (axes == nullptr) {
// `np.transpose(<array>, axes=None)`
/*
* Minor note: `np.transpose(<array>, axes=None)` is equivalent to
* `np.transpose(<array>, axes=[N-1, N-2, ..., 0])` - basically it
* is reversing the order of strides and shape.
*
* This is a fast implementation to handle this special (but very common) case.
*/
for (SizeT axis = 0; axis < ndims; axis++) {
dst_ndarray->shape[axis] = src_ndarray->shape[ndims - axis - 1];
dst_ndarray->strides[axis] = src_ndarray->strides[ndims - axis - 1];
}
} else {
// `np.transpose(<array>, <axes>)`
// Permute strides and shape according to `axes`, while resolving negative indices in `axes`
for (SizeT axis = 0; axis < ndims; axis++) {
// `i` cannot be OUT_OF_BOUNDS because of assertions
SizeT i = slice::resolve_index_in_length(ndims, axes[axis]);
dst_ndarray->shape[axis] = src_ndarray->shape[i];
dst_ndarray->strides[axis] = src_ndarray->strides[i];
}
}
}
} // namespace ndarray::transpose
} // namespace
extern "C" {
using namespace ndarray::transpose;
void __nac3_ndarray_transpose(const NDArray<int32_t>* src_ndarray,
NDArray<int32_t>* dst_ndarray,
int32_t num_axes,
const int32_t* axes) {
transpose(src_ndarray, dst_ndarray, num_axes, axes);
}
void __nac3_ndarray_transpose64(const NDArray<int64_t>* src_ndarray,
NDArray<int64_t>* dst_ndarray,
int64_t num_axes,
const int64_t* axes) {
transpose(src_ndarray, dst_ndarray, num_axes, axes);
}
}

10
nac3core/irrt/irrt/string.hpp
View File

@@ -4,20 +4,20 @@
namespace {
template<typename SizeT>
SizeT __nac3_str_eq_impl(const char* str1, SizeT len1, const char* str2, SizeT len2) {
if (len1 != len2){
bool __nac3_str_eq_impl(const char* str1, SizeT len1, const char* str2, SizeT len2) {
if (len1 != len2) {
return 0;
}
return (__builtin_memcmp(str1, str2, static_cast<SizeT>(len1)) == 0) ? 1 : 0;
return __builtin_memcmp(str1, str2, static_cast<SizeT>(len1)) == 0;
}
} // namespace
extern "C" {
uint32_t nac3_str_eq(const char* str1, uint32_t len1, const char* str2, uint32_t len2) {
bool nac3_str_eq(const char* str1, uint32_t len1, const char* str2, uint32_t len2) {
return __nac3_str_eq_impl<uint32_t>(str1, len1, str2, len2);
}
uint64_t nac3_str_eq64(const char* str1, uint64_t len1, const char* str2, uint64_t len2) {
bool nac3_str_eq64(const char* str1, uint64_t len1, const char* str2, uint64_t len2) {
return __nac3_str_eq_impl<uint64_t>(str1, len1, str2, len2);
}
}

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

File diff suppressed because it is too large Load Diff

671
nac3core/src/codegen/expr.rs
View File

@@ -11,7 +11,7 @@ use inkwell::{
values::{BasicValueEnum, CallSiteValue, FunctionValue, IntValue, PointerValue, StructValue},
AddressSpace, IntPredicate, OptimizationLevel,
};
use itertools::{chain, izip, Either, Itertools};
use itertools::{izip, Either, Itertools};
use nac3parser::ast::{
self, Boolop, Cmpop, Comprehension, Constant, Expr, ExprKind, Location, Operator, StrRef,
@@ -27,23 +27,23 @@ use super::{
call_memcpy_generic,
},
macros::codegen_unreachable,
need_sret, numpy,
need_sret,
stmt::{
gen_for_callback_incrementing, gen_if_callback, gen_if_else_expr_callback, gen_raise,
gen_var,
},
types::{ndarray::NDArrayType, ListType},
values::{
ndarray::{NDArrayValue, RustNDIndex},
ndarray::{NDArrayOut, RustNDIndex, ScalarOrNDArray},
ArrayLikeIndexer, ArrayLikeValue, ListValue, ProxyValue, RangeValue,
TypedArrayLikeAccessor, UntypedArrayLikeAccessor,
UntypedArrayLikeAccessor,
},
CodeGenContext, CodeGenTask, CodeGenerator,
};
use crate::{
symbol_resolver::{SymbolValue, ValueEnum},
toplevel::{
helper::{extract_ndims, PrimDef},
helper::{arraylike_flatten_element_type, PrimDef},
numpy::unpack_ndarray_var_tys,
DefinitionId, TopLevelDef,
},
@@ -79,7 +79,7 @@ pub fn get_subst_key(
.join(", ")
}
impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
impl<'ctx> CodeGenContext<'ctx, '_> {
/// Builds a sequence of `getelementptr` and `load` instructions which stores the value of a
/// struct field into an LLVM value.
pub fn build_gep_and_load(
@@ -1095,33 +1095,6 @@ pub fn destructure_range<'ctx>(
(start, end, step)
}
/// Allocates a List structure with the given [type][ty] and [length]. The name of the resulting
/// LLVM value is `{name}.addr`, or `list.addr` if [name] is not specified.
///
/// Setting `ty` to [`None`] implies that the list is empty **and** does not have a known element
/// type, and will therefore set the `list.data` type as `size_t*`. It is undefined behavior to
/// generate a sized list with an unknown element type.
pub fn allocate_list<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ty: Option<BasicTypeEnum<'ctx>>,
length: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ListValue<'ctx> {
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_elem_ty = ty.unwrap_or(llvm_usize.into());
// List structure; type { ty*, size_t }
let arr_ty = ListType::new(generator, ctx.ctx, llvm_elem_ty);
let list = arr_ty.alloca(generator, ctx, name);
let length = ctx.builder.build_int_z_extend(length, llvm_usize, "").unwrap();
list.store_size(ctx, generator, length);
list.create_data(ctx, llvm_elem_ty, None);
list
}
/// Generates LLVM IR for a [list comprehension expression][expr].
pub fn gen_comprehension<'ctx, G: CodeGenerator>(
generator: &mut G,
@@ -1194,12 +1167,11 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
"listcomp.alloc_size",
)
.unwrap();
list = allocate_list(
list = ListType::new(generator, ctx.ctx, elem_ty).construct(
generator,
ctx,
Some(elem_ty),
list_alloc_size.into_int_value(),
Some("listcomp.addr"),
Some("listcomp"),
);
let i = generator.gen_store_target(ctx, target, Some("i.addr"))?.unwrap();
@@ -1246,7 +1218,12 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
Some("length"),
)
.into_int_value();
list = allocate_list(generator, ctx, Some(elem_ty), length, Some("listcomp"));
list = ListType::new(generator, ctx.ctx, elem_ty).construct(
generator,
ctx,
length,
Some("listcomp"),
);
let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?;
// counter = -1
@@ -1411,7 +1388,8 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
.build_int_add(lhs.load_size(ctx, None), rhs.load_size(ctx, None), "")
.unwrap();
let new_list = allocate_list(generator, ctx, Some(llvm_elem_ty), size, None);
let new_list = ListType::new(generator, ctx.ctx, llvm_elem_ty)
.construct(generator, ctx, size, None);
let lhs_size = ctx
.builder
@@ -1498,10 +1476,9 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
let elem_llvm_ty = ctx.get_llvm_type(generator, elem_ty);
let sizeof_elem = elem_llvm_ty.size_of().unwrap();
let new_list = allocate_list(
let new_list = ListType::new(generator, ctx.ctx, elem_llvm_ty).construct(
generator,
ctx,
Some(elem_llvm_ty),
ctx.builder.build_int_mul(list_val.load_size(ctx, None), int_val, "").unwrap(),
None,
);
@@ -1554,98 +1531,77 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
} else if ty1.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())
|| ty2.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())
{
let is_ndarray1 = ty1.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id());
let is_ndarray2 = ty2.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id());
let left = ScalarOrNDArray::from_value(generator, ctx, (ty1, left_val));
let right = ScalarOrNDArray::from_value(generator, ctx, (ty2, right_val));
if is_ndarray1 && is_ndarray2 {
let (ndarray_dtype1, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty1);
let (ndarray_dtype2, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty2);
let ty1_dtype = arraylike_flatten_element_type(&mut ctx.unifier, ty1);
let ty2_dtype = arraylike_flatten_element_type(&mut ctx.unifier, ty2);
assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
// Inhomogeneous binary operations are not supported.
assert!(ctx.unifier.unioned(ty1_dtype, ty2_dtype));
let left_val = NDArrayType::from_unifier_type(generator, ctx, ty1)
.map_value(left_val.into_pointer_value(), None);
let right_val = NDArrayType::from_unifier_type(generator, ctx, ty2)
.map_value(right_val.into_pointer_value(), None);
let common_dtype = ty1_dtype;
let llvm_common_dtype = left.get_dtype();
let res = if op.base == Operator::MatMult {
// MatMult is the only binop which is not an elementwise op
numpy::ndarray_matmul_2d(
generator,
ctx,
ndarray_dtype1,
match op.variant {
BinopVariant::Normal => None,
BinopVariant::AugAssign => Some(left_val),
},
left_val,
right_val,
)?
} else {
numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ndarray_dtype1,
match op.variant {
BinopVariant::Normal => None,
BinopVariant::AugAssign => Some(left_val),
},
(ty1, left_val.as_base_value().into(), false),
(ty2, right_val.as_base_value().into(), false),
|generator, ctx, (lhs, rhs)| {
gen_binop_expr_with_values(
generator,
ctx,
(&Some(ndarray_dtype1), lhs),
op,
(&Some(ndarray_dtype2), rhs),
ctx.current_loc,
)?
.unwrap()
.to_basic_value_enum(
ctx,
generator,
ndarray_dtype1,
)
},
)?
};
let out = match op.variant {
BinopVariant::Normal => NDArrayOut::NewNDArray { dtype: llvm_common_dtype },
BinopVariant::AugAssign => {
// Augmented assignment - `left` has to be an ndarray. If it were a scalar then NAC3
// simply doesn't support it.
if let ScalarOrNDArray::NDArray(out_ndarray) = left {
NDArrayOut::WriteToNDArray { ndarray: out_ndarray }
} else {
panic!("left must be an ndarray")
}
}
};
Ok(Some(res.as_base_value().into()))
if op.base == Operator::MatMult {
let left = left.to_ndarray(generator, ctx);
let right = right.to_ndarray(generator, ctx);
let result = left
.matmul(generator, ctx, ty1, (ty2, right), (common_dtype, out))
.split_unsized(generator, ctx);
Ok(Some(result.to_basic_value_enum().into()))
} else {
let (ndarray_dtype, _) =
unpack_ndarray_var_tys(&mut ctx.unifier, if is_ndarray1 { ty1 } else { ty2 });
let ndarray_val =
NDArrayType::from_unifier_type(generator, ctx, if is_ndarray1 { ty1 } else { ty2 })
.map_value(
if is_ndarray1 { left_val } else { right_val }.into_pointer_value(),
None,
);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ndarray_dtype,
match op.variant {
BinopVariant::Normal => None,
BinopVariant::AugAssign => Some(ndarray_val),
},
(ty1, left_val, !is_ndarray1),
(ty2, right_val, !is_ndarray2),
|generator, ctx, (lhs, rhs)| {
gen_binop_expr_with_values(
generator,
ctx,
(&Some(ndarray_dtype), lhs),
op,
(&Some(ndarray_dtype), rhs),
ctx.current_loc,
)?
.unwrap()
.to_basic_value_enum(ctx, generator, ndarray_dtype)
},
)?;
// For other operations, they are all elementwise operations.
Ok(Some(res.as_base_value().into()))
// There are only three cases:
// - LHS is a scalar, RHS is an ndarray.
// - LHS is an ndarray, RHS is a scalar.
// - LHS is an ndarray, RHS is an ndarray.
//
// For all cases, the scalar operand is promoted to an ndarray,
// the two are then broadcasted, and starmapped through.
let left = left.to_ndarray(generator, ctx);
let right = right.to_ndarray(generator, ctx);
let result = NDArrayType::new_broadcast(
generator,
ctx.ctx,
llvm_common_dtype,
&[left.get_type(), right.get_type()],
)
.broadcast_starmap(generator, ctx, &[left, right], out, |generator, ctx, scalars| {
let left_value = scalars[0];
let right_value = scalars[1];
let result = gen_binop_expr_with_values(
generator,
ctx,
(&Some(ty1_dtype), left_value),
op,
(&Some(ty2_dtype), right_value),
ctx.current_loc,
)?
.unwrap()
.to_basic_value_enum(ctx, generator, common_dtype)?;
Ok(result)
})
.unwrap();
Ok(Some(result.as_base_value().into()))
}
} else {
let left_ty_enum = ctx.unifier.get_ty_immutable(left_ty.unwrap());
@@ -1808,10 +1764,10 @@ pub fn gen_unaryop_expr_with_values<'ctx, G: CodeGenerator>(
_ => val.into(),
}
} else if ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) {
let llvm_ndarray_ty = NDArrayType::from_unifier_type(generator, ctx, ty);
let (ndarray_dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
let val = llvm_ndarray_ty.map_value(val.into_pointer_value(), None);
let ndarray = NDArrayType::from_unifier_type(generator, ctx, ty)
.map_value(val.into_pointer_value(), None);
// ndarray uses `~` rather than `not` to perform elementwise inversion, convert it before
// passing it to the elementwise codegen function
@@ -1829,20 +1785,18 @@ pub fn gen_unaryop_expr_with_values<'ctx, G: CodeGenerator>(
op
};
let res = numpy::ndarray_elementwise_unaryop_impl(
let mapped_ndarray = ndarray.map(
generator,
ctx,
ndarray_dtype,
None,
val,
|generator, ctx, val| {
gen_unaryop_expr_with_values(generator, ctx, op, (&Some(ndarray_dtype), val))?
NDArrayOut::NewNDArray { dtype: ndarray.get_type().element_type() },
|generator, ctx, scalar| {
gen_unaryop_expr_with_values(generator, ctx, op, (&Some(ndarray_dtype), scalar))?
.map(|val| val.to_basic_value_enum(ctx, generator, ndarray_dtype))
.unwrap()
.to_basic_value_enum(ctx, generator, ndarray_dtype)
},
)?;
res.as_base_value().into()
mapped_ndarray.as_base_value().into()
} else {
unimplemented!()
}))
@@ -1885,87 +1839,56 @@ pub fn gen_cmpop_expr_with_values<'ctx, G: CodeGenerator>(
if left_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())
|| right_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id())
{
let (Some(left_ty), lhs) = left else { codegen_unreachable!(ctx) };
let (Some(right_ty), rhs) = comparators[0] else { codegen_unreachable!(ctx) };
let (Some(left_ty), left) = left else { codegen_unreachable!(ctx) };
let (Some(right_ty), right) = comparators[0] else { codegen_unreachable!(ctx) };
let op = ops[0];
let is_ndarray1 =
left_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id());
let is_ndarray2 =
right_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id());
let left_ty_dtype = arraylike_flatten_element_type(&mut ctx.unifier, left_ty);
let right_ty_dtype = arraylike_flatten_element_type(&mut ctx.unifier, right_ty);
return if is_ndarray1 && is_ndarray2 {
let (ndarray_dtype1, _) = unpack_ndarray_var_tys(&mut ctx.unifier, left_ty);
let (ndarray_dtype2, _) = unpack_ndarray_var_tys(&mut ctx.unifier, right_ty);
let left = ScalarOrNDArray::from_value(generator, ctx, (left_ty, left))
.to_ndarray(generator, ctx);
let right = ScalarOrNDArray::from_value(generator, ctx, (right_ty, right))
.to_ndarray(generator, ctx);
assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
let result_ndarray = NDArrayType::new_broadcast(
generator,
ctx.ctx,
ctx.ctx.i8_type().into(),
&[left.get_type(), right.get_type()],
)
.broadcast_starmap(
generator,
ctx,
&[left, right],
NDArrayOut::NewNDArray { dtype: ctx.ctx.i8_type().into() },
|generator, ctx, scalars| {
let left_scalar = scalars[0];
let right_scalar = scalars[1];
let left_val = NDArrayType::from_unifier_type(generator, ctx, left_ty)
.map_value(lhs.into_pointer_value(), None);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ctx.primitives.bool,
None,
(left_ty, left_val.as_base_value().into(), false),
(right_ty, rhs, false),
|generator, ctx, (lhs, rhs)| {
let val = gen_cmpop_expr_with_values(
generator,
ctx,
(Some(ndarray_dtype1), lhs),
&[op],
&[(Some(ndarray_dtype2), rhs)],
)?
.unwrap()
.to_basic_value_enum(
ctx,
generator,
ctx.primitives.bool,
)?;
let val = gen_cmpop_expr_with_values(
generator,
ctx,
(Some(left_ty_dtype), left_scalar),
&[op],
&[(Some(right_ty_dtype), right_scalar)],
)?
.unwrap()
.to_basic_value_enum(
ctx,
generator,
ctx.primitives.bool,
)?;
Ok(generator.bool_to_i8(ctx, val.into_int_value()).into())
},
)?;
Ok(generator.bool_to_i8(ctx, val.into_int_value()).into())
},
)?;
Ok(Some(res.as_base_value().into()))
} else {
let (ndarray_dtype, _) = unpack_ndarray_var_tys(
&mut ctx.unifier,
if is_ndarray1 { left_ty } else { right_ty },
);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ctx.primitives.bool,
None,
(left_ty, lhs, !is_ndarray1),
(right_ty, rhs, !is_ndarray2),
|generator, ctx, (lhs, rhs)| {
let val = gen_cmpop_expr_with_values(
generator,
ctx,
(Some(ndarray_dtype), lhs),
&[op],
&[(Some(ndarray_dtype), rhs)],
)?
.unwrap()
.to_basic_value_enum(
ctx,
generator,
ctx.primitives.bool,
)?;
Ok(generator.bool_to_i8(ctx, val.into_int_value()).into())
},
)?;
Ok(Some(res.as_base_value().into()))
};
return Ok(Some(result_ndarray.as_base_value().into()));
}
}
let cmp_val = izip!(chain(once(&left), comparators.iter()), comparators.iter(), ops.iter(),)
let cmp_val = izip!(once(&left).chain(comparators.iter()), comparators.iter(), ops.iter(),)
.fold(Ok(None), |prev: Result<Option<_>, String>, (lhs, rhs, op)| {
let (left_ty, lhs) = lhs;
let (right_ty, rhs) = rhs;
@@ -2444,319 +2367,6 @@ pub fn gen_cmpop_expr<'ctx, G: CodeGenerator>(
)
}
/// Generates code for a subscript expression on an `ndarray`.
///
/// * `ty` - The `Type` of the `NDArray` elements.
/// * `ndims` - The `Type` of the `NDArray` number-of-dimensions `Literal`.
/// * `v` - The `NDArray` value.
/// * `slice` - The slice expression used to subscript into the `ndarray`.
fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ty: Type,
ndims_ty: Type,
v: NDArrayValue<'ctx>,
slice: &Expr<Option<Type>>,
) -> Result<Option<ValueEnum<'ctx>>, String> {
let llvm_i1 = ctx.ctx.bool_type();
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims_ty) else {
codegen_unreachable!(ctx)
};
let ndims = values
.iter()
.map(|ndim| u64::try_from(ndim.clone()).map_err(|()| ndim.clone()))
.collect::<Result<Vec<_>, _>>()
.map_err(|val| {
format!(
"Expected non-negative literal for ndarray.ndims, got {}",
i128::try_from(val).unwrap()
)
})?;
assert!(!ndims.is_empty());
// The number of dimensions subscripted by the index expression.
// Slicing a ndarray will yield the same number of dimensions, whereas indexing into a
// dimension will remove a dimension.
let subscripted_dims = match &slice.node {
ExprKind::Tuple { elts, .. } => elts.iter().fold(0, |acc, value_subexpr| {
if let ExprKind::Slice { .. } = &value_subexpr.node {
acc
} else {
acc + 1
}
}),
ExprKind::Slice { .. } => 0,
_ => 1,
};
let llvm_ndarray_data_t = ctx.get_llvm_type(generator, ty).as_basic_type_enum();
let sizeof_elem = llvm_ndarray_data_t.size_of().unwrap();
// Check that len is non-zero
let len = v.load_ndims(ctx);
ctx.make_assert(
generator,
ctx.builder.build_int_compare(IntPredicate::SGT, len, llvm_usize.const_zero(), "").unwrap(),
"0:IndexError",
"too many indices for array: array is {0}-dimensional but 1 were indexed",
[Some(len), None, None],
slice.location,
);
// Normalizes a possibly-negative index to its corresponding positive index
let normalize_index = |generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
index: IntValue<'ctx>,
dim: u64| {
gen_if_else_expr_callback(
generator,
ctx,
|_, ctx| {
Ok(ctx
.builder
.build_int_compare(IntPredicate::SGE, index, index.get_type().const_zero(), "")
.unwrap())
},
|_, _| Ok(Some(index)),
|generator, ctx| {
let llvm_i32 = ctx.ctx.i32_type();
let len = unsafe {
v.shape().get_typed_unchecked(
ctx,
generator,
&llvm_usize.const_int(dim, true),
None,
)
};
let index = ctx
.builder
.build_int_add(
len,
ctx.builder.build_int_s_extend(index, llvm_usize, "").unwrap(),
"",
)
.unwrap();
Ok(Some(ctx.builder.build_int_truncate(index, llvm_i32, "").unwrap()))
},
)
.map(|v| v.map(BasicValueEnum::into_int_value))
};
// Converts a slice expression into a slice-range tuple
let expr_to_slice = |generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
node: &ExprKind<Option<Type>>,
dim: u64| {
match node {
ExprKind::Constant { value: Constant::Int(v), .. } => {
let Some(index) =
normalize_index(generator, ctx, llvm_i32.const_int(*v as u64, true), dim)?
else {
return Ok(None);
};
Ok(Some((index, index, llvm_i32.const_int(1, true))))
}
ExprKind::Slice { lower, upper, step } => {
let dim_sz = unsafe {
v.shape().get_typed_unchecked(
ctx,
generator,
&llvm_usize.const_int(dim, false),
None,
)
};
handle_slice_indices(lower, upper, step, ctx, generator, dim_sz)
}
_ => {
let Some(index) = generator.gen_expr(ctx, slice)? else { return Ok(None) };
let index = index
.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?
.into_int_value();
let Some(index) = normalize_index(generator, ctx, index, dim)? else {
return Ok(None);
};
Ok(Some((index, index, llvm_i32.const_int(1, true))))
}
}
};
let make_indices_arr = |generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>|
-> Result<_, String> {
Ok(if let ExprKind::Tuple { elts, .. } = &slice.node {
let llvm_int_ty = ctx.get_llvm_type(generator, elts[0].custom.unwrap());
let index_addr = generator.gen_array_var_alloc(
ctx,
llvm_int_ty,
llvm_usize.const_int(elts.len() as u64, false),
None,
)?;
for (i, elt) in elts.iter().enumerate() {
let Some(index) = generator.gen_expr(ctx, elt)? else {
return Ok(None);
};
let index = index
.to_basic_value_enum(ctx, generator, elt.custom.unwrap())?
.into_int_value();
let Some(index) = normalize_index(generator, ctx, index, 0)? else {
return Ok(None);
};
let store_ptr = unsafe {
index_addr.ptr_offset_unchecked(
ctx,
generator,
&llvm_usize.const_int(i as u64, false),
None,
)
};
ctx.builder.build_store(store_ptr, index).unwrap();
}
Some(index_addr)
} else if let Some(index) = generator.gen_expr(ctx, slice)? {
let llvm_int_ty = ctx.get_llvm_type(generator, slice.custom.unwrap());
let index_addr = generator.gen_array_var_alloc(
ctx,
llvm_int_ty,
llvm_usize.const_int(1u64, false),
None,
)?;
let index =
index.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?.into_int_value();
let Some(index) = normalize_index(generator, ctx, index, 0)? else { return Ok(None) };
let store_ptr = unsafe {
index_addr.ptr_offset_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
};
ctx.builder.build_store(store_ptr, index).unwrap();
Some(index_addr)
} else {
None
})
};
Ok(Some(if ndims.len() == 1 && ndims[0] - subscripted_dims == 0 {
let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
v.data().get(ctx, generator, &index_addr, None).into()
} else {
match &slice.node {
ExprKind::Tuple { elts, .. } => {
let slices = elts
.iter()
.enumerate()
.map(|(dim, elt)| expr_to_slice(generator, ctx, &elt.node, dim as u64))
.take_while_inclusive(|slice| slice.as_ref().is_ok_and(Option::is_some))
.collect::<Result<Vec<_>, _>>()?;
if slices.len() < elts.len() {
return Ok(None);
}
let slices = slices.into_iter().map(Option::unwrap).collect_vec();
numpy::ndarray_sliced_copy(generator, ctx, ty, v, &slices)?.as_base_value().into()
}
ExprKind::Slice { .. } => {
let Some(slice) = expr_to_slice(generator, ctx, &slice.node, 0)? else {
return Ok(None);
};
numpy::ndarray_sliced_copy(generator, ctx, ty, v, &[slice])?.as_base_value().into()
}
_ => {
// Accessing an element from a multi-dimensional `ndarray`
let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
let num_dims = extract_ndims(&ctx.unifier, ndims_ty) - 1;
// Create a new array, remove the top dimension from the dimension-size-list, and copy the
// elements over
let ndarray =
NDArrayType::new(generator, ctx.ctx, llvm_ndarray_data_t, Some(num_dims))
.construct_uninitialized(generator, ctx, None);
let ndarray_num_dims = ctx
.builder
.build_int_z_extend_or_bit_cast(
ndarray.load_ndims(ctx),
llvm_usize.size_of().get_type(),
"",
)
.unwrap();
let v_dims_src_ptr = unsafe {
v.shape().ptr_offset_unchecked(
ctx,
generator,
&llvm_usize.const_int(1, false),
None,
)
};
call_memcpy_generic(
ctx,
ndarray.shape().base_ptr(ctx, generator),
v_dims_src_ptr,
ctx.builder
.build_int_mul(ndarray_num_dims, llvm_usize.size_of(), "")
.map(Into::into)
.unwrap(),
llvm_i1.const_zero(),
);
let ndarray_num_elems = ndarray::call_ndarray_calc_size(
generator,
ctx,
&ndarray.shape().as_slice_value(ctx, generator),
(None, None),
);
let ndarray_num_elems = ctx
.builder
.build_int_z_extend_or_bit_cast(ndarray_num_elems, sizeof_elem.get_type(), "")
.unwrap();
unsafe { ndarray.create_data(generator, ctx) };
let v_data_src_ptr = v.data().ptr_offset(ctx, generator, &index_addr, None);
call_memcpy_generic(
ctx,
ndarray.data().base_ptr(ctx, generator),
v_data_src_ptr,
ctx.builder
.build_int_mul(
ndarray_num_elems,
llvm_ndarray_data_t.size_of().unwrap(),
"",
)
.map(Into::into)
.unwrap(),
llvm_i1.const_zero(),
);
ndarray.as_base_value().into()
}
}
}))
}
/// See [`CodeGenerator::gen_expr`].
pub fn gen_expr<'ctx, G: CodeGenerator>(
generator: &mut G,
@@ -2871,7 +2481,20 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
Some(elements[0].get_type())
};
let length = generator.get_size_type(ctx.ctx).const_int(elements.len() as u64, false);
let arr_str_ptr = allocate_list(generator, ctx, ty, length, Some("list"));
let arr_str_ptr = if let Some(ty) = ty {
ListType::new(generator, ctx.ctx, ty).construct(
generator,
ctx,
length,
Some("list"),
)
} else {
ListType::new_untyped(generator, ctx.ctx).construct_empty(
generator,
ctx,
Some("list"),
)
};
let arr_ptr = arr_str_ptr.data();
for (i, v) in elements.iter().enumerate() {
let elem_ptr = arr_ptr.ptr_offset(
@@ -3349,8 +2972,12 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
.unwrap(),
step,
);
let res_array_ret =
allocate_list(generator, ctx, Some(ty), length, Some("ret"));
let res_array_ret = ListType::new(generator, ctx.ctx, ty).construct(
generator,
ctx,
length,
Some("ret"),
);
let Some(res_ind) = handle_slice_indices(
&None,
&None,

1
nac3core/src/codegen/generator.rs
View File

@@ -17,6 +17,7 @@ pub trait CodeGenerator {
/// Return the module name for the code generator.
fn get_name(&self) -> &str;
/// Return an instance of [`IntType`] corresponding to the type of `size_t` for this instance.
fn get_size_type<'ctx>(&self, ctx: &'ctx Context) -> IntType<'ctx>;
/// Generate function call and returns the function return value.

52
nac3core/src/codegen/irrt/list.rs
View File

@@ -24,42 +24,52 @@ pub fn list_slice_assignment<'ctx, G: CodeGenerator + ?Sized>(
src_arr: ListValue<'ctx>,
src_idx: (IntValue<'ctx>, IntValue<'ctx>, IntValue<'ctx>),
) {
let size_ty = generator.get_size_type(ctx.ctx);
let int8_ptr = ctx.ctx.i8_type().ptr_type(AddressSpace::default());
let int32 = ctx.ctx.i32_type();
let (fun_symbol, elem_ptr_type) = ("__nac3_list_slice_assign_var_size", int8_ptr);
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pi8 = ctx.ctx.i8_type().ptr_type(AddressSpace::default());
let llvm_i32 = ctx.ctx.i32_type();
assert_eq!(dest_idx.0.get_type(), llvm_i32);
assert_eq!(dest_idx.1.get_type(), llvm_i32);
assert_eq!(dest_idx.2.get_type(), llvm_i32);
assert_eq!(src_idx.0.get_type(), llvm_i32);
assert_eq!(src_idx.1.get_type(), llvm_i32);
assert_eq!(src_idx.2.get_type(), llvm_i32);
let (fun_symbol, elem_ptr_type) = ("__nac3_list_slice_assign_var_size", llvm_pi8);
let slice_assign_fun = {
let ty_vec = vec![
int32.into(), // dest start idx
int32.into(), // dest end idx
int32.into(), // dest step
llvm_i32.into(), // dest start idx
llvm_i32.into(), // dest end idx
llvm_i32.into(), // dest step
elem_ptr_type.into(), // dest arr ptr
int32.into(), // dest arr len
int32.into(), // src start idx
int32.into(), // src end idx
int32.into(), // src step
llvm_i32.into(), // dest arr len
llvm_i32.into(), // src start idx
llvm_i32.into(), // src end idx
llvm_i32.into(), // src step
elem_ptr_type.into(), // src arr ptr
int32.into(), // src arr len
int32.into(), // size
llvm_i32.into(), // src arr len
llvm_i32.into(), // size
];
ctx.module.get_function(fun_symbol).unwrap_or_else(|| {
let fn_t = int32.fn_type(ty_vec.as_slice(), false);
let fn_t = llvm_i32.fn_type(ty_vec.as_slice(), false);
ctx.module.add_function(fun_symbol, fn_t, None)
})
};
let zero = int32.const_zero();
let one = int32.const_int(1, false);
let zero = llvm_i32.const_zero();
let one = llvm_i32.const_int(1, false);
let dest_arr_ptr = dest_arr.data().base_ptr(ctx, generator);
let dest_arr_ptr =
ctx.builder.build_pointer_cast(dest_arr_ptr, elem_ptr_type, "dest_arr_ptr_cast").unwrap();
let dest_len = dest_arr.load_size(ctx, Some("dest.len"));
let dest_len = ctx.builder.build_int_truncate_or_bit_cast(dest_len, int32, "srclen32").unwrap();
let dest_len =
ctx.builder.build_int_truncate_or_bit_cast(dest_len, llvm_i32, "srclen32").unwrap();
let src_arr_ptr = src_arr.data().base_ptr(ctx, generator);
let src_arr_ptr =
ctx.builder.build_pointer_cast(src_arr_ptr, elem_ptr_type, "src_arr_ptr_cast").unwrap();
let src_len = src_arr.load_size(ctx, Some("src.len"));
let src_len = ctx.builder.build_int_truncate_or_bit_cast(src_len, int32, "srclen32").unwrap();
let src_len =
ctx.builder.build_int_truncate_or_bit_cast(src_len, llvm_i32, "srclen32").unwrap();
// index in bound and positive should be done
// assert if dest.step == 1 then len(src) <= len(dest) else len(src) == len(dest), and
@@ -136,7 +146,7 @@ pub fn list_slice_assignment<'ctx, G: CodeGenerator + ?Sized>(
BasicTypeEnum::StructType(t) => t.size_of().unwrap(),
_ => codegen_unreachable!(ctx),
};
ctx.builder.build_int_truncate_or_bit_cast(s, int32, "size").unwrap()
ctx.builder.build_int_truncate_or_bit_cast(s, llvm_i32, "size").unwrap()
}
.into(),
];
@@ -147,6 +157,7 @@ pub fn list_slice_assignment<'ctx, G: CodeGenerator + ?Sized>(
.map(Either::unwrap_left)
.unwrap()
};
// update length
let need_update =
ctx.builder.build_int_compare(IntPredicate::NE, new_len, dest_len, "need_update").unwrap();
@@ -155,7 +166,8 @@ pub fn list_slice_assignment<'ctx, G: CodeGenerator + ?Sized>(
let cont_bb = ctx.ctx.append_basic_block(current, "cont");
ctx.builder.build_conditional_branch(need_update, update_bb, cont_bb).unwrap();
ctx.builder.position_at_end(update_bb);
let new_len = ctx.builder.build_int_z_extend_or_bit_cast(new_len, size_ty, "new_len").unwrap();
let new_len =
ctx.builder.build_int_z_extend_or_bit_cast(new_len, llvm_usize, "new_len").unwrap();
dest_arr.store_size(ctx, generator, new_len);
ctx.builder.build_unconditional_branch(cont_bb).unwrap();
ctx.builder.position_at_end(cont_bb);

20
nac3core/src/codegen/irrt/math.rs
View File

@@ -62,8 +62,13 @@ pub fn call_isinf<'ctx, G: CodeGenerator + ?Sized>(
ctx: &CodeGenContext<'ctx, '_>,
v: FloatValue<'ctx>,
) -> IntValue<'ctx> {
let llvm_i32 = ctx.ctx.i32_type();
let llvm_f64 = ctx.ctx.f64_type();
assert_eq!(v.get_type(), llvm_f64);
let intrinsic_fn = ctx.module.get_function("__nac3_isinf").unwrap_or_else(|| {
let fn_type = ctx.ctx.i32_type().fn_type(&[ctx.ctx.f64_type().into()], false);
let fn_type = llvm_i32.fn_type(&[llvm_f64.into()], false);
ctx.module.add_function("__nac3_isinf", fn_type, None)
});
@@ -84,8 +89,13 @@ pub fn call_isnan<'ctx, G: CodeGenerator + ?Sized>(
ctx: &CodeGenContext<'ctx, '_>,
v: FloatValue<'ctx>,
) -> IntValue<'ctx> {
let llvm_i32 = ctx.ctx.i32_type();
let llvm_f64 = ctx.ctx.f64_type();
assert_eq!(v.get_type(), llvm_f64);
let intrinsic_fn = ctx.module.get_function("__nac3_isnan").unwrap_or_else(|| {
let fn_type = ctx.ctx.i32_type().fn_type(&[ctx.ctx.f64_type().into()], false);
let fn_type = llvm_i32.fn_type(&[llvm_f64.into()], false);
ctx.module.add_function("__nac3_isnan", fn_type, None)
});
@@ -104,6 +114,8 @@ pub fn call_isnan<'ctx, G: CodeGenerator + ?Sized>(
pub fn call_gamma<'ctx>(ctx: &CodeGenContext<'ctx, '_>, v: FloatValue<'ctx>) -> FloatValue<'ctx> {
let llvm_f64 = ctx.ctx.f64_type();
assert_eq!(v.get_type(), llvm_f64);
let intrinsic_fn = ctx.module.get_function("__nac3_gamma").unwrap_or_else(|| {
let fn_type = llvm_f64.fn_type(&[llvm_f64.into()], false);
ctx.module.add_function("__nac3_gamma", fn_type, None)
@@ -121,6 +133,8 @@ pub fn call_gamma<'ctx>(ctx: &CodeGenContext<'ctx, '_>, v: FloatValue<'ctx>) ->
pub fn call_gammaln<'ctx>(ctx: &CodeGenContext<'ctx, '_>, v: FloatValue<'ctx>) -> FloatValue<'ctx> {
let llvm_f64 = ctx.ctx.f64_type();
assert_eq!(v.get_type(), llvm_f64);
let intrinsic_fn = ctx.module.get_function("__nac3_gammaln").unwrap_or_else(|| {
let fn_type = llvm_f64.fn_type(&[llvm_f64.into()], false);
ctx.module.add_function("__nac3_gammaln", fn_type, None)
@@ -138,6 +152,8 @@ pub fn call_gammaln<'ctx>(ctx: &CodeGenContext<'ctx, '_>, v: FloatValue<'ctx>) -
pub fn call_j0<'ctx>(ctx: &CodeGenContext<'ctx, '_>, v: FloatValue<'ctx>) -> FloatValue<'ctx> {
let llvm_f64 = ctx.ctx.f64_type();
assert_eq!(v.get_type(), llvm_f64);
let intrinsic_fn = ctx.module.get_function("__nac3_j0").unwrap_or_else(|| {
let fn_type = llvm_f64.fn_type(&[llvm_f64.into()], false);
ctx.module.add_function("__nac3_j0", fn_type, None)

11
nac3core/src/codegen/irrt/mod.rs
View File

@@ -132,10 +132,11 @@ pub fn handle_slice_indices<'ctx, G: CodeGenerator>(
generator: &mut G,
length: IntValue<'ctx>,
) -> Result<Option<(IntValue<'ctx>, IntValue<'ctx>, IntValue<'ctx>)>, String> {
let int32 = ctx.ctx.i32_type();
let zero = int32.const_zero();
let one = int32.const_int(1, false);
let length = ctx.builder.build_int_truncate_or_bit_cast(length, int32, "leni32").unwrap();
let llvm_i32 = ctx.ctx.i32_type();
let zero = llvm_i32.const_zero();
let one = llvm_i32.const_int(1, false);
let length = ctx.builder.build_int_truncate_or_bit_cast(length, llvm_i32, "leni32").unwrap();
Ok(Some(match (start, end, step) {
(s, e, None) => (
if let Some(s) = s.as_ref() {
@@ -144,7 +145,7 @@ pub fn handle_slice_indices<'ctx, G: CodeGenerator>(
None => return Ok(None),
}
} else {
int32.const_zero()
llvm_i32.const_zero()
},
{
let e = if let Some(s) = e.as_ref() {

80
nac3core/src/codegen/irrt/ndarray/array.rs Normal file
View File

@@ -0,0 +1,80 @@
use inkwell::{types::BasicTypeEnum, values::IntValue};
use crate::codegen::{
expr::infer_and_call_function,
irrt::get_usize_dependent_function_name,
values::{ndarray::NDArrayValue, ListValue, ProxyValue, TypedArrayLikeAccessor},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_array_set_and_validate_list_shape`.
///
/// Deduces the target shape of the `ndarray` from the provided `list`, raising an exception if
/// there is any issue with the resultant `shape`.
///
/// `shape` must be pre-allocated by the caller of this function to `[usize; ndims]`, and must be
/// initialized to all `-1`s.
pub fn call_nac3_ndarray_array_set_and_validate_list_shape<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
list: ListValue<'ctx>,
ndims: IntValue<'ctx>,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
assert_eq!(list.get_type().element_type().unwrap(), ctx.ctx.i8_type().into());
assert_eq!(ndims.get_type(), llvm_usize);
assert_eq!(
BasicTypeEnum::try_from(shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name = get_usize_dependent_function_name(
generator,
ctx,
"__nac3_ndarray_array_set_and_validate_list_shape",
);
infer_and_call_function(
ctx,
&name,
None,
&[list.as_base_value().into(), ndims.into(), shape.base_ptr(ctx, generator).into()],
None,
None,
);
}
/// Generates a call to `__nac3_ndarray_array_write_list_to_array`.
///
/// Copies the contents stored in `list` into `ndarray`.
///
/// The `ndarray` must fulfill the following preconditions:
///
/// - `ndarray.itemsize`: Must be initialized.
/// - `ndarray.ndims`: Must be initialized.
/// - `ndarray.shape`: Must be initialized.
/// - `ndarray.data`: Must be allocated and contiguous.
pub fn call_nac3_ndarray_array_write_list_to_array<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
list: ListValue<'ctx>,
ndarray: NDArrayValue<'ctx>,
) {
assert_eq!(list.get_type().element_type().unwrap(), ctx.ctx.i8_type().into());
let name = get_usize_dependent_function_name(
generator,
ctx,
"__nac3_ndarray_array_write_list_to_array",
);
infer_and_call_function(
ctx,
&name,
None,
&[list.as_base_value().into(), ndarray.as_base_value().into()],
None,
None,
);
}

88
nac3core/src/codegen/irrt/ndarray/basic.rs
View File

@@ -1,4 +1,5 @@
use inkwell::{
types::BasicTypeEnum,
values::{BasicValueEnum, IntValue, PointerValue},
AddressSpace,
};
@@ -7,19 +8,26 @@ use crate::codegen::{
expr::{create_and_call_function, infer_and_call_function},
irrt::get_usize_dependent_function_name,
types::ProxyType,
values::{ndarray::NDArrayValue, ProxyValue},
values::{ndarray::NDArrayValue, ProxyValue, TypedArrayLikeAccessor},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_util_assert_shape_no_negative`.
///
/// Assets that `shape` does not contain negative dimensions.
pub fn call_nac3_ndarray_util_assert_shape_no_negative<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
ndims: IntValue<'ctx>,
shape: PointerValue<'ctx>,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
assert_eq!(
BasicTypeEnum::try_from(shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name = get_usize_dependent_function_name(
generator,
ctx,
@@ -30,23 +38,37 @@ pub fn call_nac3_ndarray_util_assert_shape_no_negative<'ctx, G: CodeGenerator +
ctx,
&name,
Some(llvm_usize.into()),
&[(llvm_usize.into(), ndims.into()), (llvm_pusize.into(), shape.into())],
&[
(llvm_usize.into(), shape.size(ctx, generator).into()),
(llvm_pusize.into(), shape.base_ptr(ctx, generator).into()),
],
None,
None,
);
}
/// Generates a call to `__nac3_ndarray_util_assert_shape_output_shape_same`.
///
/// Asserts that `ndarray_shape` and `output_shape` are the same in the context of writing output to
/// an `ndarray`.
pub fn call_nac3_ndarray_util_assert_output_shape_same<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
ndarray_ndims: IntValue<'ctx>,
ndarray_shape: PointerValue<'ctx>,
output_ndims: IntValue<'ctx>,
output_shape: IntValue<'ctx>,
ndarray_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
output_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
assert_eq!(
BasicTypeEnum::try_from(ndarray_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
assert_eq!(
BasicTypeEnum::try_from(output_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name = get_usize_dependent_function_name(
generator,
ctx,
@@ -58,16 +80,20 @@ pub fn call_nac3_ndarray_util_assert_output_shape_same<'ctx, G: CodeGenerator +
&name,
Some(llvm_usize.into()),
&[
(llvm_usize.into(), ndarray_ndims.into()),
(llvm_pusize.into(), ndarray_shape.into()),
(llvm_usize.into(), output_ndims.into()),
(llvm_pusize.into(), output_shape.into()),
(llvm_usize.into(), ndarray_shape.size(ctx, generator).into()),
(llvm_pusize.into(), ndarray_shape.base_ptr(ctx, generator).into()),
(llvm_usize.into(), output_shape.size(ctx, generator).into()),
(llvm_pusize.into(), output_shape.base_ptr(ctx, generator).into()),
],
None,
None,
);
}
/// Generates a call to `__nac3_ndarray_size`.
///
/// Returns a [`usize`][CodeGenerator::get_size_type] value of the number of elements of an
/// `ndarray`, corresponding to the value of `ndarray.size`.
pub fn call_nac3_ndarray_size<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -90,6 +116,10 @@ pub fn call_nac3_ndarray_size<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_ndarray_nbytes`.
///
/// Returns a [`usize`][CodeGenerator::get_size_type] value of the number of bytes consumed by the
/// data of the `ndarray`, corresponding to the value of `ndarray.nbytes`.
pub fn call_nac3_ndarray_nbytes<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -112,6 +142,10 @@ pub fn call_nac3_ndarray_nbytes<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_ndarray_len`.
///
/// Returns a [`usize`][CodeGenerator::get_size_type] value of the size of the topmost dimension of
/// the `ndarray`, corresponding to the value of `ndarray.__len__`.
pub fn call_nac3_ndarray_len<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -134,6 +168,9 @@ pub fn call_nac3_ndarray_len<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_ndarray_is_c_contiguous`.
///
/// Returns an `i1` value indicating whether the `ndarray` is C-contiguous.
pub fn call_nac3_ndarray_is_c_contiguous<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -156,6 +193,9 @@ pub fn call_nac3_ndarray_is_c_contiguous<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_ndarray_get_nth_pelement`.
///
/// Returns a [`PointerValue`] to the `index`-th flattened element of the `ndarray`.
pub fn call_nac3_ndarray_get_nth_pelement<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -167,6 +207,8 @@ pub fn call_nac3_ndarray_get_nth_pelement<'ctx, G: CodeGenerator + ?Sized>(
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_ndarray = ndarray.get_type().as_base_type();
assert_eq!(index.get_type(), llvm_usize);
let name = get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_get_nth_pelement");
create_and_call_function(
@@ -181,11 +223,16 @@ pub fn call_nac3_ndarray_get_nth_pelement<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_ndarray_get_pelement_by_indices`.
///
/// `indices` must have the same number of elements as the number of dimensions in `ndarray`.
///
/// Returns a [`PointerValue`] to the element indexed by `indices`.
pub fn call_nac3_ndarray_get_pelement_by_indices<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
ndarray: NDArrayValue<'ctx>,
indices: PointerValue<'ctx>,
indices: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) -> PointerValue<'ctx> {
let llvm_i8 = ctx.ctx.i8_type();
let llvm_pi8 = llvm_i8.ptr_type(AddressSpace::default());
@@ -193,6 +240,11 @@ pub fn call_nac3_ndarray_get_pelement_by_indices<'ctx, G: CodeGenerator + ?Sized
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
let llvm_ndarray = ndarray.get_type().as_base_type();
assert_eq!(
BasicTypeEnum::try_from(indices.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name =
get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_get_pelement_by_indices");
@@ -202,7 +254,7 @@ pub fn call_nac3_ndarray_get_pelement_by_indices<'ctx, G: CodeGenerator + ?Sized
Some(llvm_pi8.into()),
&[
(llvm_ndarray.into(), ndarray.as_base_value().into()),
(llvm_pusize.into(), indices.into()),
(llvm_pusize.into(), indices.base_ptr(ctx, generator).into()),
],
Some("pelement"),
None,
@@ -211,6 +263,9 @@ pub fn call_nac3_ndarray_get_pelement_by_indices<'ctx, G: CodeGenerator + ?Sized
.unwrap()
}
/// Generates a call to `__nac3_ndarray_set_strides_by_shape`.
///
/// Sets `ndarray.strides` assuming that `ndarray.shape` is C-contiguous.
pub fn call_nac3_ndarray_set_strides_by_shape<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -231,6 +286,11 @@ pub fn call_nac3_ndarray_set_strides_by_shape<'ctx, G: CodeGenerator + ?Sized>(
);
}
/// Generates a call to `__nac3_ndarray_copy_data`.
///
/// Copies all elements from `src_ndarray` to `dst_ndarray` using their flattened views. The number
/// of elements in `src_ndarray` must be greater than or equal to the number of elements in
/// `dst_ndarray`.
pub fn call_nac3_ndarray_copy_data<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,

82
nac3core/src/codegen/irrt/ndarray/broadcast.rs Normal file
View File

@@ -0,0 +1,82 @@
use inkwell::values::IntValue;
use crate::codegen::{
expr::infer_and_call_function,
irrt::get_usize_dependent_function_name,
types::{ndarray::ShapeEntryType, ProxyType},
values::{
ndarray::NDArrayValue, ArrayLikeValue, ArraySliceValue, ProxyValue, TypedArrayLikeAccessor,
TypedArrayLikeMutator,
},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_broadcast_to`.
///
/// Attempts to broadcast `src_ndarray` to the new shape defined by `dst_ndarray`.
///
/// `dst_ndarray` must meet the following preconditions:
///
/// - `dst_ndarray.ndims` must be initialized and matching the length of `dst_ndarray.shape`.
/// - `dst_ndarray.shape` must be initialized and contains the target broadcast shape.
/// - `dst_ndarray.strides` must be allocated and may contain uninitialized values.
pub fn call_nac3_ndarray_broadcast_to<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
src_ndarray: NDArrayValue<'ctx>,
dst_ndarray: NDArrayValue<'ctx>,
) {
let name = get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_broadcast_to");
infer_and_call_function(
ctx,
&name,
None,
&[src_ndarray.as_base_value().into(), dst_ndarray.as_base_value().into()],
None,
None,
);
}
/// Generates a call to `__nac3_ndarray_broadcast_shapes`.
///
/// Attempts to calculate the resultant shape from broadcasting all shapes in `shape_entries`,
/// writing the result to `dst_shape`.
pub fn call_nac3_ndarray_broadcast_shapes<'ctx, G, Shape>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
num_shape_entries: IntValue<'ctx>,
shape_entries: ArraySliceValue<'ctx>,
dst_ndims: IntValue<'ctx>,
dst_shape: &Shape,
) where
G: CodeGenerator + ?Sized,
Shape: TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>
+ TypedArrayLikeMutator<'ctx, G, IntValue<'ctx>>,
{
let llvm_usize = generator.get_size_type(ctx.ctx);
assert_eq!(num_shape_entries.get_type(), llvm_usize);
assert!(ShapeEntryType::is_type(
generator,
ctx.ctx,
shape_entries.base_ptr(ctx, generator).get_type()
)
.is_ok());
assert_eq!(dst_ndims.get_type(), llvm_usize);
assert_eq!(dst_shape.element_type(ctx, generator), llvm_usize.into());
let name = get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_broadcast_shapes");
infer_and_call_function(
ctx,
&name,
None,
&[
num_shape_entries.into(),
shape_entries.base_ptr(ctx, generator).into(),
dst_ndims.into(),
dst_shape.base_ptr(ctx, generator).into(),
],
None,
None,
);
}

5
nac3core/src/codegen/irrt/ndarray/indexing.rs
View File

@@ -5,6 +5,11 @@ use crate::codegen::{
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_index`.
///
/// Performs [basic indexing](https://numpy.org/doc/stable/user/basics.indexing.html#basic-indexing)
/// on `src_ndarray` using `indices`, writing the result to `dst_ndarray`, corresponding to the
/// operation `dst_ndarray = src_ndarray[indices]`.
pub fn call_nac3_ndarray_index<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,

20
nac3core/src/codegen/irrt/ndarray/iter.rs
View File

@@ -1,4 +1,5 @@
use inkwell::{
types::BasicTypeEnum,
values::{BasicValueEnum, IntValue},
AddressSpace,
};
@@ -9,21 +10,29 @@ use crate::codegen::{
types::ProxyType,
values::{
ndarray::{NDArrayValue, NDIterValue},
ArrayLikeValue, ArraySliceValue, ProxyValue,
ProxyValue, TypedArrayLikeAccessor,
},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_nditer_initialize`.
///
/// Initializes the `iter` object.
pub fn call_nac3_nditer_initialize<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
iter: NDIterValue<'ctx>,
ndarray: NDArrayValue<'ctx>,
indices: ArraySliceValue<'ctx>,
indices: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
assert_eq!(
BasicTypeEnum::try_from(indices.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name = get_usize_dependent_function_name(generator, ctx, "__nac3_nditer_initialize");
create_and_call_function(
@@ -40,6 +49,10 @@ pub fn call_nac3_nditer_initialize<'ctx, G: CodeGenerator + ?Sized>(
);
}
/// Generates a call to `__nac3_nditer_initialize_has_element`.
///
/// Returns an `i1` value indicating whether there are elements left to traverse for the `iter`
/// object.
pub fn call_nac3_nditer_has_element<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
@@ -59,6 +72,9 @@ pub fn call_nac3_nditer_has_element<'ctx, G: CodeGenerator + ?Sized>(
.unwrap()
}
/// Generates a call to `__nac3_nditer_next`.
///
/// Moves `iter` to point to the next element.
pub fn call_nac3_nditer_next<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,

66
nac3core/src/codegen/irrt/ndarray/matmul.rs Normal file
View File

@@ -0,0 +1,66 @@
use inkwell::{types::BasicTypeEnum, values::IntValue};
use crate::codegen::{
expr::infer_and_call_function, irrt::get_usize_dependent_function_name,
values::TypedArrayLikeAccessor, CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_matmul_calculate_shapes`.
///
/// Calculates the broadcasted shapes for `a`, `b`, and the `ndarray` holding the final values of
/// `a @ b`.
#[allow(clippy::too_many_arguments)]
pub fn call_nac3_ndarray_matmul_calculate_shapes<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
a_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
b_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
final_ndims: IntValue<'ctx>,
new_a_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
new_b_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
dst_shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
assert_eq!(
BasicTypeEnum::try_from(a_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
assert_eq!(
BasicTypeEnum::try_from(b_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
assert_eq!(
BasicTypeEnum::try_from(new_a_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
assert_eq!(
BasicTypeEnum::try_from(new_b_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
assert_eq!(
BasicTypeEnum::try_from(dst_shape.element_type(ctx, generator)).unwrap(),
llvm_usize.into()
);
let name =
get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_matmul_calculate_shapes");
infer_and_call_function(
ctx,
&name,
None,
&[
a_shape.size(ctx, generator).into(),
a_shape.base_ptr(ctx, generator).into(),
b_shape.size(ctx, generator).into(),
b_shape.base_ptr(ctx, generator).into(),
final_ndims.into(),
new_a_shape.base_ptr(ctx, generator).into(),
new_b_shape.base_ptr(ctx, generator).into(),
dst_shape.base_ptr(ctx, generator).into(),
],
None,
None,
);
}

394
nac3core/src/codegen/irrt/ndarray/mod.rs
View File

@@ -1,391 +1,17 @@
use inkwell::{
types::IntType,
values::{BasicValueEnum, CallSiteValue, IntValue},
AddressSpace, IntPredicate,
};
use itertools::Either;
use crate::codegen::{
llvm_intrinsics,
macros::codegen_unreachable,
stmt::gen_for_callback_incrementing,
values::{
ndarray::NDArrayValue, ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue,
TypedArrayLikeAccessor, TypedArrayLikeAdapter, UntypedArrayLikeAccessor,
},
CodeGenContext, CodeGenerator,
};
pub use array::*;
pub use basic::*;
pub use broadcast::*;
pub use indexing::*;
pub use iter::*;
pub use matmul::*;
pub use reshape::*;
pub use transpose::*;
mod array;
mod basic;
mod broadcast;
mod indexing;
mod iter;
/// Generates a call to `__nac3_ndarray_calc_size`. Returns an [`IntValue`] representing the
/// calculated total size.
///
/// * `dims` - An [`ArrayLikeIndexer`] containing the size of each dimension.
/// * `range` - The dimension index to begin and end (exclusively) calculating the dimensions for,
/// or [`None`] if starting from the first dimension and ending at the last dimension
/// respectively.
pub fn call_ndarray_calc_size<'ctx, G, Dims>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
dims: &Dims,
(begin, end): (Option<IntValue<'ctx>>, Option<IntValue<'ctx>>),
) -> IntValue<'ctx>
where
G: CodeGenerator + ?Sized,
Dims: ArrayLikeIndexer<'ctx>,
{
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
let ndarray_calc_size_fn_name = match llvm_usize.get_bit_width() {
32 => "__nac3_ndarray_calc_size",
64 => "__nac3_ndarray_calc_size64",
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let ndarray_calc_size_fn_t = llvm_usize.fn_type(
&[llvm_pusize.into(), llvm_usize.into(), llvm_usize.into(), llvm_usize.into()],
false,
);
let ndarray_calc_size_fn =
ctx.module.get_function(ndarray_calc_size_fn_name).unwrap_or_else(|| {
ctx.module.add_function(ndarray_calc_size_fn_name, ndarray_calc_size_fn_t, None)
});
let begin = begin.unwrap_or_else(|| llvm_usize.const_zero());
let end = end.unwrap_or_else(|| dims.size(ctx, generator));
ctx.builder
.build_call(
ndarray_calc_size_fn,
&[
dims.base_ptr(ctx, generator).into(),
dims.size(ctx, generator).into(),
begin.into(),
end.into(),
],
"",
)
.map(CallSiteValue::try_as_basic_value)
.map(|v| v.map_left(BasicValueEnum::into_int_value))
.map(Either::unwrap_left)
.unwrap()
}
/// Generates a call to `__nac3_ndarray_calc_nd_indices`. Returns a [`TypeArrayLikeAdpater`]
/// containing `i32` indices of the flattened index.
///
/// * `index` - The index to compute the multidimensional index for.
/// * `ndarray` - LLVM pointer to the `NDArray`. This value must be the LLVM representation of an
/// `NDArray`.
pub fn call_ndarray_calc_nd_indices<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
index: IntValue<'ctx>,
ndarray: NDArrayValue<'ctx>,
) -> TypedArrayLikeAdapter<'ctx, IntValue<'ctx>> {
let llvm_void = ctx.ctx.void_type();
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pi32 = llvm_i32.ptr_type(AddressSpace::default());
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
let ndarray_calc_nd_indices_fn_name = match llvm_usize.get_bit_width() {
32 => "__nac3_ndarray_calc_nd_indices",
64 => "__nac3_ndarray_calc_nd_indices64",
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let ndarray_calc_nd_indices_fn =
ctx.module.get_function(ndarray_calc_nd_indices_fn_name).unwrap_or_else(|| {
let fn_type = llvm_void.fn_type(
&[llvm_usize.into(), llvm_pusize.into(), llvm_usize.into(), llvm_pi32.into()],
false,
);
ctx.module.add_function(ndarray_calc_nd_indices_fn_name, fn_type, None)
});
let ndarray_num_dims = ndarray.load_ndims(ctx);
let ndarray_dims = ndarray.shape();
let indices = ctx.builder.build_array_alloca(llvm_i32, ndarray_num_dims, "").unwrap();
ctx.builder
.build_call(
ndarray_calc_nd_indices_fn,
&[
index.into(),
ndarray_dims.base_ptr(ctx, generator).into(),
ndarray_num_dims.into(),
indices.into(),
],
"",
)
.unwrap();
TypedArrayLikeAdapter::from(
ArraySliceValue::from_ptr_val(indices, ndarray_num_dims, None),
Box::new(|_, v| v.into_int_value()),
Box::new(|_, v| v.into()),
)
}
fn call_ndarray_flatten_index_impl<'ctx, G, Indices>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
ndarray: NDArrayValue<'ctx>,
indices: &Indices,
) -> IntValue<'ctx>
where
G: CodeGenerator + ?Sized,
Indices: ArrayLikeIndexer<'ctx>,
{
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pi32 = llvm_i32.ptr_type(AddressSpace::default());
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
debug_assert_eq!(
IntType::try_from(indices.element_type(ctx, generator))
.map(IntType::get_bit_width)
.unwrap_or_default(),
llvm_i32.get_bit_width(),
"Expected i32 value for argument `indices` to `call_ndarray_flatten_index_impl`"
);
debug_assert_eq!(
indices.size(ctx, generator).get_type().get_bit_width(),
llvm_usize.get_bit_width(),
"Expected usize integer value for argument `indices_size` to `call_ndarray_flatten_index_impl`"
);
let ndarray_flatten_index_fn_name = match llvm_usize.get_bit_width() {
32 => "__nac3_ndarray_flatten_index",
64 => "__nac3_ndarray_flatten_index64",
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let ndarray_flatten_index_fn =
ctx.module.get_function(ndarray_flatten_index_fn_name).unwrap_or_else(|| {
let fn_type = llvm_usize.fn_type(
&[llvm_pusize.into(), llvm_usize.into(), llvm_pi32.into(), llvm_usize.into()],
false,
);
ctx.module.add_function(ndarray_flatten_index_fn_name, fn_type, None)
});
let ndarray_num_dims = ndarray.load_ndims(ctx);
let ndarray_dims = ndarray.shape();
let index = ctx
.builder
.build_call(
ndarray_flatten_index_fn,
&[
ndarray_dims.base_ptr(ctx, generator).into(),
ndarray_num_dims.into(),
indices.base_ptr(ctx, generator).into(),
indices.size(ctx, generator).into(),
],
"",
)
.map(CallSiteValue::try_as_basic_value)
.map(|v| v.map_left(BasicValueEnum::into_int_value))
.map(Either::unwrap_left)
.unwrap();
index
}
/// Generates a call to `__nac3_ndarray_flatten_index`. Returns the flattened index for the
/// multidimensional index.
///
/// * `ndarray` - LLVM pointer to the `NDArray`. This value must be the LLVM representation of an
/// `NDArray`.
/// * `indices` - The multidimensional index to compute the flattened index for.
pub fn call_ndarray_flatten_index<'ctx, G, Index>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
ndarray: NDArrayValue<'ctx>,
indices: &Index,
) -> IntValue<'ctx>
where
G: CodeGenerator + ?Sized,
Index: ArrayLikeIndexer<'ctx>,
{
call_ndarray_flatten_index_impl(generator, ctx, ndarray, indices)
}
/// Generates a call to `__nac3_ndarray_calc_broadcast`. Returns a tuple containing the number of
/// dimension and size of each dimension of the resultant `ndarray`.
pub fn call_ndarray_calc_broadcast<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
lhs: NDArrayValue<'ctx>,
rhs: NDArrayValue<'ctx>,
) -> TypedArrayLikeAdapter<'ctx, IntValue<'ctx>> {
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
let ndarray_calc_broadcast_fn_name = match llvm_usize.get_bit_width() {
32 => "__nac3_ndarray_calc_broadcast",
64 => "__nac3_ndarray_calc_broadcast64",
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let ndarray_calc_broadcast_fn =
ctx.module.get_function(ndarray_calc_broadcast_fn_name).unwrap_or_else(|| {
let fn_type = llvm_usize.fn_type(
&[
llvm_pusize.into(),
llvm_usize.into(),
llvm_pusize.into(),
llvm_usize.into(),
llvm_pusize.into(),
],
false,
);
ctx.module.add_function(ndarray_calc_broadcast_fn_name, fn_type, None)
});
let lhs_ndims = lhs.load_ndims(ctx);
let rhs_ndims = rhs.load_ndims(ctx);
let min_ndims = llvm_intrinsics::call_int_umin(ctx, lhs_ndims, rhs_ndims, None);
gen_for_callback_incrementing(
generator,
ctx,
None,
llvm_usize.const_zero(),
(min_ndims, false),
|generator, ctx, _, idx| {
let idx = ctx.builder.build_int_sub(min_ndims, idx, "").unwrap();
let (lhs_dim_sz, rhs_dim_sz) = unsafe {
(
lhs.shape().get_typed_unchecked(ctx, generator, &idx, None),
rhs.shape().get_typed_unchecked(ctx, generator, &idx, None),
)
};
let llvm_usize_const_one = llvm_usize.const_int(1, false);
let lhs_eqz = ctx
.builder
.build_int_compare(IntPredicate::EQ, lhs_dim_sz, llvm_usize_const_one, "")
.unwrap();
let rhs_eqz = ctx
.builder
.build_int_compare(IntPredicate::EQ, rhs_dim_sz, llvm_usize_const_one, "")
.unwrap();
let lhs_or_rhs_eqz = ctx.builder.build_or(lhs_eqz, rhs_eqz, "").unwrap();
let lhs_eq_rhs = ctx
.builder
.build_int_compare(IntPredicate::EQ, lhs_dim_sz, rhs_dim_sz, "")
.unwrap();
let is_compatible = ctx.builder.build_or(lhs_or_rhs_eqz, lhs_eq_rhs, "").unwrap();
ctx.make_assert(
generator,
is_compatible,
"0:ValueError",
"operands could not be broadcast together",
[None, None, None],
ctx.current_loc,
);
Ok(())
},
llvm_usize.const_int(1, false),
)
.unwrap();
let max_ndims = llvm_intrinsics::call_int_umax(ctx, lhs_ndims, rhs_ndims, None);
let lhs_dims = lhs.shape().base_ptr(ctx, generator);
let lhs_ndims = lhs.load_ndims(ctx);
let rhs_dims = rhs.shape().base_ptr(ctx, generator);
let rhs_ndims = rhs.load_ndims(ctx);
let out_dims = ctx.builder.build_array_alloca(llvm_usize, max_ndims, "").unwrap();
let out_dims = ArraySliceValue::from_ptr_val(out_dims, max_ndims, None);
ctx.builder
.build_call(
ndarray_calc_broadcast_fn,
&[
lhs_dims.into(),
lhs_ndims.into(),
rhs_dims.into(),
rhs_ndims.into(),
out_dims.base_ptr(ctx, generator).into(),
],
"",
)
.unwrap();
TypedArrayLikeAdapter::from(
out_dims,
Box::new(|_, v| v.into_int_value()),
Box::new(|_, v| v.into()),
)
}
/// Generates a call to `__nac3_ndarray_calc_broadcast_idx`. Returns an [`ArrayAllocaValue`]
/// containing the indices used for accessing `array` corresponding to the index of the broadcasted
/// array `broadcast_idx`.
pub fn call_ndarray_calc_broadcast_index<
'ctx,
G: CodeGenerator + ?Sized,
BroadcastIdx: UntypedArrayLikeAccessor<'ctx>,
>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
array: NDArrayValue<'ctx>,
broadcast_idx: &BroadcastIdx,
) -> TypedArrayLikeAdapter<'ctx, IntValue<'ctx>> {
let llvm_i32 = ctx.ctx.i32_type();
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_pi32 = llvm_i32.ptr_type(AddressSpace::default());
let llvm_pusize = llvm_usize.ptr_type(AddressSpace::default());
let ndarray_calc_broadcast_fn_name = match llvm_usize.get_bit_width() {
32 => "__nac3_ndarray_calc_broadcast_idx",
64 => "__nac3_ndarray_calc_broadcast_idx64",
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let ndarray_calc_broadcast_fn =
ctx.module.get_function(ndarray_calc_broadcast_fn_name).unwrap_or_else(|| {
let fn_type = llvm_usize.fn_type(
&[llvm_pusize.into(), llvm_usize.into(), llvm_pi32.into(), llvm_pi32.into()],
false,
);
ctx.module.add_function(ndarray_calc_broadcast_fn_name, fn_type, None)
});
let broadcast_size = broadcast_idx.size(ctx, generator);
let out_idx = ctx.builder.build_array_alloca(llvm_i32, broadcast_size, "").unwrap();
let array_dims = array.shape().base_ptr(ctx, generator);
let array_ndims = array.load_ndims(ctx);
let broadcast_idx_ptr = unsafe {
broadcast_idx.ptr_offset_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
};
ctx.builder
.build_call(
ndarray_calc_broadcast_fn,
&[array_dims.into(), array_ndims.into(), broadcast_idx_ptr.into(), out_idx.into()],
"",
)
.unwrap();
TypedArrayLikeAdapter::from(
ArraySliceValue::from_ptr_val(out_idx, broadcast_size, None),
Box::new(|_, v| v.into_int_value()),
Box::new(|_, v| v.into()),
)
}
mod matmul;
mod reshape;
mod transpose;

40
nac3core/src/codegen/irrt/ndarray/reshape.rs Normal file
View File

@@ -0,0 +1,40 @@
use inkwell::values::IntValue;
use crate::codegen::{
expr::infer_and_call_function,
irrt::get_usize_dependent_function_name,
values::{ArrayLikeValue, ArraySliceValue},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_reshape_resolve_and_check_new_shape`.
///
/// Resolves unknown dimensions in `new_shape` for `numpy.reshape(<ndarray>, new_shape)`, raising an
/// assertion if multiple dimensions are unknown (`-1`).
pub fn call_nac3_ndarray_reshape_resolve_and_check_new_shape<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
new_ndims: IntValue<'ctx>,
new_shape: ArraySliceValue<'ctx>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
assert_eq!(size.get_type(), llvm_usize);
assert_eq!(new_ndims.get_type(), llvm_usize);
assert_eq!(new_shape.element_type(ctx, generator), llvm_usize.into());
let name = get_usize_dependent_function_name(
generator,
ctx,
"__nac3_ndarray_reshape_resolve_and_check_new_shape",
);
infer_and_call_function(
ctx,
&name,
None,
&[size.into(), new_ndims.into(), new_shape.base_ptr(ctx, generator).into()],
None,
None,
);
}

48
nac3core/src/codegen/irrt/ndarray/transpose.rs Normal file
View File

@@ -0,0 +1,48 @@
use inkwell::{values::IntValue, AddressSpace};
use crate::codegen::{
expr::infer_and_call_function,
irrt::get_usize_dependent_function_name,
values::{ndarray::NDArrayValue, ProxyValue, TypedArrayLikeAccessor},
CodeGenContext, CodeGenerator,
};
/// Generates a call to `__nac3_ndarray_transpose`.
///
/// Creates a transpose view of `src_ndarray` and writes the result to `dst_ndarray`.
///
/// `dst_ndarray` must fulfill the following preconditions:
///
/// - `dst_ndarray.ndims` must be initialized and must be equal to `src_ndarray.ndims`.
/// - `dst_ndarray.shape` must be allocated and may contain uninitialized values.
/// - `dst_ndarray.strides` must be allocated and may contain uninitialized values.
pub fn call_nac3_ndarray_transpose<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
src_ndarray: NDArrayValue<'ctx>,
dst_ndarray: NDArrayValue<'ctx>,
axes: Option<&impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>>,
) {
let llvm_usize = generator.get_size_type(ctx.ctx);
assert!(axes.is_none_or(|axes| axes.size(ctx, generator).get_type() == llvm_usize));
assert!(axes.is_none_or(|axes| axes.element_type(ctx, generator) == llvm_usize.into()));
let name = get_usize_dependent_function_name(generator, ctx, "__nac3_ndarray_transpose");
infer_and_call_function(
ctx,
&name,
None,
&[
src_ndarray.as_base_value().into(),
dst_ndarray.as_base_value().into(),
axes.map_or(llvm_usize.const_zero(), |axes| axes.size(ctx, generator)).into(),
axes.map_or(llvm_usize.ptr_type(AddressSpace::default()).const_null(), |axes| {
axes.base_ptr(ctx, generator)
})
.into(),
],
None,
None,
);
}

18
nac3core/src/codegen/irrt/range.rs
View File

@@ -6,6 +6,13 @@ use itertools::Either;
use crate::codegen::{CodeGenContext, CodeGenerator};
/// Invokes the `__nac3_range_slice_len` in IRRT.
///
/// - `start`: The `i32` start value for the slice.
/// - `end`: The `i32` end value for the slice.
/// - `step`: The `i32` step value for the slice.
///
/// Returns an `i32` value of the length of the slice.
pub fn calculate_len_for_slice_range<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
@@ -14,9 +21,15 @@ pub fn calculate_len_for_slice_range<'ctx, G: CodeGenerator + ?Sized>(
step: IntValue<'ctx>,
) -> IntValue<'ctx> {
const SYMBOL: &str = "__nac3_range_slice_len";
let llvm_i32 = ctx.ctx.i32_type();
assert_eq!(start.get_type(), llvm_i32);
assert_eq!(end.get_type(), llvm_i32);
assert_eq!(step.get_type(), llvm_i32);
let len_func = ctx.module.get_function(SYMBOL).unwrap_or_else(|| {
let i32_t = ctx.ctx.i32_type();
let fn_t = i32_t.fn_type(&[i32_t.into(), i32_t.into(), i32_t.into()], false);
let fn_t = llvm_i32.fn_type(&[llvm_i32.into(), llvm_i32.into(), llvm_i32.into()], false);
ctx.module.add_function(SYMBOL, fn_t, None)
});
@@ -33,6 +46,7 @@ pub fn calculate_len_for_slice_range<'ctx, G: CodeGenerator + ?Sized>(
[None, None, None],
ctx.current_loc,
);
ctx.builder
.build_call(len_func, &[start.into(), end.into(), step.into()], "calc_len")
.map(CallSiteValue::try_as_basic_value)

24
nac3core/src/codegen/irrt/string.rs
View File

@@ -1,7 +1,8 @@
use inkwell::values::{BasicValueEnum, CallSiteValue, IntValue, PointerValue};
use itertools::Either;
use crate::codegen::{macros::codegen_unreachable, CodeGenContext, CodeGenerator};
use super::get_usize_dependent_function_name;
use crate::codegen::{CodeGenContext, CodeGenerator};
/// Generates a call to string equality comparison. Returns an `i1` representing whether the strings are equal.
pub fn call_string_eq<'ctx, G: CodeGenerator + ?Sized>(
@@ -12,16 +13,14 @@ pub fn call_string_eq<'ctx, G: CodeGenerator + ?Sized>(
str2_ptr: PointerValue<'ctx>,
str2_len: IntValue<'ctx>,
) -> IntValue<'ctx> {
let (func_name, return_type) = match ctx.ctx.i32_type().get_bit_width() {
32 => ("nac3_str_eq", ctx.ctx.i32_type()),
64 => ("nac3_str_eq64", ctx.ctx.i64_type()),
bw => codegen_unreachable!(ctx, "Unsupported size type bit width: {}", bw),
};
let llvm_i1 = ctx.ctx.bool_type();
let func = ctx.module.get_function(func_name).unwrap_or_else(|| {
let func_name = get_usize_dependent_function_name(generator, ctx, "nac3_str_eq");
let func = ctx.module.get_function(&func_name).unwrap_or_else(|| {
ctx.module.add_function(
func_name,
return_type.fn_type(
&func_name,
llvm_i1.fn_type(
&[
str1_ptr.get_type().into(),
str1_len.get_type().into(),
@@ -33,8 +32,8 @@ pub fn call_string_eq<'ctx, G: CodeGenerator + ?Sized>(
None,
)
});
let result = ctx
.builder
ctx.builder
.build_call(
func,
&[str1_ptr.into(), str1_len.into(), str2_ptr.into(), str2_len.into()],
@@ -43,6 +42,5 @@ pub fn call_string_eq<'ctx, G: CodeGenerator + ?Sized>(
.map(CallSiteValue::try_as_basic_value)
.map(|v| v.map_left(BasicValueEnum::into_int_value))
.map(Either::unwrap_left)
.unwrap();
generator.bool_to_i1(ctx, result)
.unwrap()
}

33
nac3core/src/codegen/llvm_intrinsics.rs
View File

@@ -1,7 +1,6 @@
use inkwell::{
context::Context,
intrinsics::Intrinsic,
types::{AnyTypeEnum::IntType, FloatType},
types::AnyTypeEnum::IntType,
values::{BasicValueEnum, CallSiteValue, FloatValue, IntValue, PointerValue},
AddressSpace,
};
@@ -9,34 +8,6 @@ use itertools::Either;
use super::CodeGenContext;
/// Returns the string representation for the floating-point type `ft` when used in intrinsic
/// functions.
fn get_float_intrinsic_repr(ctx: &Context, ft: FloatType) -> &'static str {
// Standard LLVM floating-point types
if ft == ctx.f16_type() {
return "f16";
}
if ft == ctx.f32_type() {
return "f32";
}
if ft == ctx.f64_type() {
return "f64";
}
if ft == ctx.f128_type() {
return "f128";
}
// Non-standard floating-point types
if ft == ctx.x86_f80_type() {
return "f80";
}
if ft == ctx.ppc_f128_type() {
return "ppcf128";
}
unreachable!()
}
/// Invokes the [`llvm.va_start`](https://llvm.org/docs/LangRef.html#llvm-va-start-intrinsic)
/// intrinsic.
pub fn call_va_start<'ctx>(ctx: &CodeGenContext<'ctx, '_>, arglist: PointerValue<'ctx>) {
@@ -54,7 +25,7 @@ pub fn call_va_start<'ctx>(ctx: &CodeGenContext<'ctx, '_>, arglist: PointerValue
ctx.builder.build_call(intrinsic_fn, &[arglist.into()], "").unwrap();
}
/// Invokes the [`llvm.va_start`](https://llvm.org/docs/LangRef.html#llvm-va-start-intrinsic)
/// Invokes the [`llvm.va_end`](https://llvm.org/docs/LangRef.html#llvm-va-end-intrinsic)
/// intrinsic.
pub fn call_va_end<'ctx>(ctx: &CodeGenContext<'ctx, '_>, arglist: PointerValue<'ctx>) {
const FN_NAME: &str = "llvm.va_end";

8
nac3core/src/codegen/mod.rs
View File

@@ -42,7 +42,7 @@ use crate::{
};
use concrete_type::{ConcreteType, ConcreteTypeEnum, ConcreteTypeStore};
pub use generator::{CodeGenerator, DefaultCodeGenerator};
use types::{ndarray::NDArrayType, ListType, ProxyType, RangeType};
use types::{ndarray::NDArrayType, ListType, ProxyType, RangeType, TupleType};
pub mod builtin_fns;
pub mod concrete_type;
@@ -228,7 +228,7 @@ pub struct CodeGenContext<'ctx, 'a> {
pub current_loc: Location,
}
impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
impl CodeGenContext<'_, '_> {
/// Whether the [current basic block][Builder::get_insert_block] referenced by `builder`
/// contains a [terminator statement][BasicBlock::get_terminator].
pub fn is_terminated(&self) -> bool {
@@ -520,7 +520,7 @@ fn get_llvm_type<'ctx, G: CodeGenerator + ?Sized>(
ctx, module, generator, unifier, top_level, type_cache, dtype,
);
NDArrayType::new(generator, ctx, element_type, Some(ndims)).as_base_type().into()
NDArrayType::new(generator, ctx, element_type, ndims).as_base_type().into()
}
_ => unreachable!(
@@ -574,7 +574,7 @@ fn get_llvm_type<'ctx, G: CodeGenerator + ?Sized>(
get_llvm_type(ctx, module, generator, unifier, top_level, type_cache, *ty)
})
.collect_vec();
ctx.struct_type(&fields, false).into()
TupleType::new(generator, ctx, &fields).as_base_type().into()
}
TVirtual { .. } => unimplemented!(),
_ => unreachable!("{}", ty_enum.get_type_name()),

2439
nac3core/src/codegen/numpy.rs
View File

File diff suppressed because it is too large Load Diff

53
nac3core/src/codegen/stmt.rs
View File

@@ -16,7 +16,11 @@ use super::{
gen_in_range_check,
irrt::{handle_slice_indices, list_slice_assignment},
macros::codegen_unreachable,
values::{ArrayLikeIndexer, ArraySliceValue, ListValue, RangeValue},
types::ndarray::NDArrayType,
values::{
ndarray::{RustNDIndex, ScalarOrNDArray},
ArrayLikeIndexer, ArraySliceValue, ListValue, ProxyValue, RangeValue,
},
CodeGenContext, CodeGenerator,
};
use crate::{
@@ -411,7 +415,52 @@ pub fn gen_setitem<'ctx, G: CodeGenerator>(
if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() =>
{
// Handle NDArray item assignment
todo!("ndarray subscript assignment is not yet implemented");
// Process target
let target = generator
.gen_expr(ctx, target)?
.unwrap()
.to_basic_value_enum(ctx, generator, target_ty)?;
// Process key
let key = RustNDIndex::from_subscript_expr(generator, ctx, key)?;
// Process value
let value = value.to_basic_value_enum(ctx, generator, value_ty)?;
// Reference code:
// ```python
// target = target[key]
// value = np.asarray(value)
//
// shape = np.broadcast_shape((target, value))
//
// target = np.broadcast_to(target, shape)
// value = np.broadcast_to(value, shape)
//
// # ...and finally copy 1-1 from value to target.
// ```
let target = NDArrayType::from_unifier_type(generator, ctx, target_ty)
.map_value(target.into_pointer_value(), None);
let target = target.index(generator, ctx, &key);
let value = ScalarOrNDArray::from_value(generator, ctx, (value_ty, value))
.to_ndarray(generator, ctx);
let broadcast_ndims =
[target.get_type().ndims(), value.get_type().ndims()].into_iter().max().unwrap();
let broadcast_result = NDArrayType::new(
generator,
ctx.ctx,
value.get_type().element_type(),
broadcast_ndims,
)
.broadcast(generator, ctx, &[target, value]);
let target = broadcast_result.ndarrays[0];
let value = broadcast_result.ndarrays[1];
target.copy_data_from(generator, ctx, value);
}
_ => {
panic!("encountered unknown target type: {}", ctx.unifier.stringify(target_ty));

17
nac3core/src/codegen/test.rs
View File

@@ -36,7 +36,6 @@ use crate::{
struct Resolver {
id_to_type: HashMap<StrRef, Type>,
id_to_def: RwLock<HashMap<StrRef, DefinitionId>>,
class_names: HashMap<StrRef, Type>,
}
impl Resolver {
@@ -104,11 +103,9 @@ fn test_primitives() {
let top_level = Arc::new(composer.make_top_level_context());
unifier.top_level = Some(top_level.clone());
let resolver = Arc::new(Resolver {
id_to_type: HashMap::new(),
id_to_def: RwLock::new(HashMap::new()),
class_names: HashMap::default(),
}) as Arc<dyn SymbolResolver + Send + Sync>;
let resolver =
Arc::new(Resolver { id_to_type: HashMap::new(), id_to_def: RwLock::new(HashMap::new()) })
as Arc<dyn SymbolResolver + Send + Sync>;
let threads = vec![DefaultCodeGenerator::new("test".into(), 32).into()];
let signature = FunSignature {
@@ -298,11 +295,7 @@ fn test_simple_call() {
loc: None,
})));
let resolver = Resolver {
id_to_type: HashMap::new(),
id_to_def: RwLock::new(HashMap::new()),
class_names: HashMap::default(),
};
let resolver = Resolver { id_to_type: HashMap::new(), id_to_def: RwLock::new(HashMap::new()) };
resolver.add_id_def("foo".into(), DefinitionId(foo_id));
let resolver = Arc::new(resolver) as Arc<dyn SymbolResolver + Send + Sync>;
@@ -471,6 +464,6 @@ fn test_classes_ndarray_type_new() {
let llvm_i32 = ctx.i32_type();
let llvm_usize = generator.get_size_type(&ctx);
let llvm_ndarray = NDArrayType::new(&generator, &ctx, llvm_i32.into(), None);
let llvm_ndarray = NDArrayType::new(&generator, &ctx, llvm_i32.into(), 2);
assert!(NDArrayType::is_representable(llvm_ndarray.as_base_type(), llvm_usize).is_ok());
}

317
nac3core/src/codegen/types/list.rs
View File

@@ -1,69 +1,113 @@
use inkwell::{
context::Context,
context::{AsContextRef, Context},
types::{AnyTypeEnum, BasicType, BasicTypeEnum, IntType, PointerType},
values::IntValue,
AddressSpace,
values::{IntValue, PointerValue},
AddressSpace, IntPredicate, OptimizationLevel,
};
use itertools::Itertools;
use nac3core_derive::StructFields;
use super::ProxyType;
use crate::codegen::{
values::{ArraySliceValue, ListValue, ProxyValue},
CodeGenContext, CodeGenerator,
use crate::{
codegen::{
types::structure::{
check_struct_type_matches_fields, FieldIndexCounter, StructField, StructFields,
},
values::{ListValue, ProxyValue},
CodeGenContext, CodeGenerator,
},
typecheck::typedef::{iter_type_vars, Type, TypeEnum},
};
/// Proxy type for a `list` type in LLVM.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct ListType<'ctx> {
ty: PointerType<'ctx>,
item: Option<BasicTypeEnum<'ctx>>,
llvm_usize: IntType<'ctx>,
}
#[derive(PartialEq, Eq, Clone, Copy, StructFields)]
pub struct ListStructFields<'ctx> {
/// Array pointer to content.
#[value_type(i8_type().ptr_type(AddressSpace::default()))]
pub items: StructField<'ctx, PointerValue<'ctx>>,
/// Number of items in the array.
#[value_type(usize)]
pub len: StructField<'ctx, IntValue<'ctx>>,
}
impl<'ctx> ListStructFields<'ctx> {
#[must_use]
pub fn new_typed(item: BasicTypeEnum<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
let mut counter = FieldIndexCounter::default();
ListStructFields {
items: StructField::create(
&mut counter,
"items",
item.ptr_type(AddressSpace::default()),
),
len: StructField::create(&mut counter, "len", llvm_usize),
}
}
}
impl<'ctx> ListType<'ctx> {
/// Checks whether `llvm_ty` represents a `list` type, returning [Err] if it does not.
pub fn is_representable(
llvm_ty: PointerType<'ctx>,
llvm_usize: IntType<'ctx>,
) -> Result<(), String> {
let llvm_list_ty = llvm_ty.get_element_type();
let AnyTypeEnum::StructType(llvm_list_ty) = llvm_list_ty else {
return Err(format!("Expected struct type for `list` type, got {llvm_list_ty}"));
};
if llvm_list_ty.count_fields() != 2 {
return Err(format!(
"Expected 2 fields in `list`, got {}",
llvm_list_ty.count_fields()
));
}
let ctx = llvm_ty.get_context();
let list_size_ty = llvm_list_ty.get_field_type_at_index(0).unwrap();
let Ok(_) = PointerType::try_from(list_size_ty) else {
return Err(format!("Expected pointer type for `list.0`, got {list_size_ty}"));
let llvm_ty = llvm_ty.get_element_type();
let AnyTypeEnum::StructType(llvm_ty) = llvm_ty else {
return Err(format!("Expected struct type for `list` type, got {llvm_ty}"));
};
let list_data_ty = llvm_list_ty.get_field_type_at_index(1).unwrap();
let Ok(list_data_ty) = IntType::try_from(list_data_ty) else {
return Err(format!("Expected int type for `list.1`, got {list_data_ty}"));
};
if list_data_ty.get_bit_width() != llvm_usize.get_bit_width() {
return Err(format!(
"Expected {}-bit int type for `list.1`, got {}-bit int",
llvm_usize.get_bit_width(),
list_data_ty.get_bit_width()
));
}
let fields = ListStructFields::new(ctx, llvm_usize);
Ok(())
check_struct_type_matches_fields(
fields,
llvm_ty,
"list",
&[(fields.items.name(), &|ty| {
if ty.is_pointer_type() {
Ok(())
} else {
Err(format!("Expected T* for `list.items`, got {ty}"))
}
})],
)
}
/// Returns an instance of [`StructFields`] containing all field accessors for this type.
#[must_use]
fn fields(item: BasicTypeEnum<'ctx>, llvm_usize: IntType<'ctx>) -> ListStructFields<'ctx> {
ListStructFields::new_typed(item, llvm_usize)
}
/// See [`ListType::fields`].
// TODO: Move this into e.g. StructProxyType
#[must_use]
pub fn get_fields(&self, _ctx: &impl AsContextRef<'ctx>) -> ListStructFields<'ctx> {
Self::fields(self.item.unwrap_or(self.llvm_usize.into()), self.llvm_usize)
}
/// Creates an LLVM type corresponding to the expected structure of a `List`.
#[must_use]
fn llvm_type(
ctx: &'ctx Context,
element_type: BasicTypeEnum<'ctx>,
element_type: Option<BasicTypeEnum<'ctx>>,
llvm_usize: IntType<'ctx>,
) -> PointerType<'ctx> {
// struct List { data: T*, size: size_t }
let field_tys = [element_type.ptr_type(AddressSpace::default()).into(), llvm_usize.into()];
let element_type = element_type.unwrap_or(llvm_usize.into());
let field_tys =
Self::fields(element_type, llvm_usize).into_iter().map(|field| field.1).collect_vec();
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
}
@@ -76,9 +120,50 @@ impl<'ctx> ListType<'ctx> {
element_type: BasicTypeEnum<'ctx>,
) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_list = Self::llvm_type(ctx, element_type, llvm_usize);
let llvm_list = Self::llvm_type(ctx, Some(element_type), llvm_usize);
ListType::from_type(llvm_list, llvm_usize)
Self { ty: llvm_list, item: Some(element_type), llvm_usize }
}
/// Creates an instance of [`ListType`] with an unknown element type.
#[must_use]
pub fn new_untyped<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_list = Self::llvm_type(ctx, None, llvm_usize);
Self { ty: llvm_list, item: None, llvm_usize }
}
/// Creates an [`ListType`] from a [unifier type][Type].
#[must_use]
pub fn from_unifier_type<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &mut CodeGenContext<'ctx, '_>,
ty: Type,
) -> Self {
// Check unifier type and extract `item_type`
let elem_type = match &*ctx.unifier.get_ty_immutable(ty) {
TypeEnum::TObj { obj_id, params, .. }
if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
{
iter_type_vars(params).next().unwrap().ty
}
_ => panic!("Expected `list` type, but got {}", ctx.unifier.stringify(ty)),
};
let llvm_usize = generator.get_size_type(ctx.ctx);
let llvm_elem_type = if let TypeEnum::TVar { .. } = &*ctx.unifier.get_ty_immutable(ty) {
None
} else {
Some(ctx.get_llvm_type(generator, elem_type))
};
Self {
ty: Self::llvm_type(ctx.ctx, llvm_elem_type, llvm_usize),
item: llvm_elem_type,
llvm_usize,
}
}
/// Creates an [`ListType`] from a [`PointerType`].
@@ -86,47 +171,141 @@ impl<'ctx> ListType<'ctx> {
pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
debug_assert!(Self::is_representable(ptr_ty, llvm_usize).is_ok());
ListType { ty: ptr_ty, llvm_usize }
let ctx = ptr_ty.get_context();
// We are just searching for the index off a field - Slot an arbitrary element type in.
let item_field_idx =
Self::fields(ctx.i8_type().into(), llvm_usize).index_of_field(|f| f.items);
let item = unsafe {
ptr_ty
.get_element_type()
.into_struct_type()
.get_field_type_at_index_unchecked(item_field_idx)
.into_pointer_type()
.get_element_type()
};
let item = BasicTypeEnum::try_from(item).unwrap_or_else(|()| {
panic!(
"Expected BasicTypeEnum for list element type, got {}",
ptr_ty.get_element_type().print_to_string()
)
});
ListType { ty: ptr_ty, item: Some(item), llvm_usize }
}
/// Returns the type of the `size` field of this `list` type.
#[must_use]
pub fn size_type(&self) -> IntType<'ctx> {
self.as_base_type()
.get_element_type()
.into_struct_type()
.get_field_type_at_index(1)
.map(BasicTypeEnum::into_int_type)
.unwrap()
self.llvm_usize
}
/// Returns the element type of this `list` type.
#[must_use]
pub fn element_type(&self) -> AnyTypeEnum<'ctx> {
self.as_base_type()
.get_element_type()
.into_struct_type()
.get_field_type_at_index(0)
.map(BasicTypeEnum::into_pointer_type)
.map(PointerType::get_element_type)
.unwrap()
pub fn element_type(&self) -> Option<BasicTypeEnum<'ctx>> {
self.item
}
/// Allocates an instance of [`ListValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`ListValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
self.llvm_usize,
name,
)
}
/// Allocates a [`ListValue`] on the stack using `item` of this [`ListType`] instance.
///
/// The returned list will contain:
///
/// - `data`: Allocated with `len` number of elements.
/// - `len`: Initialized to the value of `len` passed to this function.
#[must_use]
pub fn construct<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
len: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let len = ctx.builder.build_int_z_extend(len, self.llvm_usize, "").unwrap();
// Generate a runtime assertion if allocating a non-empty list with unknown element type
if ctx.registry.llvm_options.opt_level == OptimizationLevel::None && self.item.is_none() {
let len_eqz = ctx
.builder
.build_int_compare(IntPredicate::EQ, len, self.llvm_usize.const_zero(), "")
.unwrap();
ctx.make_assert(
generator,
len_eqz,
"0:AssertionError",
"Cannot allocate a non-empty list with unknown element type",
[None, None, None],
ctx.current_loc,
);
}
let plist = self.alloca_var(generator, ctx, name);
plist.store_size(ctx, generator, len);
let item = self.item.unwrap_or(self.llvm_usize.into());
plist.create_data(ctx, item, None);
plist
}
/// Convenience function for creating a list with zero elements.
///
/// This function is preferred over [`ListType::construct`] if the length is known to always be
/// 0, as this function avoids injecting an IR assertion for checking if a non-empty untyped
/// list is being allocated.
///
/// The returned list will contain:
///
/// - `data`: Initialized to `(T*) 0`.
/// - `len`: Initialized to `0`.
#[must_use]
pub fn construct_empty<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let plist = self.alloca_var(generator, ctx, name);
plist.store_size(ctx, generator, self.llvm_usize.const_zero());
plist.create_data(ctx, self.item.unwrap_or(self.llvm_usize.into()), None);
plist
}
/// Converts an existing value into a [`ListValue`].
#[must_use]
pub fn map_value(
@@ -162,36 +341,8 @@ impl<'ctx> ProxyType<'ctx> for ListType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

65
nac3core/src/codegen/types/mod.rs
View File

@@ -16,7 +16,11 @@
//! the returned object. This is similar to a `new` expression in C++ but the object is allocated
//! on the stack.
use inkwell::{context::Context, types::BasicType, values::IntValue};
use inkwell::{
context::Context,
types::BasicType,
values::{IntValue, PointerValue},
};
use super::{
values::{ArraySliceValue, ProxyValue},
@@ -24,11 +28,13 @@ use super::{
};
pub use list::*;
pub use range::*;
pub use tuple::*;
mod list;
pub mod ndarray;
mod range;
pub mod structure;
mod tuple;
pub mod utils;
/// A LLVM type that is used to represent a corresponding type in NAC3.
@@ -53,23 +59,66 @@ pub trait ProxyType<'ctx>: Into<Self::Base> {
llvm_ty: Self::Base,
) -> Result<(), String>;
/// Creates a new value of this type, returning the LLVM instance of this value.
fn raw_alloca<G: CodeGenerator + ?Sized>(
/// Returns the type that should be used in `alloca` IR statements.
fn alloca_type(&self) -> impl BasicType<'ctx>;
/// Creates a new value of this type by invoking `alloca` at the current builder location,
/// returning a [`PointerValue`] instance representing the allocated value.
fn raw_alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> PointerValue<'ctx> {
ctx.builder
.build_alloca(self.alloca_type().as_basic_type_enum(), name.unwrap_or_default())
.unwrap()
}
/// Creates a new value of this type by invoking `alloca` at the beginning of the function,
/// returning a [`PointerValue`] instance representing the allocated value.
fn raw_alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base;
) -> PointerValue<'ctx> {
generator.gen_var_alloc(ctx, self.alloca_type().as_basic_type_enum(), name).unwrap()
}
/// Creates a new array value of this type, returning an [`ArraySliceValue`] encapsulating the
/// resulting array.
fn array_alloca<G: CodeGenerator + ?Sized>(
/// Creates a new array value of this type by invoking `alloca` at the current builder location,
/// returning an [`ArraySliceValue`] encapsulating the resulting array.
fn array_alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
ArraySliceValue::from_ptr_val(
ctx.builder
.build_array_alloca(
self.alloca_type().as_basic_type_enum(),
size,
name.unwrap_or_default(),
)
.unwrap(),
size,
name,
)
}
/// Creates a new array value of this type by invoking `alloca` at the beginning of the
/// function, returning an [`ArraySliceValue`] encapsulating the resulting array.
fn array_alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx>;
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(ctx, self.alloca_type().as_basic_type_enum(), size, name)
.unwrap()
}
/// Returns the [base type][Self::Base] of this proxy.
fn as_base_type(&self) -> Self::Base;

243
nac3core/src/codegen/types/ndarray/array.rs Normal file
View File

@@ -0,0 +1,243 @@
use inkwell::{
types::BasicTypeEnum,
values::{BasicValueEnum, IntValue},
AddressSpace,
};
use crate::{
codegen::{
irrt,
stmt::gen_if_else_expr_callback,
types::{ndarray::NDArrayType, ListType, ProxyType},
values::{
ndarray::NDArrayValue, ArrayLikeValue, ArraySliceValue, ListValue, ProxyValue,
TypedArrayLikeAdapter, TypedArrayLikeMutator,
},
CodeGenContext, CodeGenerator,
},
toplevel::helper::{arraylike_flatten_element_type, arraylike_get_ndims},
typecheck::typedef::{Type, TypeEnum},
};
/// Get the expected `dtype` and `ndims` of the ndarray returned by `np_array(<list>)`.
fn get_list_object_dtype_and_ndims<'ctx, G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &mut CodeGenContext<'ctx, '_>,
list_ty: Type,
) -> (BasicTypeEnum<'ctx>, u64) {
let dtype = arraylike_flatten_element_type(&mut ctx.unifier, list_ty);
let ndims = arraylike_get_ndims(&mut ctx.unifier, list_ty);
(ctx.get_llvm_type(generator, dtype), ndims)
}
impl<'ctx> NDArrayType<'ctx> {
/// Implementation of `np_array(<list>, copy=True)`
fn construct_numpy_array_from_list_copy_true_impl<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(list_ty, list): (Type, ListValue<'ctx>),
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let (dtype, ndims_int) = get_list_object_dtype_and_ndims(generator, ctx, list_ty);
assert!(self.ndims >= ndims_int);
assert_eq!(dtype, self.dtype);
let list_value = list.as_i8_list(generator, ctx);
// Validate `list` has a consistent shape.
// Raise an exception if `list` is something abnormal like `[[1, 2], [3]]`.
// If `list` has a consistent shape, deduce the shape and write it to `shape`.
let ndims = self.llvm_usize.const_int(ndims_int, false);
let shape = ctx.builder.build_array_alloca(self.llvm_usize, ndims, "").unwrap();
let shape = ArraySliceValue::from_ptr_val(shape, ndims, None);
let shape = TypedArrayLikeAdapter::from(
shape,
|_, _, val| val.into_int_value(),
|_, _, val| val.into(),
);
irrt::ndarray::call_nac3_ndarray_array_set_and_validate_list_shape(
generator, ctx, list_value, ndims, &shape,
);
let ndarray = Self::new(generator, ctx.ctx, dtype, ndims_int)
.construct_uninitialized(generator, ctx, name);
ndarray.copy_shape_from_array(generator, ctx, shape.base_ptr(ctx, generator));
unsafe { ndarray.create_data(generator, ctx) };
// Copy all contents from the list.
irrt::ndarray::call_nac3_ndarray_array_write_list_to_array(
generator, ctx, list_value, ndarray,
);
ndarray
}
/// Implementation of `np_array(<list>, copy=None)`
fn construct_numpy_array_from_list_copy_none_impl<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(list_ty, list): (Type, ListValue<'ctx>),
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
// np_array without copying is only possible `list` is not nested.
//
// If `list` is `list[T]`, we can create an ndarray with `data` set
// to the array pointer of `list`.
//
// If `list` is `list[list[T]]` or worse, copy.
let (dtype, ndims) = get_list_object_dtype_and_ndims(generator, ctx, list_ty);
if ndims == 1 {
// `list` is not nested
assert_eq!(ndims, 1);
assert!(self.ndims >= ndims);
assert_eq!(dtype, self.dtype);
let llvm_pi8 = ctx.ctx.i8_type().ptr_type(AddressSpace::default());
let ndarray = Self::new(generator, ctx.ctx, dtype, 1)
.construct_uninitialized(generator, ctx, name);
// Set data
let data = ctx
.builder
.build_pointer_cast(list.data().base_ptr(ctx, generator), llvm_pi8, "")
.unwrap();
ndarray.store_data(ctx, data);
// ndarray->shape[0] = list->len;
let shape = ndarray.shape();
let list_len = list.load_size(ctx, None);
unsafe {
shape.set_typed_unchecked(ctx, generator, &self.llvm_usize.const_zero(), list_len);
}
// Set strides, the `data` is contiguous
ndarray.set_strides_contiguous(generator, ctx);
ndarray
} else {
// `list` is nested, copy
self.construct_numpy_array_from_list_copy_true_impl(
generator,
ctx,
(list_ty, list),
name,
)
}
}
/// Implementation of `np_array(<list>, copy=copy)`
fn construct_numpy_array_list_impl<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(list_ty, list): (Type, ListValue<'ctx>),
copy: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(copy.get_type(), ctx.ctx.bool_type());
let (dtype, ndims) = get_list_object_dtype_and_ndims(generator, ctx, list_ty);
let ndarray = gen_if_else_expr_callback(
generator,
ctx,
|_generator, _ctx| Ok(copy),
|generator, ctx| {
let ndarray = self.construct_numpy_array_from_list_copy_true_impl(
generator,
ctx,
(list_ty, list),
name,
);
Ok(Some(ndarray.as_base_value()))
},
|generator, ctx| {
let ndarray = self.construct_numpy_array_from_list_copy_none_impl(
generator,
ctx,
(list_ty, list),
name,
);
Ok(Some(ndarray.as_base_value()))
},
)
.unwrap()
.map(BasicValueEnum::into_pointer_value)
.unwrap();
NDArrayType::new(generator, ctx.ctx, dtype, ndims).map_value(ndarray, None)
}
/// Implementation of `np_array(<ndarray>, copy=copy)`.
pub fn construct_numpy_array_ndarray_impl<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ndarray: NDArrayValue<'ctx>,
copy: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(ndarray.get_type().dtype, self.dtype);
assert!(self.ndims >= ndarray.get_type().ndims);
assert_eq!(copy.get_type(), ctx.ctx.bool_type());
let ndarray_val = gen_if_else_expr_callback(
generator,
ctx,
|_generator, _ctx| Ok(copy),
|generator, ctx| {
let ndarray = ndarray.make_copy(generator, ctx); // Force copy
Ok(Some(ndarray.as_base_value()))
},
|_generator, _ctx| {
// No need to copy. Return `ndarray` itself.
Ok(Some(ndarray.as_base_value()))
},
)
.unwrap()
.map(BasicValueEnum::into_pointer_value)
.unwrap();
ndarray.get_type().map_value(ndarray_val, name)
}
/// Create a new ndarray like
/// [`np.array()`](https://numpy.org/doc/stable/reference/generated/numpy.array.html).
///
/// Note that the returned [`NDArrayValue`] may have fewer dimensions than is specified by this
/// instance. Use [`NDArrayValue::atleast_nd`] on the returned value if an `ndarray` instance
/// with the exact number of dimensions is needed.
pub fn construct_numpy_array<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
(object_ty, object): (Type, BasicValueEnum<'ctx>),
copy: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
match &*ctx.unifier.get_ty_immutable(object_ty) {
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.list.obj_id(&ctx.unifier).unwrap() =>
{
let list = ListType::from_unifier_type(generator, ctx, object_ty)
.map_value(object.into_pointer_value(), None);
self.construct_numpy_array_list_impl(generator, ctx, (object_ty, list), copy, name)
}
TypeEnum::TObj { obj_id, .. }
if *obj_id == ctx.primitives.ndarray.obj_id(&ctx.unifier).unwrap() =>
{
let ndarray = NDArrayType::from_unifier_type(generator, ctx, object_ty)
.map_value(object.into_pointer_value(), None);
self.construct_numpy_array_ndarray_impl(generator, ctx, ndarray, copy, name)
}
_ => panic!("Unrecognized object type: {}", ctx.unifier.stringify(object_ty)), // Typechecker ensures this
}
}
}

176
nac3core/src/codegen/types/ndarray/broadcast.rs Normal file
View File

@@ -0,0 +1,176 @@
use inkwell::{
context::{AsContextRef, Context},
types::{AnyTypeEnum, BasicType, BasicTypeEnum, IntType, PointerType},
values::{IntValue, PointerValue},
AddressSpace,
};
use itertools::Itertools;
use nac3core_derive::StructFields;
use crate::codegen::{
types::{
structure::{check_struct_type_matches_fields, StructField, StructFields},
ProxyType,
},
values::{ndarray::ShapeEntryValue, ProxyValue},
CodeGenContext, CodeGenerator,
};
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub struct ShapeEntryType<'ctx> {
ty: PointerType<'ctx>,
llvm_usize: IntType<'ctx>,
}
#[derive(PartialEq, Eq, Clone, Copy, StructFields)]
pub struct ShapeEntryStructFields<'ctx> {
#[value_type(usize)]
pub ndims: StructField<'ctx, IntValue<'ctx>>,
#[value_type(usize.ptr_type(AddressSpace::default()))]
pub shape: StructField<'ctx, PointerValue<'ctx>>,
}
impl<'ctx> ShapeEntryType<'ctx> {
/// Checks whether `llvm_ty` represents a [`ShapeEntryType`], returning [Err] if it does not.
pub fn is_representable(
llvm_ty: PointerType<'ctx>,
llvm_usize: IntType<'ctx>,
) -> Result<(), String> {
let ctx = llvm_ty.get_context();
let llvm_ndarray_ty = llvm_ty.get_element_type();
let AnyTypeEnum::StructType(llvm_ndarray_ty) = llvm_ndarray_ty else {
return Err(format!(
"Expected struct type for `ShapeEntry` type, got {llvm_ndarray_ty}"
));
};
check_struct_type_matches_fields(
Self::fields(ctx, llvm_usize),
llvm_ndarray_ty,
"NDArray",
&[],
)
}
/// Returns an instance of [`StructFields`] containing all field accessors for this type.
#[must_use]
fn fields(
ctx: impl AsContextRef<'ctx>,
llvm_usize: IntType<'ctx>,
) -> ShapeEntryStructFields<'ctx> {
ShapeEntryStructFields::new(ctx, llvm_usize)
}
/// See [`ShapeEntryStructFields::fields`].
// TODO: Move this into e.g. StructProxyType
#[must_use]
pub fn get_fields(&self, ctx: impl AsContextRef<'ctx>) -> ShapeEntryStructFields<'ctx> {
Self::fields(ctx, self.llvm_usize)
}
/// Creates an LLVM type corresponding to the expected structure of a `ShapeEntry`.
#[must_use]
fn llvm_type(ctx: &'ctx Context, llvm_usize: IntType<'ctx>) -> PointerType<'ctx> {
let field_tys =
Self::fields(ctx, llvm_usize).into_iter().map(|field| field.1).collect_vec();
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
}
/// Creates an instance of [`ShapeEntryType`].
#[must_use]
pub fn new<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_ty = Self::llvm_type(ctx, llvm_usize);
Self { ty: llvm_ty, llvm_usize }
}
/// Creates a [`ShapeEntryType`] from a [`PointerType`] representing an `ShapeEntry`.
#[must_use]
pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
debug_assert!(Self::is_representable(ptr_ty, llvm_usize).is_ok());
Self { ty: ptr_ty, llvm_usize }
}
/// Allocates an instance of [`ShapeEntryValue`] as if by calling `alloca` on the base type.
#[must_use]
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`ShapeEntryValue`] as if by calling `alloca` on the base type.
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca_var(generator, ctx, name),
self.llvm_usize,
name,
)
}
/// Converts an existing value into a [`ShapeEntryValue`].
#[must_use]
pub fn map_value(
&self,
value: <<Self as ProxyType<'ctx>>::Value as ProxyValue<'ctx>>::Base,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(value, self.llvm_usize, name)
}
}
impl<'ctx> ProxyType<'ctx> for ShapeEntryType<'ctx> {
type Base = PointerType<'ctx>;
type Value = ShapeEntryValue<'ctx>;
fn is_type<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
llvm_ty: impl BasicType<'ctx>,
) -> Result<(), String> {
if let BasicTypeEnum::PointerType(ty) = llvm_ty.as_basic_type_enum() {
<Self as ProxyType<'ctx>>::is_representable(generator, ctx, ty)
} else {
Err(format!("Expected pointer type, got {llvm_ty:?}"))
}
}
fn is_representable<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
llvm_ty: Self::Base,
) -> Result<(), String> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {
self.ty
}
}
impl<'ctx> From<ShapeEntryType<'ctx>> for PointerType<'ctx> {
fn from(value: ShapeEntryType<'ctx>) -> Self {
value.as_base_type()
}
}

75
nac3core/src/codegen/types/ndarray/contiguous.rs
View File

@@ -16,7 +16,7 @@ use crate::{
},
ProxyType,
},
values::{ndarray::ContiguousNDArrayValue, ArraySliceValue, ProxyValue},
values::{ndarray::ContiguousNDArrayValue, ProxyValue},
CodeGenContext, CodeGenerator,
},
toplevel::numpy::unpack_ndarray_var_tys,
@@ -31,7 +31,7 @@ pub struct ContiguousNDArrayType<'ctx> {
}
#[derive(PartialEq, Eq, Clone, Copy, StructFields)]
pub struct ContiguousNDArrayFields<'ctx> {
pub struct ContiguousNDArrayStructFields<'ctx> {
#[value_type(usize)]
pub ndims: StructField<'ctx, IntValue<'ctx>>,
#[value_type(usize.ptr_type(AddressSpace::default()))]
@@ -40,12 +40,12 @@ pub struct ContiguousNDArrayFields<'ctx> {
pub data: StructField<'ctx, PointerValue<'ctx>>,
}
impl<'ctx> ContiguousNDArrayFields<'ctx> {
impl<'ctx> ContiguousNDArrayStructFields<'ctx> {
#[must_use]
pub fn new_typed(item: BasicTypeEnum<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
let mut counter = FieldIndexCounter::default();
ContiguousNDArrayFields {
ContiguousNDArrayStructFields {
ndims: StructField::create(&mut counter, "ndims", llvm_usize),
shape: StructField::create(
&mut counter,
@@ -72,7 +72,7 @@ impl<'ctx> ContiguousNDArrayType<'ctx> {
));
};
let fields = ContiguousNDArrayFields::new(ctx, llvm_usize);
let fields = ContiguousNDArrayStructFields::new(ctx, llvm_usize);
check_struct_type_matches_fields(
fields,
@@ -93,14 +93,14 @@ impl<'ctx> ContiguousNDArrayType<'ctx> {
fn fields(
item: BasicTypeEnum<'ctx>,
llvm_usize: IntType<'ctx>,
) -> ContiguousNDArrayFields<'ctx> {
ContiguousNDArrayFields::new_typed(item, llvm_usize)
) -> ContiguousNDArrayStructFields<'ctx> {
ContiguousNDArrayStructFields::new_typed(item, llvm_usize)
}
/// See [`NDArrayType::fields`].
// TODO: Move this into e.g. StructProxyType
#[must_use]
pub fn get_fields(&self) -> ContiguousNDArrayFields<'ctx> {
pub fn get_fields(&self) -> ContiguousNDArrayStructFields<'ctx> {
Self::fields(self.item, self.llvm_usize)
}
@@ -157,16 +157,37 @@ impl<'ctx> ContiguousNDArrayType<'ctx> {
Self { ty: ptr_ty, item, llvm_usize }
}
/// Allocates an instance of [`ContiguousNDArrayValue`] as if by calling `alloca` on the base type.
/// Allocates an instance of [`ContiguousNDArrayValue`] as if by calling `alloca` on the base
/// type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.item,
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`ContiguousNDArrayValue`] as if by calling `alloca` on the base
/// type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
self.item,
self.llvm_usize,
name,
@@ -213,36 +234,8 @@ impl<'ctx> ProxyType<'ctx> for ContiguousNDArrayType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

236
nac3core/src/codegen/types/ndarray/factory.rs Normal file
View File

@@ -0,0 +1,236 @@
use inkwell::{
values::{BasicValueEnum, IntValue},
IntPredicate,
};
use super::NDArrayType;
use crate::{
codegen::{
irrt, types::ProxyType, values::TypedArrayLikeAccessor, CodeGenContext, CodeGenerator,
},
typecheck::typedef::Type,
};
/// Get the zero value in `np.zeros()` of a `dtype`.
fn ndarray_zero_value<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
) -> BasicValueEnum<'ctx> {
if [ctx.primitives.int32, ctx.primitives.uint32]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
ctx.ctx.i32_type().const_zero().into()
} else if [ctx.primitives.int64, ctx.primitives.uint64]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
ctx.ctx.i64_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.float) {
ctx.ctx.f64_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.bool) {
ctx.ctx.bool_type().const_zero().into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.str) {
ctx.gen_string(generator, "").into()
} else {
panic!("unrecognized dtype: {}", ctx.unifier.stringify(dtype));
}
}
/// Get the one value in `np.ones()` of a `dtype`.
fn ndarray_one_value<'ctx, G: CodeGenerator + ?Sized>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
) -> BasicValueEnum<'ctx> {
if [ctx.primitives.int32, ctx.primitives.uint32]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
let is_signed = ctx.unifier.unioned(dtype, ctx.primitives.int32);
ctx.ctx.i32_type().const_int(1, is_signed).into()
} else if [ctx.primitives.int64, ctx.primitives.uint64]
.iter()
.any(|ty| ctx.unifier.unioned(dtype, *ty))
{
let is_signed = ctx.unifier.unioned(dtype, ctx.primitives.int64);
ctx.ctx.i64_type().const_int(1, is_signed).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.float) {
ctx.ctx.f64_type().const_float(1.0).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.bool) {
ctx.ctx.bool_type().const_int(1, false).into()
} else if ctx.unifier.unioned(dtype, ctx.primitives.str) {
ctx.gen_string(generator, "1").into()
} else {
panic!("unrecognized dtype: {}", ctx.unifier.stringify(dtype));
}
}
impl<'ctx> NDArrayType<'ctx> {
/// Create an ndarray like
/// [`np.empty`](https://numpy.org/doc/stable/reference/generated/numpy.empty.html).
pub fn construct_numpy_empty<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let ndarray = self.construct_uninitialized(generator, ctx, name);
// Validate `shape`
irrt::ndarray::call_nac3_ndarray_util_assert_shape_no_negative(generator, ctx, shape);
ndarray.copy_shape_from_array(generator, ctx, shape.base_ptr(ctx, generator));
unsafe { ndarray.create_data(generator, ctx) };
ndarray
}
/// Create an ndarray like
/// [`np.full`](https://numpy.org/doc/stable/reference/generated/numpy.full.html).
pub fn construct_numpy_full<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
fill_value: BasicValueEnum<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let ndarray = self.construct_numpy_empty(generator, ctx, shape, name);
ndarray.fill(generator, ctx, fill_value);
ndarray
}
/// Create an ndarray like
/// [`np.zero`](https://numpy.org/doc/stable/reference/generated/numpy.zeros.html).
pub fn construct_numpy_zeros<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(
ctx.get_llvm_type(generator, dtype),
self.dtype,
"Expected LLVM dtype={} but got {}",
self.dtype.print_to_string(),
ctx.get_llvm_type(generator, dtype).print_to_string(),
);
let fill_value = ndarray_zero_value(generator, ctx, dtype);
self.construct_numpy_full(generator, ctx, shape, fill_value, name)
}
/// Create an ndarray like
/// [`np.ones`](https://numpy.org/doc/stable/reference/generated/numpy.ones.html).
pub fn construct_numpy_ones<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
shape: &impl TypedArrayLikeAccessor<'ctx, G, IntValue<'ctx>>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(
ctx.get_llvm_type(generator, dtype),
self.dtype,
"Expected LLVM dtype={} but got {}",
self.dtype.print_to_string(),
ctx.get_llvm_type(generator, dtype).print_to_string(),
);
let fill_value = ndarray_one_value(generator, ctx, dtype);
self.construct_numpy_full(generator, ctx, shape, fill_value, name)
}
/// Create an ndarray like
/// [`np.eye`](https://numpy.org/doc/stable/reference/generated/numpy.eye.html).
#[allow(clippy::too_many_arguments)]
pub fn construct_numpy_eye<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
nrows: IntValue<'ctx>,
ncols: IntValue<'ctx>,
offset: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(
ctx.get_llvm_type(generator, dtype),
self.dtype,
"Expected LLVM dtype={} but got {}",
self.dtype.print_to_string(),
ctx.get_llvm_type(generator, dtype).print_to_string(),
);
assert_eq!(nrows.get_type(), self.llvm_usize);
assert_eq!(ncols.get_type(), self.llvm_usize);
assert_eq!(offset.get_type(), self.llvm_usize);
let ndzero = ndarray_zero_value(generator, ctx, dtype);
let ndone = ndarray_one_value(generator, ctx, dtype);
let ndarray = self.construct_dyn_shape(generator, ctx, &[nrows, ncols], name);
// Create data and make the matrix like look np.eye()
unsafe {
ndarray.create_data(generator, ctx);
}
ndarray
.foreach(generator, ctx, |generator, ctx, _, nditer| {
// NOTE: rows and cols can never be zero here, since this ndarray's `np.size` would be zero
// and this loop would not execute.
let indices = nditer.get_indices();
let row_i = unsafe {
indices.get_typed_unchecked(ctx, generator, &self.llvm_usize.const_zero(), None)
};
let col_i = unsafe {
indices.get_typed_unchecked(
ctx,
generator,
&self.llvm_usize.const_int(1, false),
None,
)
};
let be_one = ctx
.builder
.build_int_compare(
IntPredicate::EQ,
ctx.builder.build_int_add(row_i, offset, "").unwrap(),
col_i,
"",
)
.unwrap();
let value = ctx.builder.build_select(be_one, ndone, ndzero, "value").unwrap();
let p = nditer.get_pointer(ctx);
ctx.builder.build_store(p, value).unwrap();
Ok(())
})
.unwrap();
ndarray
}
/// Create an ndarray like
/// [`np.identity`](https://numpy.org/doc/stable/reference/generated/numpy.identity.html).
pub fn construct_numpy_identity<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
dtype: Type,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let offset = self.llvm_usize.const_zero();
self.construct_numpy_eye(generator, ctx, dtype, size, size, offset, name)
}
}

56
nac3core/src/codegen/types/ndarray/indexing.rs
View File

@@ -90,15 +90,33 @@ impl<'ctx> NDIndexType<'ctx> {
Self { ty: ptr_ty, llvm_usize }
}
/// Allocates an instance of [`NDIndexValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`NDIndexValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
self.llvm_usize,
name,
)
@@ -114,7 +132,7 @@ impl<'ctx> NDIndexType<'ctx> {
) -> ArraySliceValue<'ctx> {
// Allocate the LLVM ndindices.
let num_ndindices = self.llvm_usize.const_int(in_ndindices.len() as u64, false);
let ndindices = self.array_alloca(generator, ctx, num_ndindices, None);
let ndindices = self.array_alloca_var(generator, ctx, num_ndindices, None);
// Initialize all of them.
for (i, in_ndindex) in in_ndindices.iter().enumerate() {
@@ -171,36 +189,8 @@ impl<'ctx> ProxyType<'ctx> for NDIndexType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

187
nac3core/src/codegen/types/ndarray/map.rs Normal file
View File

@@ -0,0 +1,187 @@
use inkwell::{types::BasicTypeEnum, values::BasicValueEnum};
use itertools::Itertools;
use crate::codegen::{
stmt::gen_for_callback,
types::{
ndarray::{NDArrayType, NDIterType},
ProxyType,
},
values::{
ndarray::{NDArrayOut, NDArrayValue, ScalarOrNDArray},
ArrayLikeValue, ProxyValue,
},
CodeGenContext, CodeGenerator,
};
impl<'ctx> NDArrayType<'ctx> {
/// Generate LLVM IR to broadcast `ndarray`s together, and starmap through them with `mapping`
/// elementwise.
///
/// `mapping` is an LLVM IR generator. The input of `mapping` is the list of elements when
/// iterating through the input `ndarrays` after broadcasting. The output of `mapping` is the
/// result of the elementwise operation.
///
/// `out` specifies whether the result should be a new ndarray or to be written an existing
/// ndarray.
pub fn broadcast_starmap<'a, G, MappingFn>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, 'a>,
ndarrays: &[NDArrayValue<'ctx>],
out: NDArrayOut<'ctx>,
mapping: MappingFn,
) -> Result<<Self as ProxyType<'ctx>>::Value, String>
where
G: CodeGenerator + ?Sized,
MappingFn: FnOnce(
&mut G,
&mut CodeGenContext<'ctx, 'a>,
&[BasicValueEnum<'ctx>],
) -> Result<BasicValueEnum<'ctx>, String>,
{
// Broadcast inputs
let broadcast_result = self.broadcast(generator, ctx, ndarrays);
let out_ndarray = match out {
NDArrayOut::NewNDArray { dtype } => {
// Create a new ndarray based on the broadcast shape.
let result_ndarray =
NDArrayType::new(generator, ctx.ctx, dtype, broadcast_result.ndims)
.construct_uninitialized(generator, ctx, None);
result_ndarray.copy_shape_from_array(
generator,
ctx,
broadcast_result.shape.base_ptr(ctx, generator),
);
unsafe {
result_ndarray.create_data(generator, ctx);
}
result_ndarray
}
NDArrayOut::WriteToNDArray { ndarray: result_ndarray } => {
// Use an existing ndarray.
// Check that its shape is compatible with the broadcast shape.
result_ndarray.assert_can_be_written_by_out(generator, ctx, broadcast_result.shape);
result_ndarray
}
};
// Map element-wise and store results into `mapped_ndarray`.
let nditer = NDIterType::new(generator, ctx.ctx).construct(generator, ctx, out_ndarray);
gen_for_callback(
generator,
ctx,
Some("broadcast_starmap"),
|generator, ctx| {
// Create NDIters for all broadcasted input ndarrays.
let other_nditers = broadcast_result
.ndarrays
.iter()
.map(|ndarray| {
NDIterType::new(generator, ctx.ctx).construct(generator, ctx, *ndarray)
})
.collect_vec();
Ok((nditer, other_nditers))
},
|generator, ctx, (out_nditer, _in_nditers)| {
// We can simply use `out_nditer`'s `has_element()`.
// `in_nditers`' `has_element()`s should return the same value.
Ok(out_nditer.has_element(generator, ctx))
},
|generator, ctx, _hooks, (out_nditer, in_nditers)| {
// Get all the scalars from the broadcasted input ndarrays, pass them to `mapping`,
// and write to `out_ndarray`.
let in_scalars =
in_nditers.iter().map(|nditer| nditer.get_scalar(ctx)).collect_vec();
let result = mapping(generator, ctx, &in_scalars)?;
let p = out_nditer.get_pointer(ctx);
ctx.builder.build_store(p, result).unwrap();
Ok(())
},
|generator, ctx, (out_nditer, in_nditers)| {
// Advance all iterators
out_nditer.next(generator, ctx);
in_nditers.iter().for_each(|nditer| nditer.next(generator, ctx));
Ok(())
},
)?;
Ok(out_ndarray)
}
}
impl<'ctx> ScalarOrNDArray<'ctx> {
/// Starmap through a list of inputs using `mapping`, where an input could be an ndarray, a
/// scalar.
///
/// This function is very helpful when implementing NumPy functions that takes on either scalars
/// or ndarrays or a mix of them as their inputs and produces either an ndarray with broadcast,
/// or a scalar if all its inputs are all scalars.
///
/// For example ,this function can be used to implement `np.add`, which has the following
/// behaviors:
///
/// - `np.add(3, 4) = 7` # (scalar, scalar) -> scalar
/// - `np.add(3, np.array([4, 5, 6]))` # (scalar, ndarray) -> ndarray; the first `scalar` is
/// converted into an ndarray and broadcasted.
/// - `np.add(np.array([[1], [2], [3]]), np.array([[4, 5, 6]]))` # (ndarray, ndarray) ->
/// ndarray; there is broadcasting.
///
/// ## Details:
///
/// If `inputs` are all [`ScalarOrNDArray::Scalar`], the output will be a
/// [`ScalarOrNDArray::Scalar`] with type `ret_dtype`.
///
/// Otherwise (if there are any [`ScalarOrNDArray::NDArray`] in `inputs`), all inputs will be
/// 'as-ndarray'-ed into ndarrays, then all inputs (now all ndarrays) will be passed to
/// [`NDArrayValue::broadcasting_starmap`] and **create** a new ndarray with dtype `ret_dtype`.
pub fn broadcasting_starmap<'a, G, MappingFn>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, 'a>,
inputs: &[ScalarOrNDArray<'ctx>],
ret_dtype: BasicTypeEnum<'ctx>,
mapping: MappingFn,
) -> Result<ScalarOrNDArray<'ctx>, String>
where
G: CodeGenerator + ?Sized,
MappingFn: FnOnce(
&mut G,
&mut CodeGenContext<'ctx, 'a>,
&[BasicValueEnum<'ctx>],
) -> Result<BasicValueEnum<'ctx>, String>,
{
// Check if all inputs are Scalars
let all_scalars: Option<Vec<_>> =
inputs.iter().map(BasicValueEnum::<'ctx>::try_from).try_collect().ok();
if let Some(scalars) = all_scalars {
let scalars = scalars.iter().copied().collect_vec();
let value = mapping(generator, ctx, &scalars)?;
Ok(ScalarOrNDArray::Scalar(value))
} else {
// Promote all input to ndarrays and map through them.
let inputs = inputs.iter().map(|input| input.to_ndarray(generator, ctx)).collect_vec();
let ndarray = NDArrayType::new_broadcast(
generator,
ctx.ctx,
ret_dtype,
&inputs.iter().map(NDArrayValue::get_type).collect_vec(),
)
.broadcast_starmap(
generator,
ctx,
&inputs,
NDArrayOut::NewNDArray { dtype: ret_dtype },
mapping,
)?;
Ok(ScalarOrNDArray::NDArray(ndarray))
}
}
}

131
nac3core/src/codegen/types/ndarray/mod.rs
View File

@@ -14,18 +14,23 @@ use super::{
};
use crate::{
codegen::{
values::{ndarray::NDArrayValue, ArraySliceValue, ProxyValue, TypedArrayLikeMutator},
values::{ndarray::NDArrayValue, ProxyValue, TypedArrayLikeMutator},
{CodeGenContext, CodeGenerator},
},
toplevel::{helper::extract_ndims, numpy::unpack_ndarray_var_tys},
typecheck::typedef::Type,
};
pub use broadcast::*;
pub use contiguous::*;
pub use indexing::*;
pub use nditer::*;
mod array;
mod broadcast;
mod contiguous;
pub mod factory;
mod indexing;
mod map;
mod nditer;
/// Proxy type for a `ndarray` type in LLVM.
@@ -33,7 +38,7 @@ mod nditer;
pub struct NDArrayType<'ctx> {
ty: PointerType<'ctx>,
dtype: BasicTypeEnum<'ctx>,
ndims: Option<u64>,
ndims: u64,
llvm_usize: IntType<'ctx>,
}
@@ -108,7 +113,7 @@ impl<'ctx> NDArrayType<'ctx> {
generator: &G,
ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>,
ndims: Option<u64>,
ndims: u64,
) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_ndarray = Self::llvm_type(ctx, llvm_usize);
@@ -116,6 +121,20 @@ impl<'ctx> NDArrayType<'ctx> {
NDArrayType { ty: llvm_ndarray, dtype, ndims, llvm_usize }
}
/// Creates an instance of [`NDArrayType`] as a result of a broadcast operation over one or more
/// `ndarray` operands.
#[must_use]
pub fn new_broadcast<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>,
inputs: &[NDArrayType<'ctx>],
) -> Self {
assert!(!inputs.is_empty());
Self::new(generator, ctx, dtype, inputs.iter().map(NDArrayType::ndims).max().unwrap())
}
/// Creates an instance of [`NDArrayType`] with `ndims` of 0.
#[must_use]
pub fn new_unsized<G: CodeGenerator + ?Sized>(
@@ -126,7 +145,7 @@ impl<'ctx> NDArrayType<'ctx> {
let llvm_usize = generator.get_size_type(ctx);
let llvm_ndarray = Self::llvm_type(ctx, llvm_usize);
NDArrayType { ty: llvm_ndarray, dtype, ndims: Some(0), llvm_usize }
NDArrayType { ty: llvm_ndarray, dtype, ndims: 0, llvm_usize }
}
/// Creates an [`NDArrayType`] from a [unifier type][Type].
@@ -145,7 +164,7 @@ impl<'ctx> NDArrayType<'ctx> {
NDArrayType {
ty: Self::llvm_type(ctx.ctx, llvm_usize),
dtype: llvm_dtype,
ndims: Some(ndims),
ndims,
llvm_usize,
}
}
@@ -155,7 +174,7 @@ impl<'ctx> NDArrayType<'ctx> {
pub fn from_type(
ptr_ty: PointerType<'ctx>,
dtype: BasicTypeEnum<'ctx>,
ndims: Option<u64>,
ndims: u64,
llvm_usize: IntType<'ctx>,
) -> Self {
debug_assert!(Self::is_representable(ptr_ty, llvm_usize).is_ok());
@@ -177,20 +196,40 @@ impl<'ctx> NDArrayType<'ctx> {
/// Returns the number of dimensions of this `ndarray` type.
#[must_use]
pub fn ndims(&self) -> Option<u64> {
pub fn ndims(&self) -> u64 {
self.ndims
}
/// Allocates an instance of [`NDArrayValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.dtype,
self.ndims,
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`NDArrayValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
self.dtype,
self.ndims,
self.llvm_usize,
@@ -214,7 +253,7 @@ impl<'ctx> NDArrayType<'ctx> {
ndims: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let ndarray = self.alloca(generator, ctx, name);
let ndarray = self.alloca_var(generator, ctx, name);
let itemsize = ctx
.builder
@@ -247,35 +286,7 @@ impl<'ctx> NDArrayType<'ctx> {
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert!(self.ndims.is_some(), "NDArrayType::construct can only be called on an instance with compile-time known ndims (self.ndims = Some(ndims))");
let Some(ndims) = self.ndims.map(|ndims| self.llvm_usize.const_int(ndims, false)) else {
unreachable!()
};
self.construct_impl(generator, ctx, ndims, name)
}
/// Allocate an [`NDArrayValue`] on the stack given its `ndims` and `dtype`.
///
/// `shape` and `strides` will be automatically allocated onto the stack.
///
/// The returned ndarray's content will be:
/// - `data`: uninitialized.
/// - `itemsize`: set to the size of `dtype`.
/// - `ndims`: set to the value of `ndims`.
/// - `shape`: allocated with an array of length `ndims` with uninitialized values.
/// - `strides`: allocated with an array of length `ndims` with uninitialized values.
#[deprecated = "Prefer construct_uninitialized or construct_*_shape."]
#[must_use]
pub fn construct_dyn_ndims<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ndims: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert!(self.ndims.is_none(), "NDArrayType::construct_dyn_ndims can only be called on an instance with compile-time unknown ndims (self.ndims = None)");
let ndims = self.llvm_usize.const_int(self.ndims, false);
self.construct_impl(generator, ctx, ndims, name)
}
@@ -291,9 +302,9 @@ impl<'ctx> NDArrayType<'ctx> {
shape: &[u64],
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert!(self.ndims.is_none_or(|ndims| shape.len() as u64 == ndims));
assert_eq!(shape.len() as u64, self.ndims);
let ndarray = Self::new(generator, ctx.ctx, self.dtype, Some(shape.len() as u64))
let ndarray = Self::new(generator, ctx.ctx, self.dtype, shape.len() as u64)
.construct_uninitialized(generator, ctx, name);
let llvm_usize = generator.get_size_type(ctx.ctx);
@@ -326,9 +337,9 @@ impl<'ctx> NDArrayType<'ctx> {
shape: &[IntValue<'ctx>],
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
assert!(self.ndims.is_none_or(|ndims| shape.len() as u64 == ndims));
assert_eq!(shape.len() as u64, self.ndims);
let ndarray = Self::new(generator, ctx.ctx, self.dtype, Some(shape.len() as u64))
let ndarray = Self::new(generator, ctx.ctx, self.dtype, shape.len() as u64)
.construct_uninitialized(generator, ctx, name);
let llvm_usize = generator.get_size_type(ctx.ctx);
@@ -368,7 +379,7 @@ impl<'ctx> NDArrayType<'ctx> {
let value = value.as_basic_value_enum();
assert_eq!(value.get_type(), self.dtype);
assert!(self.ndims.is_none_or(|ndims| ndims == 0));
assert_eq!(self.ndims, 0);
// We have to put the value on the stack to get a data pointer.
let data = ctx.builder.build_alloca(value.get_type(), "construct_unsized").unwrap();
@@ -425,36 +436,8 @@ impl<'ctx> ProxyType<'ctx> for NDArrayType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

82
nac3core/src/codegen/types/ndarray/nditer.rs
View File

@@ -14,7 +14,7 @@ use crate::codegen::{
types::structure::{check_struct_type_matches_fields, StructField, StructFields},
values::{
ndarray::{NDArrayValue, NDIterValue},
ArraySliceValue, ProxyValue,
ArrayLikeValue, ArraySliceValue, ProxyValue, TypedArrayLikeAdapter,
},
CodeGenContext, CodeGenerator,
};
@@ -109,8 +109,31 @@ impl<'ctx> NDIterType<'ctx> {
self.llvm_usize
}
/// Allocates an instance of [`NDIterValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
parent: NDArrayValue<'ctx>,
indices: ArraySliceValue<'ctx>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
parent,
indices,
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`NDIterValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
@@ -119,7 +142,7 @@ impl<'ctx> NDIterType<'ctx> {
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
parent,
indices,
self.llvm_usize,
@@ -128,6 +151,11 @@ impl<'ctx> NDIterType<'ctx> {
}
/// Allocate an [`NDIter`] that iterates through the given `ndarray`.
///
/// Note: This function allocates an array on the stack at the current builder location, which
/// may lead to stack explosion if called in a hot loop. Therefore, callers are recommended to
/// call `llvm.stacksave` before calling this function and call `llvm.stackrestore` after the
/// [`NDIter`] is no longer needed.
#[must_use]
pub fn construct<G: CodeGenerator + ?Sized>(
&self,
@@ -135,22 +163,18 @@ impl<'ctx> NDIterType<'ctx> {
ctx: &mut CodeGenContext<'ctx, '_>,
ndarray: NDArrayValue<'ctx>,
) -> <Self as ProxyType<'ctx>>::Value {
let nditer = self.raw_alloca(generator, ctx, None);
let ndims = ndarray.load_ndims(ctx);
let nditer = self.raw_alloca_var(generator, ctx, None);
let ndims = self.llvm_usize.const_int(ndarray.get_type().ndims(), false);
// The caller has the responsibility to allocate 'indices' for `NDIter`.
let indices =
generator.gen_array_var_alloc(ctx, self.llvm_usize.into(), ndims, None).unwrap();
let indices =
TypedArrayLikeAdapter::from(indices, |_, _, v| v.into_int_value(), |_, _, v| v.into());
let nditer = <Self as ProxyType<'ctx>>::Value::from_pointer_value(
nditer,
ndarray,
indices,
self.llvm_usize,
None,
);
let nditer = self.map_value(nditer, ndarray, indices.as_slice_value(ctx, generator), None);
irrt::ndarray::call_nac3_nditer_initialize(generator, ctx, nditer, ndarray, indices);
irrt::ndarray::call_nac3_nditer_initialize(generator, ctx, nditer, ndarray, &indices);
nditer
}
@@ -197,36 +221,8 @@ impl<'ctx> ProxyType<'ctx> for NDIterType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

51
nac3core/src/codegen/types/range.rs
View File

@@ -1,13 +1,12 @@
use inkwell::{
context::Context,
types::{AnyTypeEnum, BasicType, BasicTypeEnum, IntType, PointerType},
values::IntValue,
AddressSpace,
};
use super::ProxyType;
use crate::codegen::{
values::{ArraySliceValue, ProxyValue, RangeValue},
values::{ProxyValue, RangeValue},
{CodeGenContext, CodeGenerator},
};
@@ -78,15 +77,29 @@ impl<'ctx> RangeType<'ctx> {
}
/// Allocates an instance of [`RangeValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(self.raw_alloca(ctx, name), name)
}
/// Allocates an instance of [`RangeValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
name,
)
}
@@ -126,36 +139,8 @@ impl<'ctx> ProxyType<'ctx> for RangeType<'ctx> {
Self::is_representable(llvm_ty)
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

15
nac3core/src/codegen/types/structure.rs
View File

@@ -5,6 +5,7 @@ use inkwell::{
types::{BasicTypeEnum, IntType, StructType},
values::{BasicValue, BasicValueEnum, IntValue, PointerValue, StructValue},
};
use itertools::Itertools;
use crate::codegen::CodeGenContext;
@@ -55,6 +56,20 @@ pub trait StructFields<'ctx>: Eq + Copy {
{
self.into_vec().into_iter()
}
/// Returns the field index of a field in this structure.
fn index_of_field<V>(&self, name: impl FnOnce(&Self) -> StructField<'ctx, V>) -> u32
where
V: BasicValue<'ctx> + TryFrom<BasicValueEnum<'ctx>, Error = ()>,
{
let field_name = name(self).name;
self.index_of_field_name(field_name).unwrap()
}
/// Returns the field index of a field with the given name in this structure.
fn index_of_field_name(&self, field_name: &str) -> Option<u32> {
self.iter().find_position(|(name, _)| *name == field_name).map(|(idx, _)| idx as u32)
}
}
/// A single field of an LLVM structure.

184
nac3core/src/codegen/types/tuple.rs Normal file
View File

@@ -0,0 +1,184 @@
use inkwell::{
context::Context,
types::{BasicType, BasicTypeEnum, IntType, StructType},
values::BasicValueEnum,
};
use itertools::Itertools;
use super::ProxyType;
use crate::{
codegen::{
values::{ProxyValue, TupleValue},
CodeGenContext, CodeGenerator,
},
typecheck::typedef::{Type, TypeEnum},
};
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct TupleType<'ctx> {
ty: StructType<'ctx>,
llvm_usize: IntType<'ctx>,
}
impl<'ctx> TupleType<'ctx> {
/// Checks whether `llvm_ty` represents any tuple type, returning [Err] if it does not.
pub fn is_representable(_value: StructType<'ctx>) -> Result<(), String> {
Ok(())
}
/// Creates an LLVM type corresponding to the expected structure of a tuple.
#[must_use]
fn llvm_type(ctx: &'ctx Context, tys: &[BasicTypeEnum<'ctx>]) -> StructType<'ctx> {
ctx.struct_type(tys, false)
}
/// Creates an instance of [`TupleType`].
#[must_use]
pub fn new<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
tys: &[BasicTypeEnum<'ctx>],
) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_tuple = Self::llvm_type(ctx, tys);
Self { ty: llvm_tuple, llvm_usize }
}
/// Creates an [`TupleType`] from a [unifier type][Type].
#[must_use]
pub fn from_unifier_type<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &mut CodeGenContext<'ctx, '_>,
ty: Type,
) -> Self {
let llvm_usize = generator.get_size_type(ctx.ctx);
// Sanity check on object type.
let TypeEnum::TTuple { ty: tys, .. } = &*ctx.unifier.get_ty_immutable(ty) else {
panic!("Expected type to be a TypeEnum::TTuple, got {}", ctx.unifier.stringify(ty));
};
let llvm_tys = tys.iter().map(|ty| ctx.get_llvm_type(generator, *ty)).collect_vec();
Self { ty: Self::llvm_type(ctx.ctx, &llvm_tys), llvm_usize }
}
/// Creates an [`TupleType`] from a [`StructType`].
#[must_use]
pub fn from_type(struct_ty: StructType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
debug_assert!(Self::is_representable(struct_ty).is_ok());
TupleType { ty: struct_ty, llvm_usize }
}
/// Returns the number of elements present in this [`TupleType`].
#[must_use]
pub fn num_elements(&self) -> u32 {
self.ty.count_fields()
}
/// Returns the type of the tuple element at the given `index`, or [`None`] if `index` is out of
/// range.
#[must_use]
pub fn type_at_index(&self, index: u32) -> Option<BasicTypeEnum<'ctx>> {
if index < self.num_elements() {
Some(unsafe { self.type_at_index_unchecked(index) })
} else {
None
}
}
/// Returns the type of the tuple element at the given `index`.
///
/// # Safety
///
/// The caller must ensure that the index is valid.
#[must_use]
pub unsafe fn type_at_index_unchecked(&self, index: u32) -> BasicTypeEnum<'ctx> {
self.ty.get_field_type_at_index_unchecked(index)
}
/// Constructs a [`TupleValue`] from this type by zero-initializing the tuple value.
#[must_use]
pub fn construct(
&self,
ctx: &CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
self.map_value(Self::llvm_type(ctx.ctx, &self.ty.get_field_types()).const_zero(), name)
}
/// Constructs a [`TupleValue`] from `objects`. The resulting tuple preserves the order of
/// objects.
#[must_use]
pub fn construct_from_objects<I: IntoIterator<Item = BasicValueEnum<'ctx>>>(
&self,
ctx: &CodeGenContext<'ctx, '_>,
objects: I,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
let values = objects.into_iter().collect_vec();
assert_eq!(values.len(), self.num_elements() as usize);
assert!(values
.iter()
.enumerate()
.all(|(i, v)| { v.get_type() == unsafe { self.type_at_index_unchecked(i as u32) } }));
let mut value = self.construct(ctx, name);
for (i, val) in values.into_iter().enumerate() {
value.store_element(ctx, i as u32, val);
}
value
}
/// Converts an existing value into a [`ListValue`].
#[must_use]
pub fn map_value(
&self,
value: <<Self as ProxyType<'ctx>>::Value as ProxyValue<'ctx>>::Base,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_struct_value(value, self.llvm_usize, name)
}
}
impl<'ctx> ProxyType<'ctx> for TupleType<'ctx> {
type Base = StructType<'ctx>;
type Value = TupleValue<'ctx>;
fn is_type<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
llvm_ty: impl BasicType<'ctx>,
) -> Result<(), String> {
if let BasicTypeEnum::StructType(ty) = llvm_ty.as_basic_type_enum() {
<Self as ProxyType<'ctx>>::is_representable(generator, ctx, ty)
} else {
Err(format!("Expected struct type, got {llvm_ty:?}"))
}
}
fn is_representable<G: CodeGenerator + ?Sized>(
_generator: &G,
_ctx: &'ctx Context,
llvm_ty: Self::Base,
) -> Result<(), String> {
Self::is_representable(llvm_ty)
}
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type()
}
fn as_base_type(&self) -> Self::Base {
self.ty
}
}
impl<'ctx> From<TupleType<'ctx>> for StructType<'ctx> {
fn from(value: TupleType<'ctx>) -> Self {
value.as_base_type()
}
}

59
nac3core/src/codegen/types/utils/slice.rs
View File

@@ -15,7 +15,7 @@ use crate::codegen::{
},
ProxyType,
},
values::{utils::SliceValue, ArraySliceValue, ProxyValue},
values::{utils::SliceValue, ProxyValue},
CodeGenContext, CodeGenerator,
};
@@ -154,16 +154,35 @@ impl<'ctx> SliceType<'ctx> {
self.int_ty
}
/// Allocates an instance of [`ContiguousNDArrayValue`] as if by calling `alloca` on the base type.
/// Allocates an instance of [`SliceValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca`].
#[must_use]
pub fn alloca<G: CodeGenerator + ?Sized>(
pub fn alloca(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(ctx, name),
self.int_ty,
self.llvm_usize,
name,
)
}
/// Allocates an instance of [`SliceValue`] as if by calling `alloca` on the base type.
///
/// See [`ProxyType::raw_alloca_var`].
#[must_use]
pub fn alloca_var<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self as ProxyType<'ctx>>::Value {
<Self as ProxyType<'ctx>>::Value::from_pointer_value(
self.raw_alloca(generator, ctx, name),
self.raw_alloca_var(generator, ctx, name),
self.int_ty,
self.llvm_usize,
name,
@@ -210,36 +229,8 @@ impl<'ctx> ProxyType<'ctx> for SliceType<'ctx> {
Self::is_representable(llvm_ty, generator.get_size_type(ctx))
}
fn raw_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
name: Option<&'ctx str>,
) -> <Self::Value as ProxyValue<'ctx>>::Base {
generator
.gen_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
name,
)
.unwrap()
}
fn array_alloca<G: CodeGenerator + ?Sized>(
&self,
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
size: IntValue<'ctx>,
name: Option<&'ctx str>,
) -> ArraySliceValue<'ctx> {
generator
.gen_array_var_alloc(
ctx,
self.as_base_type().get_element_type().into_struct_type().into(),
size,
name,
)
.unwrap()
fn alloca_type(&self) -> impl BasicType<'ctx> {
self.as_base_type().get_element_type().into_struct_type()
}
fn as_base_type(&self) -> Self::Base {

139
nac3core/src/codegen/values/array.rs
View File

@@ -51,8 +51,8 @@ pub trait ArrayLikeIndexer<'ctx, Index = IntValue<'ctx>>: ArrayLikeValue<'ctx> {
/// This function should be called with a valid index.
unsafe fn ptr_offset_unchecked<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
ctx: &CodeGenContext<'ctx, '_>,
generator: &G,
idx: &Index,
name: Option<&str>,
) -> PointerValue<'ctx>;
@@ -76,8 +76,8 @@ pub trait UntypedArrayLikeAccessor<'ctx, Index = IntValue<'ctx>>:
/// This function should be called with a valid index.
unsafe fn get_unchecked<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
ctx: &CodeGenContext<'ctx, '_>,
generator: &G,
idx: &Index,
name: Option<&str>,
) -> BasicValueEnum<'ctx> {
@@ -107,8 +107,8 @@ pub trait UntypedArrayLikeMutator<'ctx, Index = IntValue<'ctx>>:
/// This function should be called with a valid index.
unsafe fn set_unchecked<G: CodeGenerator + ?Sized>(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
generator: &mut G,
ctx: &CodeGenContext<'ctx, '_>,
generator: &G,
idx: &Index,
value: BasicValueEnum<'ctx>,
) {
@@ -130,32 +130,33 @@ pub trait UntypedArrayLikeMutator<'ctx, Index = IntValue<'ctx>>:
}
/// An array-like value that can have its array elements accessed as an arbitrary type `T`.
pub trait TypedArrayLikeAccessor<'ctx, T, Index = IntValue<'ctx>>:
pub trait TypedArrayLikeAccessor<'ctx, G: CodeGenerator + ?Sized, T, Index = IntValue<'ctx>>:
UntypedArrayLikeAccessor<'ctx, Index>
{
/// Casts an element from [`BasicValueEnum`] into `T`.
fn downcast_to_type(
&self,
ctx: &mut CodeGenContext<'ctx, '_>,
ctx: &CodeGenContext<'ctx, '_>,
generator: &G,
value: BasicValueEnum<'ctx>,
) -> T;
/// # Safety
///
/// This function should be called with a valid index.