update kokkos example

examples/CMakeLists.txt

@@ -56,15 +56,8 @@ if(GINKGO_BUILD_MPI)
     list(APPEND EXAMPLES_LIST distributed-solver)
 endif()
 
-find_package(Kokkos QUIET)
-if(Kokkos_FOUND)
-    if(GINKGO_WITH_CCACHE)
-        message(WARNING "The CMAKE_CXX_COMPILER_LAUNCHER is set due to "
-            "GINKGO_WITH_CCACHE=ON which is known to cause issues with CUDA enabled "
-            "Kokkos (https://github.com/kokkos/kokkos/issues/4821) including compilation "
-            "failures. This can be prevented by setting GINKGO_WITH_CCACHE=OFF.")
-    endif()
-    list(APPEND EXAMPLES_LIST kokkos_assembly)
+if(GINKGO_EXTENSION_KOKKOS)
+    list(APPEND EXAMPLES_LIST kokkos-assembly)
 else()
     message(STATUS "No Kokkos found, disabling examples with Kokkos assembly.")
 endif()

examples/kokkos_assembly/CMakeLists.txt → examples/kokkos-assembly/CMakeLists.txt

@@ -9,7 +9,5 @@ if(NOT GINKGO_EXTENSION_KOKKOS)
     message(FATAL_ERROR "The Kokkos examples requires that GINKGO_EXTENSION_KOKKOS=ON is set.")
 endif()
 
-find_package(Kokkos REQUIRED)
-
-add_executable(kokkos-assembly kokkos_assembly.cpp)
-target_link_libraries(kokkos-assembly Ginkgo::ginkgo Kokkos::kokkos)
+add_executable(kokkos-assembly kokkos-assembly.cpp)
+target_link_libraries(kokkos-assembly Ginkgo::ginkgo Ginkgo::ext::kokkos)

examples/kokkos_assembly/kokkos_assembly.cpp → examples/kokkos-assembly/kokkos-assembly.cpp

@@ -4,13 +4,10 @@
 
 #include <iostream>
 #include <Kokkos_Core.hpp>
-#include <map>
 #include <string>
 
 
-#include <omp.h>
-
-
+#include <ginkgo/extensions/kokkos.hpp>
 #include <ginkgo/ginkgo.hpp>
 
 
@@ -29,12 +26,7 @@ void generate_stencil_matrix(gko::matrix::Csr<ValueType, IndexType>* matrix)
                                                      discretization_points * 3);
 
     // Create Kokkos views on Ginkgo data.
-    Kokkos::View<IndexType*> v_row_idxs(md.get_row_idxs(),
-                                        md.get_num_stored_elements());
-    Kokkos::View<IndexType*> v_col_idxs(md.get_col_idxs(),
-                                        md.get_num_stored_elements());
-    Kokkos::View<ValueType*> v_values(md.get_values(),
-                                      md.get_num_stored_elements());
+    auto k_md = gko::ext::kokkos::map_data(md);
 
     // Create the matrix entries. This also creates zero entries for the
     // first and second row to handle all rows uniformly.
@@ -51,9 +43,9 @@ void generate_stencil_matrix(gko::matrix::Csr<ValueType, IndexType>* matrix)
             auto mask =
                 static_cast<IndexType>(0 <= col && col < discretization_points);
 
-            v_row_idxs[i] = mask * row;
-            v_col_idxs[i] = mask * col;
-            v_values[i] = mask * coefs[ofs + 1];
+            k_md.row_idxs[i] = mask * row;
+            k_md.col_idxs[i] = mask * col;
+            k_md.values[i] = mask * coefs[ofs + 1];
         });
 
     // Add up duplicate (zero) entries.
@@ -70,18 +62,17 @@ void generate_rhs(Closure&& f, ValueType u0, ValueType u1,
                   gko::matrix::Dense<ValueType>* rhs)
 {
     const auto discretization_points = rhs->get_size()[0];
-    auto values = rhs->get_values();
-    Kokkos::View<ValueType*> values_view(values, discretization_points);
+    auto k_rhs = gko::ext::kokkos::map_data(rhs);
     Kokkos::parallel_for(
         "generate_rhs", discretization_points, KOKKOS_LAMBDA(int i) {
             const ValueType h = 1.0 / (discretization_points + 1);
             const ValueType xi = ValueType(i + 1) * h;
-            values_view[i] = -f(xi) * h * h;
+            k_rhs(i, 0) = -f(xi) * h * h;
             if (i == 0) {
-                values_view[i] += u0;
+                k_rhs(i, 0) += u0;
             }
             if (i == discretization_points - 1) {
-                values_view[i] += u1;
+                k_rhs(i, 0) += u1;
             }
         });
 }
@@ -94,17 +85,15 @@ double calculate_error(int discretization_points,
                        const gko::matrix::Dense<ValueType>* u,
                        Closure&& correct_u)
 {
-    Kokkos::View<const ValueType*> v_u(u->get_const_values(),
-                                       discretization_points);
+    auto k_u = gko::ext::kokkos::map_data(u);
     auto error = 0.0;
     Kokkos::parallel_reduce(
         "calculate_error", discretization_points,
         KOKKOS_LAMBDA(int i, double& lsum) {
             const auto h = 1.0 / (discretization_points + 1);
             const auto xi = (i + 1) * h;
-            lsum += Kokkos::Experimental::abs(
-                (v_u(i) - correct_u(xi)) /
-                Kokkos::Experimental::abs(correct_u(xi)));
+            lsum += Kokkos::abs((k_u(i, 0) - correct_u(xi)) /
+                                Kokkos::abs(correct_u(xi)));
         },
         error);
     return error;
@@ -113,8 +102,6 @@ double calculate_error(int discretization_points,
 
 int main(int argc, char* argv[])
 {
-    Kokkos::ScopeGuard kokkos(argc, argv);
-
     // Some shortcuts
     using ValueType = double;
     using RealValueType = gko::remove_complex<ValueType>;
@@ -130,39 +117,17 @@ int main(int argc, char* argv[])
     if (argc == 2 && (std::string(argv[1]) == "--help")) {
         std::cerr << "Usage: " << argv[0]
                   << " [discretization_points] [kokkos-options]" << std::endl;
-        std::exit(-1);
+        Kokkos::ScopeGuard kokkos(argc, argv);  // print Kokkos help
+        std::exit(1);
     }
 
-    const unsigned int discretization_points =
-        argc >= 2 ? std::atoi(argv[1]) : 100u;
+    Kokkos::ScopeGuard kokkos(argc, argv);
+
+    const auto discretization_points =
+        static_cast<gko::size_type>(argc >= 2 ? std::atoi(argv[1]) : 100u);
 
     // chooses the executor that corresponds to the Kokkos DefaultExecutionSpace
-    auto exec = []() -> std::shared_ptr<gko::Executor> {
-#ifdef KOKKOS_ENABLE_SERIAL
-        if (std::is_same<Kokkos::DefaultExecutionSpace,
-                         Kokkos::Serial>::value) {
-            return gko::ReferenceExecutor::create();
-        }
-#endif
-#ifdef KOKKOS_ENABLE_OPENMP
-        if (std::is_same<Kokkos::DefaultExecutionSpace,
-                         Kokkos::OpenMP>::value) {
-            return gko::OmpExecutor::create();
-        }
-#endif
-#ifdef KOKKOS_ENABLE_CUDA
-        if (std::is_same<Kokkos::DefaultExecutionSpace, Kokkos::Cuda>::value) {
-            return gko::CudaExecutor::create(0,
-                                             gko::ReferenceExecutor::create());
-        }
-#endif
-#ifdef KOKKOS_ENABLE_HIP
-        if (std::is_same<Kokkos::DefaultExecutionSpace, Kokkos::HIP>::value) {
-            return gko::HipExecutor::create(0,
-                                            gko::ReferenceExecutor::create());
-        }
-#endif
-    }();
+    auto exec = gko::ext::kokkos::create_default_executor();
 
     // problem:
     auto correct_u = [] KOKKOS_FUNCTION(ValueType x) { return x * x * x; };
@@ -174,9 +139,7 @@ int main(int argc, char* argv[])
     auto rhs = vec::create(exec, gko::dim<2>(discretization_points, 1));
     generate_rhs(f, u0, u1, rhs.get());
     auto u = vec::create(exec, gko::dim<2>(discretization_points, 1));
-    for (int i = 0; i < u->get_size()[0]; ++i) {
-        u->get_values()[i] = 0.0;
-    }
+    u->fill(0.0);
 
     // initialize the stencil matrix
     auto A = share(mtx::create(