Compare commits

..

9 Commits

Author SHA1 Message Date
abdul124 43991b150b nac3artiq: allow class attribute access without init function 2024-06-21 10:02:41 +08:00
abdul124 3b6a6f560f core: add support for class attributes 2024-06-21 09:42:02 +08:00
abdul124 e91d24fd10 core: add attribute field to class definition 2024-06-21 09:42:02 +08:00
lyken d89146aa02 core: use no_run on builtin_fns docs 2024-06-20 13:53:25 +08:00
David Mak 5bade81ddb standalone: Add test for multidim array index with one index 2024-06-20 12:50:30 +08:00
David Mak 0452e6de78 core: Fix codegen for tuple-index into ndarray 2024-06-20 12:50:30 +08:00
David Mak 635c944c90 core: Fix type inference for tuple-index into ndarray
Fixes #420.
2024-06-20 12:50:30 +08:00
lyken e36af3b0a3 core: reduce code duplication in codegen/builtin_fns (#422)
Used macros to generate some unary math functions.

Reviewed-on: #422
Reviewed-by: David Mak <chmakac@connect.ust.hk>
Co-authored-by: lyken <lyken@m-labs.hk>
Co-committed-by: lyken <lyken@m-labs.hk>
2024-06-20 12:48:44 +08:00
Sebastien Bourdeauducq 5b1aa812ed update dependencies 2024-06-20 10:43:55 +08:00
7 changed files with 624 additions and 1442 deletions

8
Cargo.lock generated
View File

@ -574,9 +574,9 @@ checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c"
[[package]] [[package]]
name = "memchr" name = "memchr"
version = "2.7.2" version = "2.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c8640c5d730cb13ebd907d8d04b52f55ac9a2eec55b440c8892f40d56c76c1d" checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
[[package]] [[package]]
name = "memoffset" name = "memoffset"
@ -927,9 +927,9 @@ dependencies = [
[[package]] [[package]]
name = "redox_syscall" name = "redox_syscall"
version = "0.5.1" version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "469052894dcb553421e483e4209ee581a45100d31b4018de03e5a7ad86374a7e" checksum = "c82cf8cff14456045f55ec4241383baeff27af886adb72ffb2162f99911de0fd"
dependencies = [ dependencies = [
"bitflags", "bitflags",
] ]

View File

@ -2,11 +2,11 @@
"nodes": { "nodes": {
"nixpkgs": { "nixpkgs": {
"locked": { "locked": {
"lastModified": 1717196966, "lastModified": 1718530797,
"narHash": "sha256-yZKhxVIKd2lsbOqYd5iDoUIwsRZFqE87smE2Vzf6Ck0=", "narHash": "sha256-pup6cYwtgvzDpvpSCFh1TEUjw2zkNpk8iolbKnyFmmU=",
"owner": "NixOS", "owner": "NixOS",
"repo": "nixpkgs", "repo": "nixpkgs",
"rev": "57610d2f8f0937f39dbd72251e9614b1561942d8", "rev": "b60ebf54c15553b393d144357375ea956f89e9a9",
"type": "github" "type": "github"
}, },
"original": { "original": {

View File

@ -3,41 +3,37 @@ from numpy import int32
@nac3 @nac3
class XXX: class Demo:
ATTR1: KernelInvariant[int32] = 2 attr1: KernelInvariant[int32] = 2
ATTR2:int32 = 4 attr2: int32 = 4
ATTR3: Kernel[int32] attr3: Kernel[int32]
@kernel @kernel
def __init__(self): def __init__(self):
self.ATTR3 = 8 self.attr3 = 8
@nac3 @nac3
class NAC3Devices: class NAC3Devices:
core: KernelInvariant[Core] core: KernelInvariant[Core]
attr4: KernelInvariant[int32] = 16
ATTR4: KernelInvariant[int32] = 16
def __init__(self): def __init__(self):
self.core = Core() self.core = Core()
@kernel @kernel
def run(self): def run(self):
XXX.ATTR1 # Supported Demo.attr1 # Supported
# Demo.attr2 # Field not accessible on Kernel
# XXX.ATTR2 # Field not accessible on Kernel # Demo.attr3 # Only attributes can be accessed in this way
# Demo.attr1 = 2 # Attributes are immutable
# XXX.ATTR3 # Only attributes can be accessed in this way
# XXX.ATTR1 = 2 # Attributes are immutable
self.ATTR4 # Attributes can be accessed within class self.ATTR4 # Attributes can be accessed within class
obj = XXX() obj = Demo()
obj.ATTR1 # Attributes can be accessed by class objects obj.attr1 # Attributes can be accessed by class objects
NAC3Devices.ATTR4 # Attributes accessible for classes without __init__ NAC3Devices.attr4 # Attributes accessible for classes without __init__
if __name__ == "__main__": if __name__ == "__main__":

File diff suppressed because it is too large Load Diff

View File

@ -3,8 +3,8 @@ use std::{collections::HashMap, convert::TryInto, iter::once, iter::zip};
use crate::{ use crate::{
codegen::{ codegen::{
classes::{ classes::{
ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue, ListValue, NDArrayValue, ProxyValue, ArrayLikeIndexer, ArrayLikeValue, ListValue, NDArrayValue, ProxyValue, RangeValue,
RangeValue, TypedArrayLikeAccessor, UntypedArrayLikeAccessor, TypedArrayLikeAccessor, UntypedArrayLikeAccessor,
}, },
concrete_type::{ConcreteFuncArg, ConcreteTypeEnum, ConcreteTypeStore}, concrete_type::{ConcreteFuncArg, ConcreteTypeEnum, ConcreteTypeStore},
gen_in_range_check, get_llvm_abi_type, get_llvm_type, gen_in_range_check, get_llvm_abi_type, get_llvm_type,
@ -1757,22 +1757,37 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
let ndims = values let ndims = values
.iter() .iter()
.map(|ndim| match *ndim { .map(|ndim| u64::try_from(ndim.clone()).map_err(|()| ndim.clone()))
SymbolValue::U64(v) => Ok(v), .collect::<Result<Vec<_>, _>>()
SymbolValue::U32(v) => Ok(u64::from(v)), .map_err(|val| {
SymbolValue::I32(v) => u64::try_from(v) format!(
.map_err(|_| format!("Expected non-negative literal for ndarray.ndims, got {v}")), "Expected non-negative literal for ndarray.ndims, got {}",
SymbolValue::I64(v) => u64::try_from(v) i128::try_from(val).unwrap()
.map_err(|_| format!("Expected non-negative literal for ndarray.ndims, got {v}")), )
_ => unreachable!(), })?;
})
.collect::<Result<Vec<_>, _>>()?;
assert!(!ndims.is_empty()); assert!(!ndims.is_empty());
let ndarray_ndims_ty = ctx // The number of dimensions subscripted by the index expression.
.unifier // Slicing a ndarray will yield the same number of dimensions, whereas indexing into a
.get_fresh_literal(ndims.iter().map(|v| SymbolValue::U64(v - 1)).collect(), None); // 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 ndarray_ndims_ty = ctx.unifier.get_fresh_literal(
ndims.iter().map(|v| SymbolValue::U64(v - subscripted_dims)).collect(),
None,
);
let ndarray_ty = let ndarray_ty =
make_ndarray_ty(&mut ctx.unifier, &ctx.primitives, Some(ty), Some(ndarray_ndims_ty)); make_ndarray_ty(&mut ctx.unifier, &ctx.primitives, Some(ty), Some(ndarray_ndims_ty));
let llvm_pndarray_t = ctx.get_llvm_type(generator, ndarray_ty).into_pointer_type(); let llvm_pndarray_t = ctx.get_llvm_type(generator, ndarray_ty).into_pointer_type();
@ -1875,123 +1890,165 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
} }
}; };
Ok(Some(match &slice.node { let make_indices_arr = |generator: &mut G,
ExprKind::Tuple { elts, .. } => { ctx: &mut CodeGenContext<'ctx, '_>|
let slices = elts -> Result<_, String> {
.iter() Ok(if let ExprKind::Tuple { elts, .. } = &slice.node {
.enumerate() let llvm_int_ty = ctx.get_llvm_type(generator, elts[0].custom.unwrap());
.map(|(dim, elt)| expr_to_slice(generator, ctx, &elt.node, dim as u64)) let index_addr = generator.gen_array_var_alloc(
.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()
}
_ => {
let index = if let Some(index) = generator.gen_expr(ctx, slice)? {
index.to_basic_value_enum(ctx, generator, slice.custom.unwrap())?.into_int_value()
} else {
return Ok(None);
};
let Some(index) = normalize_index(generator, ctx, index, 0)? else { return Ok(None) };
let index_addr = generator.gen_var_alloc(ctx, index.get_type().into(), None)?;
ctx.builder.build_store(index_addr, index).unwrap();
if ndims.len() == 1 && ndims[0] == 1 {
// Accessing an element from a 1-dimensional `ndarray`
return Ok(Some(
v.data()
.get(
ctx,
generator,
&ArraySliceValue::from_ptr_val(
index_addr,
llvm_usize.const_int(1, false),
None,
),
None,
)
.into(),
));
}
// Accessing an element from a multi-dimensional `ndarray`
// Create a new array, remove the top dimension from the dimension-size-list, and copy the
// elements over
let subscripted_ndarray = generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), None)?;
let ndarray = NDArrayValue::from_ptr_val(subscripted_ndarray, llvm_usize, None);
let num_dims = v.load_ndims(ctx);
ndarray.store_ndims(
ctx, ctx,
generator, llvm_int_ty,
ctx.builder.build_int_sub(num_dims, llvm_usize.const_int(1, false), "").unwrap(), llvm_usize.const_int(elts.len() as u64, false),
); None,
)?;
let ndarray_num_dims = ndarray.load_ndims(ctx); for (i, elt) in elts.iter().enumerate() {
ndarray.create_dim_sizes(ctx, llvm_usize, ndarray_num_dims); let Some(index) = generator.gen_expr(ctx, elt)? else {
return Ok(None);
};
let ndarray_num_dims = ndarray.load_ndims(ctx); let index = index
let v_dims_src_ptr = unsafe { .to_basic_value_enum(ctx, generator, elt.custom.unwrap())?
v.dim_sizes().ptr_offset_unchecked( .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) };
// Create a new array, remove the top dimension from the dimension-size-list, and copy the
// elements over
let subscripted_ndarray =
generator.gen_var_alloc(ctx, llvm_ndarray_t.into(), None)?;
let ndarray = NDArrayValue::from_ptr_val(subscripted_ndarray, llvm_usize, None);
let num_dims = v.load_ndims(ctx);
ndarray.store_ndims(
ctx, ctx,
generator, generator,
&llvm_usize.const_int(1, false), ctx.builder
None, .build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
) .unwrap(),
}; );
call_memcpy_generic(
ctx,
ndarray.dim_sizes().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 = call_ndarray_calc_size( let ndarray_num_dims = ndarray.load_ndims(ctx);
generator, ndarray.create_dim_sizes(ctx, llvm_usize, ndarray_num_dims);
ctx,
&ndarray.dim_sizes().as_slice_value(ctx, generator),
(None, None),
);
ndarray.create_data(ctx, llvm_ndarray_data_t, ndarray_num_elems);
let v_data_src_ptr = v.data().ptr_offset( let ndarray_num_dims = ndarray.load_ndims(ctx);
ctx, let v_dims_src_ptr = unsafe {
generator, v.dim_sizes().ptr_offset_unchecked(
&ArraySliceValue::from_ptr_val(index_addr, llvm_usize.const_int(1, false), None), ctx,
None, generator,
); &llvm_usize.const_int(1, false),
call_memcpy_generic( None,
ctx, )
ndarray.data().base_ptr(ctx, generator), };
v_data_src_ptr, call_memcpy_generic(
ctx.builder ctx,
.build_int_mul(ndarray_num_elems, llvm_ndarray_data_t.size_of().unwrap(), "") ndarray.dim_sizes().base_ptr(ctx, generator),
.map(Into::into) v_dims_src_ptr,
.unwrap(), ctx.builder
llvm_i1.const_zero(), .build_int_mul(ndarray_num_dims, llvm_usize.size_of(), "")
); .map(Into::into)
.unwrap(),
llvm_i1.const_zero(),
);
ndarray.as_base_value().into() let ndarray_num_elems = call_ndarray_calc_size(
generator,
ctx,
&ndarray.dim_sizes().as_slice_value(ctx, generator),
(None, None),
);
ndarray.create_data(ctx, llvm_ndarray_data_t, ndarray_num_elems);
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()
}
} }
})) }))
} }

View File

@ -1646,6 +1646,7 @@ impl<'a> Inferencer<'a> {
fn infer_subscript_ndarray( fn infer_subscript_ndarray(
&mut self, &mut self,
value: &ast::Expr<Option<Type>>, value: &ast::Expr<Option<Type>>,
slice: &ast::Expr<Option<Type>>,
dummy_tvar: Type, dummy_tvar: Type,
ndims: Type, ndims: Type,
) -> InferenceResult { ) -> InferenceResult {
@ -1664,48 +1665,66 @@ impl<'a> Inferencer<'a> {
let ndims = values let ndims = values
.iter() .iter()
.map(|ndim| match *ndim { .map(|ndim| u64::try_from(ndim.clone()).map_err(|()| ndim.clone()))
SymbolValue::U64(v) => Ok(v), .collect::<Result<Vec<_>, _>>()
SymbolValue::U32(v) => Ok(u64::from(v)), .map_err(|val| {
SymbolValue::I32(v) => u64::try_from(v).map_err(|_| { HashSet::from([format!(
HashSet::from([format!( "Expected non-negative literal for ndarray.ndims, got {}",
"Expected non-negative literal for ndarray.ndims, got {v}" i128::try_from(val).unwrap()
)]) )])
}), })?;
SymbolValue::I64(v) => u64::try_from(v).map_err(|_| {
HashSet::from([format!(
"Expected non-negative literal for ndarray.ndims, got {v}"
)])
}),
_ => unreachable!(),
})
.collect::<Result<Vec<_>, _>>()?;
assert!(!ndims.is_empty()); assert!(!ndims.is_empty());
if ndims.len() == 1 && ndims[0] == 1 { // The number of dimensions subscripted by the index expression.
// ndarray[T, Literal[1]] - Index always returns an object of type T // 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,
};
if ndims.len() == 1 && ndims[0] - subscripted_dims == 0 {
// ndarray[T, Literal[1]] - Non-Slice index always returns an object of type T
assert_ne!(ndims[0], 0); assert_ne!(ndims[0], 0);
Ok(dummy_tvar) Ok(dummy_tvar)
} else { } else {
// ndarray[T, Literal[N]] where N != 1 - Index returns an object of type ndarray[T, Literal[N - 1]] // Otherwise - Index returns an object of type ndarray[T, Literal[N - subscripted_dims]]
if ndims.iter().any(|v| *v == 0) { // Disallow subscripting if any Literal value will subscript on an element
let new_ndims = ndims
.into_iter()
.map(|v| {
let v = i128::from(v) - i128::from(subscripted_dims);
u64::try_from(v)
})
.collect::<Result<Vec<_>, _>>()
.map_err(|_| {
HashSet::from([format!(
"Cannot subscript {} by {subscripted_dims} dimensions",
self.unifier.stringify(value.custom.unwrap()),
)])
})?;
if new_ndims.iter().any(|v| *v == 0) {
unimplemented!("Inference for ndarray subscript operator with Literal[0, ...] bound unimplemented") unimplemented!("Inference for ndarray subscript operator with Literal[0, ...] bound unimplemented")
} }
let ndims_min_one_ty = self.unifier.get_fresh_literal( let ndims_ty = self
ndims.into_iter().map(|v| SymbolValue::U64(v - 1)).collect(), .unifier
None, .get_fresh_literal(new_ndims.into_iter().map(SymbolValue::U64).collect(), None);
); let subscripted_ty =
let subscripted_ty = make_ndarray_ty( make_ndarray_ty(self.unifier, self.primitives, Some(dummy_tvar), Some(ndims_ty));
self.unifier,
self.primitives,
Some(dummy_tvar),
Some(ndims_min_one_ty),
);
Ok(subscripted_ty) Ok(subscripted_ty)
} }
@ -1742,7 +1761,7 @@ impl<'a> Inferencer<'a> {
TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { TypeEnum::TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
let (_, ndims) = let (_, ndims) =
unpack_ndarray_var_tys(self.unifier, value.custom.unwrap()); unpack_ndarray_var_tys(self.unifier, value.custom.unwrap());
self.infer_subscript_ndarray(value, ty, ndims) self.infer_subscript_ndarray(value, slice, ty, ndims)
} }
_ => { _ => {
// the index is a constant, so value can be a sequence. // the index is a constant, so value can be a sequence.
@ -1785,10 +1804,7 @@ impl<'a> Inferencer<'a> {
} }
let (_, ndims) = unpack_ndarray_var_tys(self.unifier, value.custom.unwrap()); let (_, ndims) = unpack_ndarray_var_tys(self.unifier, value.custom.unwrap());
let ndarray_ty = self.infer_subscript_ndarray(value, slice, ty, ndims)
make_ndarray_ty(self.unifier, self.primitives, Some(ty), Some(ndims));
self.constrain(value.custom.unwrap(), ndarray_ty, &value.location)?;
Ok(ndarray_ty)
} }
_ => { _ => {
if let TypeEnum::TTuple { .. } = &*self.unifier.get_ty(value.custom.unwrap()) { if let TypeEnum::TTuple { .. } = &*self.unifier.get_ty(value.custom.unwrap()) {
@ -1823,7 +1839,7 @@ impl<'a> Inferencer<'a> {
.get_fresh_var_with_range(valid_index_tys.as_slice(), None, None) .get_fresh_var_with_range(valid_index_tys.as_slice(), None, None)
.ty; .ty;
self.constrain(slice.custom.unwrap(), valid_index_ty, &slice.location)?; self.constrain(slice.custom.unwrap(), valid_index_ty, &slice.location)?;
self.infer_subscript_ndarray(value, ty, ndims) self.infer_subscript_ndarray(value, slice, ty, ndims)
} }
_ => unreachable!(), _ => unreachable!(),
} }

View File

@ -150,6 +150,15 @@ def test_ndarray_slices():
x2 = x[0::2, 0::2] x2 = x[0::2, 0::2]
output_ndarray_float_2(x2) output_ndarray_float_2(x2)
def test_ndarray_nd_idx():
x = np_identity(2)
x0: float = x[0, 0]
output_float64(x0)
output_float64(x[0, 1])
output_float64(x[1, 0])
output_float64(x[1, 1])
def test_ndarray_add(): def test_ndarray_add():
x = np_identity(2) x = np_identity(2)
y = x + np_ones([2, 2]) y = x + np_ones([2, 2])
@ -1393,6 +1402,7 @@ def run() -> int32:
test_ndarray_neg_idx() test_ndarray_neg_idx()
test_ndarray_slices() test_ndarray_slices()
test_ndarray_nd_idx()
test_ndarray_add() test_ndarray_add()
test_ndarray_add_broadcast() test_ndarray_add_broadcast()