static inline ndt::type get_tagged_dims_from_type(intptr_t ndim, const ndt::type &tp, intptr_t *out_tagged_dims)
{
ndt::type dtp = tp.without_memory_type();
for (int i = 0; i < ndim; ++i) {
switch (dtp.get_id()) {
case fixed_dim_id:
if (dtp.is_symbolic()) {
out_tagged_dims[i] = -2;
}
else {
out_tagged_dims[i] = dtp.extended<ndt::fixed_dim_type>()->get_fixed_dim_size();
}
break;
case var_dim_id:
out_tagged_dims[i] = -1;
break;
default: {
stringstream ss;
ss << "dim_fragment_type failed to get shape from type " << tp;
throw type_error(ss.str());
}
}
dtp = dtp.extended<ndt::base_dim_type>()->get_element_type();
}
return dtp;
}
intptr_t dynd::make_cuda_from_device_builtin_type_assignment_kernel(
const callable_type_data *DYND_UNUSED(self),
const ndt::callable_type *DYND_UNUSED(af_tp), char *DYND_UNUSED(data),
void *ckb, intptr_t ckb_offset, const ndt::type &dst_tp,
const char *DYND_UNUSED(dst_arrmeta), intptr_t DYND_UNUSED(nsrc),
const ndt::type *src_tp, const char *const *DYND_UNUSED(src_arrmeta),
kernel_request_t kernreq, const eval::eval_context *ectx,
const nd::array &DYND_UNUSED(kwds),
const std::map<std::string, ndt::type> &DYND_UNUSED(tp_vars))
{
assign_error_mode errmode = ectx->errmode;
if (errmode != assign_error_nocheck &&
is_lossless_assignment(dst_tp, *src_tp)) {
errmode = assign_error_nocheck;
}
if (!dst_tp.is_builtin() || !src_tp->is_builtin() ||
errmode == assign_error_default) {
stringstream ss;
ss << "cannot assign from CUDA device with types " << *src_tp << " to "
<< dst_tp;
throw type_error(ss.str());
}
nd::cuda_device_to_host_assign_ck::make(ckb, kernreq, ckb_offset,
src_tp->get_data_size());
return make_builtin_type_assignment_kernel(
ckb, ckb_offset, dst_tp.without_memory_type().get_type_id(),
src_tp->without_memory_type().get_type_id(), kernel_request_single,
errmode);
}
static inline bool broadcast_tagged_dims_from_type(intptr_t ndim, ndt::type tp, const intptr_t *tagged_dims,
intptr_t *out_tagged_dims) {
tp = tp.without_memory_type();
for (intptr_t i = 0; i < ndim; ++i) {
intptr_t tagged_dim = tagged_dims[i], dim_size;
switch (tp.get_id()) {
case fixed_dim_kind_id: {
if (tagged_dim < 0) {
out_tagged_dims[i] = -2;
}
} break;
case fixed_dim_id: {
dim_size = tp.extended<ndt::fixed_dim_type>()->get_fixed_dim_size();
if (tagged_dim < 0 || tagged_dim == 1) {
out_tagged_dims[i] = dim_size;
} else if (tagged_dim != dim_size && dim_size != 1) {
return false;
}
} break;
case var_dim_id:
// All broadcasting is done dynamically for var
break;
default: {
stringstream ss;
ss << "dim_fragment_type failed to get shape from type " << tp;
throw type_error(ss.str());
}
}
tp = tp.extended<ndt::base_dim_type>()->get_element_type();
}
return true;
}
static void resolve_dst_type(char *static_data, char *DYND_UNUSED(data), ndt::type &dst_tp, intptr_t nsrc,
const ndt::type *src_tp, intptr_t nkwd, const dynd::nd::array *kwds,
const std::map<std::string, ndt::type> &tp_vars)
{
base_callable *child = reinterpret_cast<callable *>(static_data)->get();
const ndt::callable_type *child_tp = reinterpret_cast<callable *>(static_data)->get_type();
if (child->resolve_dst_type != NULL) {
child->resolve_dst_type(child->static_data(), NULL, dst_tp, nsrc, src_tp, nkwd, kwds, tp_vars);
} else {
dst_tp = ndt::substitute(child_tp->get_return_type(), tp_vars, false);
}
ndt::type tp = dst_tp.without_memory_type();
for (intptr_t i = nsrc - 1; i >= 0; --i) {
if (!src_tp[i].without_memory_type().is_scalar()) {
tp = src_tp[i].without_memory_type().with_replaced_dtype(tp);
}
}
if (dst_tp.get_kind() == memory_kind) {
dst_tp = dst_tp.extended<ndt::base_memory_type>()->with_replaced_storage_type(tp);
} else {
dst_tp = tp;
}
}
size_t dynd::ndt::get_cuda_device_data_alignment(const ndt::type &tp) {
if (tp.is_symbolic()) {
return 0;
}
const ndt::type &dtp = tp.without_memory_type().get_dtype();
if (dtp.is_builtin()) {
return dtp.get_data_size();
} else {
// TODO: Return the data size of the largest built-in component
return 0;
}
}
