| 34 | } |
| 35 | |
| 36 | std::string ElmwiseKernel::GetKernelBody(TContext* context) const { |
| 37 | auto op0 = context->getAttrOprand("operand:0"); |
| 38 | auto src_dtype = op0.dtype; |
| 39 | auto specifier = Utils::cvt_dtype_specifier(src_dtype); |
| 40 | |
| 41 | auto dst_operand = Utils::get_last_operand(context); |
| 42 | int nr_dim = dst_operand.shape.size(); |
| 43 | |
| 44 | std::stringstream writer; |
| 45 | writer << "#include <math.h> \n"; |
| 46 | writer << "#include <stdbool.h> \n"; |
| 47 | auto auto_kern = codegen::GenCode(KernelPack::KernType::ElemwiseKernel); |
| 48 | auto func_name = auto_kern->GetKernelSymbol(context); |
| 49 | writer << StringTemplate::StringTemplateArgs() |
| 50 | .add("func_name", func_name) |
| 51 | .add("nr_dim", nr_dim) |
| 52 | .add("specifier", specifier) |
| 53 | .render(R"( |
| 54 | typedef struct MemRef_descriptor_* MemRef_descriptor; |
| 55 | typedef struct MemRef_descriptor_ { |
| 56 | ${specifier}* allocated; |
| 57 | ${specifier}* aligned; |
| 58 | size_t offset; |
| 59 | size_t sizes[${nr_dim}]; |
| 60 | size_t strides[${nr_dim}]; |
| 61 | } Memref; |
| 62 | #define FUNC_NAME _mlir_ciface_${func_name} |
| 63 | extern void FUNC_NAME(Memref* a, Memref* c); |
| 64 | )"); |
| 65 | |
| 66 | writer << GenCommonRet() << " "; |
| 67 | writer << GetKernelSignature(context); |
| 68 | |
| 69 | if (context->getAttrInt("nr_operands") == 2) { |
| 70 | std::string unary_str = R"({ |
| 71 | ${specifier}* input_data = (${specifier}*)inputs[0]->ptr; |
| 72 | TINYNN_ASSERT(input_data); |
| 73 | ${specifier}* output_data = (${specifier}*)outputs[0]->ptr; |
| 74 | TINYNN_ASSERT(output_data); |
| 75 | |
| 76 | Layout in_layout = inputs[0]->layout; |
| 77 | size_t nr_elem = 1; |
| 78 | for (size_t i = 0; i < in_layout.nr_dim; ++i) { |
| 79 | nr_elem *= in_layout.dims[i]; |
| 80 | } |
| 81 | |
| 82 | Memref ref_a; |
| 83 | const Layout src_layout = inputs[0]->layout; |
| 84 | ${fill_memref(ref_a, src_layout, input_data)} |
| 85 | Memref ref_c; |
| 86 | const Layout dst_layout = outputs[0]->layout; |
| 87 | ${fill_memref(ref_c, dst_layout, output_data)} |
| 88 | |
| 89 | FUNC_NAME(&ref_a, &ref_c); |
| 90 | return TinyNN_SUCCESS; |
| 91 | } |
| 92 | )"; |
| 93 | writer << StringTemplate::StringTemplateArgs() |
nothing calls this directly
no test coverage detected