Can't reuse dpex.func device functions with different signatures

[fcharras]

This used to work with numba_dpex<=0.18.1 but fails with >=0.19:

import numba_dpex as dpex
import dpctl.tensor as dpt
import numpy as np


@dpex.func
def g(array_in, idx, const):
    array_in[idx] = const

@dpex.kernel
def kernel_a(array_in):
    idx = dpex.get_global_id(0)
    g(array_in, idx, np.int64(0))

@dpex.kernel
def kernel_b(array_in):
    idx = dpex.get_global_id(0)
    g(array_in, idx, np.int32(0))  # NB: call with inputs of different types than in kernel_a

   
dtype = np.float32
size = 16
array_in = dpt.zeros(sh=(size,), dtype=dtype)

kernel_a[size, size](array_in)
kernel_b[size, size](array_in)

with numba_dpex>=0.19 this snippet gives the following exception:

<...traceback elided...>
LoweringError: Failed in dpex_nopython mode pipeline (step: Custom Lowerer with auto-offload support)
No definition for lowering <function get_global_id at 0x7f47c2edf700>(uint32,) -> int64

File "<ipython-input-14-e32b2bd5f6a0>", line 17:
def kernel_b(array_in):
    idx = dpex.get_global_id(0)
    ^

During: lowering "idx = call $4load_method.1($const6.2, func=$4load_method.1, args=[Var($const6.2, <ipython-input-14-e32b2bd5f6a0>:17)], kws=(), vararg=None, varkwarg=None, target=None)" at <ipython-input-14-e32b2bd5f6a0> (17)

There's a workaround that consists in ensuring that in both kernels the device function is at different memory locations with a different name, e.g.:

import numba_dpex as dpex
import dpctl.tensor as dpt
import numpy as np


def make_g():
    @dpex.func
    def g(array_in, idx, const):
        array_in[idx] = const
    return g

g = make_g()
@dpex.kernel
def kernel_a(array_in):
    idx = dpex.get_global_id(0)
    g(array_in, idx, np.int64(0))

g_ = make_g()
@dpex.kernel
def kernel_b(array_in):
    idx = dpex.get_global_id(0)
    g_(array_in, idx, np.int32(0))

   
dtype = np.float32
size = 16
array_in = dpt.zeros(sh=(size,), dtype=dtype)

kernel_a[size, size](array_in)
kernel_b[size, size](array_in)

[diptorupd]

@fcharras #877 fixes the issue along with overall improvements to how we cache and specialize func decorated functions. Can you please test the branch and confirm that the issue you were seeing is addressed?

[fcharras]

The cache in #877 looks like it's working. I've checked that there are cache hits here

numba-dpex/numba_dpex/core/kernel_interface/dispatcher.py

Line 580 in 5c16478

device_driver_ir_module, kernel_module_name = artifact

when it's expected.

I have another issue but I don't think it's related to #877 since I checked for cache hit and misses. My user code works fine with numba_dpex==0.19.0. But now it's super slow and some kernels seem to output wrong values, I think it comes from more recent commits on main. I can try to bisect..

[fcharras]

The new problems are not related to #877 but to #876 , reported in #816

[diptorupd]

Thanks @fcharras for your review. I am going ahead and merging #877 and closing this ticket.

I am opening a separate issue to track the performance regression introduced by #816.

[chudur-budur]

But now it's super slow and some kernels seem to output wrong values,

Which kernels are running slow and outputting wrong values?