Method GetKernelBody

compiler/lib/KernelGen/BareMetal/Reduce.cpp:92–134

}

std::string ReduceKernel::GetKernelBody(TContext* context) const {
    auto mode = context->getAttrStr("mode");
    size_t axis = context->getAttrInt("axis");
    auto input = context->getAttrOprand("operand:0");
    std::stringstream writer;
    auto dtype = context->getAttrOprand("operand:0").dtype;
    auto specifier = Utils::cvt_dtype_specifier(dtype);
    writer << gen_helper(mode, dtype);
    writer << GenCommonRet() << " ";
    writer << GetKernelSignature(context) << "{\n";
    writer << "const size_t axis = " << axis << ";\n";
    writer << StringTemplate::StringTemplateArgs()
                      .add("specifier", specifier)
                      .add("gen_init", gen_init(mode, dtype))
                      .add("gen_apply",
                           gen_apply(mode, "acc", "input_data[i * B * C + j * C + k]"))
                      .add("gen_write", gen_write(mode, "acc", "B"))
                      .render(R"(
        ${specifier}* input_data = (${specifier}*)inputs[0]->ptr;
        ${specifier}* output_data = (${specifier}*)outputs[0]->ptr;

        Layout in_layout = inputs[0]->layout;
        int A = 1, B, C = 1;
        for (int i = 0; i < axis; ++ i)
            A *= in_layout.dims[i];
        B = in_layout.dims[axis];
        for (int i = axis + 1; i < in_layout.nr_dim; ++ i)
            C *= in_layout.dims[i];

        for (int i = 0; i < A; ++ i) {
            for (int k = 0; k < C; ++ k) {
                ${specifier} acc = ${gen_init};
                for (int j = 0; j < B; ++ j) {
                    acc = ${gen_apply};
                }
                
                output_data[i * C + k] = ${gen_write};
            }
        }
        return TinyNN_SUCCESS;
    })");
    return writer.str();
}

// vim: syntax=cpp.doxygen