[WIP] Add CUDAArray type and implementation with addresspace information

[ZzEeKkAa]

Replace internal usage of types.Array towards cuda specific array type CUDAArray that can handle addrespace. The idea is to help the nvvm compiler recognize address space for the memory load/store operation. The issue is that in complex workloads compiler loses track of address space and produce general purpose instructions instead of memory specific one. One such example is device GEMM. This PR results in LDS shared memory specific instruction generated instead of general purpose LD.E.

List of changes happened after transitioning to CUDAArray:

• introduce CUDAArray type and model that supports address space pointer;
• all arrays inside cuda.jit are now using CUDAArray instead of types.Array. That breaks some api like requesting implementation with general purpose array signature. Thats why many tests had to be updated;
• atomics was updated to use nvvm intrinsics specific to the memory address: https://docs.nvidia.com/cuda/nvvm-ir-spec/#atomic .

Future improvements that unblocked:

• get rid of memory_info data field for cuda array. The purpose was to have it like a smart pointer when going back and force between python and compiled code. For kernels we have only compiled code, so we don't need this field ;
• implement address space validation ;
• implement static size arrays, so that size could be resolved at compile time .

TODO:

• Fix rest of the tests.
• Add auto type casting when passing specific address space pointer array to a function that accepts generic address space pointer array.

[copy-pr-bot]

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

Contributors can view more details about this message here.

[ZzEeKkAa]

/ok to test

[gmarkall]

/ok to test

[gmarkall]

I haven't fully reviewed this yet but want to discuss thoughts so far - some of them are marked on the diff.

I think the main design concern I have at the moment is with trying to ensure that all array types in kernels are a CUDAArray type instead of an Array type - I think this might impact launch latency and have a lot of edge cases we need to find. Is an alternative path to keep Array types coexisting with CUDAArray types in kernels, but treat Array types as being in the generic address space? The idea here is to leave the decorator and dispatcher logic unchanged so we don't have to try and make CUDAArray types in the critical path of a launch.

[gmarkall]

As the address spaces in nvvm.py are just integers, might it be worth converting them to be an enum class so that it's easier to get the array type to print like

array(int64, 1, 'C', SHARED)

instead of

array(int64, 1, 'C', addrspace(3))

?

It might make interactive debug / development a little easier without having to mentally translate the address space numbers to names - there aren't to many uses of the address spaces so I would hope that updating the uses (if necessary) wouldn't be too burdensome.

[gmarkall]

Might the mangling_args property need implementing as well? Two methods that differ only in the address space of an array could end up mangling to the same name and potentially creating a symbol clash.

For example:

from numba.core.itanium_mangler import mangle_type
from numba.cuda.types import CUDAArray
from numba import types

shared_array = CUDAArray(types.int64, 1, 'C', addrspace=3)
generic_array = CUDAArray(types.int64, 1, 'C', addrspace=0)

shared_mangled = mangle_type(shared_array)
generic_mangled = mangle_type(generic_array)

print(shared_mangled)
print(generic_mangled)
assert shared_mangled != generic_mangled

gives

9CUDAArrayIxLi1E1C7mutable7alignedE
9CUDAArrayIxLi1E1C7mutable7alignedE
Traceback (most recent call last):
  File "/home/gmarkall/numbadev/issues/numba-cuda-236/mangle_test.py", line 13, in <module>
    assert shared_mangled != generic_mangled
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[gmarkall]

Similarly, what about the unification (the unify() method) and conversion (can_convert_to())? If unify() is not implemented, then all CUDA arrays will end up unifying to Array types instead, even if the set of types to unify were all in the same address space.

Conversions will also lose address space information, or perhaps even allow invalid conversions - I think we should not allow conversion from shared to local address space, for example, but conversions to the generic address space should always be OK.

[gmarkall]

I'm concerned this could have a non-trivial impact on kernel launch time. Can you do a microbenchmark to check how much this impacts the latency of launches with various numbers of array arguments?

[gmarkall]

all arrays inside cuda.jit are now using CUDAArray instead of types.Array. That breaks some api like requesting implementation with general purpose array signature. Thats why many tests had to be updated;

I think this will break user code too - if we can find a way to avoid doing that I'd strongly prefer to.

[ZzEeKkAa]

all arrays inside cuda.jit are now using CUDAArray instead of types.Array. That breaks some api like requesting implementation with general purpose array signature. Thats why many tests had to be updated;

I think this will break user code too - if we can find a way to avoid doing that I'd strongly prefer to.

Yes, that's what I'm worry about. Still thinking about potential ways to avoid it

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}