| 34 | } |
| 35 | |
| 36 | std::string MatmulKernel::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::MatrixMulKernel); |
| 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* b, Memref* c); |
| 64 | |
| 65 | )"); |
| 66 | |
| 67 | writer << GenCommonRet() << " "; |
| 68 | writer << GetKernelSignature(context); |
| 69 | |
| 70 | std::string unary_str = R"({ |
| 71 | ${specifier}* input_a = (${specifier}*)inputs[0]->ptr; |
| 72 | TINYNN_ASSERT(input_a); |
| 73 | ${specifier}* input_b = (${specifier}*)inputs[1]->ptr; |
| 74 | TINYNN_ASSERT(input_b); |
| 75 | ${specifier}* output_data = (${specifier}*)outputs[0]->ptr; |
| 76 | TINYNN_ASSERT(output_data); |
| 77 | |
| 78 | |
| 79 | Memref ref_a; |
| 80 | const Layout src_layout_a = inputs[0]->layout; |
| 81 | ${fill_memref(ref_a, src_layout_a, input_a)} |
| 82 | |
| 83 | Memref ref_b; |
| 84 | const Layout src_layout_b = inputs[1]->layout; |
| 85 | ${fill_memref(ref_b, src_layout_b, input_b)} |
| 86 | |
| 87 | Memref ref_c; |
| 88 | const Layout dst_layout = outputs[0]->layout; |
| 89 | ${fill_memref(ref_c, dst_layout, output_data)} |
| 90 | |
| 91 | Layout out_layout = outputs[0]->layout; |
| 92 | size_t nr_elem = 1; |
| 93 | for (size_t i = 0; i < out_layout.nr_dim; ++i) { |
nothing calls this directly
no test coverage detected