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

core: reap numpy

Browse Source
This commit is contained in:
lyken 2024-07-08 10:19:29 +08:00
parent eeb474f9e6
commit 6233f84ee9
6 changed files with 2908 additions and 2848 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

File diff suppressed because it is too large Load Diff

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

@ -1362,100 +1362,101 @@ 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 llvm_usize = generator.get_size_type(ctx.ctx);
todo!()
// let llvm_usize = generator.get_size_type(ctx.ctx);
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 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());
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);
// 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);
assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
// assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
let left_val =
NDArrayValue::from_ptr_val(left_val.into_pointer_value(), llvm_usize, None);
let right_val =
NDArrayValue::from_ptr_val(right_val.into_pointer_value(), llvm_usize, None);
// let left_val =
// NDArrayValue::from_ptr_val(left_val.into_pointer_value(), llvm_usize, None);
// let right_val =
// NDArrayValue::from_ptr_val(right_val.into_pointer_value(), llvm_usize, None);
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),
},
(left_val.as_base_value().into(), false),
(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 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),
// },
// (left_val.as_base_value().into(), false),
// (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,
// )
// },
// )?
// };
Ok(Some(res.as_base_value().into()))
} else {
let (ndarray_dtype, _) =
unpack_ndarray_var_tys(&mut ctx.unifier, if is_ndarray1 { ty1 } else { ty2 });
let ndarray_val = NDArrayValue::from_ptr_val(
if is_ndarray1 { left_val } else { right_val }.into_pointer_value(),
llvm_usize,
None,
);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ndarray_dtype,
match op.variant {
BinopVariant::Normal => None,
BinopVariant::AugAssign => Some(ndarray_val),
},
(left_val, !is_ndarray1),
(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)
},
)?;
// Ok(Some(res.as_base_value().into()))
// } else {
// let (ndarray_dtype, _) =
// unpack_ndarray_var_tys(&mut ctx.unifier, if is_ndarray1 { ty1 } else { ty2 });
// let ndarray_val = NDArrayValue::from_ptr_val(
// if is_ndarray1 { left_val } else { right_val }.into_pointer_value(),
// llvm_usize,
// None,
// );
// let res = numpy::ndarray_elementwise_binop_impl(
// generator,
// ctx,
// ndarray_dtype,
// match op.variant {
// BinopVariant::Normal => None,
// BinopVariant::AugAssign => Some(ndarray_val),
// },
// (left_val, !is_ndarray1),
// (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)
// },
// )?;
Ok(Some(res.as_base_value().into()))
}
// Ok(Some(res.as_base_value().into()))
// }
} else {
let left_ty_enum = ctx.unifier.get_ty_immutable(left_ty.unwrap());
let TypeEnum::TObj { fields, obj_id, .. } = left_ty_enum.as_ref() else {
@ -1612,40 +1613,41 @@ 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_usize = generator.get_size_type(ctx.ctx);
let (ndarray_dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
todo!()
// let llvm_usize = generator.get_size_type(ctx.ctx);
// let (ndarray_dtype, _) = unpack_ndarray_var_tys(&mut ctx.unifier, ty);
let val = NDArrayValue::from_ptr_val(val.into_pointer_value(), llvm_usize, None);
// let val = NDArrayValue::from_ptr_val(val.into_pointer_value(), llvm_usize, None);
// ndarray uses `~` rather than `not` to perform elementwise inversion, convert it before
// passing it to the elementwise codegen function
let op = if ndarray_dtype.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::Bool.id()) {
if op == ast::Unaryop::Invert {
ast::Unaryop::Not
} else {
unreachable!(
"ufunc {} not supported for ndarray[bool, N]",
op.op_info().method_name,
)
}
} else {
op
};
// // ndarray uses `~` rather than `not` to perform elementwise inversion, convert it before
// // passing it to the elementwise codegen function
// let op = if ndarray_dtype.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::Bool.id()) {
// if op == ast::Unaryop::Invert {
// ast::Unaryop::Not
// } else {
// unreachable!(
// "ufunc {} not supported for ndarray[bool, N]",
// op.op_info().method_name,
// )
// }
// } else {
// op
// };
let res = numpy::ndarray_elementwise_unaryop_impl(
generator,
ctx,
ndarray_dtype,
None,
val,
|generator, ctx, val| {
gen_unaryop_expr_with_values(generator, ctx, op, (&Some(ndarray_dtype), val))?
.unwrap()
.to_basic_value_enum(ctx, generator, ndarray_dtype)
},
)?;
// let res = numpy::ndarray_elementwise_unaryop_impl(
// generator,
// ctx,
// ndarray_dtype,
// None,
// val,
// |generator, ctx, val| {
// gen_unaryop_expr_with_values(generator, ctx, op, (&Some(ndarray_dtype), val))?
// .unwrap()
// .to_basic_value_enum(ctx, generator, ndarray_dtype)
// },
// )?;
res.as_base_value().into()
// res.as_base_value().into()
} else {
unimplemented!()
}))
@ -1688,85 +1690,86 @@ 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 llvm_usize = generator.get_size_type(ctx.ctx);
todo!()
// let llvm_usize = generator.get_size_type(ctx.ctx);
let (Some(left_ty), lhs) = left else { unreachable!() };
let (Some(right_ty), rhs) = comparators[0] else { unreachable!() };
let op = ops[0];
// let (Some(left_ty), lhs) = left else { unreachable!() };
// let (Some(right_ty), rhs) = comparators[0] else { unreachable!() };
// 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 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());
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);
// 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);
assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
// assert!(ctx.unifier.unioned(ndarray_dtype1, ndarray_dtype2));
let left_val =
NDArrayValue::from_ptr_val(lhs.into_pointer_value(), llvm_usize, None);
let res = numpy::ndarray_elementwise_binop_impl(
generator,
ctx,
ctx.primitives.bool,
None,
(left_val.as_base_value().into(), false),
(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 left_val =
// NDArrayValue::from_ptr_val(lhs.into_pointer_value(), llvm_usize, None);
// let res = numpy::ndarray_elementwise_binop_impl(
// generator,
// ctx,
// ctx.primitives.bool,
// None,
// (left_val.as_base_value().into(), false),
// (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,
// )?;
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,
(lhs, !is_ndarray1),
(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(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,
// (lhs, !is_ndarray1),
// (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(generator.bool_to_i8(ctx, val.into_int_value()).into())
// },
// )?;
Ok(Some(res.as_base_value().into()))
};
// Ok(Some(res.as_base_value().into()))
// };
}
}
@ -2102,310 +2105,312 @@ fn gen_ndarray_subscript_expr<'ctx, G: CodeGenerator>(
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);
todo!()
let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) else {
unreachable!()
};
// let llvm_i1 = ctx.ctx.bool_type();
// let llvm_i32 = ctx.ctx.i32_type();
// let llvm_usize = generator.get_size_type(ctx.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()
)
})?;
// let TypeEnum::TLiteral { values, .. } = &*ctx.unifier.get_ty_immutable(ndims) else {
// unreachable!()
// };
assert!(!ndims.is_empty());
// 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()
// )
// })?;
// 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
}
}),
// assert!(!ndims.is_empty());
ExprKind::Slice { .. } => 0,
_ => 1,
};
// // 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
// }
// }),
let ndarray_ndims_ty = ctx.unifier.get_fresh_literal(
ndims.iter().map(|v| SymbolValue::U64(v - subscripted_dims)).collect(),
None,
);
let ndarray_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_ndarray_t = llvm_pndarray_t.get_element_type().into_struct_type();
let llvm_ndarray_data_t = ctx.get_llvm_type(generator, ty).as_basic_type_enum();
// ExprKind::Slice { .. } => 0,
// _ => 1,
// };
// 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,
);
// let ndarray_ndims_ty = ctx.unifier.get_fresh_literal(
// ndims.iter().map(|v| SymbolValue::U64(v - subscripted_dims)).collect(),
// None,
// );
// let ndarray_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_ndarray_t = llvm_pndarray_t.get_element_type().into_struct_type();
// let llvm_ndarray_data_t = ctx.get_llvm_type(generator, ty).as_basic_type_enum();
// 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();
// // 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,
// );
let len = unsafe {
v.dim_sizes().get_typed_unchecked(
ctx,
generator,
&llvm_usize.const_int(dim, true),
None,
)
};
// // 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 index = ctx
.builder
.build_int_add(
len,
ctx.builder.build_int_s_extend(index, llvm_usize, "").unwrap(),
"",
)
.unwrap();
// let len = unsafe {
// v.dim_sizes().get_typed_unchecked(
// ctx,
// generator,
// &llvm_usize.const_int(dim, true),
// None,
// )
// };
Ok(Some(ctx.builder.build_int_truncate(index, llvm_i32, "").unwrap()))
},
)
.map(|v| v.map(BasicValueEnum::into_int_value))
};
// let index = ctx
// .builder
// .build_int_add(
// len,
// ctx.builder.build_int_s_extend(index, llvm_usize, "").unwrap(),
// "",
// )
// .unwrap();
// 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(ctx.builder.build_int_truncate(index, llvm_i32, "").unwrap()))
// },
// )
// .map(|v| v.map(BasicValueEnum::into_int_value))
// };
Ok(Some((index, index, llvm_i32.const_int(1, true))))
}
// // 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);
// };
ExprKind::Slice { lower, upper, step } => {
let dim_sz = unsafe {
v.dim_sizes().get_typed_unchecked(
ctx,
generator,
&llvm_usize.const_int(dim, false),
None,
)
};
// Ok(Some((index, index, llvm_i32.const_int(1, true))))
// }
handle_slice_indices(lower, upper, step, ctx, generator, dim_sz)
}
// ExprKind::Slice { lower, upper, step } => {
// let dim_sz = unsafe {
// v.dim_sizes().get_typed_unchecked(
// ctx,
// generator,
// &llvm_usize.const_int(dim, false),
// None,
// )
// };
_ => {
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);
};
// handle_slice_indices(lower, upper, step, ctx, generator, dim_sz)
// }
Ok(Some((index, index, llvm_i32.const_int(1, true))))
}
}
};
// _ => {
// 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);
// };
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,
)?;
// Ok(Some((index, index, llvm_i32.const_int(1, true))))
// }
// }
// };
for (i, elt) in elts.iter().enumerate() {
let Some(index) = generator.gen_expr(ctx, elt)? else {
return Ok(None);
};
// 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,
// )?;
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);
};
// for (i, elt) in elts.iter().enumerate() {
// let Some(index) = generator.gen_expr(ctx, elt)? 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();
}
// 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);
// };
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 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();
// }
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) };
// 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 store_ptr = unsafe {
index_addr.ptr_offset_unchecked(ctx, generator, &llvm_usize.const_zero(), None)
};
ctx.builder.build_store(store_ptr, index).unwrap();
// 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) };
Some(index_addr)
} else {
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();
Ok(Some(if ndims.len() == 1 && ndims[0] - subscripted_dims == 0 {
let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
// Some(index_addr)
// } else {
// 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);
}
// Ok(Some(if ndims.len() == 1 && ndims[0] - subscripted_dims == 0 {
// let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
let slices = slices.into_iter().map(Option::unwrap).collect_vec();
// 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);
// }
numpy::ndarray_sliced_copy(generator, ctx, ty, v, &slices)?.as_base_value().into()
}
// let slices = slices.into_iter().map(Option::unwrap).collect_vec();
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, &slices)?.as_base_value().into()
// }
numpy::ndarray_sliced_copy(generator, ctx, ty, v, &[slice])?.as_base_value().into()
}
// ExprKind::Slice { .. } => {
// let Some(slice) = expr_to_slice(generator, ctx, &slice.node, 0)? else {
// return Ok(None);
// };
_ => {
// Accessing an element from a multi-dimensional `ndarray`
// numpy::ndarray_sliced_copy(generator, ctx, ty, v, &[slice])?.as_base_value().into()
// }
let Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
// _ => {
// // 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 Some(index_addr) = make_indices_arr(generator, ctx)? else { return Ok(None) };
let num_dims = v.load_ndims(ctx);
ndarray.store_ndims(
ctx,
generator,
ctx.builder
.build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
.unwrap(),
);
// // 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 ndarray_num_dims = ndarray.load_ndims(ctx);
ndarray.create_dim_sizes(ctx, llvm_usize, ndarray_num_dims);
// let num_dims = v.load_ndims(ctx);
// ndarray.store_ndims(
// ctx,
// generator,
// ctx.builder
// .build_int_sub(num_dims, llvm_usize.const_int(1, false), "")
// .unwrap(),
// );
let ndarray_num_dims = ndarray.load_ndims(ctx);
let v_dims_src_ptr = unsafe {
v.dim_sizes().ptr_offset_unchecked(
ctx,
generator,
&llvm_usize.const_int(1, false),
None,
)
};
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_dims = ndarray.load_ndims(ctx);
// ndarray.create_dim_sizes(ctx, llvm_usize, ndarray_num_dims);
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 ndarray_num_dims = ndarray.load_ndims(ctx);
// let v_dims_src_ptr = unsafe {
// v.dim_sizes().ptr_offset_unchecked(
// ctx,
// generator,
// &llvm_usize.const_int(1, false),
// None,
// )
// };
// 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 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(),
);
// 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);
ndarray.as_base_value().into()
}
}
}))
// 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`].

14
nac3core/src/codegen/irrt/irrt.cpp
View File

@ -137,6 +137,20 @@ void __nac3_ndarray_calc_broadcast_idx_impl(
}
} // namespace
template<typename SizeT>
static void __nac3_ndarray_strides_from_shape_impl(
SizeT ndims,
SizeT *shape,
SizeT *dst_strides
) {
SizeT stride_product = 1;
for (SizeT i = 0; i < ndims; i++) {
int dim_i = ndims - i - 1;
dst_strides[dim_i] = stride_product;
stride_product *= shape[dim_i];
}
}
extern "C" {
#define DEF_nac3_int_exp_(T) \
T __nac3_int_exp_##T(T base, T exp) {\

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

@ -2,8 +2,7 @@ use crate::typecheck::typedef::Type;
use super::{
classes::{
ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue, ListValue, NDArrayValue,
TypedArrayLikeAdapter, UntypedArrayLikeAccessor,
ArrayLikeIndexer, ArrayLikeValue, ArraySliceValue, ListValue, NDArrayValue, NpArrayValue, TypedArrayLikeAdapter, UntypedArrayLikeAccessor
},
llvm_intrinsics, CodeGenContext, CodeGenerator,
};
@ -15,7 +14,7 @@ use inkwell::{
memory_buffer::MemoryBuffer,
module::Module,
types::{BasicTypeEnum, IntType},
values::{BasicValueEnum, CallSiteValue, FloatValue, IntValue},
values::{BasicValueEnum, CallSiteValue, FloatValue, IntValue, PointerValue},
AddressSpace, IntPredicate,
};
use itertools::Either;
@ -927,3 +926,13 @@ pub fn call_ndarray_calc_broadcast_index<
Box::new(|_, v| v.into()),
)
}
fn call_ndarray_strides_from_shape<'ctx, G, Index>(
generator: &mut G,
ctx: &mut CodeGenContext<'ctx, '_>,
ndims: IntValue<'ctx>,
shape: PointerValue<'ctx>,
dst_strides: PointerValue<'ctx>,
) -> IntValue<'ctx> {
todo!()
}

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

File diff suppressed because it is too large Load Diff

45
nac3core/src/toplevel/builtins.rs
View File

@ -961,8 +961,9 @@ impl<'a> BuiltinBuilder<'a> {
resolver: None,
codegen_callback: Some(Arc::new(GenCall::new(Box::new(
|ctx, obj, fun, args, generator| {
gen_ndarray_copy(ctx, &obj, fun, &args, generator)
.map(|val| Some(val.as_basic_value_enum()))
todo!()
// gen_ndarray_copy(ctx, &obj, fun, &args, generator)
// .map(|val| Some(val.as_basic_value_enum()))
},
)))),
loc: None,
@ -978,8 +979,9 @@ impl<'a> BuiltinBuilder<'a> {
resolver: None,
codegen_callback: Some(Arc::new(GenCall::new(Box::new(
|ctx, obj, fun, args, generator| {
gen_ndarray_fill(ctx, &obj, fun, &args, generator)?;
Ok(None)
todo!()
// gen_ndarray_fill(ctx, &obj, fun, &args, generator)?;
// Ok(None)
},
)))),
loc: None,
@ -1199,13 +1201,14 @@ impl<'a> BuiltinBuilder<'a> {
self.ndarray_float,
&[(self.ndarray_factory_fn_shape_arg_tvar.ty, "shape")],
Box::new(move |ctx, obj, fun, args, generator| {
let func = match prim {
PrimDef::FunNpNDArray | PrimDef::FunNpEmpty => gen_ndarray_empty,
PrimDef::FunNpZeros => gen_ndarray_zeros,
PrimDef::FunNpOnes => gen_ndarray_ones,
_ => unreachable!(),
};
func(ctx, &obj, fun, &args, generator).map(|val| Some(val.as_basic_value_enum()))
todo!()
// let func = match prim {
// PrimDef::FunNpNDArray | PrimDef::FunNpEmpty => gen_ndarray_empty,
// PrimDef::FunNpZeros => gen_ndarray_zeros,
// PrimDef::FunNpOnes => gen_ndarray_ones,
// _ => unreachable!(),
// };
// func(ctx, &obj, fun, &args, generator).map(|val| Some(val.as_basic_value_enum()))
}),
)
}
@ -1251,8 +1254,9 @@ impl<'a> BuiltinBuilder<'a> {
resolver: None,
codegen_callback: Some(Arc::new(GenCall::new(Box::new(
|ctx, obj, fun, args, generator| {
gen_ndarray_array(ctx, &obj, fun, &args, generator)
.map(|val| Some(val.as_basic_value_enum()))
todo!()
// gen_ndarray_array(ctx, &obj, fun, &args, generator)
// .map(|val| Some(val.as_basic_value_enum()))
},
)))),
loc: None,
@ -1270,8 +1274,9 @@ impl<'a> BuiltinBuilder<'a> {
// type variable
&[(self.list_int32, "shape"), (tv.ty, "fill_value")],
Box::new(move |ctx, obj, fun, args, generator| {
gen_ndarray_full(ctx, &obj, fun, &args, generator)
.map(|val| Some(val.as_basic_value_enum()))
todo!()
// gen_ndarray_full(ctx, &obj, fun, &args, generator)
// .map(|val| Some(val.as_basic_value_enum()))
}),
)
}
@ -1303,8 +1308,9 @@ impl<'a> BuiltinBuilder<'a> {
resolver: None,
codegen_callback: Some(Arc::new(GenCall::new(Box::new(
|ctx, obj, fun, args, generator| {
gen_ndarray_eye(ctx, &obj, fun, &args, generator)
.map(|val| Some(val.as_basic_value_enum()))
todo!()
// gen_ndarray_eye(ctx, &obj, fun, &args, generator)
// .map(|val| Some(val.as_basic_value_enum()))
},
)))),
loc: None,
@ -1317,8 +1323,9 @@ impl<'a> BuiltinBuilder<'a> {
self.ndarray_float_2d,
&[(int32, "n")],
Box::new(|ctx, obj, fun, args, generator| {
gen_ndarray_identity(ctx, &obj, fun, &args, generator)
.map(|val| Some(val.as_basic_value_enum()))
todo!()
// gen_ndarray_identity(ctx, &obj, fun, &args, generator)
// .map(|val| Some(val.as_basic_value_enum()))
}),
),
_ => unreachable!(),