JIT: illegal Python function (with undefined variable) gives correct kernel by accident

On the current master, the following code compiles and runs:

import cupy as cp
from cupyx import jit

@jit.rawkernel()
def f(b):
    laneId = jit.threadIdx.x & 0x1f
    if laneId == 0:
        value = cp.int64(100)
    b[laneId] = value


b = cp.empty((32,), dtype=cp.int64)
f[1, 32](b)
print(b)
print(f.cached_code)

However, a compilation error (better be NameError or UnboundLocalError) should be raised by the JIT compiler, because the variable value is not well defined for all code paths.

[leofang]

Thanks, @asi1024!

Current CuPy JIT already supports y = numpy.float32(12.) to specify the dtype of a local variable. We can support more additional styles like y = x.dtype.type(1.) or y: numpy.float32 = 1.

Sorry for my slow action…(still on vacation) I just opened an issue to track this feature request (#5417).

For example, do we pass a CuPy array or a NumPy array to the kernel under the simulation mode? Are we able to not generate any CUDA context initialization request when using JIT?

This feature is still in the planning stage, but it will probably not rely on any CUDA context.

Sounds cool! Not that I am a big fan of such simulators, but I am definitely interested in learning how it’s implemented.

[asi1024]

@leofang Sorry for my late response.

Having the ability of accessing array types (ex: b.dtype) in a kernel has additional benefits, such as supporting C++ template-like kernels and allowing user-managed explicit casting. For example, Numba supports statements like y = x.dtype.type(1.) and y = numba.float32(12.) in a kernel. I think it’s best to be discussed separately. Let me open a new issue.

Current CuPy JIT already supports y = numpy.float32(12.) to specify the dtype of a local variable. We can support more additional styles like y = x.dtype.type(1.) or y: numpy.float32 = 1.

For example, do we pass a CuPy array or a NumPy array to the kernel under the simulation mode? Are we able to not generate any CUDA context initialization request when using JIT?

This feature is still in the planning stage, but it will probably not rely on any CUDA context.