Minimal reproducer for regression in #804

[fcharras]

Performances issues #804 have been reported in #891 and #886

In fact, what I have reported as a performance issue is really a JIT issue, maybe related to caching. Sorry for misleading you into thinking that might be only related to performance. The way it surfaced on the benchmarks was a general slowdown on GPU because rather than segfaulting it returns corrupted results that trigger a slower code branch in our KMeans (cluster relocation). On CPU it gives segfaults, meaning the wrong code was executed, so I could have thought sooner about errors in the compiled code.

I think it might be related to caching, but I have looked for the error in the python depth of numba_dpex (post- #804 ) without any luck, the bug must be well hidden.

I finally came up with a minimal reproducer of what cause the regression in my code:

import numba_dpex as dpex
import dpctl.tensor as dpt


def make_write_values_kernel(n_rows):

    write_values = make_write_values_kernel_func()

    @dpex.kernel
    def write_values_kernel(array_in):
        for row_idx in range(n_rows):
            is_even = (row_idx % 2) == 0
            write_values(array_in, row_idx, is_even)

    return write_values_kernel[1, 1]


def make_write_values_kernel_func():
    write_when_even = _make_write_values_kernel_func(4)
    write_when_odd = _make_write_values_kernel_func(2)

    @dpex.func
    def write_values(array_in, row_idx, is_even):
        if is_even:
            write_when_even(array_in, row_idx)
        else:
            write_when_odd(array_in, row_idx)

    return write_values


def _make_write_values_kernel_func(n_cols):
    @dpex.func
    def write_values(array_in, row_idx):
        for idx in range(n_cols):
            array_in[row_idx, idx] = 1

    return write_values


kernel = make_write_values_kernel(10)
array_in = dpt.empty(sh=(10, 10), dtype=dpt.int64)
kernel(array_in)
print(array_in)

Read it attentively, it should output the following:

[[1 1 1 1 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0 0 0]]

but instead the output is:

[[1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]
 [1 1 1 1 0 0 0 0 0 0]]

(basically, when write_when_odd should be called, it's write_when_even that is called instead, since those two funcs have the same code and the same name it suggests a caching issue)

Important notes:

• it works with numba_dpex == 0.19.0 when using this patch that we apply by monkey-patching in numba_dpex: https://github.com/soda-inria/sklearn-numba-dpex/blob/main/sklearn_numba_dpex/patches/load_numba_dpex.py#L45
• it (probably - haven't tested this specific case but our KMeans didn't have issues then) works with numba_dpex < 0.19.0

In case you wonder why one would come up with such convoluted cascade of dpex.funcs: basically in sklearn_numba_dpex we factor some code using dpex.func kernel functions that are defined as closures and re-instantiated with different parameters every time it's needed. See https://github.com/soda-inria/sklearn-numba-dpex/blob/main/sklearn_numba_dpex/kmeans/kernels/_base_kmeans_kernel_funcs.py . This is the best solution for factoring redundant code that I've thought of (and it only works as long as there are no local memory allocation or barriers in the factored code) ).

Note that #892 doesn't seem to be related at all, since it's already wrong for numba_dpex == 0.19.0.

[chudur-budur]

You can disable caching with NUMBA_DPEX_ENABLE_CACHE=0 python <code>. I did a quick check by disabling the cache and the results are still wrong (all rows 4 zeros). Which means the bug is not coming from the caching mechanism.

Looks like the issue is with @func decorator. We are looking into it.

[diptorupd]

The problem here is caused by the way we generate the function name for any JIT compiled function. The write_values is compiled into multiple versions based on the value of the "free variable" n_cols that is used inside the function (welcome to template specialization in Numba 😄 ).

def _make_write_values_kernel_func(n_cols):
    @dpex.func
    def write_values(array_in, row_idx):
        for idx in range(n_cols):
            array_in[row_idx, idx] = 1

    return write_values

In your example program, we "specialize" the write_values function on two n_cols values. Numba correctly generates two functions and in the Numba IR assigns them separate unique names. However, when the lowering from Numba to LLVM IR happens Numba/numba-dpex is not using the unique name and generates two LLVM functions with the same name. Due to the same names, there is no way to distinguish between them and the linker randomly picks one. Thus, the wrong results.

We have a workaround that will fix it in numba-dpex.

@stuartarchibald In NativeLowering, the FunctionDescriptor uses the qualname instead of the unique_name for the Numba function IR function. Is there a reason for it? To fix the present issue, we are overriding the behaviour and forcing the FunctionDescriptor to use the unique_name by setting the qualname to the unique_name in the state variable. Is there a better way to handle the issue?
I am referring to the following function:
https://github.com/numba/numba/blob/28573bb985cf4816e1c6f54ce2f4bd1e44617bcd/numba/core/typed_passes.py#L397

[fcharras]

Great news and great debug 🎉 , TY for following up. Ready to test the patch when it's out 😀

the linker randomly picks one

Now that you mention it, the bug was maybe already there in numba == 0.19.0 and this accidentally fixed it ? I noticed __name__ or __qualname__ were used somewhere, but I thought it was for caching. When bumping to post- #804 commits I removed this hack that might explain why the bug surfaced.

I never noticed it prior to 0.19.0 though.

[diptorupd]

@fcharras yes, the bug was there in 0.19 as well. Your patch must have been hiding it.

[diptorupd]

@fcharras Can you confirm if #905 fixes the issue?

[fcharras]

I've run a quick try with #905 . Outputs seem correct now but the benchmarks still show a huge slow down. I'd need a bit more time to understand what is going on.

[diptorupd]

@fcharras the slowdown issue should be addressed by #896.

[fcharras]

Unfortunately, the kmeans is still almost 100 times slower despite fixes in #896 😿 it seems that #905 was not enough. (will investigate, maybe there are issues with the user code)

[fcharras]

Issues cleared, there was some mistake in my setup that caused the 100x slowdown (apparently, activating the oneapi basekit interfers with numba_dpex in my environment 🤷 )

[stuartarchibald]

@diptorupd #898 (comment) sounds like a bug (I also saw something similar to this problem with derived IRs the other day). Do you think it will be reproducible in Numba alone or is this numba-dpex specific?