bool try_take_( callable & ca)
{
if ( empty_() )
return false;
callable tmp( queue_.top().ca);
queue_.pop();
ca.swap( tmp);
return ! ca.empty();
}
ndt::type resolve(base_callable *caller, char *DYND_UNUSED(data), call_graph &cg, const ndt::type &dst_tp,
size_t DYND_UNUSED(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) {
size_t self_offset = kb.size();
kb.emplace_back<forward_na_kernel<I...>>(kernreq);
kb(kernel_request_single, nullptr, dst_arrmeta, nsrc, src_arrmeta);
for (intptr_t i : std::array<index_t, sizeof...(I)>({I...})) {
size_t is_na_offset = kb.size() - self_offset;
kb(kernel_request_single, nullptr, nullptr, 1, src_arrmeta + i);
kb.get_at<forward_na_kernel<I...>>(self_offset)->is_na_offset[i] = is_na_offset;
}
size_t assign_na_offset = kb.size() - self_offset;
kb(kernel_request_single, nullptr, nullptr, 0, nullptr);
kb.get_at<forward_na_kernel<I...>>(self_offset)->assign_na_offset = assign_na_offset;
});
ndt::type src_value_tp[2];
for (intptr_t i = 0; i < 2; ++i) {
src_value_tp[i] = src_tp[i];
}
for (intptr_t i : std::array<index_t, sizeof...(I)>({I...})) {
src_value_tp[i] = src_value_tp[i].extended<ndt::option_type>()->get_value_type();
}
base_callable *child;
if (m_child.is_null()) {
child = caller;
} else {
child = m_child.get();
}
ndt::type res_value_tp =
child->resolve(this, nullptr, cg, dst_tp.is_symbolic() ? child->get_ret_type() : dst_tp, 2, src_value_tp,
nkwd, kwds, tp_vars);
for (index_t i : std::array<index_t, sizeof...(I)>({I...})) {
is_na->resolve(this, nullptr, cg, ndt::make_type<bool>(), 1, src_tp + i, 0, nullptr, tp_vars);
}
return assign_na->resolve(this, nullptr, cg, ndt::make_type<ndt::option_type>(res_value_tp), 0, nullptr, nkwd,
kwds, tp_vars);
}
static callable make()
{
const callable self = functional::call<F>(ndt::type("(Any) -> Any"));
for (const std::pair<const type_id_t, callable> &pair :
callable::make_all<K, numeric_type_ids>(0)) {
children[pair.first] = pair.second;
}
for (type_id_t i0 : dim_type_ids()) {
const ndt::type child_tp = ndt::callable_type::make(
self.get_type()->get_return_type(), ndt::type(i0));
children[i0] = functional::elwise(child_tp, self);
}
return functional::multidispatch(self.get_array_type(), children,
default_child);
}
nd::callable nd::functional::neighborhood(const callable &neighborhood_op, const callable &boundary_child)
{
const ndt::callable_type *funcproto_tp = neighborhood_op.get_array_type().extended<ndt::callable_type>();
intptr_t nh_ndim = funcproto_tp->get_pos_type(0).get_ndim();
std::vector<ndt::type> arg_tp(2);
arg_tp[0] = ndt::type("?" + std::to_string(nh_ndim) + " * int");
arg_tp[1] = ndt::type("?" + std::to_string(nh_ndim) + " * int");
return callable::make<neighborhood_kernel<1>>(
ndt::callable_type::make(funcproto_tp->get_pos_type(0).with_replaced_dtype(funcproto_tp->get_return_type()),
funcproto_tp->get_pos_tuple(), ndt::struct_type::make({"shape", "offset"}, arg_tp)),
neighborhood_kernel<1>::static_data_type(neighborhood_op, boundary_child));
}
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());
}
nd::callable nd::functional::outer(const callable &child)
{
return callable::make<outer_ck>(outer_make_type(child.get_type()), child, 0);
}
void swap( value_type & other)
{
ca.swap( other.ca);
std::swap( attr, other.attr);
}
value_type(
callable const& ca_,
attribute_type const& attr_) :
ca( ca_), attr( attr_)
{ BOOST_ASSERT( ! ca.empty() ); }