Method codegen_cuda

src/jit/impl/nvrtc/codegen_cuda.cpp:117–252

}  // anonymous namespace

std::pair<std::string, std::string> mgb::jit::codegen_cuda(
        const InternalGraph& internal_graph, const JITExecutor::Args& args,
        bool copy_param_to_dev) {
    std::string cuda_kernel =
            R"(
#include <cuda_fp16.h>

struct Uint32Fastdiv {
    unsigned int m_mul, m_divisor, m_divisor_is_not_1, m_inc_dividend, m_shift;

    static const unsigned int MAX_DIVIDEND = ~0u - 1;
};

template <int ndim>
struct ParamElemVisitor {
    int m_stride[ndim];

    //! m_shape_highdim[i] = original_shape[i + 1]
    Uint32Fastdiv m_shape_highdim[ndim > 1 ? ndim - 1 : 1];
    static const int NDIM = ndim;
};

struct Data {
    void* inputs[{{NR_INPS}}];
    {{OUTPUT_DTYPE}}* output;
};

struct PEVisitors {
    ParamElemVisitor<{{NDIM}}> m[{{NR_INPS}}];
};

template<typename T>
static __forceinline__ __device__ T jit_log_sum_exp(T x, T y) {
    T a, b;
    a = x < y ? x : y;
    b = x < y ? y : x;
    return T(b + log1pf(expf(a - b)));
}

)";

    cuda_kernel += copy_param_to_dev ? R"(
extern "C" __global__ void {{KERNEL_NAME}} (Data* data_ptr, size_t num_elements, PEVisitors* visitors_ptr) {
    Data data = *data_ptr;
    PEVisitors visitors = *visitors_ptr;
)"
                                     : R"(
extern "C" __global__ void {{KERNEL_NAME}} (Data data, size_t num_elements,
 PEVisitors visitors) { )";

    cuda_kernel += R"(
    unsigned int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned int delta = blockDim.x * gridDim.x;
    unsigned int tmp_idx;

    {{DECL_EXPRS}}
    {{INTERNAL_DECL_EXPRS}}
    {{DECL_fastdiv_offset}}