Draft: feat: implement `__cuda_array_interface__`
https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html
This is supported by:
- pycuda
- numba
- cupy
- torch
- nvcv https://github.com/CvCuda/CV-CUDA
- maybe by tensorflow in future: https://github.com/tensorflow/tensorflow/issues/29039
Also allow to execute with cupy (https://docs.cupy.dev/en/stable/index.html) instead of pycuda
TODO:
-
check that pointers in correct CUDA context and if not import into current -
make execution with pycuda aware of __cuda_array_interface__
-
what/how to test
Edited by Stephan Seitz
Merge request reports
Activity
32 34 Returns: 33 35 compiled kernel as Python function 34 36 """ 35 import pycuda.autoinit # NOQA 36 from pycuda.compiler import SourceModule 37 if USE_PYCUDA: should be an enum or maybe and autoselect if cupy installed but pycuda not
Edited by Stephan Seitzchanged this line in version 4 of the diff
2 2 3 3 from pystencils.backends.cbackend import get_headers 4 4 from pystencils.backends.cuda_backend import generate_cuda 5 from pystencils.typing import StructType 6 5 from pystencils.field import FieldType 7 6 from pystencils.include import get_pycuda_include_path, get_pystencils_include_path 8 7 from pystencils.kernel_wrapper import KernelWrapper 8 from pystencils.typing import StructType 9 9 from pystencils.typing.typed_sympy import FieldPointerSymbol 10 from pystencils.utils import DotDict 10 11 11 12 USE_FAST_MATH = True 13 USE_PYCUDA = False 12 14 13 15 14 16 def get_cubic_interpolation_include_paths(): 4 4 from pystencils.backends.cuda_backend import generate_cuda 5 from pystencils.typing import StructType 6 5 from pystencils.field import FieldType 7 6 from pystencils.include import get_pycuda_include_path, get_pystencils_include_path 8 7 from pystencils.kernel_wrapper import KernelWrapper 8 from pystencils.typing import StructType 9 9 from pystencils.typing.typed_sympy import FieldPointerSymbol 10 from pystencils.utils import DotDict 10 11 11 12 USE_FAST_MATH = True 13 USE_PYCUDA = False 12 14 13 15 14 16 def get_cubic_interpolation_include_paths(): 15 from os.path import join, dirname 17 from os.path import dirname, join this is
isort
https://pypi.org/project/isort/. Matrin added aisort
config to this repo.
1 import numpy as np 2 import pytest 3 import sympy as sp 4 from scipy.ndimage import convolve 5 6 from pystencils import Assignment, CreateKernelConfig, Field, Target, create_kernel 7 from pystencils.simp import sympy_cse_on_assignment_list 8 from pystencils.slicing import add_ghost_layers, remove_ghost_layers 9 10 11 @pytest.mark.parametrize("framework", ("cupy", "torch", "numba")) 12 def test_cuda_array_interface(framework): changed this line in version 8 of the diff
50 53 if USE_FAST_MATH: 51 54 nvcc_options.append("-use_fast_math") 52 55 53 mod = SourceModule(code, options=nvcc_options, include_dirs=[ 54 get_pystencils_include_path(), get_pycuda_include_path()]) 55 func = mod.get_function(kernel_function_node.function_name) 56 if USE_PYCUDA: 57 nvcc_options.append("-Wno-deprecated-gpu-targets") 58 mod = SourceModule(code, options=nvcc_options, include_dirs=[ 59 get_pystencils_include_path(), get_pycuda_include_path()]) 60 func = mod.get_function(kernel_function_node.function_name) 61 else: 62 import cupy 63 nvcc_options.append("-I" + get_pystencils_include_path()) 64 nvcc_options.append("-I" + get_pystencils_include_path()) 65 func = cupy.RawKernel(code, kernel_function_node.function_name, options=tuple(nvcc_options), jitify=True) 147 176 symbolic_field_strides = tuple(int(i) * field_arr.dtype.itemsize for i in symbolic_field.strides) 148 177 if isinstance(symbolic_field.dtype, StructType): 149 178 symbolic_field_strides = symbolic_field_strides[:-1] 150 if symbolic_field_strides != field_arr.strides: 179 if field_arr.strides and symbolic_field_strides != field_arr.strides: 42 41 includes = "\n".join([f"#include {include_file}" for include_file in header_list]) 43 42 44 43 code = includes + "\n" 45 code += "#define FUNC_PREFIX __global__\n" 44 code += "#define FUNC_PREFIX extern \"C\" __global__\n"
Please register or sign in to reply