Method codegen_opencl

src/jit/impl/tiny_opencl/codegen_opencl.cpp:95–194

}  // anonymous namespace

std::pair<std::string, std::string> mgb::jit::codegen_opencl(
        const InternalGraph& internal_graph, const JITExecutor::Args& args) {
    std::string opencl_kernel = R"(
__kernel void {{KERNEL_NAME}} (
        {{KERNEL_SRC_ARGS}}
        __write_only image2d_t dst,
        __private const int global_size_dim0,
        __private const int global_size_dim1,
        __private const int wc_size,
        __private const int hb_size,
        __private const int h,
        __private const uint w_size
        ) {
    #if OPENCL_ENABLE_FP16
    #pragma OPENCL EXTENSION cl_khr_fp16 : enable
    #endif

    const sampler_t SAMPLER = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP;

    int wc = get_global_id(0);
    int hb = get_global_id(1);

#ifndef NON_UNIFORM_WORK_GROUP
    if (wc >= global_size_dim0 || hb >= global_size_dim1)
        return;
#endif

    for (; hb < hb_size; hb += global_size_dim1) {
        for (; wc < wc_size; wc += global_size_dim0) {
            int2 coord = (int2)(wc, hb);
            int2 coord_b = (int2)(wc / w_size, hb/h);
            {{INTERNAL_DECL_EXPRS}}
            {{ASSIGN_EXPRS}}
            {{INTERNAL_ASSIGN_EXPRS}}
            {{WRITE_IMAGE}}(dst, coord, {{EXP}});
        }
        wc = get_global_id(0);
    }
}
    )";

    auto input_dtype = args.inputs[0].layout.dtype;
    for (size_t i = 0; i < args.inputs.size(); i++) {
        mgb_assert(
                args.inputs[i].layout.dtype == input_dtype,
                "OpenCL jit all oprs should have same dtype");
    }
    mgb_assert(
            args.outputs.size() == 1 && args.outputs[0].layout.dtype == input_dtype,
            "output size should be 1 and output dtype should be same with input");
    mgb_assert(
            dtype::Float16() == input_dtype || dtype::Float32() == input_dtype,
            "OpenCL jit dtype only support float32 or float16, %s not support",
            input_dtype.name());
    auto is_half = dtype::Float16() == input_dtype;

    VarNode2AST var2ast;
    str_util::StrReplaceMap source_replace_map;