@memoize
def select_block_grid_sizes(dev, data_shape, threads_per_block=None):
    """
    Determine CUDA block and grid dimensions given device constraints.

    Determine the CUDA block and grid dimensions allowed by a GPU
    device that are sufficient for processing every element of an
    array in a separate thread.

    Parameters
    ----------
    d : pycuda.driver.Device
        Device object to be used.
    data_shape : tuple
        Shape of input data array. Must be of length 2.
    threads_per_block : int, optional
        Number of threads to execute in each block. If this is None,
        the maximum number of threads per block allowed by device `d`
        is used.

    Returns
    -------
    block_dim : tuple
        X, Y, and Z dimensions of minimal required thread block.
    grid_dim : tuple
        X and Y dimensions of minimal required block grid.

    Notes
    -----
    Using the scheme in this function, all of the threads in the grid can be enumerated
    as `i = blockIdx.y*max_threads_per_block*max_blocks_per_grid+
    blockIdx.x*max_threads_per_block+threadIdx.x`.

    For 2D shapes, the subscripts of the element `data[a, b]` where `data.shape == (A, B)`
    can be computed as
    `a = i/B`
    `b = mod(i,B)`.

    For 3D shapes, the subscripts of the element `data[a, b, c]` where
    `data.shape == (A, B, C)` can be computed as
    `a = i/(B*C)`
    `b = mod(i, B*C)/C`
    `c = mod(mod(i, B*C), C)`.

    For 4D shapes, the subscripts of the element `data[a, b, c, d]`
    where `data.shape == (A, B, C, D)` can be computed as
    `a = i/(B*C*D)`
    `b = mod(i, B*C*D)/(C*D)`
    `c = mod(mod(i, B*C*D)%(C*D))/D`
    `d = mod(mod(mod(i, B*C*D)%(C*D)), D)`

    It is advisable that the number of threads per block be a multiple
    of the warp size to fully utilize a device's computing resources.
    """

    # Sanity checks:
    if np.isscalar(data_shape):
        data_shape = (data_shape,)

    # Number of elements to process; we need to cast the result of
    # np.prod to a Python int to prevent PyCUDA's kernel execution
    # framework from getting confused when
    N = int(np.prod(data_shape))

    # Get device constraints:
    max_threads_per_block, max_block_dim, max_grid_dim = get_dev_attrs(dev)

    if threads_per_block is not None:
        if max_threads_per_block > threads_per_block:
            raise ValueError('threads per block exceeds device maximum')
        else:
            max_threads_per_block = threads_per_block

    # Actual number of thread blocks needed:
    blocks_needed = iceil(N/float(max_threads_per_block))

    if blocks_needed <= max_grid_dim[0]:
        return (max_threads_per_block, 1, 1), (blocks_needed, 1, 1)
    elif blocks_needed > max_grid_dim[0] and \
         blocks_needed <= max_grid_dim[0]*max_grid_dim[1]:
        return (max_threads_per_block, 1, 1), \
            (max_grid_dim[0], iceil(blocks_needed/float(max_grid_dim[0])), 1)
    elif blocks_needed > max_grid_dim[0]*max_grid_dim[1] and \
         blocks_needed <= max_grid_dim[0]*max_grid_dim[1]*max_grid_dim[2]:
        return (max_threads_per_block, 1, 1), \
            (max_grid_dim[0], max_grid_dim[1],
             iceil(blocks_needed/float(max_grid_dim[0]*max_grid_dim[1])))
    else:
        raise ValueError('array size too large')
dataset/ETHPy150Open lebedov/scikit-cuda/skcuda/misc.py/select_block_grid_sizes/swapped operands of `>`