inline void init(const ndt::type& tp0, const char *metadata0, char *data0)
{
m_array_tp = tp0;
m_iter_ndim = m_array_tp.get_ndim();
m_itersize = 1;
if (m_iter_ndim != 0) {
m_iterindex.init(m_iter_ndim);
memset(m_iterindex.get(), 0, sizeof(intptr_t) * m_iter_ndim);
m_itershape.init(m_iter_ndim);
m_array_tp.extended()->get_shape(m_iter_ndim, 0, m_itershape.get(), metadata0);
size_t iterdata_size = m_array_tp.extended()->get_iterdata_size(m_iter_ndim);
m_iterdata = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata) {
throw std::bad_alloc();
}
m_metadata = metadata0;
m_array_tp.iterdata_construct(m_iterdata,
&m_metadata, m_iter_ndim, m_itershape.get(), m_uniform_tp);
m_data = m_iterdata->reset(m_iterdata, data0, m_iter_ndim);
for (size_t i = 0, i_end = m_iter_ndim; i != i_end; ++i) {
m_itersize *= m_itershape[i];
}
} else {
m_iterdata = NULL;
m_uniform_tp = m_array_tp;
m_data = data0;
m_metadata = metadata0;
}
}
// Constructor which creates the output based on the input's broadcast shape
array_iter(const ndt::type& op0_dtype, nd::array& out_op0, const nd::array& op1, const nd::array& op2, const nd::array& op3) {
create_broadcast_result(op0_dtype, op1, op2, op3, out_op0, m_iter_ndim[0], m_itershape);
nd::array ops[4] = {out_op0, op1, op2, op3};
m_array_tp[0] = out_op0.get_type();
m_array_tp[1] = op1.get_type();
m_array_tp[2] = op2.get_type();
m_array_tp[3] = op3.get_type();
m_itersize = 1;
m_iter_ndim[1] = m_array_tp[1].get_ndim();
m_iter_ndim[2] = m_array_tp[2].get_ndim();
m_iter_ndim[3] = m_array_tp[3].get_ndim();
// Allocate and initialize the iterdata
if (m_iter_ndim[0] != 0) {
m_iterindex.init(m_iter_ndim[0]);
memset(m_iterindex.get(), 0, sizeof(intptr_t) * m_iter_ndim[0]);
// The destination iterdata
size_t iterdata_size = m_array_tp[0].get_iterdata_size(m_iter_ndim[0]);
m_iterdata[0] = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata[0]) {
throw std::bad_alloc();
}
m_metadata[0] = out_op0.get_ndo_meta();
m_array_tp[0].iterdata_construct(m_iterdata[0],
&m_metadata[0], m_iter_ndim[0], m_itershape.get(), m_uniform_tp[0]);
m_data[0] = m_iterdata[0]->reset(m_iterdata[0], out_op0.get_readwrite_originptr(), m_iter_ndim[0]);
// The op iterdata
for (int i = 1; i < 4; ++i) {
iterdata_size = m_array_tp[i].get_broadcasted_iterdata_size(m_iter_ndim[i]);
m_iterdata[i] = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata[i]) {
throw std::bad_alloc();
}
m_metadata[i] = ops[i].get_ndo_meta();
m_array_tp[i].broadcasted_iterdata_construct(m_iterdata[i],
&m_metadata[i], m_iter_ndim[i],
m_itershape.get() + (m_iter_ndim[0] - m_iter_ndim[i]), m_uniform_tp[i]);
m_data[i] = m_iterdata[i]->reset(m_iterdata[i], ops[i].get_ndo()->m_data_pointer, m_iter_ndim[0]);
}
for (size_t i = 0, i_end = m_iter_ndim[0]; i != i_end; ++i) {
m_itersize *= m_itershape[i];
}
} else {
for (size_t i = 0; i < 4; ++i) {
m_iterdata[i] = NULL;
m_uniform_tp[i] = m_array_tp[i];
m_data[i] = ops[i].get_ndo()->m_data_pointer;
m_metadata[i] = ops[i].get_ndo_meta();
}
}
}
void dynd::create_broadcast_result(const ndt::type& result_inner_tp,
const nd::array& op0, const nd::array& op1, const nd::array& op2,
nd::array &out, intptr_t& out_ndim, dimvector& out_shape)
{
// Get the shape of the result
shortvector<int> axis_perm;
nd::array ops[3] = {op0, op1, op2};
broadcast_input_shapes(3, ops, out_ndim, out_shape, axis_perm);
out = nd::make_strided_array(result_inner_tp, out_ndim, out_shape.get(),
nd::read_access_flag|nd::write_access_flag, axis_perm.get());
}
array_iter(const nd::array& op0, const nd::array& op1) {
nd::array ops[2] = {op0, op1};
m_array_tp[0] = op0.get_type();
m_array_tp[1] = op1.get_type();
m_itersize = 1;
shortvector<int> axis_perm; // TODO: Use this to affect the iteration order
broadcast_input_shapes(2, ops, m_iter_ndim, m_itershape, axis_perm);
// Allocate and initialize the iterdata
if (m_iter_ndim != 0) {
m_iterindex.init(m_iter_ndim);
memset(m_iterindex.get(), 0, sizeof(intptr_t) * m_iter_ndim);
// The op iterdata
for (int i = 0; i < 2; ++i) {
size_t iter_ndim_i = m_array_tp[i].get_ndim();
size_t iterdata_size = m_array_tp[i].get_broadcasted_iterdata_size(iter_ndim_i);
m_iterdata[i] = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata[i]) {
throw std::bad_alloc();
}
m_metadata[i] = ops[i].get_ndo_meta();
m_array_tp[i].broadcasted_iterdata_construct(m_iterdata[i],
&m_metadata[i], iter_ndim_i,
m_itershape.get() + (m_iter_ndim - iter_ndim_i), m_uniform_tp[i]);
m_data[i] = m_iterdata[i]->reset(m_iterdata[i], ops[i].get_ndo()->m_data_pointer, m_iter_ndim);
}
for (size_t i = 0, i_end = m_iter_ndim; i != i_end; ++i) {
m_itersize *= m_itershape[i];
}
} else {
for (size_t i = 0; i < 2; ++i) {
m_iterdata[i] = NULL;
m_uniform_tp[i] = m_array_tp[i];
m_data[i] = ops[i].get_ndo()->m_data_pointer;
m_metadata[i] = ops[i].get_ndo_meta();
}
}
}
void dynd::broadcast_input_shapes(intptr_t ninputs, const nd::array* inputs,
intptr_t& out_undim, dimvector& out_shape, shortvector<int>& out_axis_perm)
{
// Get the number of broadcast dimensions
intptr_t undim = inputs[0].get_ndim();
for (intptr_t i = 0; i < ninputs; ++i) {
intptr_t candidate_undim = inputs[i].get_ndim();
if (candidate_undim > undim) {
undim = candidate_undim;
}
}
out_undim = undim;
out_shape.init(undim);
out_axis_perm.init(undim);
intptr_t *shape = out_shape.get();
// Fill in the broadcast shape
for (intptr_t k = 0; k < undim; ++k) {
shape[k] = 1;
}
dimvector tmpshape(undim);
for (intptr_t i = 0; i < ninputs; ++i) {
intptr_t input_undim = inputs[i].get_ndim();
inputs[i].get_shape(tmpshape.get());
intptr_t dimdelta = undim - input_undim;
for (intptr_t k = dimdelta; k < undim; ++k) {
intptr_t size = tmpshape[k - dimdelta];
intptr_t itershape_size = shape[k];
if (itershape_size == 1) {
shape[k] = size;
} else if (size < 0) {
// A negative shape value means variable-sized
if (itershape_size > 0) {
shape[k] = -itershape_size;
} else {
shape[k] = -1;
}
} else if (itershape_size >= 0) {
if (size != 1 && itershape_size != size) {
//cout << "operand " << i << ", comparing size " << itershape_size << " vs " << size << "\n";
throw broadcast_error(ninputs, inputs);
}
} else { // itershape_size < 0
if (itershape_size == -1 && size > 0) {
shape[k] = -size;
} else if (size > 1 && itershape_size != -size) {
throw broadcast_error(ninputs, inputs);
}
}
}
}
// Fill in the axis permutation
if (undim > 1) {
int *axis_perm = out_axis_perm.get();
// TODO: keeporder behavior, currently always C order
for (intptr_t i = 0; i < undim; ++i) {
axis_perm[i] = int(undim - i - 1);
}
} else if (undim == 1) {
out_axis_perm[0] = 0;
}
}
inline void init(const ndt::type& tp0, const char *metadata0, char *data0,
const ndt::type& tp1, const char *metadata1, const char *data1)
{
m_array_tp[0] = tp0;
m_array_tp[1] = tp1;
m_itersize = 1;
// The destination shape
m_iter_ndim[0] = m_array_tp[0].get_ndim();
m_itershape.init(m_iter_ndim[0]);
if (m_iter_ndim[0] > 0) {
m_array_tp[0].extended()->get_shape(m_iter_ndim[0], 0, m_itershape.get(), metadata0);
}
// The source shape
dimvector src_shape;
m_iter_ndim[1] = m_array_tp[1].get_ndim();
src_shape.init(m_iter_ndim[1]);
if (m_iter_ndim[1] > 0) {
m_array_tp[1].extended()->get_shape(m_iter_ndim[1], 0, src_shape.get(), metadata1);
}
// Check that the source shape broadcasts ok
if (!shape_can_broadcast(m_iter_ndim[0], m_itershape.get(),
m_iter_ndim[1], src_shape.get())) {
throw broadcast_error(m_iter_ndim[0], m_itershape.get(),
m_iter_ndim[1], src_shape.get());
}
// Allocate and initialize the iterdata
if (m_iter_ndim[0] != 0) {
m_iterindex.init(m_iter_ndim[0]);
memset(m_iterindex.get(), 0, sizeof(intptr_t) * m_iter_ndim[0]);
// The destination iterdata
size_t iterdata_size = m_array_tp[0].get_iterdata_size(m_iter_ndim[0]);
m_iterdata[0] = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata[0]) {
throw std::bad_alloc();
}
m_metadata[0] = metadata0;
m_array_tp[0].iterdata_construct(m_iterdata[0],
&m_metadata[0], m_iter_ndim[0], m_itershape.get(), m_uniform_tp[0]);
m_data[0] = m_iterdata[0]->reset(m_iterdata[0], data0, m_iter_ndim[0]);
// The source iterdata
iterdata_size = m_array_tp[1].get_broadcasted_iterdata_size(m_iter_ndim[1]);
m_iterdata[1] = reinterpret_cast<iterdata_common *>(malloc(iterdata_size));
if (!m_iterdata[1]) {
throw std::bad_alloc();
}
m_metadata[1] = metadata1;
m_array_tp[1].broadcasted_iterdata_construct(m_iterdata[1],
&m_metadata[1], m_iter_ndim[1],
m_itershape.get() + (m_iter_ndim[0] - m_iter_ndim[1]), m_uniform_tp[1]);
m_data[1] = m_iterdata[1]->reset(m_iterdata[1], const_cast<char *>(data1), m_iter_ndim[0]);
for (size_t i = 0, i_end = m_iter_ndim[0]; i != i_end; ++i) {
m_itersize *= m_itershape[i];
}
} else {
m_iterdata[0] = NULL;
m_iterdata[1] = NULL;
m_uniform_tp[0] = m_array_tp[0];
m_uniform_tp[1] = m_array_tp[1];
m_data[0] = data0;
m_data[1] = const_cast<char *>(data1);
m_metadata[0] = metadata0;
m_metadata[1] = metadata1;
}
}
/**
* The tagged_dims should be interpreted as an array of
* size get_ndim() containing:
* -1 : var
* -2 : strided
* N >= 0 : fixed[N]
*/
inline const intptr_t *get_tagged_dims() const
{
return m_tagged_dims.get();
}
