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

Method GetKernelBody

compiler/lib/KernelGen/BareMetal/Relayout.cpp:33–101  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

31}
32
33std::string RelayoutKernel::GetKernelBody(TContext* context) const {
34 std::stringstream ss;
35 auto src_dtype_str = context->getAttrOprand("operand:0").dtype;
36 std::string specifier = Utils::cvt_dtype_specifier(src_dtype_str);
37 int data_size = Utils::get_dtype_size(src_dtype_str);
38 ss << R"(
39 #include <stdbool.h>
40 #include <string.h>
41
42 )";
43 ss << RelayoutHelper::GetLayoutHelper();
44 ss << RelayoutHelper::GetTransposeModule(specifier, data_size);
45 ss << RelayoutHelper::GetNonconMemcpyModule(specifier);
46
47 ss << GenCommonRet() << " " << GetKernelSignature(context);
48 std::string body_temp = R"({
49 const Tensor* src_tensor = inputs[0];
50 const Tensor* dst_tensor = outputs[0];
51 TINYNN_ASSERT(src_tensor->dtype.type_enum == dst_tensor->dtype.type_enum);
52
53 ${specifier}* src_data = (${specifier}*)(src_tensor->ptr);
54 ${specifier}* dst_data = (${specifier}*)(dst_tensor->ptr);
55 TINYNN_ASSERT(src_data);
56 TINYNN_ASSERT(dst_data);
57
58 Layout src_layout = src_tensor->layout;
59 Layout dst_layout = dst_tensor->layout;
60 src_layout = collapse_contiguous(src_layout);
61 dst_layout = collapse_contiguous(dst_layout);
62 size_t nr_elem = 1;
63 for (int i = 0; i < src_layout.nr_dim; ++i) {
64 nr_elem *= src_layout.dims[i];
65 }
66 bool src_contig = is_contiguous(src_layout);
67 bool dst_contig = is_contiguous(dst_layout);
68
69 if (src_contig && dst_contig){
70 memcpy(dst_data, src_data, nr_elem * sizeof(${specifier}));
71 } else {
72 //! compress layout to retain no contig shape
73 //! try transpose opt
74 ${do_transpose}
75 //! try partial copy
76 if(src_contig && copy_check(dst_layout)){
77 memcpy_cont2nocont(dst_data, src_data, src_layout, dst_layout, postive_memcpy);
78 return TinyNN_SUCCESS;
79 }
80 if(dst_contig && copy_check(src_layout)){
81 memcpy_cont2nocont(src_data, dst_data, dst_layout, src_layout, reverse_memcpy);
82 return TinyNN_SUCCESS;
83 }
84 //! fallback to naive
85 NoconIter src_iter = init_iter(src_layout);
86 NoconIter dst_iter = init_iter(dst_layout);
87 for (size_t i = 0; i < nr_elem; ++i){
88 dst_data[dst_iter.offset] = src_data[src_iter.offset];
89 inc_iter(src_layout, &src_iter, src_layout.nr_dim-1);
90 inc_iter(dst_layout, &dst_iter, dst_layout.nr_dim-1);

Callers

nothing calls this directly

Calls 7

cvt_dtype_specifierFunction · 0.85
get_dtype_sizeFunction · 0.85
GenCommonRetFunction · 0.85
StringTemplateArgsClass · 0.85
getAttrOprandMethod · 0.80
renderMethod · 0.45
addMethod · 0.45

Tested by

no test coverage detected