From 002584c0b92124af476d1f8013a0ba9df65567f7 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 14 Oct 2024 16:39:46 -0500 Subject: [PATCH 1/3] Views: add support for complex numbers --- pykokkos/interface/__init__.py | 1 + pykokkos/interface/data_types.py | 102 +++++++++++++++++++++++++++++++ pykokkos/interface/views.py | 31 ++++++++-- 3 files changed, 130 insertions(+), 4 deletions(-) diff --git a/pykokkos/interface/__init__.py b/pykokkos/interface/__init__.py index bf6f9210..04489e8e 100644 --- a/pykokkos/interface/__init__.py +++ b/pykokkos/interface/__init__.py @@ -18,6 +18,7 @@ uint16, uint32, uint64, float, double, real, float32, float64, bool, + complex32, complex64 ) from .decorators import ( callback, classtype, Decorator, function, functor, main, diff --git a/pykokkos/interface/data_types.py b/pykokkos/interface/data_types.py index 5e88fad2..468f795f 100644 --- a/pykokkos/interface/data_types.py +++ b/pykokkos/interface/data_types.py @@ -1,6 +1,8 @@ from enum import Enum from pykokkos.bindings import kokkos +import pykokkos.kokkos_manager as km + import numpy as np @@ -26,6 +28,8 @@ class DataType(Enum): float64 = kokkos.double real = None bool = np.bool_ + complex32 = kokkos.complex_float32_dtype + complex64 = kokkos.complex_float64_dtype class DataTypeClass: @@ -91,3 +95,101 @@ class float64(DataTypeClass): class bool(DataTypeClass): value = kokkos.uint8 np_equiv = np.bool_ + + +class complex(DataTypeClass): + def __add__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot add '{}' and '{}'".format(type(other), type(self))) + + if isinstance(self, complex32): + return complex32(self.kokkos_complex + other.kokkos_complex) + elif isinstance(self, complex64): + return complex64(self.kokkos_complex + other.kokkos_complex) + + def __iadd__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot add '{}' and '{}'".format(type(other), type(self))) + + self.kokkos_complex += other.kokkos_complex + return self + + def __sub__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot subtract '{}' and '{}'".format(type(other), type(self))) + + if isinstance(self, complex32): + return complex32(self.kokkos_complex - other.kokkos_complex) + elif isinstance(self, complex64): + return complex64(self.kokkos_complex - other.kokkos_complex) + + def __isub__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot subtract '{}' and '{}'".format(type(other), type(self))) + + self.kokkos_complex -= other.kokkos_complex + return self + + def __mul__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot multiply '{}' and '{}'".format(type(other), type(self))) + + if isinstance(self, complex32): + return complex32(self.kokkos_complex * other.kokkos_complex) + elif isinstance(self, complex64): + return complex64(self.kokkos_complex * other.kokkos_complex) + + def __imul__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot multiply '{}' and '{}'".format(type(other), type(self))) + + self.kokkos_complex *= other.kokkos_complex + return self + + def __truediv__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot divide '{}' and '{}'".format(type(other), type(self))) + + if isinstance(self, complex32): + return complex32(self.kokkos_complex / other.kokkos_complex) + elif isinstance(self, complex64): + return complex64(self.kokkos_complex / other.kokkos_complex) + + def __itruediv__(self, other): + if not isinstance(other, type(self)): + raise TypeError("cannot divide '{}' and '{}'".format(type(other), type(self))) + + self.kokkos_complex /= other.kokkos_complex + return self + + def __repr__(self): + return f"({self.kokkos_complex.real_const()}, {self.kokkos_complex.imag_const()})" + + @property + def real(self): + return self.kokkos_complex.real_const() + + @property + def imag(self): + return self.kokkos_complex.imag_const() + +class complex32(complex): + value = kokkos.complex_float32_dtype + np_equiv = np.complex64 # 32 bits from real + 32 from imaginary + + def __init__(self, real: float, imaginary: float = 0.0): + if isinstance(real, kokkos.complex_float32): + self.kokkos_complex = real + else: + self.kokkos_complex = km.get_kokkos_module(is_cpu=True).complex_float32(real, imaginary) + + +class complex64(complex): + value = kokkos.complex_float64_dtype + np_equiv = np.complex128 # 64 bits from real + 64 from imaginary + + def __init__(self, real: float, imaginary: float = 0.0): + if isinstance(real, kokkos.complex_float64): + self.kokkos_complex = real + else: + self.kokkos_complex = km.get_kokkos_module(is_cpu=True).complex_float64(real, imaginary) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 4493a8f7..3b165753 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -25,6 +25,7 @@ uint8, uint16, uint32, uint64, double, float32, float64, + complex32, complex64 ) from .data_types import float as pk_float from .layout import get_default_layout, Layout @@ -98,13 +99,16 @@ def extent(self, dimension: int) -> int: return self.shape[dimension] - def fill(self, value: Union[int, float]) -> None: + def fill(self, value: Union[int, float, complex, complex32, complex64]) -> None: """ Sets all elements to a scalar value :param value: the scalar value """ + if isinstance(value, (complex, complex32, complex64)): + value = np.complex64(value.real, value.imag) if self.dtype is complex32 else np.complex128(value.real, value.imag) + if self.trait is Trait.Unmanaged: self.xp_array.fill(value) else: @@ -126,8 +130,16 @@ def __getitem__(self, key: Union[int, TeamMember, slice, Tuple]) -> Union[int, f if isinstance(key, int) or isinstance(key, TeamMember): if self.trait is Trait.Unmanaged: - return self.xp_array[key] - return self.data[key] + return_val = self.xp_array[key] + else: + return_val = self.data[key] + + if self.dtype is complex32: + return_val = complex32(return_val.real, return_val.imag) + elif self.dtype is complex64: + return_val = complex64(return_val.real, return_val.imag) + + return return_val length: int = 1 if isinstance(key, slice) else len(key) if length != self.rank(): @@ -137,7 +149,7 @@ def __getitem__(self, key: Union[int, TeamMember, slice, Tuple]) -> Union[int, f return subview - def __setitem__(self, key: Union[int, TeamMember], value: Union[int, float]) -> None: + def __setitem__(self, key: Union[int, TeamMember], value: Union[int, float, complex, complex32, complex64]) -> None: """ Overloads the indexing operator setting an item in the View. @@ -148,6 +160,9 @@ def __setitem__(self, key: Union[int, TeamMember], value: Union[int, float]) -> if "PK_FUSION" in os.environ: runtime_singleton.runtime.flush_data(self) + if isinstance(value, (complex, complex32, complex64)): + value = np.complex64(value.real, value.imag) if self.dtype is complex32 else np.complex128(value.real, value.imag) + if self.trait is Trait.Unmanaged: self.xp_array[key] = value else: @@ -674,6 +689,10 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O dtype = float64 elif np_dtype is np.bool_: dtype = uint8 + elif np_dtype is np.complex64: + dtype = complex32 + elif np_dtype is np.complex128: + dtype = complex64 else: raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}") @@ -736,6 +755,10 @@ def from_array(array) -> ViewType: ctype = ctypes.c_double elif np_dtype is np.bool_: ctype = ctypes.c_uint8 + elif np_dtype is np.complex64: + dtype = complex32 + elif np_dtype is np.complex128: + dtype = complex64 else: raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}") From 98bcedf723cd8539cf0bd7581645c652f645607d Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 14 Oct 2024 16:40:07 -0500 Subject: [PATCH 2/3] Visitors: add support for translating complex numbers --- pykokkos/core/visitors/pykokkos_visitor.py | 8 +++++++- pykokkos/core/visitors/visitors_util.py | 7 +++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pykokkos/core/visitors/pykokkos_visitor.py b/pykokkos/core/visitors/pykokkos_visitor.py index 8a8606b9..dff00a98 100644 --- a/pykokkos/core/visitors/pykokkos_visitor.py +++ b/pykokkos/core/visitors/pykokkos_visitor.py @@ -420,11 +420,17 @@ def visit_Call(self, node: ast.Call) -> cppast.CallExpr: ) elif name in ["PerTeam", "PerThread", "fence"]: name = "Kokkos::" + name + elif name in {"complex32", "complex64"}: + name = "Kokkos::complex" + if "32" in name: + name += "" + else: + name += "" function = cppast.DeclRefExpr(name) args: List[cppast.Expr] = [self.visit(a) for a in node.args] - if visitors_util.is_math_function(name) or name in ["printf", "abs", "Kokkos::PerTeam", "Kokkos::PerThread", "Kokkos::fence"]: + if visitors_util.is_math_function(name) or name in ["printf", "abs", "Kokkos::PerTeam", "Kokkos::PerThread", "Kokkos::fence", "Kokkos::complex", "Kokkos::complex"]: return cppast.CallExpr(function, args) if function in self.kokkos_functions: diff --git a/pykokkos/core/visitors/visitors_util.py b/pykokkos/core/visitors/visitors_util.py index bd4c9078..e2db3536 100644 --- a/pykokkos/core/visitors/visitors_util.py +++ b/pykokkos/core/visitors/visitors_util.py @@ -19,7 +19,9 @@ def pretty_print(node): "double": "double", "bool": "bool", "TeamMember": f"Kokkos::TeamPolicy<{Keywords.DefaultExecSpace.value}>::member_type", - "cpp_auto": "auto" + "cpp_auto": "auto", + "complex32": "Kokkos::complex", + "complex64": "Kokkos::complex" } # Maps from the DataType enum to cppast @@ -307,7 +309,8 @@ def parse_view_template_params( if parameter in ("int", "double", "float", "int8_t", "int16_t", "int32_t", "int64_t", - "uint8_t", "uint16_t", "uint32_t", "uint64_t"): + "uint8_t", "uint16_t", "uint32_t", "uint64_t", + "Kokkos::complex", "Kokkos::complex"): datatype: str = parameter + "*" * rank params["dtype"] = datatype From 7eadd914f87bfa817dfd4e43085b2cc4bcf0e884 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Mon, 28 Oct 2024 19:03:44 -0500 Subject: [PATCH 3/3] Views: fix issue with creating views from cupy complex number arrays --- pykokkos/interface/views.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/pykokkos/interface/views.py b/pykokkos/interface/views.py index 3b165753..83273809 100644 --- a/pykokkos/interface/views.py +++ b/pykokkos/interface/views.py @@ -667,6 +667,13 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O dtype: DataTypeClass np_dtype = array.dtype.type + if np_dtype is np.void and cp_array is not None: + # This means that this is a cupy array passed through + # from_array(). When this happens, if the cupy array was + # originally a complex number dtype, np_dtype will be void. We + # should therefore retreive the data type from cp_array. + np_dtype = cp_array.dtype.type + if np_dtype is np.int8: dtype = int8 elif np_dtype is np.int16: @@ -724,6 +731,12 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O return View(ret_list, dtype, space=space, trait=Trait.Unmanaged, array=array, layout=layout, cp_array=cp_array) +class ctypes_complex32(ctypes.Structure): + _fields_ = [("real", ctypes.c_float), ("imag", ctypes.c_float)] + +class ctypes_complex64(ctypes.Structure): + _fields_ = [("real", ctypes.c_double), ("imag", ctypes.c_double)] + def from_array(array) -> ViewType: """ Create a PyKokkos View from a cupy (or other numpy-like) array @@ -756,9 +769,9 @@ def from_array(array) -> ViewType: elif np_dtype is np.bool_: ctype = ctypes.c_uint8 elif np_dtype is np.complex64: - dtype = complex32 + ctype = ctypes_complex32 elif np_dtype is np.complex128: - dtype = complex64 + ctype = ctypes_complex64 else: raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}") @@ -769,6 +782,11 @@ def from_array(array) -> ViewType: ptr = ctypes.cast(ptr, ctypes.POINTER(ctype)) np_array = np.ctypeslib.as_array(ptr, shape=array.shape) + if np_dtype in {np.complex64, np.complex128}: + # This sets the arrays dtype to numpy's complex number types. + # Without this the type would be np.void. + np_array = np_array.view(np_dtype) + # need to select the layout here since the np_array flags do not # preserve the original flags layout: Layout