Function mean_dim_kernel

src/backend/cuda/kernel/mean.hpp:53–166

template<typename Ti, typename Tw, typename To, uint dim, uint DIMY>
__global__ static void mean_dim_kernel(Param<To> out, Param<Tw> owt,
                                       CParam<Ti> in, CParam<Tw> iwt,
                                       uint blocks_x, uint blocks_y,
                                       uint offset_dim) {
    const uint tidx = threadIdx.x;
    const uint tidy = threadIdx.y;
    const uint tid  = tidy * THREADS_X + tidx;

    const uint zid        = blockIdx.x / blocks_x;
    const uint wid        = blockIdx.y / blocks_y;
    const uint blockIdx_x = blockIdx.x - (blocks_x)*zid;
    const uint blockIdx_y = blockIdx.y - (blocks_y)*wid;
    const uint xid        = blockIdx_x * blockDim.x + tidx;
    const uint yid = blockIdx_y;  // yid  of output. updated for input later.

    uint ids[4] = {xid, yid, zid, wid};

    const Ti *iptr  = in.ptr;
    const Tw *iwptr = iwt.ptr;
    To *optr        = out.ptr;
    Tw *owptr       = owt.ptr;

    int ooffset = ids[3] * out.strides[3] + ids[2] * out.strides[2] +
                  ids[1] * out.strides[1] + ids[0];
    // There is only one element per block for out
    // There are blockDim.y elements per block for in
    // Hence increment ids[dim] just after offseting out and before offsetting
    // in
    optr += ooffset;
    if (owptr != NULL) owptr += ooffset;

    const uint blockIdx_dim = ids[dim];

    ids[dim] = ids[dim] * blockDim.y + tidy;

    int ioffset = ids[3] * in.strides[3] + ids[2] * in.strides[2] +
                  ids[1] * in.strides[1] + ids[0];
    iptr += ioffset;
    if (iwptr != NULL) iwptr += ioffset;

    const uint id_dim_in   = ids[dim];
    const uint istride_dim = in.strides[dim];

    bool is_valid = (ids[0] < in.dims[0]) && (ids[1] < in.dims[1]) &&
                    (ids[2] < in.dims[2]) && (ids[3] < in.dims[3]);

    common::Transform<Ti, compute_t<To>, af_add_t> transform;

    compute_t<To> val    = common::Binary<compute_t<To>, af_add_t>::init();
    compute_t<Tw> weight = common::Binary<compute_t<Tw>, af_add_t>::init();

    if (is_valid && id_dim_in < in.dims[dim]) {
        val = transform(*iptr);
        if (iwptr != NULL) {
            weight = *iwptr;
        } else {
            weight = (Tw)1;
        }