Implement a thread-local means to access kernel launch config.

[tpn]

This allows downstream passes, such as rewriting, to access information about the kernel launch for which they have been enlisted to participate.

Posting this as a PR now to get feedback on the overall approach. Assuming this solution is acceptable, I'll follow up with tests and docs.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[tpn]

Will close #280.

[gmarkall]

/ok to test

[gmarkall]

Using the simple benchmark from numba/numba#3003 (comment):

from numba import cuda
import numpy as np


@cuda.jit('void()')
def some_kernel_1():
    return

@cuda.jit('void(float32[:])')
def some_kernel_2(arr1):
    return

@cuda.jit('void(float32[:],float32[:])')
def some_kernel_3(arr1,arr2):
    return

@cuda.jit('void(float32[:],float32[:],float32[:])')
def some_kernel_4(arr1,arr2,arr3):
    return

@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')
def some_kernel_5(arr1,arr2,arr3,arr4):
    return

arr = cuda.device_array(10000, dtype=np.float32)

%timeit some_kernel_1[1, 1]()
%timeit some_kernel_2[1, 1](arr)
%timeit some_kernel_3[1, 1](arr,arr)
%timeit some_kernel_4[1, 1](arr,arr,arr)
%timeit some_kernel_5[1, 1](arr,arr,arr,arr)

On main:

6.6 μs ± 11.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
10.6 μs ± 79.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
14.1 μs ± 65.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.4 μs ± 236 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 38.5 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

On this branch:

8.6 μs ± 70.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
13 μs ± 41.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.2 μs ± 152 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 35.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
24.4 μs ± 44.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

From this crude benchmark, it appears to add 2-4μs per launch, or 17-23% overhead. I don't yet know how we should consider this (or whether the benchmark is appropriate) but I think it's a consideration to keep in mind when thinking about this approach.

[gmarkall]

Specialized kernels cannot be recompiled, so a new launch configuration would not be able to affect the compilation of a new version - so this check could be kept outside the context manager.

[gmarkall]

There seems to be quite some overlap with dispatcher._LaunchConfiguration in this class (an observation at this point - I don't know whether it makes sense to combine them)

[kkraus14]

Using the simple benchmark from numba/numba#3003 (comment):

from numba import cuda
import numpy as np


@cuda.jit('void()')
def some_kernel_1():
    return

@cuda.jit('void(float32[:])')
def some_kernel_2(arr1):
    return

@cuda.jit('void(float32[:],float32[:])')
def some_kernel_3(arr1,arr2):
    return

@cuda.jit('void(float32[:],float32[:],float32[:])')
def some_kernel_4(arr1,arr2,arr3):
    return

@cuda.jit('void(float32[:],float32[:],float32[:],float32[:])')
def some_kernel_5(arr1,arr2,arr3,arr4):
    return

arr = cuda.device_array(10000, dtype=np.float32)

%timeit some_kernel_1[1, 1]()
%timeit some_kernel_2[1, 1](arr)
%timeit some_kernel_3[1, 1](arr,arr)
%timeit some_kernel_4[1, 1](arr,arr,arr)
%timeit some_kernel_5[1, 1](arr,arr,arr,arr)

On main:

6.6 μs ± 11.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
10.6 μs ± 79.9 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
14.1 μs ± 65.8 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.4 μs ± 236 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 38.5 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

On this branch:

8.6 μs ± 70.7 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
13 μs ± 41.3 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
17.2 μs ± 152 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
20.3 μs ± 35.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
24.4 μs ± 44.4 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

From this crude benchmark, it appears to add 2-4μs per launch, or 17-23% overhead. I don't yet know how we should consider this (or whether the benchmark is appropriate) but I think it's a consideration to keep in mind when thinking about this approach.

Kernel launch latency is something we need to care about moving forward and is becoming more and more of a bottleneck in important workloads. Adding 2-4μs per launch for this is probably unacceptable, but that being said our launch latency right now is much higher than we'd like in general where we probably need to rework the entire launch path at some point in the not too distant future.

[gmarkall]

we probably need to rework the entire launch path at some point in the not too distant future.

I think that's another thing that concerns me - if we implement something like this, then it constrains how we can rework the launch path in future if we have to go on supporting it.