forked from M-Labs/nac3
core: use C++ for irrt source
This commit is contained in:
parent
f062ef5f59
commit
650f354b74
|
|
@ -12,7 +12,7 @@
|
||||||
llvm-tools-irrt = pkgs.runCommandNoCC "llvm-tools-irrt" {}
|
llvm-tools-irrt = pkgs.runCommandNoCC "llvm-tools-irrt" {}
|
||||||
''
|
''
|
||||||
mkdir -p $out/bin
|
mkdir -p $out/bin
|
||||||
ln -s ${pkgs.llvmPackages_14.clang-unwrapped}/bin/clang $out/bin/clang-irrt
|
ln -s ${pkgs.llvmPackages_14.clang-unwrapped}/bin/clang++ $out/bin/clang++-irrt
|
||||||
ln -s ${pkgs.llvmPackages_14.llvm.out}/bin/llvm-as $out/bin/llvm-as-irrt
|
ln -s ${pkgs.llvmPackages_14.llvm.out}/bin/llvm-as $out/bin/llvm-as-irrt
|
||||||
'';
|
'';
|
||||||
nac3artiq = pkgs.python3Packages.toPythonModule (
|
nac3artiq = pkgs.python3Packages.toPythonModule (
|
||||||
|
|
|
||||||
|
|
@ -8,7 +8,7 @@ use std::{
|
||||||
};
|
};
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
const FILE: &str = "src/codegen/irrt/irrt.c";
|
const FILE: &str = "src/codegen/irrt/irrt.cpp";
|
||||||
|
|
||||||
/*
|
/*
|
||||||
* HACK: Sadly, clang doesn't let us emit generic LLVM bitcode.
|
* HACK: Sadly, clang doesn't let us emit generic LLVM bitcode.
|
||||||
|
|
@ -18,6 +18,8 @@ fn main() {
|
||||||
"--target=wasm32",
|
"--target=wasm32",
|
||||||
FILE,
|
FILE,
|
||||||
"-fno-discard-value-names",
|
"-fno-discard-value-names",
|
||||||
|
"-fno-exceptions",
|
||||||
|
"-fno-rtti",
|
||||||
match env::var("PROFILE").as_deref() {
|
match env::var("PROFILE").as_deref() {
|
||||||
Ok("debug") => "-O0",
|
Ok("debug") => "-O0",
|
||||||
Ok("release") => "-O3",
|
Ok("release") => "-O3",
|
||||||
|
|
@ -35,7 +37,7 @@ fn main() {
|
||||||
let out_dir = env::var("OUT_DIR").unwrap();
|
let out_dir = env::var("OUT_DIR").unwrap();
|
||||||
let out_path = Path::new(&out_dir);
|
let out_path = Path::new(&out_dir);
|
||||||
|
|
||||||
let output = Command::new("clang-irrt")
|
let output = Command::new("clang++-irrt")
|
||||||
.args(flags)
|
.args(flags)
|
||||||
.output()
|
.output()
|
||||||
.map(|o| {
|
.map(|o| {
|
||||||
|
|
|
||||||
|
|
@ -1,389 +0,0 @@
|
||||||
typedef _BitInt(8) int8_t;
|
|
||||||
typedef unsigned _BitInt(8) uint8_t;
|
|
||||||
typedef _BitInt(32) int32_t;
|
|
||||||
typedef unsigned _BitInt(32) uint32_t;
|
|
||||||
typedef _BitInt(64) int64_t;
|
|
||||||
typedef unsigned _BitInt(64) uint64_t;
|
|
||||||
|
|
||||||
# define MAX(a, b) (a > b ? a : b)
|
|
||||||
# define MIN(a, b) (a > b ? b : a)
|
|
||||||
|
|
||||||
# define NULL ((void *) 0)
|
|
||||||
|
|
||||||
// adapted from GNU Scientific Library: https://git.savannah.gnu.org/cgit/gsl.git/tree/sys/pow_int.c
|
|
||||||
// need to make sure `exp >= 0` before calling this function
|
|
||||||
#define DEF_INT_EXP(T) T __nac3_int_exp_##T( \
|
|
||||||
T base, \
|
|
||||||
T exp \
|
|
||||||
) { \
|
|
||||||
T res = (T)1; \
|
|
||||||
/* repeated squaring method */ \
|
|
||||||
do { \
|
|
||||||
if (exp & 1) res *= base; /* for n odd */ \
|
|
||||||
exp >>= 1; \
|
|
||||||
base *= base; \
|
|
||||||
} while (exp); \
|
|
||||||
return res; \
|
|
||||||
} \
|
|
||||||
|
|
||||||
DEF_INT_EXP(int32_t)
|
|
||||||
DEF_INT_EXP(int64_t)
|
|
||||||
DEF_INT_EXP(uint32_t)
|
|
||||||
DEF_INT_EXP(uint64_t)
|
|
||||||
|
|
||||||
|
|
||||||
int32_t __nac3_slice_index_bound(int32_t i, const int32_t len) {
|
|
||||||
if (i < 0) {
|
|
||||||
i = len + i;
|
|
||||||
}
|
|
||||||
if (i < 0) {
|
|
||||||
return 0;
|
|
||||||
} else if (i > len) {
|
|
||||||
return len;
|
|
||||||
}
|
|
||||||
return i;
|
|
||||||
}
|
|
||||||
|
|
||||||
int32_t __nac3_range_slice_len(const int32_t start, const int32_t end, const int32_t step) {
|
|
||||||
int32_t diff = end - start;
|
|
||||||
if (diff > 0 && step > 0) {
|
|
||||||
return ((diff - 1) / step) + 1;
|
|
||||||
} else if (diff < 0 && step < 0) {
|
|
||||||
return ((diff + 1) / step) + 1;
|
|
||||||
} else {
|
|
||||||
return 0;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Handle list assignment and dropping part of the list when
|
|
||||||
// both dest_step and src_step are +1.
|
|
||||||
// - All the index must *not* be out-of-bound or negative,
|
|
||||||
// - The end index is *inclusive*,
|
|
||||||
// - The length of src and dest slice size should already
|
|
||||||
// be checked: if dest.step == 1 then len(src) <= len(dest) else len(src) == len(dest)
|
|
||||||
int32_t __nac3_list_slice_assign_var_size(
|
|
||||||
int32_t dest_start,
|
|
||||||
int32_t dest_end,
|
|
||||||
int32_t dest_step,
|
|
||||||
uint8_t *dest_arr,
|
|
||||||
int32_t dest_arr_len,
|
|
||||||
int32_t src_start,
|
|
||||||
int32_t src_end,
|
|
||||||
int32_t src_step,
|
|
||||||
uint8_t *src_arr,
|
|
||||||
int32_t src_arr_len,
|
|
||||||
const int32_t size
|
|
||||||
) {
|
|
||||||
/* if dest_arr_len == 0, do nothing since we do not support extending list */
|
|
||||||
if (dest_arr_len == 0) return dest_arr_len;
|
|
||||||
/* if both step is 1, memmove directly, handle the dropping of the list, and shrink size */
|
|
||||||
if (src_step == dest_step && dest_step == 1) {
|
|
||||||
const int32_t src_len = (src_end >= src_start) ? (src_end - src_start + 1) : 0;
|
|
||||||
const int32_t dest_len = (dest_end >= dest_start) ? (dest_end - dest_start + 1) : 0;
|
|
||||||
if (src_len > 0) {
|
|
||||||
__builtin_memmove(
|
|
||||||
dest_arr + dest_start * size,
|
|
||||||
src_arr + src_start * size,
|
|
||||||
src_len * size
|
|
||||||
);
|
|
||||||
}
|
|
||||||
if (dest_len > 0) {
|
|
||||||
/* dropping */
|
|
||||||
__builtin_memmove(
|
|
||||||
dest_arr + (dest_start + src_len) * size,
|
|
||||||
dest_arr + (dest_end + 1) * size,
|
|
||||||
(dest_arr_len - dest_end - 1) * size
|
|
||||||
);
|
|
||||||
}
|
|
||||||
/* shrink size */
|
|
||||||
return dest_arr_len - (dest_len - src_len);
|
|
||||||
}
|
|
||||||
/* if two range overlaps, need alloca */
|
|
||||||
uint8_t need_alloca =
|
|
||||||
(dest_arr == src_arr)
|
|
||||||
&& !(
|
|
||||||
MAX(dest_start, dest_end) < MIN(src_start, src_end)
|
|
||||||
|| MAX(src_start, src_end) < MIN(dest_start, dest_end)
|
|
||||||
);
|
|
||||||
if (need_alloca) {
|
|
||||||
uint8_t *tmp = __builtin_alloca(src_arr_len * size);
|
|
||||||
__builtin_memcpy(tmp, src_arr, src_arr_len * size);
|
|
||||||
src_arr = tmp;
|
|
||||||
}
|
|
||||||
int32_t src_ind = src_start;
|
|
||||||
int32_t dest_ind = dest_start;
|
|
||||||
for (;
|
|
||||||
(src_step > 0) ? (src_ind <= src_end) : (src_ind >= src_end);
|
|
||||||
src_ind += src_step, dest_ind += dest_step
|
|
||||||
) {
|
|
||||||
/* for constant optimization */
|
|
||||||
if (size == 1) {
|
|
||||||
__builtin_memcpy(dest_arr + dest_ind, src_arr + src_ind, 1);
|
|
||||||
} else if (size == 4) {
|
|
||||||
__builtin_memcpy(dest_arr + dest_ind * 4, src_arr + src_ind * 4, 4);
|
|
||||||
} else if (size == 8) {
|
|
||||||
__builtin_memcpy(dest_arr + dest_ind * 8, src_arr + src_ind * 8, 8);
|
|
||||||
} else {
|
|
||||||
/* memcpy for var size, cannot overlap after previous alloca */
|
|
||||||
__builtin_memcpy(dest_arr + dest_ind * size, src_arr + src_ind * size, size);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
/* only dest_step == 1 can we shrink the dest list. */
|
|
||||||
/* size should be ensured prior to calling this function */
|
|
||||||
if (dest_step == 1 && dest_end >= dest_start) {
|
|
||||||
__builtin_memmove(
|
|
||||||
dest_arr + dest_ind * size,
|
|
||||||
dest_arr + (dest_end + 1) * size,
|
|
||||||
(dest_arr_len - dest_end - 1) * size
|
|
||||||
);
|
|
||||||
return dest_arr_len - (dest_end - dest_ind) - 1;
|
|
||||||
}
|
|
||||||
return dest_arr_len;
|
|
||||||
}
|
|
||||||
|
|
||||||
int32_t __nac3_isinf(double x) {
|
|
||||||
return __builtin_isinf(x);
|
|
||||||
}
|
|
||||||
|
|
||||||
int32_t __nac3_isnan(double x) {
|
|
||||||
return __builtin_isnan(x);
|
|
||||||
}
|
|
||||||
|
|
||||||
double tgamma(double arg);
|
|
||||||
|
|
||||||
double __nac3_gamma(double z) {
|
|
||||||
// Handling for denormals
|
|
||||||
// | x | Python gamma(x) | C tgamma(x) |
|
|
||||||
// --- | ----------------- | --------------- | ----------- |
|
|
||||||
// (1) | nan | nan | nan |
|
|
||||||
// (2) | -inf | -inf | inf |
|
|
||||||
// (3) | inf | inf | inf |
|
|
||||||
// (4) | 0.0 | inf | inf |
|
|
||||||
// (5) | {-1.0, -2.0, ...} | inf | nan |
|
|
||||||
|
|
||||||
// (1)-(3)
|
|
||||||
if (__builtin_isinf(z) || __builtin_isnan(z)) {
|
|
||||||
return z;
|
|
||||||
}
|
|
||||||
|
|
||||||
double v = tgamma(z);
|
|
||||||
|
|
||||||
// (4)-(5)
|
|
||||||
return __builtin_isinf(v) || __builtin_isnan(v) ? __builtin_inf() : v;
|
|
||||||
}
|
|
||||||
|
|
||||||
double lgamma(double arg);
|
|
||||||
|
|
||||||
double __nac3_gammaln(double x) {
|
|
||||||
// libm's handling of value overflows differs from scipy:
|
|
||||||
// - scipy: gammaln(-inf) -> -inf
|
|
||||||
// - libm : lgamma(-inf) -> inf
|
|
||||||
|
|
||||||
if (__builtin_isinf(x)) {
|
|
||||||
return x;
|
|
||||||
}
|
|
||||||
|
|
||||||
return lgamma(x);
|
|
||||||
}
|
|
||||||
|
|
||||||
double j0(double x);
|
|
||||||
|
|
||||||
double __nac3_j0(double x) {
|
|
||||||
// libm's handling of value overflows differs from scipy:
|
|
||||||
// - scipy: j0(inf) -> nan
|
|
||||||
// - libm : j0(inf) -> 0.0
|
|
||||||
|
|
||||||
if (__builtin_isinf(x)) {
|
|
||||||
return __builtin_nan("");
|
|
||||||
}
|
|
||||||
|
|
||||||
return j0(x);
|
|
||||||
}
|
|
||||||
|
|
||||||
uint32_t __nac3_ndarray_calc_size(
|
|
||||||
const uint32_t *list_data,
|
|
||||||
uint32_t list_len,
|
|
||||||
uint32_t begin_idx,
|
|
||||||
uint32_t end_idx
|
|
||||||
) {
|
|
||||||
__builtin_assume(end_idx <= list_len);
|
|
||||||
|
|
||||||
uint32_t num_elems = 1;
|
|
||||||
for (uint32_t i = begin_idx; i < end_idx; ++i) {
|
|
||||||
uint32_t val = list_data[i];
|
|
||||||
__builtin_assume(val > 0);
|
|
||||||
num_elems *= val;
|
|
||||||
}
|
|
||||||
return num_elems;
|
|
||||||
}
|
|
||||||
|
|
||||||
uint64_t __nac3_ndarray_calc_size64(
|
|
||||||
const uint64_t *list_data,
|
|
||||||
uint64_t list_len,
|
|
||||||
uint64_t begin_idx,
|
|
||||||
uint64_t end_idx
|
|
||||||
) {
|
|
||||||
__builtin_assume(end_idx <= list_len);
|
|
||||||
|
|
||||||
uint64_t num_elems = 1;
|
|
||||||
for (uint64_t i = begin_idx; i < end_idx; ++i) {
|
|
||||||
uint64_t val = list_data[i];
|
|
||||||
__builtin_assume(val > 0);
|
|
||||||
num_elems *= val;
|
|
||||||
}
|
|
||||||
return num_elems;
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_nd_indices(
|
|
||||||
uint32_t index,
|
|
||||||
const uint32_t* dims,
|
|
||||||
uint32_t num_dims,
|
|
||||||
uint32_t* idxs
|
|
||||||
) {
|
|
||||||
uint32_t stride = 1;
|
|
||||||
for (uint32_t dim = 0; dim < num_dims; dim++) {
|
|
||||||
uint32_t i = num_dims - dim - 1;
|
|
||||||
__builtin_assume(dims[i] > 0);
|
|
||||||
idxs[i] = (index / stride) % dims[i];
|
|
||||||
stride *= dims[i];
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_nd_indices64(
|
|
||||||
uint64_t index,
|
|
||||||
const uint64_t* dims,
|
|
||||||
uint64_t num_dims,
|
|
||||||
uint32_t* idxs
|
|
||||||
) {
|
|
||||||
uint64_t stride = 1;
|
|
||||||
for (uint64_t dim = 0; dim < num_dims; dim++) {
|
|
||||||
uint64_t i = num_dims - dim - 1;
|
|
||||||
__builtin_assume(dims[i] > 0);
|
|
||||||
idxs[i] = (uint32_t) ((index / stride) % dims[i]);
|
|
||||||
stride *= dims[i];
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
uint32_t __nac3_ndarray_flatten_index(
|
|
||||||
const uint32_t* dims,
|
|
||||||
uint32_t num_dims,
|
|
||||||
const uint32_t* indices,
|
|
||||||
uint32_t num_indices
|
|
||||||
) {
|
|
||||||
uint32_t idx = 0;
|
|
||||||
uint32_t stride = 1;
|
|
||||||
for (uint32_t i = 0; i < num_dims; ++i) {
|
|
||||||
uint32_t ri = num_dims - i - 1;
|
|
||||||
if (ri < num_indices) {
|
|
||||||
idx += (stride * indices[ri]);
|
|
||||||
}
|
|
||||||
|
|
||||||
__builtin_assume(dims[i] > 0);
|
|
||||||
stride *= dims[ri];
|
|
||||||
}
|
|
||||||
return idx;
|
|
||||||
}
|
|
||||||
|
|
||||||
uint64_t __nac3_ndarray_flatten_index64(
|
|
||||||
const uint64_t* dims,
|
|
||||||
uint64_t num_dims,
|
|
||||||
const uint32_t* indices,
|
|
||||||
uint64_t num_indices
|
|
||||||
) {
|
|
||||||
uint64_t idx = 0;
|
|
||||||
uint64_t stride = 1;
|
|
||||||
for (uint64_t i = 0; i < num_dims; ++i) {
|
|
||||||
uint64_t ri = num_dims - i - 1;
|
|
||||||
if (ri < num_indices) {
|
|
||||||
idx += (stride * indices[ri]);
|
|
||||||
}
|
|
||||||
|
|
||||||
__builtin_assume(dims[i] > 0);
|
|
||||||
stride *= dims[ri];
|
|
||||||
}
|
|
||||||
return idx;
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_broadcast(
|
|
||||||
const uint32_t *lhs_dims,
|
|
||||||
uint32_t lhs_ndims,
|
|
||||||
const uint32_t *rhs_dims,
|
|
||||||
uint32_t rhs_ndims,
|
|
||||||
uint32_t *out_dims
|
|
||||||
) {
|
|
||||||
uint32_t max_ndims = lhs_ndims > rhs_ndims ? lhs_ndims : rhs_ndims;
|
|
||||||
|
|
||||||
for (uint32_t i = 0; i < max_ndims; ++i) {
|
|
||||||
uint32_t *lhs_dim_sz = i < lhs_ndims ? &lhs_dims[lhs_ndims - i - 1] : NULL;
|
|
||||||
uint32_t *rhs_dim_sz = i < rhs_ndims ? &rhs_dims[rhs_ndims - i - 1] : NULL;
|
|
||||||
uint32_t *out_dim = &out_dims[max_ndims - i - 1];
|
|
||||||
|
|
||||||
if (lhs_dim_sz == NULL) {
|
|
||||||
*out_dim = *rhs_dim_sz;
|
|
||||||
} else if (rhs_dim_sz == NULL) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else if (*lhs_dim_sz == 1) {
|
|
||||||
*out_dim = *rhs_dim_sz;
|
|
||||||
} else if (*rhs_dim_sz == 1) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else if (*lhs_dim_sz == *rhs_dim_sz) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else {
|
|
||||||
__builtin_unreachable();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_broadcast64(
|
|
||||||
const uint64_t *lhs_dims,
|
|
||||||
uint64_t lhs_ndims,
|
|
||||||
const uint64_t *rhs_dims,
|
|
||||||
uint64_t rhs_ndims,
|
|
||||||
uint64_t *out_dims
|
|
||||||
) {
|
|
||||||
uint64_t max_ndims = lhs_ndims > rhs_ndims ? lhs_ndims : rhs_ndims;
|
|
||||||
|
|
||||||
for (uint64_t i = 0; i < max_ndims; ++i) {
|
|
||||||
uint64_t *lhs_dim_sz = i < lhs_ndims ? &lhs_dims[lhs_ndims - i - 1] : NULL;
|
|
||||||
uint64_t *rhs_dim_sz = i < rhs_ndims ? &rhs_dims[rhs_ndims - i - 1] : NULL;
|
|
||||||
uint64_t *out_dim = &out_dims[max_ndims - i - 1];
|
|
||||||
|
|
||||||
if (lhs_dim_sz == NULL) {
|
|
||||||
*out_dim = *rhs_dim_sz;
|
|
||||||
} else if (rhs_dim_sz == NULL) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else if (*lhs_dim_sz == 1) {
|
|
||||||
*out_dim = *rhs_dim_sz;
|
|
||||||
} else if (*rhs_dim_sz == 1) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else if (*lhs_dim_sz == *rhs_dim_sz) {
|
|
||||||
*out_dim = *lhs_dim_sz;
|
|
||||||
} else {
|
|
||||||
__builtin_unreachable();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_broadcast_idx(
|
|
||||||
const uint32_t *src_dims,
|
|
||||||
uint32_t src_ndims,
|
|
||||||
const uint32_t *in_idx,
|
|
||||||
uint32_t *out_idx
|
|
||||||
) {
|
|
||||||
for (uint32_t i = 0; i < src_ndims; ++i) {
|
|
||||||
uint32_t src_i = src_ndims - i - 1;
|
|
||||||
out_idx[src_i] = src_dims[src_i] == 1 ? 0 : in_idx[src_i];
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
void __nac3_ndarray_calc_broadcast_idx64(
|
|
||||||
const uint64_t *src_dims,
|
|
||||||
uint64_t src_ndims,
|
|
||||||
const uint32_t *in_idx,
|
|
||||||
uint32_t *out_idx
|
|
||||||
) {
|
|
||||||
for (uint64_t i = 0; i < src_ndims; ++i) {
|
|
||||||
uint64_t src_i = src_ndims - i - 1;
|
|
||||||
out_idx[src_i] = src_dims[src_i] == 1 ? 0 : (uint32_t) in_idx[src_i];
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -0,0 +1,412 @@
|
||||||
|
typedef _BitInt(8) int8_t;
|
||||||
|
typedef unsigned _BitInt(8) uint8_t;
|
||||||
|
typedef _BitInt(32) int32_t;
|
||||||
|
typedef unsigned _BitInt(32) uint32_t;
|
||||||
|
typedef _BitInt(64) int64_t;
|
||||||
|
typedef unsigned _BitInt(64) uint64_t;
|
||||||
|
|
||||||
|
// NDArray indices are always `uint32_t`.
|
||||||
|
typedef uint32_t NDIndex;
|
||||||
|
// The type of an index or a value describing the length of a range/slice is always `int32_t`.
|
||||||
|
typedef int32_t SliceIndex;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
static T max(T a, T b) {
|
||||||
|
return a > b ? a : b;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
static T min(T a, T b) {
|
||||||
|
return a > b ? b : a;
|
||||||
|
}
|
||||||
|
|
||||||
|
// adapted from GNU Scientific Library: https://git.savannah.gnu.org/cgit/gsl.git/tree/sys/pow_int.c
|
||||||
|
// need to make sure `exp >= 0` before calling this function
|
||||||
|
template <typename T>
|
||||||
|
static T __nac3_int_exp_impl(T base, T exp) {
|
||||||
|
T res = 1;
|
||||||
|
/* repeated squaring method */
|
||||||
|
do {
|
||||||
|
if (exp & 1) {
|
||||||
|
res *= base; /* for n odd */
|
||||||
|
}
|
||||||
|
exp >>= 1;
|
||||||
|
base *= base;
|
||||||
|
} while (exp);
|
||||||
|
return res;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename SizeT>
|
||||||
|
static SizeT __nac3_ndarray_calc_size_impl(
|
||||||
|
const SizeT *list_data,
|
||||||
|
SizeT list_len,
|
||||||
|
SizeT begin_idx,
|
||||||
|
SizeT end_idx
|
||||||
|
) {
|
||||||
|
__builtin_assume(end_idx <= list_len);
|
||||||
|
|
||||||
|
SizeT num_elems = 1;
|
||||||
|
for (SizeT i = begin_idx; i < end_idx; ++i) {
|
||||||
|
SizeT val = list_data[i];
|
||||||
|
__builtin_assume(val > 0);
|
||||||
|
num_elems *= val;
|
||||||
|
}
|
||||||
|
return num_elems;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename SizeT>
|
||||||
|
static void __nac3_ndarray_calc_nd_indices_impl(
|
||||||
|
SizeT index,
|
||||||
|
const SizeT *dims,
|
||||||
|
SizeT num_dims,
|
||||||
|
NDIndex *idxs
|
||||||
|
) {
|
||||||
|
SizeT stride = 1;
|
||||||
|
for (SizeT dim = 0; dim < num_dims; dim++) {
|
||||||
|
SizeT i = num_dims - dim - 1;
|
||||||
|
__builtin_assume(dims[i] > 0);
|
||||||
|
idxs[i] = (index / stride) % dims[i];
|
||||||
|
stride *= dims[i];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename SizeT>
|
||||||
|
static SizeT __nac3_ndarray_flatten_index_impl(
|
||||||
|
const SizeT *dims,
|
||||||
|
SizeT num_dims,
|
||||||
|
const NDIndex *indices,
|
||||||
|
SizeT num_indices
|
||||||
|
) {
|
||||||
|
SizeT idx = 0;
|
||||||
|
SizeT stride = 1;
|
||||||
|
for (SizeT i = 0; i < num_dims; ++i) {
|
||||||
|
SizeT ri = num_dims - i - 1;
|
||||||
|
if (ri < num_indices) {
|
||||||
|
idx += stride * indices[ri];
|
||||||
|
}
|
||||||
|
|
||||||
|
__builtin_assume(dims[i] > 0);
|
||||||
|
stride *= dims[ri];
|
||||||
|
}
|
||||||
|
return idx;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename SizeT>
|
||||||
|
static void __nac3_ndarray_calc_broadcast_impl(
|
||||||
|
const SizeT *lhs_dims,
|
||||||
|
SizeT lhs_ndims,
|
||||||
|
const SizeT *rhs_dims,
|
||||||
|
SizeT rhs_ndims,
|
||||||
|
SizeT *out_dims
|
||||||
|
) {
|
||||||
|
SizeT max_ndims = lhs_ndims > rhs_ndims ? lhs_ndims : rhs_ndims;
|
||||||
|
|
||||||
|
for (SizeT i = 0; i < max_ndims; ++i) {
|
||||||
|
const SizeT *lhs_dim_sz = i < lhs_ndims ? &lhs_dims[lhs_ndims - i - 1] : nullptr;
|
||||||
|
const SizeT *rhs_dim_sz = i < rhs_ndims ? &rhs_dims[rhs_ndims - i - 1] : nullptr;
|
||||||
|
SizeT *out_dim = &out_dims[max_ndims - i - 1];
|
||||||
|
|
||||||
|
if (lhs_dim_sz == nullptr) {
|
||||||
|
*out_dim = *rhs_dim_sz;
|
||||||
|
} else if (rhs_dim_sz == nullptr) {
|
||||||
|
*out_dim = *lhs_dim_sz;
|
||||||
|
} else if (*lhs_dim_sz == 1) {
|
||||||
|
*out_dim = *rhs_dim_sz;
|
||||||
|
} else if (*rhs_dim_sz == 1) {
|
||||||
|
*out_dim = *lhs_dim_sz;
|
||||||
|
} else if (*lhs_dim_sz == *rhs_dim_sz) {
|
||||||
|
*out_dim = *lhs_dim_sz;
|
||||||
|
} else {
|
||||||
|
__builtin_unreachable();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename SizeT>
|
||||||
|
static void __nac3_ndarray_calc_broadcast_idx_impl(
|
||||||
|
const SizeT *src_dims,
|
||||||
|
SizeT src_ndims,
|
||||||
|
const NDIndex *in_idx,
|
||||||
|
NDIndex *out_idx
|
||||||
|
) {
|
||||||
|
for (SizeT i = 0; i < src_ndims; ++i) {
|
||||||
|
SizeT src_i = src_ndims - i - 1;
|
||||||
|
out_idx[src_i] = src_dims[src_i] == 1 ? 0 : in_idx[src_i];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
extern "C" {
|
||||||
|
#define DEF_nac3_int_exp_(T) \
|
||||||
|
T __nac3_int_exp_##T(T base, T exp) {\
|
||||||
|
return __nac3_int_exp_impl(base, exp);\
|
||||||
|
}
|
||||||
|
|
||||||
|
DEF_nac3_int_exp_(int32_t)
|
||||||
|
DEF_nac3_int_exp_(int64_t)
|
||||||
|
DEF_nac3_int_exp_(uint32_t)
|
||||||
|
DEF_nac3_int_exp_(uint64_t)
|
||||||
|
|
||||||
|
SliceIndex __nac3_slice_index_bound(SliceIndex i, const SliceIndex len) {
|
||||||
|
if (i < 0) {
|
||||||
|
i = len + i;
|
||||||
|
}
|
||||||
|
if (i < 0) {
|
||||||
|
return 0;
|
||||||
|
} else if (i > len) {
|
||||||
|
return len;
|
||||||
|
}
|
||||||
|
return i;
|
||||||
|
}
|
||||||
|
|
||||||
|
SliceIndex __nac3_range_slice_len(
|
||||||
|
const SliceIndex start,
|
||||||
|
const SliceIndex end,
|
||||||
|
const SliceIndex step
|
||||||
|
) {
|
||||||
|
SliceIndex diff = end - start;
|
||||||
|
if (diff > 0 && step > 0) {
|
||||||
|
return ((diff - 1) / step) + 1;
|
||||||
|
} else if (diff < 0 && step < 0) {
|
||||||
|
return ((diff + 1) / step) + 1;
|
||||||
|
} else {
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle list assignment and dropping part of the list when
|
||||||
|
// both dest_step and src_step are +1.
|
||||||
|
// - All the index must *not* be out-of-bound or negative,
|
||||||
|
// - The end index is *inclusive*,
|
||||||
|
// - The length of src and dest slice size should already
|
||||||
|
// be checked: if dest.step == 1 then len(src) <= len(dest) else len(src) == len(dest)
|
||||||
|
SliceIndex __nac3_list_slice_assign_var_size(
|
||||||
|
SliceIndex dest_start,
|
||||||
|
SliceIndex dest_end,
|
||||||
|
SliceIndex dest_step,
|
||||||
|
uint8_t *dest_arr,
|
||||||
|
SliceIndex dest_arr_len,
|
||||||
|
SliceIndex src_start,
|
||||||
|
SliceIndex src_end,
|
||||||
|
SliceIndex src_step,
|
||||||
|
uint8_t *src_arr,
|
||||||
|
SliceIndex src_arr_len,
|
||||||
|
const SliceIndex size
|
||||||
|
) {
|
||||||
|
/* if dest_arr_len == 0, do nothing since we do not support extending list */
|
||||||
|
if (dest_arr_len == 0) return dest_arr_len;
|
||||||
|
/* if both step is 1, memmove directly, handle the dropping of the list, and shrink size */
|
||||||
|
if (src_step == dest_step && dest_step == 1) {
|
||||||
|
const SliceIndex src_len = (src_end >= src_start) ? (src_end - src_start + 1) : 0;
|
||||||
|
const SliceIndex dest_len = (dest_end >= dest_start) ? (dest_end - dest_start + 1) : 0;
|
||||||
|
if (src_len > 0) {
|
||||||
|
__builtin_memmove(
|
||||||
|
dest_arr + dest_start * size,
|
||||||
|
src_arr + src_start * size,
|
||||||
|
src_len * size
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (dest_len > 0) {
|
||||||
|
/* dropping */
|
||||||
|
__builtin_memmove(
|
||||||
|
dest_arr + (dest_start + src_len) * size,
|
||||||
|
dest_arr + (dest_end + 1) * size,
|
||||||
|
(dest_arr_len - dest_end - 1) * size
|
||||||
|
);
|
||||||
|
}
|
||||||
|
/* shrink size */
|
||||||
|
return dest_arr_len - (dest_len - src_len);
|
||||||
|
}
|
||||||
|
/* if two range overlaps, need alloca */
|
||||||
|
uint8_t need_alloca =
|
||||||
|
(dest_arr == src_arr)
|
||||||
|
&& !(
|
||||||
|
max(dest_start, dest_end) < min(src_start, src_end)
|
||||||
|
|| max(src_start, src_end) < min(dest_start, dest_end)
|
||||||
|
);
|
||||||
|
if (need_alloca) {
|
||||||
|
uint8_t *tmp = reinterpret_cast<uint8_t *>(__builtin_alloca(src_arr_len * size));
|
||||||
|
__builtin_memcpy(tmp, src_arr, src_arr_len * size);
|
||||||
|
src_arr = tmp;
|
||||||
|
}
|
||||||
|
SliceIndex src_ind = src_start;
|
||||||
|
SliceIndex dest_ind = dest_start;
|
||||||
|
for (;
|
||||||
|
(src_step > 0) ? (src_ind <= src_end) : (src_ind >= src_end);
|
||||||
|
src_ind += src_step, dest_ind += dest_step
|
||||||
|
) {
|
||||||
|
/* for constant optimization */
|
||||||
|
if (size == 1) {
|
||||||
|
__builtin_memcpy(dest_arr + dest_ind, src_arr + src_ind, 1);
|
||||||
|
} else if (size == 4) {
|
||||||
|
__builtin_memcpy(dest_arr + dest_ind * 4, src_arr + src_ind * 4, 4);
|
||||||
|
} else if (size == 8) {
|
||||||
|
__builtin_memcpy(dest_arr + dest_ind * 8, src_arr + src_ind * 8, 8);
|
||||||
|
} else {
|
||||||
|
/* memcpy for var size, cannot overlap after previous alloca */
|
||||||
|
__builtin_memcpy(dest_arr + dest_ind * size, src_arr + src_ind * size, size);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
/* only dest_step == 1 can we shrink the dest list. */
|
||||||
|
/* size should be ensured prior to calling this function */
|
||||||
|
if (dest_step == 1 && dest_end >= dest_start) {
|
||||||
|
__builtin_memmove(
|
||||||
|
dest_arr + dest_ind * size,
|
||||||
|
dest_arr + (dest_end + 1) * size,
|
||||||
|
(dest_arr_len - dest_end - 1) * size
|
||||||
|
);
|
||||||
|
return dest_arr_len - (dest_end - dest_ind) - 1;
|
||||||
|
}
|
||||||
|
return dest_arr_len;
|
||||||
|
}
|
||||||
|
|
||||||
|
int32_t __nac3_isinf(double x) {
|
||||||
|
return __builtin_isinf(x);
|
||||||
|
}
|
||||||
|
|
||||||
|
int32_t __nac3_isnan(double x) {
|
||||||
|
return __builtin_isnan(x);
|
||||||
|
}
|
||||||
|
|
||||||
|
double tgamma(double arg);
|
||||||
|
|
||||||
|
double __nac3_gamma(double z) {
|
||||||
|
// Handling for denormals
|
||||||
|
// | x | Python gamma(x) | C tgamma(x) |
|
||||||
|
// --- | ----------------- | --------------- | ----------- |
|
||||||
|
// (1) | nan | nan | nan |
|
||||||
|
// (2) | -inf | -inf | inf |
|
||||||
|
// (3) | inf | inf | inf |
|
||||||
|
// (4) | 0.0 | inf | inf |
|
||||||
|
// (5) | {-1.0, -2.0, ...} | inf | nan |
|
||||||
|
|
||||||
|
// (1)-(3)
|
||||||
|
if (__builtin_isinf(z) || __builtin_isnan(z)) {
|
||||||
|
return z;
|
||||||
|
}
|
||||||
|
|
||||||
|
double v = tgamma(z);
|
||||||
|
|
||||||
|
// (4)-(5)
|
||||||
|
return __builtin_isinf(v) || __builtin_isnan(v) ? __builtin_inf() : v;
|
||||||
|
}
|
||||||
|
|
||||||
|
double lgamma(double arg);
|
||||||
|
|
||||||
|
double __nac3_gammaln(double x) {
|
||||||
|
// libm's handling of value overflows differs from scipy:
|
||||||
|
// - scipy: gammaln(-inf) -> -inf
|
||||||
|
// - libm : lgamma(-inf) -> inf
|
||||||
|
|
||||||
|
if (__builtin_isinf(x)) {
|
||||||
|
return x;
|
||||||
|
}
|
||||||
|
|
||||||
|
return lgamma(x);
|
||||||
|
}
|
||||||
|
|
||||||
|
double j0(double x);
|
||||||
|
|
||||||
|
double __nac3_j0(double x) {
|
||||||
|
// libm's handling of value overflows differs from scipy:
|
||||||
|
// - scipy: j0(inf) -> nan
|
||||||
|
// - libm : j0(inf) -> 0.0
|
||||||
|
|
||||||
|
if (__builtin_isinf(x)) {
|
||||||
|
return __builtin_nan("");
|
||||||
|
}
|
||||||
|
|
||||||
|
return j0(x);
|
||||||
|
}
|
||||||
|
|
||||||
|
uint32_t __nac3_ndarray_calc_size(
|
||||||
|
const uint32_t *list_data,
|
||||||
|
uint32_t list_len,
|
||||||
|
uint32_t begin_idx,
|
||||||
|
uint32_t end_idx
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_calc_size_impl(list_data, list_len, begin_idx, end_idx);
|
||||||
|
}
|
||||||
|
|
||||||
|
uint64_t __nac3_ndarray_calc_size64(
|
||||||
|
const uint64_t *list_data,
|
||||||
|
uint64_t list_len,
|
||||||
|
uint64_t begin_idx,
|
||||||
|
uint64_t end_idx
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_calc_size_impl(list_data, list_len, begin_idx, end_idx);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_nd_indices(
|
||||||
|
uint32_t index,
|
||||||
|
const uint32_t* dims,
|
||||||
|
uint32_t num_dims,
|
||||||
|
NDIndex* idxs
|
||||||
|
) {
|
||||||
|
__nac3_ndarray_calc_nd_indices_impl(index, dims, num_dims, idxs);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_nd_indices64(
|
||||||
|
uint64_t index,
|
||||||
|
const uint64_t* dims,
|
||||||
|
uint64_t num_dims,
|
||||||
|
NDIndex* idxs
|
||||||
|
) {
|
||||||
|
__nac3_ndarray_calc_nd_indices_impl(index, dims, num_dims, idxs);
|
||||||
|
}
|
||||||
|
|
||||||
|
uint32_t __nac3_ndarray_flatten_index(
|
||||||
|
const uint32_t* dims,
|
||||||
|
uint32_t num_dims,
|
||||||
|
const NDIndex* indices,
|
||||||
|
uint32_t num_indices
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_flatten_index_impl(dims, num_dims, indices, num_indices);
|
||||||
|
}
|
||||||
|
|
||||||
|
uint64_t __nac3_ndarray_flatten_index64(
|
||||||
|
const uint64_t* dims,
|
||||||
|
uint64_t num_dims,
|
||||||
|
const NDIndex* indices,
|
||||||
|
uint64_t num_indices
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_flatten_index_impl(dims, num_dims, indices, num_indices);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_broadcast(
|
||||||
|
const uint32_t *lhs_dims,
|
||||||
|
uint32_t lhs_ndims,
|
||||||
|
const uint32_t *rhs_dims,
|
||||||
|
uint32_t rhs_ndims,
|
||||||
|
uint32_t *out_dims
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_calc_broadcast_impl(lhs_dims, lhs_ndims, rhs_dims, rhs_ndims, out_dims);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_broadcast64(
|
||||||
|
const uint64_t *lhs_dims,
|
||||||
|
uint64_t lhs_ndims,
|
||||||
|
const uint64_t *rhs_dims,
|
||||||
|
uint64_t rhs_ndims,
|
||||||
|
uint64_t *out_dims
|
||||||
|
) {
|
||||||
|
return __nac3_ndarray_calc_broadcast_impl(lhs_dims, lhs_ndims, rhs_dims, rhs_ndims, out_dims);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_broadcast_idx(
|
||||||
|
const uint32_t *src_dims,
|
||||||
|
uint32_t src_ndims,
|
||||||
|
const NDIndex *in_idx,
|
||||||
|
NDIndex *out_idx
|
||||||
|
) {
|
||||||
|
__nac3_ndarray_calc_broadcast_idx_impl(src_dims, src_ndims, in_idx, out_idx);
|
||||||
|
}
|
||||||
|
|
||||||
|
void __nac3_ndarray_calc_broadcast_idx64(
|
||||||
|
const uint64_t *src_dims,
|
||||||
|
uint64_t src_ndims,
|
||||||
|
const NDIndex *in_idx,
|
||||||
|
NDIndex *out_idx
|
||||||
|
) {
|
||||||
|
__nac3_ndarray_calc_broadcast_idx_impl(src_dims, src_ndims, in_idx, out_idx);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -80,7 +80,7 @@ in rec {
|
||||||
llvm-tools-irrt = pkgs.runCommandNoCC "llvm-tools-irrt" {}
|
llvm-tools-irrt = pkgs.runCommandNoCC "llvm-tools-irrt" {}
|
||||||
''
|
''
|
||||||
mkdir -p $out/bin
|
mkdir -p $out/bin
|
||||||
ln -s ${llvm-nac3}/bin/clang.exe $out/bin/clang-irrt.exe
|
ln -s ${llvm-nac3}/bin/clang++.exe $out/bin/clang++-irrt.exe
|
||||||
ln -s ${llvm-nac3}/bin/llvm-as.exe $out/bin/llvm-as-irrt.exe
|
ln -s ${llvm-nac3}/bin/llvm-as.exe $out/bin/llvm-as-irrt.exe
|
||||||
'';
|
'';
|
||||||
nac3artiq = pkgs.rustPlatform.buildRustPackage {
|
nac3artiq = pkgs.rustPlatform.buildRustPackage {
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue