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.

[csskene]

The cylinder calculus tests were failing because the distributor was created without specifying the dtype. Then in operators the matrices were cast to the default type of None. I've pushed a quick fix, but may want to do something better to catch this in future.

[kburns]

Hmm maybe we should just make float64 the default dtype in dist.

[csskene]

I think we need to be more careful how operators are cast. The reason to cast the matrices was to go from double precision to single precision (for real and complex). They are always created in double, I used dist.dtype to get them to single, but this doesn't work if the dtype here doesn't match the bases.