Add support for Kokkos::complex

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 += "<float>"
+            else:
+                name += "<double>"
 
         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<float>", "Kokkos::complex<double>"]:
             return cppast.CallExpr(function, args)
 
         if function in self.kokkos_functions:

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<float>",
+    "complex64": "Kokkos::complex<double>"
 }
 
 # 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<float>", "Kokkos::complex<double>"):
             datatype: str = parameter + "*" * rank
             params["dtype"] = datatype
 

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,

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)

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:
@@ -652,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:
@@ -674,6 +696,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}")
 
@@ -705,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
@@ -736,6 +768,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:
+        ctype = ctypes_complex32
+    elif np_dtype is np.complex128:
+        ctype = ctypes_complex64
     else:
         raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}")
 
@@ -746,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