MCPcopy Create free account
hub / github.com/MegEngine/MegCC / execute_setsubtensor

Function execute_setsubtensor

runtime/src/vm/subtensor.c:241–311  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

239 return TinyNN_SUCCESS;
240}
241static TinyNNStatus execute_setsubtensor(Instruction* inst, VM* vm) {
242 Tensor **inputs = inst->workload.set_subtensor.inputs,
243 *output = inst->workload.set_subtensor.output;
244 SetSubTensor* set_subtensor = &inst->workload.set_subtensor;
245 Tensor dymmy_src = *(inputs[0]);
246 Tensor* src = inputs[0];
247 Tensor* value = inputs[1];
248 const int elem_bytes = dtype_length(value->dtype.type_enum, NULL);
249 output->layout = inputs[0]->layout;
250 //! alloc output
251 size_t length_in_byte = tensor_length_in_byte(src);
252 alloc_tensor(output, vm);
253 //! if src_len == value_len and contig, fast forward
254 bool is_src_contig = is_contiguous(src->layout);
255 bool is_value_contig = is_contiguous(value->layout);
256 size_t length_of_value = tensor_length_in_byte(value);
257 if (is_src_contig && is_value_contig && length_of_value == length_in_byte) {
258 memcpy(output->ptr, value->ptr, length_of_value);
259 return TinyNN_SUCCESS;
260 }
261 //! copy all memory to dst
262 memcpy(output->ptr, src->ptr, length_in_byte);
263 //! deduce output shape, and modify the input stride
264 uint32_t offset = update_layout(
265 inputs, output, &dymmy_src, set_subtensor->descs, set_subtensor->flags,
266 set_subtensor->nr_descs);
267 //! do set_subtensor
268 size_t nr_elem = 1;
269 for (int i = 0; i < output->layout.nr_dim; ++i) {
270 nr_elem *= output->layout.dims[i];
271 }
272
273 NoconIter src_iter = init_iter(value->layout);
274 NoconIter dst_iter = init_iter(output->layout);
275 if (elem_bytes == 1) {
276 char* dst_data = output->ptr;
277 char* src_data = value->ptr;
278 for (size_t i = 0; i < nr_elem; ++i) {
279 dst_data[dst_iter.offset] = src_data[src_iter.offset];
280 inc_iter(value->layout, &src_iter);
281 inc_iter(output->layout, &dst_iter);
282 }
283 } else if (elem_bytes == 2) {
284 int16_t* dst_data = output->ptr;
285 int16_t* src_data = value->ptr;
286 for (size_t i = 0; i < nr_elem; ++i) {
287 dst_data[dst_iter.offset] = src_data[src_iter.offset];
288 inc_iter(value->layout, &src_iter);
289 inc_iter(output->layout, &dst_iter);
290 }
291 } else if (elem_bytes == 4) {
292 int32_t* dst_data = output->ptr;
293 int32_t* src_data = value->ptr;
294 for (size_t i = 0; i < nr_elem; ++i) {
295 dst_data[dst_iter.offset] = src_data[src_iter.offset];
296 inc_iter(value->layout, &src_iter);
297 inc_iter(output->layout, &dst_iter);
298 }

Callers

nothing calls this directly

Calls 8

dtype_lengthFunction · 0.85
tensor_length_in_byteFunction · 0.85
alloc_tensorFunction · 0.85
is_contiguousFunction · 0.85
update_layoutFunction · 0.85
init_iterFunction · 0.85
inc_iterFunction · 0.85
log_tensorFunction · 0.85

Tested by

no test coverage detected