GPU support for Cartesian problems

[kburns]

This PR adds GPU support for one dimensional bases and cartesian problems. Remaining rough edges include good defaults for subproblem coupling, and raising or warning for unsupported features (GPU+MPI, GPU+curvilinear).

[csskene]

Just noting here that the trace operator doesn't work on GPU's as cupy's einsum doesn't take an out argument. A similar fix to dot product (commit fb9b3d6) fixes things.
Also problem.add_equation will not work for float32 as unify attributes for dtype raises an 'Objects are not all equal' error.

[csskene]

Interpolate doesn't currently work either, with error
...
File "~/Packages/dedalus_gpu/dedalus/core/operators.py", line 975, in _subspace_matrix
return cls._full_matrix(input_basis, output_basis, *args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: InterpolateRealFourier._full_matrix() missing 1 required positional argument: 'position'

[csskene]

Just noting here that the trace operator doesn't work on GPU's as cupy's einsum doesn't take an out argument. A similar fix to dot product (commit fb9b3d6) fixes things. Also problem.add_equation will not work for float32 as unify attributes for dtype raises an 'Objects are not all equal' error.

trace is fixed as of [6e32312] .
I've traced the dtype error to MultiplyNumberField. Specifically, this line

self.dtype = np.result_type(type(arg0), arg1.dtype)

If arg0 is a python int, and arg1 is np.float32, this function returns np.float64, leading to the error. This error is unavoidable, since subtract invokes a multiplication by the int (-1). For example f-g will do f + ( (-1)*g). I am not sure what the best fix is. I've tried

self.dtype = np.result_type(arg0, arg1.dtype)

which fixes things, i.e. use arg0 rather than type(arg0). This gives np.float32 which is correct behaviour here. I wanted to check this is a good fix before pushing.

[kburns]

We may want to use cupy's jit-rawkernel approach to handle type generality: https://docs.cupy.dev/en/stable/user_guide/kernel.html#jit-kernel-definition

[csskene]

We may want to use cupy's jit-rawkernel approach to handle type generality: https://docs.cupy.dev/en/stable/user_guide/kernel.html#jit-kernel-definition

This seems to work for here

@jit.rawkernel()
def apply_csr_mid_kernel(data, indices, indptr, x3, y3, N1, N2i, N2o, N3):
    n1 = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x # batch index
    n3 = jit.blockIdx.y * jit.blockDim.y + jit.threadIdx.y # output column index
    if n1 >= N1 or n3 >= N3:
        return
    # Loop over output rows = CSR matrix rows
    for i in range(N2o):
        y3[n1, i, n3] = 0
        start = indptr[i]
        end = indptr[i + 1]
        for k in range(start, end):
            j = indices[k]
            y3[n1, i, n3] += data[k] * x3[n1, j, n3]

def cupy_apply_csr_mid(matrix, array, out):
    N1, N2i, N3 = array.shape
    N2o = matrix.shape[0]
    N1 =  cp.uint32(N1)
    N2i =  cp.uint32(N2i)
    N3 =  cp.uint32(N3)
    N2o =  cp.uint32(N2o)
    # Choose thread/block config
    threads_y = min(1024, N3) # maximize concurrency along n3
    threads_x = 1024 // threads_y # make block have 1024 threads
    blockdim = (threads_x, threads_y)
    blocks_x = (N1 + threads_x - 1) // threads_x
    blocks_y = (N3 + threads_y - 1) // threads_y
    griddim = (blocks_x, blocks_y)
    # Launch kernel
    apply_csr_mid_kernel(griddim, blockdim, (matrix.data, matrix.indices, matrix.indptr, array, out, N1, N2i, N2o, N3))

[kburns]

I've changed the Chebyshev transform default from the cupy DCT to matrix transforms for now, since even up to sizes of 1024 that seems much much faster. The cupy transform seems quite slow for some reason in my tests.

Note there is no CUDA-native DCT, so cupy implements the DCT as an extended FFT.