ndt::type resolve(base_callable *DYND_UNUSED(caller), char *DYND_UNUSED(data), call_graph &cg,
const ndt::type &dst_tp, size_t nsrc, const ndt::type *src_tp, size_t nkwd, const array *kwds,
const std::map<std::string, ndt::type> &tp_vars) {
cg.emplace_back([](kernel_builder &kb, kernel_request_t kernreq, char *DYND_UNUSED(data),
const char *dst_arrmeta, size_t nsrc, const char *const *src_arrmeta) {
kb.emplace_back<right_compound_kernel>(kernreq);
const char *child_src_arrmeta[2] = {src_arrmeta[0], dst_arrmeta};
kb(kernreq | kernel_request_data_only, nullptr, dst_arrmeta, nsrc + 1, child_src_arrmeta);
});
ndt::type child_src_tp[2] = {src_tp[0], dst_tp};
m_child->resolve(this, nullptr, cg, dst_tp, nsrc + 1, child_src_tp, nkwd, kwds, tp_vars);
return dst_tp;
}
ndt::type resolve(base_callable *DYND_UNUSED(caller), char *DYND_UNUSED(data), call_graph &cg,
const ndt::type &dst_tp, size_t nsrc, const ndt::type *src_tp, size_t nkwd, const array *kwds,
const std::map<std::string, ndt::type> &tp_vars) {
cg.emplace_back([](kernel_builder &kb, kernel_request_t kernreq, char *data, const char *dst_arrmeta, size_t nsrc,
const char *const *src_arrmeta) {
kb.emplace_back<where_kernel>(
kernreq, data, reinterpret_cast<const ndt::var_dim_type::metadata_type *>(dst_arrmeta)->stride,
reinterpret_cast<const ndt::var_dim_type::metadata_type *>(dst_arrmeta)->blockref);
kb(kernel_request_single, nullptr, nullptr, nsrc - 1, src_arrmeta);
});
m_child->resolve(this, nullptr, cg, ndt::make_type<bool>(), nsrc - 1, src_tp, nkwd, kwds, tp_vars);
return dst_tp;
}
ndt::type resolve(base_callable *DYND_UNUSED(caller), char *data, call_graph &cg, const ndt::type &dst_tp,
size_t nsrc, const ndt::type *src_tp, size_t nkwd, const array *kwds,
const std::map<std::string, ndt::type> &tp_vars) {
const ndt::callable_type *child_tp =
reinterpret_cast<data_type *>(data)->child->get_type().template extended<ndt::callable_type>();
bool first = reinterpret_cast<data_type *>(data)->first;
reinterpret_cast<data_type *>(data)->first = false;
bool state = reinterpret_cast<data_type *>(data)->state;
bool res_ignore = reinterpret_cast<data_type *>(data)->res_ignore;
bool dst_variadic = dst_tp.is_variadic();
// Do a pass through the src types to classify them
bool src_all_strided = true, src_all_strided_or_var = true;
for (size_t i = 0; i < nsrc; ++i) {
intptr_t src_ndim = src_tp[i].get_ndim() - child_tp->get_pos_type(i).get_ndim();
switch (src_tp[i].get_id()) {
case fixed_dim_id:
break;
case var_dim_id:
src_all_strided = false;
break;
default:
// If it's a scalar, allow it to broadcast like
// a strided dimension
if (src_ndim > 0) {
src_all_strided_or_var = false;
}
break;
}
}
bool var_broadcast = !src_all_strided;
for (size_t i = 0; i < nsrc; ++i) {
var_broadcast &= src_tp[i].get_id() == var_dim_id ||
(src_tp[i].get_id() == fixed_dim_id &&
src_tp[i].extended<ndt::fixed_dim_type>()->get_fixed_dim_size() == 1);
}
if ((dst_variadic || (dst_tp.get_id() == fixed_dim_id || res_ignore)) && src_all_strided) {
static callable f = make_callable<elwise_callable<fixed_dim_id, fixed_dim_id, no_traits, N>>();
static callable g = make_callable<elwise_callable<fixed_dim_id, fixed_dim_id, state_traits, N>>();
if (!first && state) {
return g->resolve(this, data, cg, dst_tp, nsrc, src_tp, nkwd, kwds, tp_vars);
} else {
return f->resolve(this, data, cg, dst_tp, nsrc, src_tp, nkwd, kwds, tp_vars);
}
} else if (((dst_variadic) || dst_tp.get_id() == var_dim_id) && (var_broadcast || src_all_strided)) {
static callable f = make_callable<elwise_callable<var_dim_id, fixed_dim_id, no_traits, N>>();
return f->resolve(this, data, cg, dst_tp, nsrc, src_tp, nkwd, kwds, tp_vars);
} else if (src_all_strided_or_var) {
static callable f = make_callable<elwise_callable<fixed_dim_id, var_dim_id, no_traits, N>>();
return f->resolve(this, data, cg, dst_tp, nsrc, src_tp, nkwd, kwds, tp_vars);
}
std::stringstream ss;
ss << "Cannot process lifted elwise expression from (";
for (size_t i = 0; i < nsrc; ++i) {
ss << src_tp[i];
if (i != nsrc - 1) {
ss << ", ";
}
}
ss << ") to " << dst_tp;
throw std::runtime_error(ss.str());
}