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

Method GetKernelBody

compiler/lib/KernelGen/BareMetal/ElemwiseMultiType.cpp:73–164  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

71}
72
73std::string ElemwiseMultiTypeKernel::GetKernelBody(TContext* context) const {
74 auto mode = context->getAttrStr("mode");
75 std::stringstream writer;
76 writer << "#include <math.h> \n";
77 writer << gen_dep(mode);
78 writer << GenCommonRet() << " ";
79 writer << GetKernelSignature(context);
80 if (context->getAttrInt("nr_operands") == 2) {
81 auto op0 = context->getAttrOprand("operand:0");
82 auto dst = context->getAttrOprand("operand:1");
83 CC_ASSERT(
84 Utils::is_quant_dtype(op0.dtype, 8) &&
85 Utils::is_quant_dtype(dst.dtype, 8));
86 auto op0_specifier = Utils::cvt_dtype_specifier(op0.dtype);
87 auto dst_specifier = Utils::cvt_dtype_specifier(dst.dtype);
88 std::string binary_str = R"({
89 ${op0_specifier}* input_0 = (${op0_specifier}*)inputs[0]->ptr;
90 float scale_0 = inputs[0]->dtype.param.scale;
91 TINYNN_ASSERT(input_0);
92 ${dst_specifier}* output_data = (${dst_specifier}*)outputs[0]->ptr;
93 float scale_dst = outputs[0]->dtype.param.scale;
94 TINYNN_ASSERT(output_data);
95 float scale_div = 1.f / scale_dst;
96
97 Layout in_layout = inputs[0]->layout;
98 size_t nr_elem = 1;
99 for (int i = 0; i < in_layout.nr_dim; ++i) {
100 nr_elem *= in_layout.dims[i];
101 }
102 for(size_t i = 0; i < nr_elem; ++i){
103 ${op0_specifier} val_0 = input_0[i];
104 ${act};
105 output_data[i] = out_val;
106 }
107 return TinyNN_SUCCESS;
108 }
109 )";
110 writer << StringTemplate::StringTemplateArgs()
111 .add("op0_specifier", op0_specifier)
112 .add("dst_specifier", dst_specifier)
113 .add("act", gen_unary(mode))
114 .render(binary_str);
115 } else if (context->getAttrInt("nr_operands") == 3) {
116 auto op0 = context->getAttrOprand("operand:0");
117 auto op1 = context->getAttrOprand("operand:1");
118 auto dst = context->getAttrOprand("operand:2");
119 CC_ASSERT(
120 Utils::is_quant_dtype(op0.dtype, 8) &&
121 Utils::is_quant_dtype(op1.dtype, 8) &&
122 Utils::is_quant_dtype(dst.dtype, 8));
123 CC_ASSERT(op0.shape == op1.shape) << "no support broadcast\n";
124 auto op0_specifier = Utils::cvt_dtype_specifier(op0.dtype);
125 auto op1_specifier = Utils::cvt_dtype_specifier(op1.dtype);
126 auto dst_specifier = Utils::cvt_dtype_specifier(dst.dtype);
127 std::string binary_str = R"({
128 ${op0_specifier}* input_0 = (${op0_specifier}*)inputs[0]->ptr;
129 float scale_0 = inputs[0]->dtype.param.scale;
130 TINYNN_ASSERT(input_0);

Callers

nothing calls this directly

Calls 11

GenCommonRetFunction · 0.85
is_quant_dtypeFunction · 0.85
cvt_dtype_specifierFunction · 0.85
StringTemplateArgsClass · 0.85
getAttrStrMethod · 0.80
getAttrOprandMethod · 0.80
gen_depFunction · 0.70
gen_unaryFunction · 0.70
gen_binaryFunction · 0.70
renderMethod · 0.45
addMethod · 0.45

Tested by

no test coverage detected