Method GetKernelBody

compiler/lib/KernelGen/BareMetal/Pooling.cpp:166–273

}  // namespace

std::string PoolingKernel::GetKernelBody(TContext* context) const {
    auto format_str = context->getAttrStr("format");
    auto mode_str = context->getAttrStr("mode");
    auto src_dtype = context->getAttrOprand("operand:0").dtype;
    auto dst_dtype = context->getAttrOprand("operand:1").dtype;
    auto src_specifier = Utils::cvt_dtype_specifier(src_dtype);
    auto dst_specifier = Utils::cvt_dtype_specifier(dst_dtype);
    auto acc_specifier = get_acc_dtype_specifier(src_dtype, mode_str);

    std::stringstream ss;
    const uint32_t window_h = context->getAttrInt("window_h");
    const uint32_t window_w = context->getAttrInt("window_w");
    Pooler pooler(mode_str, window_h * window_w, acc_specifier);
    ss << R"(
#include <math.h>
#include <stdbool.h>
)";
    if (src_specifier == "gi_float16_t")
        ss << gen_fp16_define();
    ss << pooler.gen_dep();
    ss << GenFormatIter::gen_inline_format_iter_body(format_str);
    auto format_iter_symbol = GenFormatIter::gen_inline_format_iter_symbol(format_str);
    ss << GenCommonRet() << " " << GetKernelSignature(context) << "{\n";

    std::string body_temp = R"(
    const int oc_ratio = ${oc_ratio};
    const int ic_ratio = ${ic_ratio};
    const int batch_pos = ${batch_pos};
    const int spatial_start = ${spatial_start};
    const int dst_oc_idx = ${dst_oc_idx};
    const int src_ic_idx = ${src_ic_idx};

    const uint32_t window_h = ${window_h};
    const uint32_t window_w = ${window_w};
    const uint32_t ph = ${pad_h};
    const uint32_t pw = ${pad_w};
    const uint32_t sh = ${stride_h};
    const uint32_t sw = ${stride_w};
    ${src_specifier}* input_data = (${src_specifier}*)inputs[0]->ptr;
    TINYNN_ASSERT(input_data);
    ${dst_specifier}* output_data = (${dst_specifier}*)outputs[0]->ptr;
    TINYNN_ASSERT(output_data);
    const Tensor* src_tensor = inputs[0];
    TINYNN_ASSERT(src_tensor);
    const Tensor* dst_tensor = outputs[0];
    TINYNN_ASSERT(dst_tensor);
    Layout src_layout = inputs[0]->layout;
    const Layout dst_layout = dst_tensor->layout;
    const uint32_t batch = src_layout.dims[batch_pos];
    const uint32_t ih = src_layout.dims[spatial_start];
    const uint32_t iw = src_layout.dims[spatial_start + 1];
    const uint32_t oc = dst_layout.dims[dst_oc_idx] * oc_ratio;
    const uint32_t ic = src_layout.dims[src_ic_idx] * ic_ratio;
    const uint32_t oh = dst_layout.dims[spatial_start];
    const uint32_t ow = dst_layout.dims[spatial_start + 1];
    for (uint32_t batch_idx = 0; batch_idx < batch; ++batch_idx) {
        for (uint32_t oc_idx = 0; oc_idx < oc; ++oc_idx) {
            for (uint32_t oh_idx = 0; oh_idx < oh; ++oh_idx) {