Migrate to numba-cuda-mlir

[seanlaw]

Sadly, numba.cuda no longer works and numba-cuda has gone into maintenance mode and is being replaced by numa-cuda-mlir:

numba/numba#10753
nzxyin/sdtw-cuda-torch#2

The solution would be to do:

try:
    from numba_cuda_mlir import cuda
except ModuleNotFoundError:
    from numba import cuda

To Do:

• Fix dependencies (especially for targeting linux-64) to install numba-cuda-mlir on specific platforms
• Update conda.sh
• Add numba-cuda-mlir to requirements.txt
• Update import as shown above
• Update docstring examples (i.e., see gpu_mpdist, gpu_stimp, gpu_ostinato) and replace cuda import

Extra Notes

When you runtests.sh gpu and encounter:

TypeError: Signature mismatch: 2 argument types given, but function takes 1 arguments

This is because numba.cuda introduced a bug that does not support *args as inputs to the min and max builtin functions. To resolve this, you must (minimally) use

numba-cuda = "*"

However, this is insufficient on its own and you may then encounter:

AttributeError: module 'numpy' has no attribute 'row_stack'

This is because row_stack was deprecated in numpy v2.0 and fully removed in v2.5! Thus, the only way to overcome this is to use an older version of numpy with the no longer maintained numba-cuda:

numba-cuda = "*"
numpy = "<2.5"

and then this combination should allow the tests to pass! However, the real question is whether numba-cuda-mlir allows us to just do:

numba-cuda-mlir = "*"

[seanlaw]

The simplest short term solution is to install numba-cuda but note that this only works for numpy <= 2.5 (see NVIDIA/numba-cuda#907)