50 const ITensorInfo *output)
60 if (output->total_size() != 0)
72 shift_from_axis_and_replace_coordinate(
const Coordinates &
id, uint32_t axis, uint32_t idx_input, uint32_t num_dims)
74 Coordinates id_out = id;
75 for (uint32_t i = num_dims; i > axis; --i)
77 id_out.set(i,
id[i - 1]);
79 id_out.set(axis, idx_input);
83 void elementwise_stack(
const std::vector<ITensor *> &
input, ITensor *output, uint32_t axis,
const Window &window)
86 window_out.use_tensor_dimensions(output->info()->tensor_shape());
88 const int32_t num_tensors =
input.size();
89 const size_t element_size =
input[0]->info()->element_size();
90 const uint32_t num_dims =
static_cast<uint32_t
>(
input[0]->info()->num_dimensions());
92 for (int32_t idx_input = 0; idx_input < num_tensors; ++idx_input)
94 Iterator input_it(
input[idx_input], window);
98 [&](
const Coordinates &
id)
100 Coordinates id_out = shift_from_axis_and_replace_coordinate(
id, axis, idx_input, num_dims);
101 std::memcpy(output->ptr_to_element(id_out), input_it.ptr(), element_size);
107 void memcpy_stack(
const std::vector<ITensor *> &
input, ITensor *output, uint32_t axis,
const Window &window)
109 const int32_t element_size =
input[0]->info()->element_size();
110 const int32_t chunk_size =
input[0]->info()->tensor_shape().total_size_lower(axis) * element_size;
111 const int32_t num_tensors =
input.size();
112 const int32_t out_chunk_step = chunk_size * num_tensors;
114 const int32_t start_x = window.x().start();
115 const int32_t end_x = window.x().end();
116 const int32_t start_y = window.y().start();
117 const int32_t end_y = window.y().end();
119 uint8_t *out_ptr_base = output->buffer() + output->info()->offset_first_element_in_bytes() + start_x * chunk_size;
121 for (int32_t x = start_x; x < end_x; ++x)
123 const uint8_t *in_ptr =
124 input[x]->buffer() +
input[x]->info()->offset_first_element_in_bytes() + start_y * chunk_size;
125 uint8_t *out_ptr = out_ptr_base + start_y * out_chunk_step;
127 for (int32_t y = start_y; y < end_y; ++y)
129 std::memcpy(out_ptr, in_ptr, chunk_size);
131 in_ptr += chunk_size;
132 out_ptr += out_chunk_step;
135 out_ptr_base += chunk_size;
149 const int32_t num_tensors =
input.size();
152 const uint32_t rank =
input[0]->info()->num_dimensions();
155 for (int32_t i = 0; i < num_tensors; ++i)
171 const int32_t num_tensors =
input.size();
172 const size_t rank =
input[0]->num_dimensions();
174 for (int32_t i = 0; i < num_tensors; ++i)
187 const int32_t num_dims =
input_info->num_dimensions();
188 const int32_t num_tensors = _input.size();
191 bool has_padding =
false;
192 for (
const ITensor *in : _input)
194 if (
has_holes(*in->info(), num_dims - 1))
201 has_padding = has_padding ||
has_holes(*_output->
info(), num_dims);
206 _stack_fn = memcpy_stack;
214 _stack_fn = elementwise_stack;
218 INEKernel::configure(win);
227 _stack_fn(_input, _output, _axis,
window);