diff --git a/pyccel/ast/cudaext.py b/pyccel/ast/cudaext.py index 63b837be9d..b29e537cc4 100644 --- a/pyccel/ast/cudaext.py +++ b/pyccel/ast/cudaext.py @@ -306,4 +306,4 @@ def __new__(cls, dim=0): } cuda_mod = Module('cuda', variables = cuda_constants.values(), - funcs = cuda_funcs.values()) \ No newline at end of file + funcs = cuda_funcs.values()) diff --git a/pyccel/ast/cupyext.py b/pyccel/ast/cupyext.py index 68a5a07391..4ef323b79f 100644 --- a/pyccel/ast/cupyext.py +++ b/pyccel/ast/cupyext.py @@ -56,6 +56,7 @@ 'CupyOnesLike', 'CupyZeros', 'CupyZerosLike', + 'CupyRavel' 'Shape' ) @@ -83,7 +84,7 @@ class CupyArray(CupyNewArray): arg : list, tuple, PythonList """ - __slots__ = ('_arg','_dtype','_precision','_shape','_rank','_order') + __slots__ = ('_arg','_dtype','_precision','_shape','_rank','_order', '_memory_location') _attribute_nodes = ('_arg',) name = 'array' @@ -131,6 +132,7 @@ def __init__(self, arg, dtype=None, order='C'): self._dtype = dtype self._order = order self._precision = prec + self._memory_location = 'device' super().__init__() def __str__(self): @@ -140,6 +142,12 @@ def __str__(self): def arg(self): return self._arg + @property + def memory_location(self): + """ Indicate if the array is in the host or device memory + """ + return self._memory_location + #============================================================================== class CupyArange(CupyNewArray): """ @@ -454,6 +462,35 @@ def __str__(self): #============================================================================== +class CupyRavel(CupyArray): + """ + Class representing a call to the cupy ravel function which + returns flattened version of the passed array + + Parameters + ========== + arg : PyccelAstNode + A PyccelAstNode of unknown shape + memory_location : string + The location where the new array memory should be allocated + """ + name = 'ravel' + __slots__ = () + def __new__(cls, arg): + if not isinstance(arg, (list, tuple, PyccelAstNode)): + raise TypeError('Unknown type of %s.' % type(arg)) + if arg.rank == 0: + raise TypeError('Unknown type of %s.' % type(arg)) + return super().__new__(cls) + + def __init__(self, arg): + super().__init__(arg = arg) + shape = reduce((lambda x, y: x.python_value * y.python_value), self.shape) + self._shape = [shape if isinstance(shape, (LiteralInteger, PyccelArraySize)) else LiteralInteger(shape)] + self._rank = 1 + + +#============================================================================== cupy_funcs = { # ... array creation routines 'full' : PyccelFunctionDef('full' , CupyFull), @@ -466,6 +503,7 @@ def __str__(self): 'ones_like' : PyccelFunctionDef('ones_like' , CupyOnesLike), 'array' : PyccelFunctionDef('array' , CupyArray), 'arange' : PyccelFunctionDef('arange' , CupyArange), + 'ravel' : PyccelFunctionDef('ravel' , CupyRavel), # ... 'shape' : PyccelFunctionDef('shape' , Shape), 'size' : PyccelFunctionDef('size' , CupyArraySize), diff --git a/pyccel/codegen/printing/ccudacode.py b/pyccel/codegen/printing/ccudacode.py index 2d6ac5ad43..128db6e81e 100644 --- a/pyccel/codegen/printing/ccudacode.py +++ b/pyccel/codegen/printing/ccudacode.py @@ -40,7 +40,7 @@ from pyccel.ast.numpyext import NumpyFull, NumpyArray, NumpyArange from pyccel.ast.numpyext import NumpyReal, NumpyImag, NumpyFloat -from pyccel.ast.cupyext import CupyFull, CupyArray, CupyArange +from pyccel.ast.cupyext import CupyFull, CupyArray, CupyArange, CupyRavel from pyccel.ast.cudaext import CudaCopy, cuda_Internal_Var, CudaArray @@ -424,6 +424,37 @@ def _print_Assign(self, expr): rhs = self._print(expr.rhs) return prefix_code+'{} = {};\n'.format(lhs, rhs) + def _print_AliasAssign(self, expr): + lhs_var = expr.lhs + rhs_var = expr.rhs + + lhs_address = ObjectAddress(lhs_var) + rhs_address = ObjectAddress(rhs_var) + + # the below condition handles the case of reassinging a pointer to an array view. + # setting the pointer's is_view attribute to false so it can be ignored by the free_pointer function. + + if not self.stored_in_c_pointer(lhs_var) and \ + isinstance(lhs_var, Variable) and lhs_var.is_ndarray: + if isinstance(rhs_var, CupyRavel): + lhs = self._print(lhs_address) + return f'cupy_ravel({lhs}, {rhs_var});\n' + rhs = self._print(rhs_var) + if isinstance(rhs_var, Variable) and rhs_var.is_ndarray: + lhs = self._print(lhs_address) + if lhs_var.order == rhs_var.order: + return f'alias_assign({lhs}, {rhs});\n' + else: + return f'transpose_alias_assign({lhs}, {rhs});\n' + else: + lhs = self._print(lhs_var) + return f'{lhs} = {rhs};\n' + else: + lhs = self._print(lhs_address) + rhs = self._print(rhs_address) + + return f'{lhs} = {rhs};\n' + def arrayFill(self, expr): """ print the assignment of a NdArray @@ -521,19 +552,18 @@ def copy_CudaArray_Data(self, expr): declare_dtype = self.find_in_dtype_registry(self._print(rhs.dtype), rhs.precision) dtype = self.find_in_ndarray_type_registry(self._print(rhs.dtype), rhs.precision) arg = rhs.arg if isinstance(rhs, (CudaArray, CupyArray)) else rhs - if rhs.rank > 1: - # flattening the args to use them in C initialization. - arg = self._flatten_list(arg) - self.add_import(c_imports['string']) if isinstance(arg, Variable): arg = self._print(arg) - cpy_data = "cudaMemcpy({0}.raw_data, {1}.{2}, {0}.buffer_size, cudaMemcpyHostToDevice);".format(lhs, arg, dtype) + cpy_data = f"cudaMemcpy({lhs}.{dtype}, {arg}.{dtype}, {lhs}.buffer_size, cudaMemcpyHostToDevice);" return '%s\n' % (cpy_data) else : + if arg.rank > 1: + arg = self._flatten_list(arg) arg = ', '.join(self._print(i) for i in arg) dummy_array = "%s %s[] = {%s};\n" % (declare_dtype, dummy_array_name, arg) - cpy_data = "cudaMemcpy({0}.raw_data, {1}, {0}.buffer_size, cudaMemcpyHostToDevice);".format(self._print(lhs), dummy_array_name, dtype) + target_array_name = self._print(lhs) + cpy_data = f"cudaMemcpy({target_array_name}.{dtype}, {dummy_array_name}, {target_array_name}.buffer_size, cudaMemcpyHostToDevice);" return '%s%s\n' % (dummy_array, cpy_data) def _print_CudaSynchronize(self, expr): diff --git a/pyccel/parser/semantic.py b/pyccel/parser/semantic.py index d69a880955..0201fc2f31 100644 --- a/pyccel/parser/semantic.py +++ b/pyccel/parser/semantic.py @@ -61,6 +61,7 @@ from pyccel.ast.core import Assert from pyccel.ast.class_defs import NumpyArrayClass, TupleClass, get_cls_base, CudaArrayClass +from pyccel.ast.cupyext import CupyRavel from pyccel.ast.datatypes import NativeRange, str_dtype from pyccel.ast.datatypes import NativeSymbol @@ -493,6 +494,20 @@ def _infere_type(self, expr, **settings): d_var['cls_base' ] = NumpyArrayClass return d_var + elif isinstance(expr, CupyRavel): + if isinstance(expr.arg, Variable) and expr.arg.memory_location == "device": + d_var['memory_handling'] = 'alias' + else: + d_var['memory_handling'] = 'heap' + d_var['memory_location'] = expr.memory_location + d_var['datatype' ] = expr.dtype + d_var['shape' ] = expr.shape + d_var['rank' ] = expr.rank + d_var['precision' ] = expr.precision + d_var['order' ] = None + d_var['cls_base' ] = CudaNewArray + return d_var + elif isinstance(expr, CupyNewArray): d_var['datatype' ] = expr.dtype d_var['memory_handling'] = 'heap' if expr.rank > 0 else 'stack' diff --git a/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.cu b/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.cu index 2c0d517e19..e284971c35 100644 --- a/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.cu +++ b/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.cu @@ -132,6 +132,17 @@ t_ndarray cuda_array_create(int32_t nd, int64_t *shape, return (arr); } +void cupy_ravel(t_ndarray *dest, t_ndarray src) +{ + *dest = src; + dest->nd = 1; + cudaMallocManaged(&(dest->shape), sizeof(int64_t)); + cudaMallocManaged(&(dest->strides), sizeof(int64_t)); + *(dest->shape) = src.length; + *(dest->strides) = 1; + dest->is_view = true; +} + int32_t cuda_free_array(t_ndarray arr) { if (arr.shape == NULL) diff --git a/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.h b/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.h index d0cce3b5f0..c928d859e7 100644 --- a/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.h +++ b/pyccel/stdlib/cuda_ndarrays/cuda_ndarrays.h @@ -21,7 +21,9 @@ void _cuda_array_fill_int64(int64_t c, t_ndarray arr); __global__ void _cuda_array_fill_double(double c, t_ndarray arr); -t_ndarray cuda_array_create(int32_t nd, int64_t *shape, enum e_types type, bool is_view, enum e_memory_locations location); +void cupy_ravel(t_ndarray *dest, t_ndarray src); + +t_ndarray cuda_array_create(int32_t nd, int64_t *shape, enum e_types type, bool is_view, enum e_memory_locations location); int32_t cuda_free_array(t_ndarray dump); int32_t cuda_free_pointer(t_ndarray dump); #endif diff --git a/tests/internal/scripts/ccuda/cupy_ravel.py b/tests/internal/scripts/ccuda/cupy_ravel.py new file mode 100644 index 0000000000..97e3342df1 --- /dev/null +++ b/tests/internal/scripts/ccuda/cupy_ravel.py @@ -0,0 +1,11 @@ +# pylint: disable=missing-function-docstring, disable=unused-variable, missing-module-docstring + +import cupy as cp +from pyccel.decorators import kernel, types +from pyccel import cuda + +if __name__ == '__main__': + threads_per_block = 32 + n_blocks = 1 + arr1 = cp.ravel([[1,2],[1,3]]) + arr2 = cp.ravel([1,2,3,4]) diff --git a/tests/internal/scripts/ccuda/cupy_ravel_variable.py b/tests/internal/scripts/ccuda/cupy_ravel_variable.py new file mode 100644 index 0000000000..1062c262ba --- /dev/null +++ b/tests/internal/scripts/ccuda/cupy_ravel_variable.py @@ -0,0 +1,12 @@ +# pylint: disable=missing-function-docstring, disable=unused-variable, missing-module-docstring + +import cupy as cp +from pyccel import cuda + +if __name__ == '__main__': + c = ((1, 2), (1, 3)) + host_arr = cuda.array(c, dtype=int) + device_arr = cuda.array(c, dtype=int, memory_location='device') + arr1 = cp.ravel(host_arr) + arr2 = cp.ravel(device_arr) + arr3 = cp.ravel(c)