Function where

src/backend/opencl/kernel/where.hpp:65–133

template<typename T>
static void where(Param &out, Param &in) {
    uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0]));
    threads_x      = std::min(threads_x, THREADS_PER_GROUP);
    uint threads_y = THREADS_PER_GROUP / threads_x;

    uint groups_x = divup(in.info.dims[0], threads_x * REPEAT);
    uint groups_y = divup(in.info.dims[1], threads_y);

    Param rtmp;
    Param otmp;

    rtmp.info.dims[0] = groups_x;
    otmp.info.dims[0] = in.info.dims[0];

    rtmp.info.strides[0] = 1;
    otmp.info.strides[0] = 1;

    rtmp.info.offset = 0;
    otmp.info.offset = 0;

    for (int k = 1; k < 4; k++) {
        rtmp.info.dims[k]    = in.info.dims[k];
        rtmp.info.strides[k] = rtmp.info.strides[k - 1] * rtmp.info.dims[k - 1];

        otmp.info.dims[k]    = in.info.dims[k];
        otmp.info.strides[k] = otmp.info.strides[k - 1] * otmp.info.dims[k - 1];
    }

    int rtmp_elements = rtmp.info.strides[3] * rtmp.info.dims[3];
    rtmp.data         = bufferAlloc(rtmp_elements * sizeof(uint));

    int otmp_elements = otmp.info.strides[3] * otmp.info.dims[3];
    otmp.data         = bufferAlloc(otmp_elements * sizeof(uint));

    scanFirstLauncher<T, uint, af_notzero_t>(otmp, rtmp, in, false, groups_x,
                                             groups_y, threads_x);

    // Linearize the dimensions and perform scan
    Param ltmp        = rtmp;
    ltmp.info.offset  = 0;
    ltmp.info.dims[0] = rtmp_elements;
    for (int k = 1; k < 4; k++) {
        ltmp.info.dims[k]    = 1;
        ltmp.info.strides[k] = rtmp_elements;
    }

    scanFirst<uint, uint, af_add_t>(ltmp, ltmp);

    // Get output size and allocate output
    uint total;
    getQueue().enqueueReadBuffer(*rtmp.data, CL_TRUE,
                                 sizeof(uint) * (rtmp_elements - 1),
                                 sizeof(uint), &total);

    out.data = bufferAlloc(total * sizeof(uint));

    out.info.dims[0]    = total;
    out.info.strides[0] = 1;