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

Method GetKernelBody

compiler/lib/KernelGen/BareMetal/ElemwiseKernel.cpp:441–590  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

439}
440
441std::string ElmwiseKernel::GetKernelBody(TContext* context) const {
442 auto op0 = context->getAttrOprand("operand:0");
443 auto src_dtype = op0.dtype;
444 auto specifier = Utils::cvt_dtype_specifier(src_dtype);
445
446 auto mode = context->getAttrStr("mode");
447 std::stringstream writer;
448 writer << "#include <math.h> \n";
449 writer << "#include <stdbool.h> \n";
450 if (specifier == "gi_float16_t")
451 writer << gen_fp16_define();
452 writer << gen_dep_func(mode);
453 writer << GenCommonRet() << " ";
454 writer << GetKernelSignature(context);
455 if (context->getAttrInt("nr_operands") == 2) {
456 std::string unary_str = R"({
457 ${specifier}* input_data = (${specifier}*)inputs[0]->ptr;
458 TINYNN_ASSERT(input_data);
459 ${specifier}* output_data = (${specifier}*)outputs[0]->ptr;
460 TINYNN_ASSERT(output_data);
461
462 Layout in_layout = inputs[0]->layout;
463 size_t nr_elem = 1;
464 for (size_t i = 0; i < in_layout.nr_dim; ++i) {
465 nr_elem *= in_layout.dims[i];
466 }
467 for(size_t i = 0; i < nr_elem; ++i){
468 ${specifier} val = input_data[i];
469 output_data[i] = ${act};
470 }
471 return TinyNN_SUCCESS;
472 }
473 )";
474 writer << StringTemplate::StringTemplateArgs()
475 .add("specifier", specifier)
476 .add("act", gen_unary(mode))
477 .render(unary_str);
478 } else if (context->getAttrInt("nr_operands") == 3) {
479 std::string binary_str = R"({
480 ${specifier}* input_data = (${specifier}*)inputs[0]->ptr;
481 TINYNN_ASSERT(input_data);
482 ${specifier}* input_data_2 = (${specifier}*)inputs[1]->ptr;
483 TINYNN_ASSERT(input_data_2);
484 ${specifier}* output_data = (${specifier}*)outputs[0]->ptr;
485 TINYNN_ASSERT(output_data);
486
487 Layout in_layout_a = inputs[0]->layout;
488 Layout in_layout_b = inputs[1]->layout;
489 Layout out_layout = outputs[0]->layout;
490 size_t nr_elem_a = 1, nr_elem_b = 1, nr_elem = 1;
491
492 for (size_t i = 0; i < in_layout_a.nr_dim; ++i) {
493 nr_elem_a *= in_layout_a.dims[i];
494 }
495 for (size_t i = 0; i < in_layout_b.nr_dim; ++i) {
496 nr_elem_b *= in_layout_b.dims[i];
497 }
498 for (size_t i = 0; i < out_layout.nr_dim; ++i) {

Callers

nothing calls this directly

Calls 13

cvt_dtype_specifierFunction · 0.85
gen_fp16_defineFunction · 0.85
GenCommonRetFunction · 0.85
StringTemplateArgsClass · 0.85
compute_tenary_bodyFunction · 0.85
compute_quater_bodyFunction · 0.85
getAttrOprandMethod · 0.80
getAttrStrMethod · 0.80
gen_dep_funcFunction · 0.70
gen_unaryFunction · 0.70
gen_binaryFunction · 0.70
renderMethod · 0.45

Tested by

no test coverage detected