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

[core] codegen: Add ProxyType::new overloads and refactor to use them

Browse Source
This commit is contained in:
David Mak 2025-01-14 18:20:09 +08:00
parent bd66fe48d8
commit 8e614d83de
25 changed files with 320 additions and 228 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
nac3artiq/src/codegen.rs
View File

@ -476,8 +476,8 @@ fn format_rpc_arg<'ctx>(
let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, arg_ty); let (elem_ty, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, arg_ty);
let ndims = extract_ndims(&ctx.unifier, ndims); let ndims = extract_ndims(&ctx.unifier, ndims);
let dtype = ctx.get_llvm_type(generator, elem_ty); let dtype = ctx.get_llvm_type(generator, elem_ty);
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype, ndims) let ndarray =
.map_value(arg.into_pointer_value(), None); NDArrayType::new(ctx, dtype, ndims).map_value(arg.into_pointer_value(), None);
let ndims = llvm_usize.const_int(ndims, false); let ndims = llvm_usize.const_int(ndims, false);
@ -609,7 +609,7 @@ fn format_rpc_ret<'ctx>(
let (dtype, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty); let (dtype, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ret_ty);
let dtype_llvm = ctx.get_llvm_type(generator, dtype); let dtype_llvm = ctx.get_llvm_type(generator, dtype);
let ndims = extract_ndims(&ctx.unifier, ndims); let ndims = extract_ndims(&ctx.unifier, ndims);
let ndarray = NDArrayType::new(generator, ctx.ctx, dtype_llvm, ndims) let ndarray = NDArrayType::new(ctx, dtype_llvm, ndims)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
// NOTE: Current content of `ndarray`: // NOTE: Current content of `ndarray`:

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

@ -752,12 +752,8 @@ pub fn call_numpy_minimum<'ctx, G: CodeGenerator + ?Sized>(
debug_assert!(ctx.unifier.unioned(x1_dtype, x2_dtype)); debug_assert!(ctx.unifier.unioned(x1_dtype, x2_dtype));
let llvm_common_dtype = x1.get_type().element_type(); let llvm_common_dtype = x1.get_type().element_type();
let result = NDArrayType::new_broadcast( let result =
generator, NDArrayType::new_broadcast(ctx, llvm_common_dtype, &[x1.get_type(), x2.get_type()])
ctx.ctx,
llvm_common_dtype,
&[x1.get_type(), x2.get_type()],
)
.broadcast_starmap( .broadcast_starmap(
generator, generator,
ctx, ctx,
@ -1015,12 +1011,8 @@ pub fn call_numpy_maximum<'ctx, G: CodeGenerator + ?Sized>(
debug_assert!(ctx.unifier.unioned(x1_dtype, x2_dtype)); debug_assert!(ctx.unifier.unioned(x1_dtype, x2_dtype));
let llvm_common_dtype = x1.get_type().element_type(); let llvm_common_dtype = x1.get_type().element_type();
let result = NDArrayType::new_broadcast( let result =
generator, NDArrayType::new_broadcast(ctx, llvm_common_dtype, &[x1.get_type(), x2.get_type()])
ctx.ctx,
llvm_common_dtype,
&[x1.get_type(), x2.get_type()],
)
.broadcast_starmap( .broadcast_starmap(
generator, generator,
ctx, ctx,
@ -1652,7 +1644,7 @@ pub fn call_np_linalg_cholesky<'ctx, G: CodeGenerator + ?Sized>(
unsupported_type(ctx, FN_NAME, &[x1_ty]); unsupported_type(ctx, FN_NAME, &[x1_ty]);
} }
let out = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2) let out = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
out.copy_shape_from_ndarray(generator, ctx, x1); out.copy_shape_from_ndarray(generator, ctx, x1);
unsafe { out.create_data(generator, ctx) }; unsafe { out.create_data(generator, ctx) };
@ -1694,7 +1686,7 @@ pub fn call_np_linalg_qr<'ctx, G: CodeGenerator + ?Sized>(
}; };
let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None); let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None);
let out_ndarray_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2); let out_ndarray_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2);
let q = out_ndarray_ty.construct_dyn_shape(generator, ctx, &[d0, dk], None); let q = out_ndarray_ty.construct_dyn_shape(generator, ctx, &[d0, dk], None);
unsafe { q.create_data(generator, ctx) }; unsafe { q.create_data(generator, ctx) };
@ -1715,8 +1707,11 @@ pub fn call_np_linalg_qr<'ctx, G: CodeGenerator + ?Sized>(
let q = q.as_base_value().as_basic_value_enum(); let q = q.as_base_value().as_basic_value_enum();
let r = r.as_base_value().as_basic_value_enum(); let r = r.as_base_value().as_basic_value_enum();
let tuple = TupleType::new(generator, ctx.ctx, &[q.get_type(), r.get_type()]) let tuple = TupleType::new(ctx, &[q.get_type(), r.get_type()]).construct_from_objects(
.construct_from_objects(ctx, [q, r], None); ctx,
[q, r],
None,
);
Ok(tuple.as_base_value().into()) Ok(tuple.as_base_value().into())
} }
@ -1746,8 +1741,8 @@ pub fn call_np_linalg_svd<'ctx, G: CodeGenerator + ?Sized>(
}; };
let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None); let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None);
let out_ndarray1_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 1); let out_ndarray1_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 1);
let out_ndarray2_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2); let out_ndarray2_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2);
let u = out_ndarray2_ty.construct_dyn_shape(generator, ctx, &[d0, d0], None); let u = out_ndarray2_ty.construct_dyn_shape(generator, ctx, &[d0, d0], None);
unsafe { u.create_data(generator, ctx) }; unsafe { u.create_data(generator, ctx) };
@ -1775,7 +1770,7 @@ pub fn call_np_linalg_svd<'ctx, G: CodeGenerator + ?Sized>(
let u = u.as_base_value().as_basic_value_enum(); let u = u.as_base_value().as_basic_value_enum();
let s = s.as_base_value().as_basic_value_enum(); let s = s.as_base_value().as_basic_value_enum();
let vh = vh.as_base_value().as_basic_value_enum(); let vh = vh.as_base_value().as_basic_value_enum();
let tuple = TupleType::new(generator, ctx.ctx, &[u.get_type(), s.get_type(), vh.get_type()]) let tuple = TupleType::new(ctx, &[u.get_type(), s.get_type(), vh.get_type()])
.construct_from_objects(ctx, [u, s, vh], None); .construct_from_objects(ctx, [u, s, vh], None);
Ok(tuple.as_base_value().into()) Ok(tuple.as_base_value().into())
} }
@ -1796,7 +1791,7 @@ pub fn call_np_linalg_inv<'ctx, G: CodeGenerator + ?Sized>(
unsupported_type(ctx, FN_NAME, &[x1_ty]); unsupported_type(ctx, FN_NAME, &[x1_ty]);
} }
let out = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2) let out = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
out.copy_shape_from_ndarray(generator, ctx, x1); out.copy_shape_from_ndarray(generator, ctx, x1);
unsafe { out.create_data(generator, ctx) }; unsafe { out.create_data(generator, ctx) };
@ -1838,8 +1833,12 @@ pub fn call_np_linalg_pinv<'ctx, G: CodeGenerator + ?Sized>(
x1_shape.get_typed_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None) x1_shape.get_typed_unchecked(ctx, generator, &llvm_usize.const_int(1, false), None)
}; };
let out = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2) let out = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2).construct_dyn_shape(
.construct_dyn_shape(generator, ctx, &[d0, d1], None); generator,
ctx,
&[d0, d1],
None,
);
unsafe { out.create_data(generator, ctx) }; unsafe { out.create_data(generator, ctx) };
let x1_c = x1.make_contiguous_ndarray(generator, ctx); let x1_c = x1.make_contiguous_ndarray(generator, ctx);
@ -1880,7 +1879,7 @@ pub fn call_sp_linalg_lu<'ctx, G: CodeGenerator + ?Sized>(
}; };
let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None); let dk = llvm_intrinsics::call_int_smin(ctx, d0, d1, None);
let out_ndarray_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2); let out_ndarray_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2);
let l = out_ndarray_ty.construct_dyn_shape(generator, ctx, &[d0, dk], None); let l = out_ndarray_ty.construct_dyn_shape(generator, ctx, &[d0, dk], None);
unsafe { l.create_data(generator, ctx) }; unsafe { l.create_data(generator, ctx) };
@ -1901,8 +1900,11 @@ pub fn call_sp_linalg_lu<'ctx, G: CodeGenerator + ?Sized>(
let l = l.as_base_value().as_basic_value_enum(); let l = l.as_base_value().as_basic_value_enum();
let u = u.as_base_value().as_basic_value_enum(); let u = u.as_base_value().as_basic_value_enum();
let tuple = TupleType::new(generator, ctx.ctx, &[l.get_type(), u.get_type()]) let tuple = TupleType::new(ctx, &[l.get_type(), u.get_type()]).construct_from_objects(
.construct_from_objects(ctx, [l, u], None); ctx,
[l, u],
None,
);
Ok(tuple.as_base_value().into()) Ok(tuple.as_base_value().into())
} }
@ -1936,11 +1938,11 @@ pub fn call_np_linalg_matrix_power<'ctx, G: CodeGenerator + ?Sized>(
unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty]) unsupported_type(ctx, FN_NAME, &[x1_ty, x2_ty])
}; };
let x2 = NDArrayType::new_unsized(generator, ctx.ctx, ctx.ctx.f64_type().into()) let x2 = NDArrayType::new_unsized(ctx, ctx.ctx.f64_type().into())
.construct_unsized(generator, ctx, &x2, None); // x2.shape == [] .construct_unsized(generator, ctx, &x2, None); // x2.shape == []
let x2 = x2.atleast_nd(generator, ctx, 1); // x2.shape == [1] let x2 = x2.atleast_nd(generator, ctx, 1); // x2.shape == [1]
let out = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2) let out = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
out.copy_shape_from_ndarray(generator, ctx, x1); out.copy_shape_from_ndarray(generator, ctx, x1);
unsafe { out.create_data(generator, ctx) }; unsafe { out.create_data(generator, ctx) };
@ -1979,8 +1981,12 @@ pub fn call_np_linalg_det<'ctx, G: CodeGenerator + ?Sized>(
} }
// The output is a float64, but we are using an ndarray (shape == [1]) for uniformity in function call. // The output is a float64, but we are using an ndarray (shape == [1]) for uniformity in function call.
let det = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 1) let det = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 1).construct_const_shape(
.construct_const_shape(generator, ctx, &[1], None); generator,
ctx,
&[1],
None,
);
unsafe { det.create_data(generator, ctx) }; unsafe { det.create_data(generator, ctx) };
let x1_c = x1.make_contiguous_ndarray(generator, ctx); let x1_c = x1.make_contiguous_ndarray(generator, ctx);
@ -2014,7 +2020,7 @@ pub fn call_sp_linalg_schur<'ctx, G: CodeGenerator + ?Sized>(
unsupported_type(ctx, FN_NAME, &[x1_ty]); unsupported_type(ctx, FN_NAME, &[x1_ty]);
} }
let out_ndarray_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2); let out_ndarray_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2);
let t = out_ndarray_ty.construct_uninitialized(generator, ctx, None); let t = out_ndarray_ty.construct_uninitialized(generator, ctx, None);
t.copy_shape_from_ndarray(generator, ctx, x1); t.copy_shape_from_ndarray(generator, ctx, x1);
@ -2037,8 +2043,11 @@ pub fn call_sp_linalg_schur<'ctx, G: CodeGenerator + ?Sized>(
let t = t.as_base_value().as_basic_value_enum(); let t = t.as_base_value().as_basic_value_enum();
let z = z.as_base_value().as_basic_value_enum(); let z = z.as_base_value().as_basic_value_enum();
let tuple = TupleType::new(generator, ctx.ctx, &[t.get_type(), z.get_type()]) let tuple = TupleType::new(ctx, &[t.get_type(), z.get_type()]).construct_from_objects(
.construct_from_objects(ctx, [t, z], None); ctx,
[t, z],
None,
);
Ok(tuple.as_base_value().into()) Ok(tuple.as_base_value().into())
} }
@ -2059,7 +2068,7 @@ pub fn call_sp_linalg_hessenberg<'ctx, G: CodeGenerator + ?Sized>(
unsupported_type(ctx, FN_NAME, &[x1_ty]); unsupported_type(ctx, FN_NAME, &[x1_ty]);
} }
let out_ndarray_ty = NDArrayType::new(generator, ctx.ctx, ctx.ctx.f64_type().into(), 2); let out_ndarray_ty = NDArrayType::new(ctx, ctx.ctx.f64_type().into(), 2);
let h = out_ndarray_ty.construct_uninitialized(generator, ctx, None); let h = out_ndarray_ty.construct_uninitialized(generator, ctx, None);
h.copy_shape_from_ndarray(generator, ctx, x1); h.copy_shape_from_ndarray(generator, ctx, x1);
@ -2082,7 +2091,10 @@ pub fn call_sp_linalg_hessenberg<'ctx, G: CodeGenerator + ?Sized>(
let h = h.as_base_value().as_basic_value_enum(); let h = h.as_base_value().as_basic_value_enum();
let q = q.as_base_value().as_basic_value_enum(); let q = q.as_base_value().as_basic_value_enum();
let tuple = TupleType::new(generator, ctx.ctx, &[h.get_type(), q.get_type()]) let tuple = TupleType::new(ctx, &[h.get_type(), q.get_type()]).construct_from_objects(
.construct_from_objects(ctx, [h, q], None); ctx,
[h, q],
None,
);
Ok(tuple.as_base_value().into()) Ok(tuple.as_base_value().into())
} }

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

@ -1167,7 +1167,7 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
"listcomp.alloc_size", "listcomp.alloc_size",
) )
.unwrap(); .unwrap();
list = ListType::new(generator, ctx.ctx, elem_ty).construct( list = ListType::new(ctx, &elem_ty).construct(
generator, generator,
ctx, ctx,
list_alloc_size.into_int_value(), list_alloc_size.into_int_value(),
@ -1218,12 +1218,7 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
Some("length"), Some("length"),
) )
.into_int_value(); .into_int_value();
list = ListType::new(generator, ctx.ctx, elem_ty).construct( list = ListType::new(ctx, &elem_ty).construct(generator, ctx, length, Some("listcomp"));
generator,
ctx,
length,
Some("listcomp"),
);
let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?; let counter = generator.gen_var_alloc(ctx, size_t.into(), Some("counter.addr"))?;
// counter = -1 // counter = -1
@ -1386,8 +1381,8 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
.build_int_add(lhs.load_size(ctx, None), rhs.load_size(ctx, None), "") .build_int_add(lhs.load_size(ctx, None), rhs.load_size(ctx, None), "")
.unwrap(); .unwrap();
let new_list = ListType::new(generator, ctx.ctx, llvm_elem_ty) let new_list =
.construct(generator, ctx, size, None); ListType::new(ctx, &llvm_elem_ty).construct(generator, ctx, size, None);
let lhs_size = ctx let lhs_size = ctx
.builder .builder
@ -1474,7 +1469,7 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
let elem_llvm_ty = ctx.get_llvm_type(generator, elem_ty); let elem_llvm_ty = ctx.get_llvm_type(generator, elem_ty);
let sizeof_elem = elem_llvm_ty.size_of().unwrap(); let sizeof_elem = elem_llvm_ty.size_of().unwrap();
let new_list = ListType::new(generator, ctx.ctx, elem_llvm_ty).construct( let new_list = ListType::new(ctx, &elem_llvm_ty).construct(
generator, generator,
ctx, ctx,
ctx.builder.build_int_mul(list_val.load_size(ctx, None), int_val, "").unwrap(), ctx.builder.build_int_mul(list_val.load_size(ctx, None), int_val, "").unwrap(),
@ -1576,8 +1571,7 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
let right = right.to_ndarray(generator, ctx); let right = right.to_ndarray(generator, ctx);
let result = NDArrayType::new_broadcast( let result = NDArrayType::new_broadcast(
generator, ctx,
ctx.ctx,
llvm_common_dtype, llvm_common_dtype,
&[left.get_type(), right.get_type()], &[left.get_type(), right.get_type()],
) )
@ -1850,8 +1844,7 @@ pub fn gen_cmpop_expr_with_values<'ctx, G: CodeGenerator>(
.to_ndarray(generator, ctx); .to_ndarray(generator, ctx);
let result_ndarray = NDArrayType::new_broadcast( let result_ndarray = NDArrayType::new_broadcast(
generator, ctx,
ctx.ctx,
ctx.ctx.i8_type().into(), ctx.ctx.i8_type().into(),
&[left.get_type(), right.get_type()], &[left.get_type(), right.get_type()],
) )
@ -2480,18 +2473,9 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
}; };
let length = ctx.get_size_type().const_int(elements.len() as u64, false); let length = ctx.get_size_type().const_int(elements.len() as u64, false);
let arr_str_ptr = if let Some(ty) = ty { let arr_str_ptr = if let Some(ty) = ty {
ListType::new(generator, ctx.ctx, ty).construct( ListType::new(ctx, &ty).construct(generator, ctx, length, Some("list"))
generator,
ctx,
length,
Some("list"),
)
} else { } else {
ListType::new_untyped(generator, ctx.ctx).construct_empty( ListType::new_untyped(ctx).construct_empty(generator, ctx, Some("list"))
generator,
ctx,
Some("list"),
)
}; };
let arr_ptr = arr_str_ptr.data(); let arr_ptr = arr_str_ptr.data();
for (i, v) in elements.iter().enumerate() { for (i, v) in elements.iter().enumerate() {
@ -2970,12 +2954,8 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
.unwrap(), .unwrap(),
step, step,
); );
let res_array_ret = ListType::new(generator, ctx.ctx, ty).construct( let res_array_ret =
generator, ListType::new(ctx, &ty).construct(generator, ctx, length, Some("ret"));
ctx,
length,
Some("ret"),
);
let Some(res_ind) = handle_slice_indices( let Some(res_ind) = handle_slice_indices(
&None, &None,
&None, &None,

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

@ -530,7 +530,7 @@ fn get_llvm_type<'ctx, G: CodeGenerator + ?Sized>(
*params.iter().next().unwrap().1, *params.iter().next().unwrap().1,
); );
ListType::new(generator, ctx, element_type).as_base_type().into() ListType::new_with_generator(generator, ctx, element_type).as_base_type().into()
} }
TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => { TObj { obj_id, .. } if *obj_id == PrimDef::NDArray.id() => {
@ -540,7 +540,7 @@ fn get_llvm_type<'ctx, G: CodeGenerator + ?Sized>(
ctx, module, generator, unifier, top_level, type_cache, dtype, ctx, module, generator, unifier, top_level, type_cache, dtype,
); );
NDArrayType::new(generator, ctx, element_type, ndims).as_base_type().into() NDArrayType::new_with_generator(generator, ctx, element_type, ndims).as_base_type().into()
} }
_ => unreachable!( _ => unreachable!(
@ -594,7 +594,7 @@ fn get_llvm_type<'ctx, G: CodeGenerator + ?Sized>(
get_llvm_type(ctx, module, generator, unifier, top_level, type_cache, *ty) get_llvm_type(ctx, module, generator, unifier, top_level, type_cache, *ty)
}) })
.collect_vec(); .collect_vec();
TupleType::new(generator, ctx, &fields).as_base_type().into() TupleType::new_with_generator(generator, ctx, &fields).as_base_type().into()
} }
TVirtual { .. } => unimplemented!(), TVirtual { .. } => unimplemented!(),
_ => unreachable!("{}", ty_enum.get_type_name()), _ => unreachable!("{}", ty_enum.get_type_name()),

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

@ -42,7 +42,7 @@ pub fn gen_ndarray_empty<'ctx>(
let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg)); let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg));
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, ndims) let ndarray = NDArrayType::new(context, llvm_dtype, ndims)
.construct_numpy_empty(generator, context, &shape, None); .construct_numpy_empty(generator, context, &shape, None);
Ok(ndarray.as_base_value()) Ok(ndarray.as_base_value())
} }
@ -67,7 +67,7 @@ pub fn gen_ndarray_zeros<'ctx>(
let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg)); let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg));
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, ndims) let ndarray = NDArrayType::new(context, llvm_dtype, ndims)
.construct_numpy_zeros(generator, context, dtype, &shape, None); .construct_numpy_zeros(generator, context, dtype, &shape, None);
Ok(ndarray.as_base_value()) Ok(ndarray.as_base_value())
} }
@ -92,7 +92,7 @@ pub fn gen_ndarray_ones<'ctx>(
let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg)); let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg));
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, ndims) let ndarray = NDArrayType::new(context, llvm_dtype, ndims)
.construct_numpy_ones(generator, context, dtype, &shape, None); .construct_numpy_ones(generator, context, dtype, &shape, None);
Ok(ndarray.as_base_value()) Ok(ndarray.as_base_value())
} }
@ -120,7 +120,7 @@ pub fn gen_ndarray_full<'ctx>(
let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg)); let shape = parse_numpy_int_sequence(generator, context, (shape_ty, shape_arg));
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, ndims).construct_numpy_full( let ndarray = NDArrayType::new(context, llvm_dtype, ndims).construct_numpy_full(
generator, generator,
context, context,
&shape, &shape,
@ -223,7 +223,7 @@ pub fn gen_ndarray_eye<'ctx>(
.build_int_s_extend_or_bit_cast(offset_arg.into_int_value(), llvm_usize, "") .build_int_s_extend_or_bit_cast(offset_arg.into_int_value(), llvm_usize, "")
.unwrap(); .unwrap();
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, 2) let ndarray = NDArrayType::new(context, llvm_dtype, 2)
.construct_numpy_eye(generator, context, dtype, nrows, ncols, offset, None); .construct_numpy_eye(generator, context, dtype, nrows, ncols, offset, None);
Ok(ndarray.as_base_value()) Ok(ndarray.as_base_value())
} }
@ -251,7 +251,7 @@ pub fn gen_ndarray_identity<'ctx>(
.builder .builder
.build_int_s_extend_or_bit_cast(n_arg.into_int_value(), llvm_usize, "") .build_int_s_extend_or_bit_cast(n_arg.into_int_value(), llvm_usize, "")
.unwrap(); .unwrap();
let ndarray = NDArrayType::new(generator, context.ctx, llvm_dtype, 2) let ndarray = NDArrayType::new(context, llvm_dtype, 2)
.construct_numpy_identity(generator, context, dtype, n, None); .construct_numpy_identity(generator, context, dtype, n, None);
Ok(ndarray.as_base_value()) Ok(ndarray.as_base_value())
} }
@ -349,8 +349,8 @@ pub fn ndarray_dot<'ctx, G: CodeGenerator + ?Sized>(
ctx, ctx,
Some("np_dot"), Some("np_dot"),
|generator, ctx| { |generator, ctx| {
let a_iter = NDIterType::new(generator, ctx.ctx).construct(generator, ctx, a); let a_iter = NDIterType::new(ctx).construct(generator, ctx, a);
let b_iter = NDIterType::new(generator, ctx.ctx).construct(generator, ctx, b); let b_iter = NDIterType::new(ctx).construct(generator, ctx, b);
Ok((a_iter, b_iter)) Ok((a_iter, b_iter))
}, },
|_, ctx, (a_iter, _b_iter)| { |_, ctx, (a_iter, _b_iter)| {

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

@ -448,8 +448,7 @@ pub fn gen_setitem<'ctx, G: CodeGenerator>(
let broadcast_ndims = let broadcast_ndims =
[target.get_type().ndims(), value.get_type().ndims()].into_iter().max().unwrap(); [target.get_type().ndims(), value.get_type().ndims()].into_iter().max().unwrap();
let broadcast_result = NDArrayType::new( let broadcast_result = NDArrayType::new(
generator, ctx,
ctx.ctx,
value.get_type().element_type(), value.get_type().element_type(),
broadcast_ndims, broadcast_ndims,
) )

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

@ -446,7 +446,7 @@ fn test_classes_list_type_new() {
let llvm_i32 = ctx.i32_type(); let llvm_i32 = ctx.i32_type();
let llvm_usize = generator.get_size_type(&ctx); let llvm_usize = generator.get_size_type(&ctx);
let llvm_list = ListType::new(&generator, &ctx, llvm_i32.into()); let llvm_list = ListType::new_with_generator(&generator, &ctx, llvm_i32.into());
assert!(ListType::is_representable(llvm_list.as_base_type(), llvm_usize).is_ok()); assert!(ListType::is_representable(llvm_list.as_base_type(), llvm_usize).is_ok());
} }
@ -466,6 +466,6 @@ fn test_classes_ndarray_type_new() {
let llvm_i32 = ctx.i32_type(); let llvm_i32 = ctx.i32_type();
let llvm_usize = generator.get_size_type(&ctx); let llvm_usize = generator.get_size_type(&ctx);
let llvm_ndarray = NDArrayType::new(&generator, &ctx, llvm_i32.into(), 2); let llvm_ndarray = NDArrayType::new_with_generator(&generator, &ctx, llvm_i32.into(), 2);
assert!(NDArrayType::is_representable(llvm_ndarray.as_base_type(), llvm_usize).is_ok()); assert!(NDArrayType::is_representable(llvm_ndarray.as_base_type(), llvm_usize).is_ok());
} }

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

@ -104,7 +104,7 @@ impl<'ctx> ListType<'ctx> {
element_type: Option<BasicTypeEnum<'ctx>>, element_type: Option<BasicTypeEnum<'ctx>>,
llvm_usize: IntType<'ctx>, llvm_usize: IntType<'ctx>,
) -> PointerType<'ctx> { ) -> PointerType<'ctx> {
let element_type = element_type.unwrap_or(llvm_usize.into()); let element_type = element_type.map_or(llvm_usize.into(), |ty| ty.as_basic_type_enum());
let field_tys = let field_tys =
Self::fields(element_type, llvm_usize).into_iter().map(|field| field.1).collect_vec(); Self::fields(element_type, llvm_usize).into_iter().map(|field| field.1).collect_vec();
@ -112,26 +112,45 @@ impl<'ctx> ListType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
fn new_impl(
ctx: &'ctx Context,
element_type: Option<BasicTypeEnum<'ctx>>,
llvm_usize: IntType<'ctx>,
) -> Self {
let llvm_list = Self::llvm_type(ctx, element_type, llvm_usize);
Self { ty: llvm_list, item: element_type, llvm_usize }
}
/// Creates an instance of [`ListType`]. /// Creates an instance of [`ListType`].
#[must_use] #[must_use]
pub fn new<G: CodeGenerator + ?Sized>( pub fn new(ctx: &CodeGenContext<'ctx, '_>, element_type: &impl BasicType<'ctx>) -> Self {
Self::new_impl(ctx.ctx, Some(element_type.as_basic_type_enum()), ctx.get_size_type())
}
/// Creates an instance of [`ListType`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G, generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
element_type: BasicTypeEnum<'ctx>, element_type: BasicTypeEnum<'ctx>,
) -> Self { ) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx, Some(element_type.as_basic_type_enum()), generator.get_size_type(ctx))
let llvm_list = Self::llvm_type(ctx, Some(element_type), llvm_usize);
Self { ty: llvm_list, item: Some(element_type), llvm_usize }
} }
/// Creates an instance of [`ListType`] with an unknown element type. /// Creates an instance of [`ListType`] with an unknown element type.
#[must_use] #[must_use]
pub fn new_untyped<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self { pub fn new_untyped(ctx: &CodeGenContext<'ctx, '_>) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx.ctx, None, ctx.get_size_type())
let llvm_list = Self::llvm_type(ctx, None, llvm_usize); }
Self { ty: llvm_list, item: None, llvm_usize } /// Creates an instance of [`ListType`] with an unknown element type.
#[must_use]
pub fn new_untyped_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
) -> Self {
Self::new_impl(ctx, None, generator.get_size_type(ctx))
} }
/// Creates an [`ListType`] from a [unifier type][Type]. /// Creates an [`ListType`] from a [unifier type][Type].
@ -159,11 +178,7 @@ impl<'ctx> ListType<'ctx> {
Some(ctx.get_llvm_type(generator, elem_type)) Some(ctx.get_llvm_type(generator, elem_type))
}; };
Self { Self::new_impl(ctx.ctx, llvm_elem_type, llvm_usize)
ty: Self::llvm_type(ctx.ctx, llvm_elem_type, llvm_usize),
item: llvm_elem_type,
llvm_usize,
}
} }
/// Creates an [`ListType`] from a [`PointerType`]. /// Creates an [`ListType`] from a [`PointerType`].

11
nac3core/src/codegen/types/ndarray/array.rs
View File

@ -44,7 +44,7 @@ impl<'ctx> NDArrayType<'ctx> {
assert!(self.ndims >= ndims_int); assert!(self.ndims >= ndims_int);
assert_eq!(dtype, self.dtype); assert_eq!(dtype, self.dtype);
let list_value = list.as_i8_list(generator, ctx); let list_value = list.as_i8_list(ctx);
// Validate `list` has a consistent shape. // Validate `list` has a consistent shape.
// Raise an exception if `list` is something abnormal like `[[1, 2], [3]]`. // Raise an exception if `list` is something abnormal like `[[1, 2], [3]]`.
@ -61,8 +61,8 @@ impl<'ctx> NDArrayType<'ctx> {
generator, ctx, list_value, ndims, &shape, generator, ctx, list_value, ndims, &shape,
); );
let ndarray = Self::new(generator, ctx.ctx, dtype, ndims_int) let ndarray =
.construct_uninitialized(generator, ctx, name); Self::new(ctx, dtype, ndims_int).construct_uninitialized(generator, ctx, name);
ndarray.copy_shape_from_array(generator, ctx, shape.base_ptr(ctx, generator)); ndarray.copy_shape_from_array(generator, ctx, shape.base_ptr(ctx, generator));
unsafe { ndarray.create_data(generator, ctx) }; unsafe { ndarray.create_data(generator, ctx) };
@ -96,8 +96,7 @@ impl<'ctx> NDArrayType<'ctx> {
let llvm_pi8 = ctx.ctx.i8_type().ptr_type(AddressSpace::default()); let llvm_pi8 = ctx.ctx.i8_type().ptr_type(AddressSpace::default());
let ndarray = Self::new(generator, ctx.ctx, dtype, 1) let ndarray = Self::new(ctx, dtype, 1).construct_uninitialized(generator, ctx, name);
.construct_uninitialized(generator, ctx, name);
// Set data // Set data
let data = ctx let data = ctx
@ -168,7 +167,7 @@ impl<'ctx> NDArrayType<'ctx> {
.map(BasicValueEnum::into_pointer_value) .map(BasicValueEnum::into_pointer_value)
.unwrap(); .unwrap();
NDArrayType::new(generator, ctx.ctx, dtype, ndims).map_value(ndarray, None) NDArrayType::new(ctx, dtype, ndims).map_value(ndarray, None)
} }
/// Implementation of `np_array(<ndarray>, copy=copy)`. /// Implementation of `np_array(<ndarray>, copy=copy)`.

20
nac3core/src/codegen/types/ndarray/broadcast.rs
View File

@ -79,15 +79,27 @@ impl<'ctx> ShapeEntryType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
/// Creates an instance of [`ShapeEntryType`]. fn new_impl(ctx: &'ctx Context, llvm_usize: IntType<'ctx>) -> Self {
#[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); let llvm_ty = Self::llvm_type(ctx, llvm_usize);
Self { ty: llvm_ty, llvm_usize } Self { ty: llvm_ty, llvm_usize }
} }
/// Creates an instance of [`ShapeEntryType`].
#[must_use]
pub fn new(ctx: &CodeGenContext<'ctx, '_>) -> Self {
Self::new_impl(ctx.ctx, ctx.get_size_type())
}
/// Creates an instance of [`ShapeEntryType`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
) -> Self {
Self::new_impl(ctx, generator.get_size_type(ctx))
}
/// Creates a [`ShapeEntryType`] from a [`PointerType`] representing an `ShapeEntry`. /// Creates a [`ShapeEntryType`] from a [`PointerType`] representing an `ShapeEntry`.
#[must_use] #[must_use]
pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self { pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {

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

@ -117,17 +117,26 @@ impl<'ctx> ContiguousNDArrayType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
fn new_impl(ctx: &'ctx Context, item: BasicTypeEnum<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
let llvm_cndarray = Self::llvm_type(ctx, item, llvm_usize);
Self { ty: llvm_cndarray, item, llvm_usize }
}
/// Creates an instance of [`ContiguousNDArrayType`]. /// Creates an instance of [`ContiguousNDArrayType`].
#[must_use] #[must_use]
pub fn new<G: CodeGenerator + ?Sized>( pub fn new(ctx: &CodeGenContext<'ctx, '_>, item: &impl BasicType<'ctx>) -> Self {
Self::new_impl(ctx.ctx, item.as_basic_type_enum(), ctx.get_size_type())
}
/// Creates an instance of [`ContiguousNDArrayType`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G, generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
item: BasicTypeEnum<'ctx>, item: BasicTypeEnum<'ctx>,
) -> Self { ) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx, item, generator.get_size_type(ctx))
let llvm_cndarray = Self::llvm_type(ctx, item, llvm_usize);
Self { ty: llvm_cndarray, item, llvm_usize }
} }
/// Creates an [`ContiguousNDArrayType`] from a [unifier type][Type]. /// Creates an [`ContiguousNDArrayType`] from a [unifier type][Type].
@ -140,9 +149,8 @@ impl<'ctx> ContiguousNDArrayType<'ctx> {
let (dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty); let (dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
let llvm_dtype = ctx.get_llvm_type(generator, dtype); let llvm_dtype = ctx.get_llvm_type(generator, dtype);
let llvm_usize = ctx.get_size_type();
Self { ty: Self::llvm_type(ctx.ctx, llvm_dtype, llvm_usize), item: llvm_dtype, llvm_usize } Self::new_impl(ctx.ctx, llvm_dtype, ctx.get_size_type())
} }
/// Creates an [`ContiguousNDArrayType`] from a [`PointerType`] representing an `NDArray`. /// Creates an [`ContiguousNDArrayType`] from a [`PointerType`] representing an `NDArray`.

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

@ -75,14 +75,25 @@ impl<'ctx> NDIndexType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
#[must_use] fn new_impl(ctx: &'ctx Context, llvm_usize: IntType<'ctx>) -> Self {
pub fn new<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_ndindex = Self::llvm_type(ctx, llvm_usize); let llvm_ndindex = Self::llvm_type(ctx, llvm_usize);
Self { ty: llvm_ndindex, llvm_usize } Self { ty: llvm_ndindex, llvm_usize }
} }
#[must_use]
pub fn new(ctx: &CodeGenContext<'ctx, '_>) -> Self {
Self::new_impl(ctx.ctx, ctx.get_size_type())
}
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
) -> Self {
Self::new_impl(ctx, generator.get_size_type(ctx))
}
#[must_use] #[must_use]
pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self { pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
debug_assert!(Self::is_representable(ptr_ty, llvm_usize).is_ok()); debug_assert!(Self::is_representable(ptr_ty, llvm_usize).is_ok());

12
nac3core/src/codegen/types/ndarray/map.rs
View File

@ -46,8 +46,7 @@ impl<'ctx> NDArrayType<'ctx> {
let out_ndarray = match out { let out_ndarray = match out {
NDArrayOut::NewNDArray { dtype } => { NDArrayOut::NewNDArray { dtype } => {
// Create a new ndarray based on the broadcast shape. // Create a new ndarray based on the broadcast shape.
let result_ndarray = let result_ndarray = NDArrayType::new(ctx, dtype, broadcast_result.ndims)
NDArrayType::new(generator, ctx.ctx, dtype, broadcast_result.ndims)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
result_ndarray.copy_shape_from_array( result_ndarray.copy_shape_from_array(
generator, generator,
@ -70,7 +69,7 @@ impl<'ctx> NDArrayType<'ctx> {
}; };
// Map element-wise and store results into `mapped_ndarray`. // Map element-wise and store results into `mapped_ndarray`.
let nditer = NDIterType::new(generator, ctx.ctx).construct(generator, ctx, out_ndarray); let nditer = NDIterType::new(ctx).construct(generator, ctx, out_ndarray);
gen_for_callback( gen_for_callback(
generator, generator,
ctx, ctx,
@ -80,9 +79,7 @@ impl<'ctx> NDArrayType<'ctx> {
let other_nditers = broadcast_result let other_nditers = broadcast_result
.ndarrays .ndarrays
.iter() .iter()
.map(|ndarray| { .map(|ndarray| NDIterType::new(ctx).construct(generator, ctx, *ndarray))
NDIterType::new(generator, ctx.ctx).construct(generator, ctx, *ndarray)
})
.collect_vec(); .collect_vec();
Ok((nditer, other_nditers)) Ok((nditer, other_nditers))
}, },
@ -169,8 +166,7 @@ impl<'ctx> ScalarOrNDArray<'ctx> {
// Promote all input to ndarrays and map through them. // Promote all input to ndarrays and map through them.
let inputs = inputs.iter().map(|input| input.to_ndarray(generator, ctx)).collect_vec(); let inputs = inputs.iter().map(|input| input.to_ndarray(generator, ctx)).collect_vec();
let ndarray = NDArrayType::new_broadcast( let ndarray = NDArrayType::new_broadcast(
generator, ctx,
ctx.ctx,
ret_dtype, ret_dtype,
&inputs.iter().map(NDArrayValue::get_type).collect_vec(), &inputs.iter().map(NDArrayValue::get_type).collect_vec(),
) )

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

@ -107,24 +107,56 @@ impl<'ctx> NDArrayType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
/// Creates an instance of [`NDArrayType`]. fn new_impl(
#[must_use]
pub fn new<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>, dtype: BasicTypeEnum<'ctx>,
ndims: u64, ndims: u64,
llvm_usize: IntType<'ctx>,
) -> Self { ) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_ndarray = Self::llvm_type(ctx, llvm_usize); let llvm_ndarray = Self::llvm_type(ctx, llvm_usize);
NDArrayType { ty: llvm_ndarray, dtype, ndims, llvm_usize } NDArrayType { ty: llvm_ndarray, dtype, ndims, llvm_usize }
} }
/// Creates an instance of [`NDArrayType`].
#[must_use]
pub fn new(ctx: &CodeGenContext<'ctx, '_>, dtype: BasicTypeEnum<'ctx>, ndims: u64) -> Self {
Self::new_impl(ctx.ctx, dtype, ndims, ctx.get_size_type())
}
/// Creates an instance of [`NDArrayType`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>,
ndims: u64,
) -> Self {
Self::new_impl(ctx, dtype, ndims, generator.get_size_type(ctx))
}
/// Creates an instance of [`NDArrayType`] as a result of a broadcast operation over one or more /// Creates an instance of [`NDArrayType`] as a result of a broadcast operation over one or more
/// `ndarray` operands. /// `ndarray` operands.
#[must_use] #[must_use]
pub fn new_broadcast<G: CodeGenerator + ?Sized>( pub fn new_broadcast(
ctx: &CodeGenContext<'ctx, '_>,
dtype: BasicTypeEnum<'ctx>,
inputs: &[NDArrayType<'ctx>],
) -> Self {
assert!(!inputs.is_empty());
Self::new_impl(
ctx.ctx,
dtype,
inputs.iter().map(NDArrayType::ndims).max().unwrap(),
ctx.get_size_type(),
)
}
/// Creates an instance of [`NDArrayType`] as a result of a broadcast operation over one or more
/// `ndarray` operands.
#[must_use]
pub fn new_broadcast_with_generator<G: CodeGenerator + ?Sized>(
generator: &G, generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>, dtype: BasicTypeEnum<'ctx>,
@ -132,20 +164,28 @@ impl<'ctx> NDArrayType<'ctx> {
) -> Self { ) -> Self {
assert!(!inputs.is_empty()); assert!(!inputs.is_empty());
Self::new(generator, ctx, dtype, inputs.iter().map(NDArrayType::ndims).max().unwrap()) Self::new_impl(
ctx,
dtype,
inputs.iter().map(NDArrayType::ndims).max().unwrap(),
generator.get_size_type(ctx),
)
} }
/// Creates an instance of [`NDArrayType`] with `ndims` of 0. /// Creates an instance of [`NDArrayType`] with `ndims` of 0.
#[must_use] #[must_use]
pub fn new_unsized<G: CodeGenerator + ?Sized>( pub fn new_unsized(ctx: &CodeGenContext<'ctx, '_>, dtype: BasicTypeEnum<'ctx>) -> Self {
Self::new_impl(ctx.ctx, dtype, 0, ctx.get_size_type())
}
/// Creates an instance of [`NDArrayType`] with `ndims` of 0.
#[must_use]
pub fn new_unsized_with_generator<G: CodeGenerator + ?Sized>(
generator: &G, generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
dtype: BasicTypeEnum<'ctx>, dtype: BasicTypeEnum<'ctx>,
) -> Self { ) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx, dtype, 0, generator.get_size_type(ctx))
let llvm_ndarray = Self::llvm_type(ctx, llvm_usize);
NDArrayType { ty: llvm_ndarray, dtype, ndims: 0, llvm_usize }
} }
/// Creates an [`NDArrayType`] from a [unifier type][Type]. /// Creates an [`NDArrayType`] from a [unifier type][Type].
@ -158,15 +198,9 @@ impl<'ctx> NDArrayType<'ctx> {
let (dtype, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ty); let (dtype, ndims) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
let llvm_dtype = ctx.get_llvm_type(generator, dtype); let llvm_dtype = ctx.get_llvm_type(generator, dtype);
let llvm_usize = ctx.get_size_type();
let ndims = extract_ndims(&ctx.unifier, ndims); let ndims = extract_ndims(&ctx.unifier, ndims);
NDArrayType { Self::new_impl(ctx.ctx, llvm_dtype, ndims, ctx.get_size_type())
ty: Self::llvm_type(ctx.ctx, llvm_usize),
dtype: llvm_dtype,
ndims,
llvm_usize,
}
} }
/// Creates an [`NDArrayType`] from a [`PointerType`] representing an `NDArray`. /// Creates an [`NDArrayType`] from a [`PointerType`] representing an `NDArray`.
@ -304,7 +338,7 @@ impl<'ctx> NDArrayType<'ctx> {
) -> <Self as ProxyType<'ctx>>::Value { ) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(shape.len() as u64, self.ndims); assert_eq!(shape.len() as u64, self.ndims);
let ndarray = Self::new(generator, ctx.ctx, self.dtype, shape.len() as u64) let ndarray = Self::new(ctx, self.dtype, shape.len() as u64)
.construct_uninitialized(generator, ctx, name); .construct_uninitialized(generator, ctx, name);
let llvm_usize = ctx.get_size_type(); let llvm_usize = ctx.get_size_type();
@ -339,7 +373,7 @@ impl<'ctx> NDArrayType<'ctx> {
) -> <Self as ProxyType<'ctx>>::Value { ) -> <Self as ProxyType<'ctx>>::Value {
assert_eq!(shape.len() as u64, self.ndims); assert_eq!(shape.len() as u64, self.ndims);
let ndarray = Self::new(generator, ctx.ctx, self.dtype, shape.len() as u64) let ndarray = Self::new(ctx, self.dtype, shape.len() as u64)
.construct_uninitialized(generator, ctx, name); .construct_uninitialized(generator, ctx, name);
let llvm_usize = ctx.get_size_type(); let llvm_usize = ctx.get_size_type();
@ -389,8 +423,8 @@ impl<'ctx> NDArrayType<'ctx> {
.build_pointer_cast(data, ctx.ctx.i8_type().ptr_type(AddressSpace::default()), "") .build_pointer_cast(data, ctx.ctx.i8_type().ptr_type(AddressSpace::default()), "")
.unwrap(); .unwrap();
let ndarray = Self::new_unsized(generator, ctx.ctx, value.get_type()) let ndarray =
.construct_uninitialized(generator, ctx, name); Self::new_unsized(ctx, value.get_type()).construct_uninitialized(generator, ctx, name);
ctx.builder.build_store(ndarray.ptr_to_data(ctx), data).unwrap(); ctx.builder.build_store(ndarray.ptr_to_data(ctx), data).unwrap();
ndarray ndarray
} }

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

@ -86,15 +86,27 @@ impl<'ctx> NDIterType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
/// Creates an instance of [`NDIter`]. fn new_impl(ctx: &'ctx Context, llvm_usize: IntType<'ctx>) -> Self {
#[must_use]
pub fn new<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self {
let llvm_usize = generator.get_size_type(ctx);
let llvm_nditer = Self::llvm_type(ctx, llvm_usize); let llvm_nditer = Self::llvm_type(ctx, llvm_usize);
Self { ty: llvm_nditer, llvm_usize } Self { ty: llvm_nditer, llvm_usize }
} }
/// Creates an instance of [`NDIter`].
#[must_use]
pub fn new(ctx: &CodeGenContext<'ctx, '_>) -> Self {
Self::new_impl(ctx.ctx, ctx.get_size_type())
}
/// Creates an instance of [`NDIter`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
) -> Self {
Self::new_impl(ctx, generator.get_size_type(ctx))
}
/// Creates an [`NDIterType`] from a [`PointerType`] representing an `NDIter`. /// Creates an [`NDIterType`] from a [`PointerType`] representing an `NDIter`.
#[must_use] #[must_use]
pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self { pub fn from_type(ptr_ty: PointerType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {

27
nac3core/src/codegen/types/tuple.rs
View File

@ -32,17 +32,34 @@ impl<'ctx> TupleType<'ctx> {
ctx.struct_type(tys, false) ctx.struct_type(tys, false)
} }
fn new_impl(
ctx: &'ctx Context,
tys: &[BasicTypeEnum<'ctx>],
llvm_usize: IntType<'ctx>,
) -> Self {
let llvm_tuple = Self::llvm_type(ctx, tys);
Self { ty: llvm_tuple, llvm_usize }
}
/// Creates an instance of [`TupleType`]. /// Creates an instance of [`TupleType`].
#[must_use] #[must_use]
pub fn new<G: CodeGenerator + ?Sized>( pub fn new(ctx: &CodeGenContext<'ctx, '_>, tys: &[impl BasicType<'ctx>]) -> Self {
Self::new_impl(
ctx.ctx,
&tys.iter().map(BasicType::as_basic_type_enum).collect_vec(),
ctx.get_size_type(),
)
}
/// Creates an instance of [`TupleType`].
#[must_use]
pub fn new_with_generator<G: CodeGenerator + ?Sized>(
generator: &G, generator: &G,
ctx: &'ctx Context, ctx: &'ctx Context,
tys: &[BasicTypeEnum<'ctx>], tys: &[BasicTypeEnum<'ctx>],
) -> Self { ) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx, tys, 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]. /// Creates an [`TupleType`] from a [unifier type][Type].

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

@ -122,19 +122,31 @@ impl<'ctx> SliceType<'ctx> {
ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default()) ctx.struct_type(&field_tys, false).ptr_type(AddressSpace::default())
} }
/// Creates an instance of [`SliceType`] with `int_ty` as its backing integer type. fn new_impl(ctx: &'ctx Context, int_ty: IntType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
#[must_use]
pub fn new(ctx: &'ctx Context, int_ty: IntType<'ctx>, llvm_usize: IntType<'ctx>) -> Self {
let llvm_ty = Self::llvm_type(ctx, int_ty); let llvm_ty = Self::llvm_type(ctx, int_ty);
Self { ty: llvm_ty, int_ty, llvm_usize } Self { ty: llvm_ty, int_ty, llvm_usize }
} }
/// Creates an instance of [`SliceType`] with `int_ty` as its backing integer type.
#[must_use]
pub fn new(ctx: &CodeGenContext<'ctx, '_>, int_ty: IntType<'ctx>) -> Self {
Self::new_impl(ctx.ctx, int_ty, ctx.get_size_type())
}
/// Creates an instance of [`SliceType`] with `usize` as its backing integer type. /// Creates an instance of [`SliceType`] with `usize` as its backing integer type.
#[must_use] #[must_use]
pub fn new_usize<G: CodeGenerator + ?Sized>(generator: &G, ctx: &'ctx Context) -> Self { pub fn new_usize(ctx: &CodeGenContext<'ctx, '_>) -> Self {
let llvm_usize = generator.get_size_type(ctx); Self::new_impl(ctx.ctx, ctx.get_size_type(), ctx.get_size_type())
Self::new(ctx, llvm_usize, llvm_usize) }
/// Creates an instance of [`SliceType`] with `usize` as its backing integer type.
#[must_use]
pub fn new_usize_with_generator<G: CodeGenerator + ?Sized>(
generator: &G,
ctx: &'ctx Context,
) -> Self {
Self::new_impl(ctx, generator.get_size_type(ctx), generator.get_size_type(ctx))
} }
/// Creates an [`SliceType`] from a [`PointerType`] representing a `slice`. /// Creates an [`SliceType`] from a [`PointerType`] representing a `slice`.

8
nac3core/src/codegen/values/list.rs
View File

@ -114,13 +114,9 @@ impl<'ctx> ListValue<'ctx> {
/// Returns an instance of [`ListValue`] with the `items` pointer cast to `i8*`. /// Returns an instance of [`ListValue`] with the `items` pointer cast to `i8*`.
#[must_use] #[must_use]
pub fn as_i8_list<G: CodeGenerator + ?Sized>( pub fn as_i8_list(&self, ctx: &CodeGenContext<'ctx, '_>) -> ListValue<'ctx> {
&self,
generator: &G,
ctx: &CodeGenContext<'ctx, '_>,
) -> ListValue<'ctx> {
let llvm_i8 = ctx.ctx.i8_type(); let llvm_i8 = ctx.ctx.i8_type();
let llvm_list_i8 = <Self as ProxyValue>::Type::new(generator, ctx.ctx, llvm_i8.into()); let llvm_list_i8 = <Self as ProxyValue>::Type::new(ctx, &llvm_i8);
Self::from_pointer_value( Self::from_pointer_value(
ctx.builder.build_pointer_cast(self.value, llvm_list_i8.as_base_type(), "").unwrap(), ctx.builder.build_pointer_cast(self.value, llvm_list_i8.as_base_type(), "").unwrap(),

4
nac3core/src/codegen/values/ndarray/broadcast.rs
View File

@ -104,7 +104,7 @@ impl<'ctx> NDArrayValue<'ctx> {
assert!(self.ndims <= target_ndims); assert!(self.ndims <= target_ndims);
assert_eq!(target_shape.element_type(ctx, generator), self.llvm_usize.into()); assert_eq!(target_shape.element_type(ctx, generator), self.llvm_usize.into());
let broadcast_ndarray = NDArrayType::new(generator, ctx.ctx, self.dtype, target_ndims) let broadcast_ndarray = NDArrayType::new(ctx, self.dtype, target_ndims)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
broadcast_ndarray.copy_shape_from_array( broadcast_ndarray.copy_shape_from_array(
generator, generator,
@ -147,7 +147,7 @@ fn broadcast_shapes<'ctx, G, Shape>(
+ TypedArrayLikeMutator<'ctx, G, IntValue<'ctx>>, + TypedArrayLikeMutator<'ctx, G, IntValue<'ctx>>,
{ {
let llvm_usize = ctx.get_size_type(); let llvm_usize = ctx.get_size_type();
let llvm_shape_ty = ShapeEntryType::new(generator, ctx.ctx); let llvm_shape_ty = ShapeEntryType::new(ctx);
assert!(in_shape_entries assert!(in_shape_entries
.iter() .iter()

11
nac3core/src/codegen/values/ndarray/contiguous.rs
View File

@ -117,8 +117,8 @@ impl<'ctx> NDArrayValue<'ctx> {
generator: &mut G, generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>, ctx: &mut CodeGenContext<'ctx, '_>,
) -> ContiguousNDArrayValue<'ctx> { ) -> ContiguousNDArrayValue<'ctx> {
let result = ContiguousNDArrayType::new(generator, ctx.ctx, self.dtype) let result =
.alloca_var(generator, ctx, self.name); ContiguousNDArrayType::new(ctx, &self.dtype).alloca_var(generator, ctx, self.name);
// Set ndims and shape. // Set ndims and shape.
let ndims = self.llvm_usize.const_int(self.ndims, false); let ndims = self.llvm_usize.const_int(self.ndims, false);
@ -178,8 +178,11 @@ impl<'ctx> NDArrayValue<'ctx> {
// TODO: Debug assert `ndims == carray.ndims` to catch bugs. // TODO: Debug assert `ndims == carray.ndims` to catch bugs.
// Allocate the resulting ndarray. // Allocate the resulting ndarray.
let ndarray = NDArrayType::new(generator, ctx.ctx, carray.item, ndims) let ndarray = NDArrayType::new(ctx, carray.item, ndims).construct_uninitialized(
.construct_uninitialized(generator, ctx, carray.name); generator,
ctx,
carray.name,
);
// Copy shape and update strides // Copy shape and update strides
let shape = carray.load_shape(ctx); let shape = carray.load_shape(ctx);

8
nac3core/src/codegen/values/ndarray/indexing.rs
View File

@ -128,11 +128,10 @@ impl<'ctx> NDArrayValue<'ctx> {
indices: &[RustNDIndex<'ctx>], indices: &[RustNDIndex<'ctx>],
) -> Self { ) -> Self {
let dst_ndims = self.deduce_ndims_after_indexing_with(indices); let dst_ndims = self.deduce_ndims_after_indexing_with(indices);
let dst_ndarray = NDArrayType::new(generator, ctx.ctx, self.dtype, dst_ndims) let dst_ndarray = NDArrayType::new(ctx, self.dtype, dst_ndims)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
let indices = let indices = NDIndexType::new(ctx).construct_ndindices(generator, ctx, indices);
NDIndexType::new(generator, ctx.ctx).construct_ndindices(generator, ctx, indices);
irrt::ndarray::call_nac3_ndarray_index(generator, ctx, indices, *self, dst_ndarray); irrt::ndarray::call_nac3_ndarray_index(generator, ctx, indices, *self, dst_ndarray);
dst_ndarray dst_ndarray
@ -245,8 +244,7 @@ impl<'ctx> RustNDIndex<'ctx> {
} }
RustNDIndex::Slice(in_rust_slice) => { RustNDIndex::Slice(in_rust_slice) => {
let user_slice_ptr = let user_slice_ptr =
SliceType::new(ctx.ctx, ctx.ctx.i32_type(), ctx.get_size_type()) SliceType::new(ctx, ctx.ctx.i32_type()).alloca_var(generator, ctx, None);
.alloca_var(generator, ctx, None);
in_rust_slice.write_to_slice(ctx, user_slice_ptr); in_rust_slice.write_to_slice(ctx, user_slice_ptr);
dst_ndindex.store_data( dst_ndindex.store_data(

2
nac3core/src/codegen/values/ndarray/matmul.rs
View File

@ -108,7 +108,7 @@ fn matmul_at_least_2d<'ctx, G: CodeGenerator>(
let lhs = in_a.broadcast_to(generator, ctx, ndims_int, &lhs_shape); let lhs = in_a.broadcast_to(generator, ctx, ndims_int, &lhs_shape);
let rhs = in_b.broadcast_to(generator, ctx, ndims_int, &rhs_shape); let rhs = in_b.broadcast_to(generator, ctx, ndims_int, &rhs_shape);
let dst = NDArrayType::new(generator, ctx.ctx, llvm_dst_dtype, ndims_int) let dst = NDArrayType::new(ctx, llvm_dst_dtype, ndims_int)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
dst.copy_shape_from_array(generator, ctx, dst_shape.base_ptr(ctx, generator)); dst.copy_shape_from_array(generator, ctx, dst_shape.base_ptr(ctx, generator));
unsafe { unsafe {

18
nac3core/src/codegen/values/ndarray/mod.rs
View File

@ -377,11 +377,7 @@ impl<'ctx> NDArrayValue<'ctx> {
.map(|obj| obj.as_basic_value_enum()) .map(|obj| obj.as_basic_value_enum())
.collect_vec(); .collect_vec();
TupleType::new( TupleType::new(ctx, &repeat_n(llvm_i32, self.ndims as usize).collect_vec())
generator,
ctx.ctx,
&repeat_n(llvm_i32.into(), self.ndims as usize).collect_vec(),
)
.construct_from_objects(ctx, objects, None) .construct_from_objects(ctx, objects, None)
} }
@ -411,11 +407,7 @@ impl<'ctx> NDArrayValue<'ctx> {
.map(|obj| obj.as_basic_value_enum()) .map(|obj| obj.as_basic_value_enum())
.collect_vec(); .collect_vec();
TupleType::new( TupleType::new(ctx, &repeat_n(llvm_i32, self.ndims as usize).collect_vec())
generator,
ctx.ctx,
&repeat_n(llvm_i32.into(), self.ndims as usize).collect_vec(),
)
.construct_from_objects(ctx, objects, None) .construct_from_objects(ctx, objects, None)
} }
@ -998,10 +990,8 @@ impl<'ctx> ScalarOrNDArray<'ctx> {
) -> NDArrayValue<'ctx> { ) -> NDArrayValue<'ctx> {
match self { match self {
ScalarOrNDArray::NDArray(ndarray) => *ndarray, ScalarOrNDArray::NDArray(ndarray) => *ndarray,
ScalarOrNDArray::Scalar(scalar) => { ScalarOrNDArray::Scalar(scalar) => NDArrayType::new_unsized(ctx, scalar.get_type())
NDArrayType::new_unsized(generator, ctx.ctx, scalar.get_type()) .construct_unsized(generator, ctx, scalar, None),
.construct_unsized(generator, ctx, scalar, None)
}
} }
} }

4
nac3core/src/codegen/values/ndarray/nditer.rs
View File

@ -160,9 +160,7 @@ impl<'ctx> NDArrayValue<'ctx> {
generator, generator,
ctx, ctx,
Some("ndarray_foreach"), Some("ndarray_foreach"),
|generator, ctx| { |generator, ctx| Ok(NDIterType::new(ctx).construct(generator, ctx, *self)),
Ok(NDIterType::new(generator, ctx.ctx).construct(generator, ctx, *self))
},
|_, ctx, nditer| Ok(nditer.has_element(ctx)), |_, ctx, nditer| Ok(nditer.has_element(ctx)),
|generator, ctx, hooks, nditer| body(generator, ctx, hooks, nditer), |generator, ctx, hooks, nditer| body(generator, ctx, hooks, nditer),
|_, ctx, nditer| { |_, ctx, nditer| {

2
nac3core/src/codegen/values/ndarray/view.rs
View File

@ -65,7 +65,7 @@ impl<'ctx> NDArrayValue<'ctx> {
// not contiguous but could be reshaped without copying data. Look into how numpy does // not contiguous but could be reshaped without copying data. Look into how numpy does
// it. // it.
let dst_ndarray = NDArrayType::new(generator, ctx.ctx, self.dtype, new_ndims) let dst_ndarray = NDArrayType::new(ctx, self.dtype, new_ndims)
.construct_uninitialized(generator, ctx, None); .construct_uninitialized(generator, ctx, None);
dst_ndarray.copy_shape_from_array(generator, ctx, new_shape.base_ptr(ctx, generator)); dst_ndarray.copy_shape_from_array(generator, ctx, new_shape.base_ptr(ctx, generator));