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

Method GetKernelBody

compiler/lib/KernelGen/AutoBareMetal/MatmulKernel.cpp:36–127  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

34}
35
36std::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) {

Callers

nothing calls this directly

Calls 8

cvt_dtype_specifierFunction · 0.85
get_last_operandFunction · 0.85
StringTemplateArgsClass · 0.85
GenCommonRetFunction · 0.85
getAttrOprandMethod · 0.80
GetKernelSymbolMethod · 0.45
renderMethod · 0.45
addMethod · 0.45

Tested by

no test coverage detected