| 44 | } |
| 45 | |
| 46 | std::string ArgmaxKernel::GetKernelBody(TContext* context) const { |
| 47 | std::stringstream writer; |
| 48 | size_t axis = context->getAttrInt("axis"); |
| 49 | auto input_dtype = context->getAttrOprand("operand:0").dtype; |
| 50 | auto output_dtype = context->getAttrOprand("operand:1").dtype; |
| 51 | auto input_specifier = Utils::cvt_dtype_specifier(input_dtype); |
| 52 | auto output_specifier = Utils::cvt_dtype_specifier(output_dtype); |
| 53 | writer << "#include <string.h>\n"; |
| 54 | writer << GenCommonRet() << " "; |
| 55 | writer << GetKernelSignature(context) << "{\n"; |
| 56 | writer << "const size_t axis = " << axis << ";\n"; |
| 57 | // clang-format off |
| 58 | writer << StringTemplate::StringTemplateArgs() |
| 59 | .add("input_specifier", input_specifier) |
| 60 | .add("output_specifier", output_specifier) |
| 61 | .add("gen_init", gen_init("MAX", input_dtype)) |
| 62 | .add("gen_comp", gen_comp("MAX")) |
| 63 | .render(R"( |
| 64 | ${input_specifier}* input_data = (${input_specifier}*)inputs[0]->ptr; |
| 65 | ${output_specifier}* output_data = (${output_specifier}*)outputs[0]->ptr; |
| 66 | |
| 67 | Layout in_layout = inputs[0]->layout; |
| 68 | int A = 1, B, C = 1; |
| 69 | for (int i = 0; i < axis; ++ i) |
| 70 | A *= in_layout.dims[i]; |
| 71 | B = in_layout.dims[axis]; |
| 72 | for (int i = axis + 1; i < in_layout.nr_dim; ++ i) |
| 73 | C *= in_layout.dims[i]; |
| 74 | |
| 75 | for (int a = 0; a < A; ++ a) { |
| 76 | for (int c = 0; c < C; ++ c) { |
| 77 | ${input_specifier} best_val = ${gen_init}; |
| 78 | size_t best_arg = 0; |
| 79 | for (int b = 0; b < B; ++ b) { |
| 80 | ${input_specifier} curr_val = input_data[(a * B + b) * C + c]; |
| 81 | if(${gen_comp}){ |
| 82 | best_val = curr_val; |
| 83 | best_arg = b; |
| 84 | } |
| 85 | } |
| 86 | |
| 87 | output_data[a * C + c] = best_arg; |
| 88 | } |
| 89 | } |
| 90 | return TinyNN_SUCCESS; |
| 91 | })"); |
| 92 | // clang-format on |
| 93 | return writer.str(); |
| 94 | } |
| 95 | // vim: syntax=cpp.doxygen |
nothing calls this directly
no test coverage detected