Method GetKernelBody

compiler/lib/KernelGen/Arm/ArmCommon/Reduce.cpp:447–486

}

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;
    writer << "#include<arm_neon.h>\n";
    writer << "static inline float max(float a,float b){return a>b?a:b;}\n";
    writer << "static inline float min(float a,float b){return a<b?a:b;}\n";
    writer << gen_neon_intrin_compat();
    writer << GenCommonRet() << " ";
    writer << GetKernelSignature(context) << "{\n";
    writer << "const size_t axis = " << axis << ";\n";
    // clang-format off
    writer << R"(
    float* src = (float*)inputs[0]->ptr;
    float* dst = (float*)outputs[0]->ptr;
    Layout in_layout = inputs[0]->layout;
    size_t A = 1, B, C = 1;
    for (size_t i = 0; i < axis; ++ i)
        A *= in_layout.dims[i];
    B = in_layout.dims[axis];
    for (size_t i = axis + 1; i < in_layout.nr_dim; ++ i)
        C *= in_layout.dims[i];
    )";
    writer << R"(
    if(C == 1){
        )"
         << generate_reducer<true>(mode) << R"(
    })"
         << R"(else{
        )"
         << generate_reducer<false>(mode) << R"(

    })";
    writer << R"(
        return TinyNN_SUCCESS;
    })";
    // clang-format on
    return writer.str();
}

}  // namespace ArmCommon
}  // namespace KernelGen