[BUG]: cuda.cooperative passes string dtype parameters through to C++

[brycelelbach]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this bug and that I agree to the Code of Conduct

Type of Bug

Compile-time Error

Component

cuda.cooperative

Describe the bug

If you pass a string as dtype to cuda.cooperative things, it passes the string through to the C++ code. Is this intended?

At the very least, we shouldn't just pass through for builtin Numba/Numpy dtypes. We should make dtype="int32" work as it does in NumPy, not fail horribly as it does today.

How to Reproduce

import cuda.cooperative.experimental as cudax
import numba
from pynvjitlink import patch
patch.patch_numba_linker(lto=True)

load = cudax.block.load(dtype=numba.int32, threads_in_block=32, items_per_thread=4, algorithm="striped")
store = cudax.block.store(dtype="int32", threads_in_block=32, items_per_thread=4, algorithm="striped")

@numba.cuda.jit(link=load.files+store.files)
def kernel():
  pass

kernel[1, 32]()

Today this leads to:

Traceback (most recent call last):
  File "/raid/blelbach/dev/sandbox/py/bugs/cuda_cooperative_string_dtype.py", line 7, in <module>
    store = cudax.block.store(dtype="int32", threads_in_block=32, items_per_thread=4, algorithm="striped")
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/block/_block_load_store.py", line 204, in store
    for ltoir in specialization.get_lto_ir()
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/_types.py", line 726, in get_lto_ir
    _, lto_fn = nvrtc.compile(cpp=src, cc=cc, rdc=True, code="lto")
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/_nvrtc.py", line 89, in compile
    return nvrtc_version, compile_impl(
                          ^^^^^^^^^^^^^
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/_caching.py", line 46, in cacher
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/_nvrtc.py", line 55, in compile_impl
    CHECK_NVRTC(err, prog)
  File "/raid/blelbach/dev/cccl/python/cuda_cooperative/cuda/cooperative/experimental/_nvrtc.py", line 17, in CHECK_NVRTC
    raise RuntimeError(f"NVRTC error: {log.decode('ascii')}")
RuntimeError: NVRTC error: code.cu(12): error: identifier "int32" is undefined
              using algorithm_t = cub::BlockStore<int32, 32, 4, ::cub::BLOCK_STORE_STRIPED, 1, 1>;
                                                  ^

code.cu(16): error: incomplete type is not allowed
                  extern "C" __device__ void block_storeint32__32__4____cub__BLOCK_STORE_STRIPED__1__1_Puint8_Pint32_Pint32_alloc(storage_t* param_0, storage_t (&param_1)[4])
                                             ^

code.cu(16): error: identifier "storage_t" is undefined
                  extern "C" __device__ void block_storeint32__32__4____cub__BLOCK_STORE_STRIPED__1__1_Puint8_Pint32_Pint32_alloc(storage_t* param_0, storage_t (&param_1)[4])
                                                                                                                                  ^

code.cu(16): error: identifier "param_0" is undefined
                  extern "C" __device__ void block_storeint32__32__4____cub__BLOCK_STORE_STRIPED__1__1_Puint8_Pint32_Pint32_alloc(storage_t* param_0, storage_t (&param_1)[4])
                                                                                                                                             ^

code.cu(16): error: identifier "param_1" is undefined
                  extern "C" __device__ void block_storeint32__32__4____cub__BLOCK_STORE_STRIPED__1__1_Puint8_Pint32_Pint32_alloc(storage_t* param_0, storage_t (&param_1)[4])
                                                                                                                                                                  ^

code.cu(17): error: expected a ";"
                  {
                  ^

6 errors detected in the compilation of "code.cu".

Expected behavior

N/A

Reproduction link

No response

Operating System

No response

nvidia-smi output

No response

NVCC version

No response