[FEA] Add support for CPU-GPU agnostic kernel programming

[cchapmanbird]

Is your feature request related to a problem? Please describe.
Currently, top-level kernels cannot be written in a CPU/GPU agnostic format (in particular, the CPU kernel will fail to compile due to references to numba.cuda). This slows development as near-identical CPU and GPU kernels must be written and separately maintained.

Describe the solution you'd like
I would like support for compiling kernels written with GPU logic on the CPU (guarded appropriately with if statements). Similarly, a way to poll whether the kernel is being executed on CPU or GPU would make this even easier, rather than needing to pass an additional argument to indicate to the internal logic if a CPU or GPU is being used.

Describe alternatives you've considered
I have not found any work-arounds for this issue.

[gmarkall]

Numba-CUDA supports calling @jit functions from @cuda.jit functions - when using this mechanism the CUDA-specific logic is guarded from the CPU side not by using if statements, but by the function scope. See https://nvidia.github.io/numba-cuda/user/examples.html#jit-function-cpu-gpu-compatibility for an example.

Does this mechanism provide a solution for your use case? If not, can you give some example code to illustrate the limitation?

[cchapmanbird]

Indeed, this is useful but only if you can divide things on the level of an individual thread, which is not always possible. For example, I am not sure how one would handle use of shared memory in this context, as we cannot offload the shared memory call to the per-thread level (as it is inherently block-level). For example,

def fill_mem(arr):
    if <cpu>:
        input_arr=np.zeros(5000)
        start = 0
        end = 5000
        step = 1
    elif <gpu>:
        input_arr = cuda.shared.array(5000)   # also, what if this is dynamically-sized?
        start = cuda.threadIdx.x
        end = 5000
        step = cuda.blockDim.x

    for i in range(start, end, step):
        input_arr[i] = arr[i]        
    
    ...... # rest of the code 

[gmarkall]

This example is not exactly based on your pseudocode as I'm not sure where the boundaries between the host and kernel and/or jitted code are intended to be, but this demonstrates a way of keeping the shared memory allocation inside the kernel and passing it into the function containing the common logic:

from numba import cuda, njit, np
import numpy as np


N = 128
dtype = np.float32

@njit
def reverse_array(dest, src, tid):
    rtid = len(dest) - tid - 1
    dest[rtid] = src[tid]


@cuda.jit
def reverse_kernel(dest, src):
    staged = cuda.shared.array(N, dtype)

    tid = cuda.threadIdx.x

    if tid >= len(dest):
        return

    staged[tid] = src[tid]
    cuda.syncthreads()

    reverse_array(dest, src, tid)


@njit
def reverse_cpu(dest, src):
    for i in range(len(dest)):
        reverse_array(dest, src, i)


original = np.arange(N, dtype=dtype)
dest_gpu = np.empty_like(original)
dest_cpu = np.empty_like(original)

reverse_kernel[1, N](dest_gpu, original)
reverse_cpu(dest_cpu, original)


print(original)
print(dest_gpu)
print(dest_cpu)

np.testing.assert_equal(dest_gpu, original[::-1])
np.testing.assert_equal(dest_cpu, dest_gpu)

This gives:

[  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17
  18  19  20  21  22  23  24  25  26  27  28  29  30  31  32  33  34  35
  36  37  38  39  40  41  42  43  44  45  46  47  48  49  50  51  52  53
  54  55  56  57  58  59  60  61  62  63  64  65  66  67  68  69  70  71
  72  73  74  75  76  77  78  79  80  81  82  83  84  85  86  87  88  89
  90  91  92  93  94  95  96  97  98  99 100 101 102 103 104 105 106 107
 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125
 126 127]
[127 126 125 124 123 122 121 120 119 118 117 116 115 114 113 112 111 110
 109 108 107 106 105 104 103 102 101 100  99  98  97  96  95  94  93  92
  91  90  89  88  87  86  85  84  83  82  81  80  79  78  77  76  75  74
  73  72  71  70  69  68  67  66  65  64  63  62  61  60  59  58  57  56
  55  54  53  52  51  50  49  48  47  46  45  44  43  42  41  40  39  38
  37  36  35  34  33  32  31  30  29  28  27  26  25  24  23  22  21  20
  19  18  17  16  15  14  13  12  11  10   9   8   7   6   5   4   3   2
   1   0]
[127 126 125 124 123 122 121 120 119 118 117 116 115 114 113 112 111 110
 109 108 107 106 105 104 103 102 101 100  99  98  97  96  95  94  93  92
  91  90  89  88  87  86  85  84  83  82  81  80  79  78  77  76  75  74
  73  72  71  70  69  68  67  66  65  64  63  62  61  60  59  58  57  56
  55  54  53  52  51  50  49  48  47  46  45  44  43  42  41  40  39  38
  37  36  35  34  33  32  31  30  29  28  27  26  25  24  23  22  21  20
  19  18  17  16  15  14  13  12  11  10   9   8   7   6   5   4   3   2
   1   0]

I used only one thread block to keep things simple here, but multiple thread blocks could be used.

For the CPU equivalent, I skipped the shared memory and passed directly the source in, but it would also be possible to allocate a temporary array using a NumPy constructor function. Instead of using threads I kept the loop sequential - it may be that a prange() would be possible here but I'm not very experienced with that so I haven't attempted to use it.