| 21 | } |
| 22 | |
| 23 | std::string ConcatKernel::GetKernelBody(TContext* context) const { |
| 24 | std::stringstream writer; |
| 25 | int axis = context->getAttrInt("axis"); |
| 26 | int dtype_size = Utils::get_dtype_size(context->getAttrOprand("operand:0").dtype); |
| 27 | std::string ctype = Utils::get_common_dtype_specifier(dtype_size); |
| 28 | writer << R"( |
| 29 | #include "tensor_util.h" |
| 30 | )"; |
| 31 | writer << GenCommonRet() << " "; |
| 32 | writer << GetKernelSignature(context); |
| 33 | // clang-format off |
| 34 | auto temp_body = R"({ |
| 35 | Tensor* output = outputs[0]; |
| 36 | int axis = ${axis}; |
| 37 | uint32_t dtype_len = ${dtype_size}; |
| 38 | int32_t axis_index = 0; |
| 39 | for (int i = 0; i < nr_input; i++) { |
| 40 | Tensor slice = *output; |
| 41 | slice.layout.dims[axis]= inputs[i]->layout.dims[axis]; |
| 42 | uint32_t offset = axis_index * output->layout.stride[axis]; |
| 43 | slice.ptr += offset * dtype_len; |
| 44 | axis_index += inputs[i]->layout.dims[axis]; |
| 45 | size_t nr_elem = 1; |
| 46 | for (int j = 0; j < inputs[i]->layout.nr_dim; ++j) { |
| 47 | nr_elem *= inputs[i]->layout.dims[j]; |
| 48 | } |
| 49 | if (is_contiguous(slice.layout)) { |
| 50 | memcpy(slice.ptr, inputs[i]->ptr, nr_elem * dtype_len); |
| 51 | } else { |
| 52 | NoconIter src_iter = init_iter(inputs[i]->layout); |
| 53 | NoconIter dst_iter = init_iter(slice.layout); |
| 54 | |
| 55 | ${ctype}* dst_data = slice.ptr; |
| 56 | ${ctype}* src_data = inputs[i]->ptr; |
| 57 | for (size_t j = 0; j < nr_elem; ++j) { |
| 58 | dst_data[dst_iter.offset] = src_data[src_iter.offset]; |
| 59 | inc_iter(inputs[i]->layout, &src_iter); |
| 60 | inc_iter(slice.layout, &dst_iter); |
| 61 | } |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | return TinyNN_SUCCESS; |
| 66 | |
| 67 | })"; |
| 68 | // clang-format on |
| 69 | writer << StringTemplate::StringTemplateArgs() |
| 70 | .add("axis", axis) |
| 71 | .add("ctype", ctype) |
| 72 | .add("dtype_size", dtype_size) |
| 73 | .render(temp_body); |
| 74 | return writer.str(); |
| 75 | } |
| 76 | |
| 77 | } // namespace BareMetal |
| 78 | } // namespace KernelGen |
nothing calls this directly
no test coverage detected