ENH: add suppport for Kokkos::complex Views

CMakeLists.txt

@@ -44,7 +44,8 @@ SET(libpykokkos_SOURCES
     ${CMAKE_CURRENT_LIST_DIR}/src/available.cpp
     ${CMAKE_CURRENT_LIST_DIR}/src/common.cpp
     ${CMAKE_CURRENT_LIST_DIR}/src/tools.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/src/execution_spaces.cpp)
+    ${CMAKE_CURRENT_LIST_DIR}/src/execution_spaces.cpp
+    ${CMAKE_CURRENT_LIST_DIR}/src/complex_dtypes.cpp)
 
 SET(libpykokkos_HEADERS
     ${CMAKE_CURRENT_LIST_DIR}/include/libpykokkos.hpp

include/fwd.hpp

@@ -134,6 +134,8 @@ enum KokkosViewDataType {
   Uint64,
   Float32,
   Float64,
+  ComplexFloat32,
+  ComplexFloat64,
   ViewDataTypesEnd
 };
 

include/libpykokkos.hpp

@@ -58,4 +58,5 @@ void generate_atomic_variants(py::module& kokkos);
 void generate_backend_versions(py::module& kokkos);
 void generate_pool_variants(py::module& kokkos);
 void generate_execution_spaces(py::module& kokkos);
+void generate_complex_dtypes(py::module& kokkos);
 void destroy_callbacks();

include/traits.hpp

@@ -85,6 +85,10 @@ VIEW_DATA_TYPE(uint32_t, Uint32, "uint32", "unsigned", "unsigned_int")
 VIEW_DATA_TYPE(uint64_t, Uint64, "uint64", "unsigned_long")
 VIEW_DATA_TYPE(float, Float32, "float32", "float")
 VIEW_DATA_TYPE(double, Float64, "float64", "double")
+VIEW_DATA_TYPE(Kokkos::complex<float>, ComplexFloat32, "complex_float32_dtype",
+               "complex_float_dtype")
+VIEW_DATA_TYPE(Kokkos::complex<double>, ComplexFloat64, "complex_float64_dtype",
+               "complex_double_dtype")
 
 //----------------------------------------------------------------------------//
 // <data-type> <enum> <string identifiers>

include/views.hpp

@@ -51,6 +51,7 @@
 #include "fwd.hpp"
 #include "traits.hpp"
 
+#include <pybind11/numpy.h>
 #include <Kokkos_Core.hpp>
 #include <Kokkos_DynRankView.hpp>
 #include <iostream>

src/complex_dtypes.cpp

@@ -0,0 +1,82 @@
+/*
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 3.0
+//       Copyright (2020) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+// 1. Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright
+// notice, this list of conditions and the following disclaimer in the
+// documentation and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the Corporation nor the names of the
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
+// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
+// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
+// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+//
+// Questions? Contact Christian R. Trott ([email protected])
+//
+// ************************************************************************
+//@HEADER
+*/
+
+#include "common.hpp"
+
+#include <pybind11/operators.h>
+#include <Kokkos_Core.hpp>
+
+//----------------------------------------------------------------------------//
+//
+//        The Kokkos::complex dtypes
+//
+//----------------------------------------------------------------------------//
+
+template <typename Tp>
+void generate_complex_dtype(py::module& kokkos, const std::string& _name) {
+  using ComplexTp = Kokkos::complex<Tp>;
+
+  py::class_<ComplexTp>(kokkos, _name.c_str())
+      .def(py::init<Tp>())      // Constructor for real part only
+      .def(py::init<Tp, Tp>())  // Constructor for real and imaginary parts
+      .def("imag_mutable", py::overload_cast<>(&ComplexTp::imag))
+      .def("imag_const", py::overload_cast<>(&ComplexTp::imag, py::const_))
+      .def("imag_set", py::overload_cast<Tp>(&ComplexTp::imag))
+      .def("real_mutable", py::overload_cast<>(&ComplexTp::real))
+      .def("real_const", py::overload_cast<>(&ComplexTp::real, py::const_))
+      .def("real_set", py::overload_cast<Tp>(&ComplexTp::real))
+      .def(py::self += py::self)
+      .def(py::self += Tp())
+      .def(py::self -= py::self)
+      .def(py::self -= Tp())
+      .def(py::self *= py::self)
+      .def(py::self *= Tp())
+      .def(py::self /= py::self)
+      .def(py::self /= Tp());
+}
+
+void generate_complex_dtypes(py::module& kokkos) {
+  generate_complex_dtype<float>(kokkos, "complex_float32");
+  generate_complex_dtype<double>(kokkos, "complex_float64");
+}

src/libpykokkos.cpp

@@ -116,4 +116,5 @@ PYBIND11_MODULE(libpykokkos, kokkos) {
   generate_backend_versions(kokkos);
   generate_pool_variants(kokkos);
   generate_execution_spaces(kokkos);
+  generate_complex_dtypes(kokkos);
 }

src/variants/CMakeLists.txt

@@ -26,7 +26,7 @@ TARGET_LINK_LIBRARIES(libpykokkos-variants PUBLIC
 
 SET(_types              concrete dynamic)
 SET(_variants           layout memory_trait)
-SET(_data_types         Int8 Int16 Int32 Int64 Uint8 Uint16 Uint32 Uint64 Float32 Float64)
+SET(_data_types         Int8 Int16 Int32 Int64 Uint8 Uint16 Uint32 Uint64 Float32 Float64 ComplexFloat32 ComplexFloat64)
 
 SET(layout_enums        Right)
 SET(memory_trait_enums  Managed)