Add KokkosKernels::eager_initialize() to common

blas/CMakeLists.txt

@@ -36,7 +36,7 @@ IF (KOKKOSKERNELS_ENABLE_TPL_BLAS OR KOKKOSKERNELS_ENABLE_TPL_MKL OR KOKKOSKERNE
 ENDIF()
 
 # Include cuda blas TPL source file
-IF (KOKKOSKERNELS_ENABLE_TPL_CUBLAS)
+IF (KOKKOSKERNELS_ENABLE_TPL_CUBLAS OR KOKKOSKERNELS_ENABLE_TPL_MAGMA)
   LIST(APPEND SOURCES
    blas/tpls/KokkosBlas_Cuda_tpl.cpp
   )

blas/tpls/KokkosBlas_Cuda_tpl.hpp

@@ -25,12 +25,24 @@ namespace Impl {
 CudaBlasSingleton::CudaBlasSingleton() {
   cublasStatus_t stat = cublasCreate(&handle);
   if (stat != CUBLAS_STATUS_SUCCESS) Kokkos::abort("CUBLAS initialization failed\n");
-
-  Kokkos::push_finalize_hook([&]() { cublasDestroy(handle); });
 }
 
 CudaBlasSingleton& CudaBlasSingleton::singleton() {
-  static CudaBlasSingleton s;
+  std::unique_ptr<CudaBlasSingleton>& instance = get_instance();
+  if (!instance) {
+    instance = std::make_unique<CudaBlasSingleton>();
+    Kokkos::push_finalize_hook([&]() {
+      cublasDestroy(instance->handle);
+      instance.reset();
+    });
+  }
+  return *instance;
+}
+
+bool CudaBlasSingleton::is_initialized() { return get_instance() != nullptr; }
+
+std::unique_ptr<CudaBlasSingleton>& CudaBlasSingleton::get_instance() {
+  static std::unique_ptr<CudaBlasSingleton> s;
   return s;
 }
 

blas/tpls/KokkosBlas_Magma_tpl.hpp

@@ -25,12 +25,24 @@ namespace Impl {
 MagmaSingleton::MagmaSingleton() {
   magma_int_t stat = magma_init();
   if (stat != MAGMA_SUCCESS) Kokkos::abort("MAGMA initialization failed\n");
-
-  Kokkos::push_finalize_hook([&]() { magma_finalize(); });
 }
 
 MagmaSingleton& MagmaSingleton::singleton() {
-  static MagmaSingleton s;
+  std::unique_ptr<MagmaSingleton>& instance = get_instance();
+  if (!instance) {
+    instance = std::make_unique<MagmaSingleton>();
+    Kokkos::push_finalize_hook([&]() {
+      magma_finalize();
+      instance.reset();
+    });
+  }
+  return *instance;
+}
+
+bool MagmaSingleton::is_initialized() { return get_instance() != nullptr; }
+
+std::unique_ptr<MagmaSingleton>& MagmaSingleton::get_instance() {
+  static std::unique_ptr<MagmaSingleton> s;
   return s;
 }
 

blas/tpls/KokkosBlas_Rocm_tpl.hpp

@@ -22,14 +22,24 @@
 namespace KokkosBlas {
 namespace Impl {
 
-RocBlasSingleton::RocBlasSingleton() {
-  KOKKOS_ROCBLAS_SAFE_CALL_IMPL(rocblas_create_handle(&handle));
+RocBlasSingleton::RocBlasSingleton() { KOKKOS_ROCBLAS_SAFE_CALL_IMPL(rocblas_create_handle(&handle)); }
 
-  Kokkos::push_finalize_hook([&]() { KOKKOS_ROCBLAS_SAFE_CALL_IMPL(rocblas_destroy_handle(handle)); });
+RocBlasSingleton& RocBlasSingleton::singleton() {
+  std::unique_ptr<RocBlasSingleton>& instance = get_instance();
+  if (!instance) {
+    instance = std::make_unique<RocBlasSingleton>();
+    Kokkos::push_finalize_hook([&]() {
+      KOKKOS_ROCBLAS_SAFE_CALL_IMPL(rocblas_destroy_handle(instance->handle));
+      instance.reset();
+    });
+  }
+  return *instance;
 }
 
-RocBlasSingleton& RocBlasSingleton::singleton() {
-  static RocBlasSingleton s;
+bool RocBlasSingleton::is_initialized() { return get_instance() != nullptr; }
+
+std::unique_ptr<RocBlasSingleton>& RocBlasSingleton::get_instance() {
+  static std::unique_ptr<RocBlasSingleton> s;
   return s;
 }
 

blas/tpls/KokkosBlas_magma.hpp

@@ -27,7 +27,11 @@ namespace Impl {
 struct MagmaSingleton {
   MagmaSingleton();
 
+  static bool is_initialized();
   static MagmaSingleton& singleton();
+
+ private:
+  static std::unique_ptr<MagmaSingleton>& get_instance();
 };
 
 }  // namespace Impl

blas/tpls/KokkosBlas_tpl_spec.hpp

@@ -29,7 +29,11 @@ struct CudaBlasSingleton {
 
   CudaBlasSingleton();
 
+  static bool is_initialized();
   static CudaBlasSingleton& singleton();
+
+ private:
+  static std::unique_ptr<CudaBlasSingleton>& get_instance();
 };
 
 inline void cublas_internal_error_throw(cublasStatus_t cublasState, const char* name, const char* file,
@@ -111,7 +115,12 @@ struct RocBlasSingleton {
 
   RocBlasSingleton();
 
+  static bool is_initialized();
+
   static RocBlasSingleton& singleton();
+
+ private:
+  static std::unique_ptr<RocBlasSingleton>& get_instance();
 };
 
 inline void rocblas_internal_error_throw(rocblas_status rocblasState, const char* name, const char* file,

cmake/KokkosKernels_config.h.in

@@ -174,4 +174,12 @@
 #define KOKKOSKERNELS_IMPL_COMPILE_LIBRARY false
 #endif
 
+/* Enabled components */
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_BATCHED
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_BLAS
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_LAPACK
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_SPARSE
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_GRAPH
+#cmakedefine KOKKOSKERNELS_ENABLE_COMPONENT_ODE
+
 #endif  // KOKKOSKERNELS_CONFIG_H

cmake/kokkoskernels_components.cmake

@@ -102,4 +102,13 @@ IF (    KokkosKernels_ENABLE_COMPONENT_BATCHED
 ELSE()
   SET(KOKKOSKERNELS_ALL_COMPONENTS_ENABLED OFF CACHE BOOL "" FORCE)
 ENDIF()
-mark_as_advanced(FORCE KOKKOSKERNELS_ALL_COMPONENTS_ENABLED)
+mark_as_advanced(FORCE KOKKOSKERNELS_ALL_COMPONENTS_ENABLED)
+# Now that component enables are finalized, also set upper-case
+# versions of component enables for the config.h
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_BATCHED KokkosKernels_ENABLE_COMPONENT_BATCHED)
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_BLAS KokkosKernels_ENABLE_COMPONENT_BLAS)
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_LAPACK KokkosKernels_ENABLE_COMPONENT_LAPACK)
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_GRAPH KokkosKernels_ENABLE_COMPONENT_GRAPH)
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_SPARSE KokkosKernels_ENABLE_COMPONENT_SPARSE)
+SET(KOKKOSKERNELS_ENABLE_COMPONENT_ODE KokkosKernels_ENABLE_COMPONENT_ODE)
+

common/CMakeLists.txt

@@ -2,3 +2,5 @@
 LIST(APPEND KK_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/common/src)
 LIST(APPEND KK_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/common/impl)
 LIST(APPEND KK_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/common/unit_test)
+
+LIST(APPEND SOURCES common/src/KokkosKernels_EagerInitialize.cpp)

common/src/KokkosKernels_EagerInitialize.cpp

@@ -0,0 +1,78 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) 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.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#include "KokkosKernels_EagerInitialize.hpp"
+#include "KokkosKernels_config.h"
+#include "Kokkos_Core.hpp"
+
+// Include the minimal set of headers that declare all TPL singletons
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_BLAS
+#include "KokkosBlas_tpl_spec.hpp"  //cuBLAS, rocBLAS
+#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA
+#include "KokkosBlas_magma.hpp"
+#endif
+#endif
+
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_SPARSE
+// note: this file declares both cuSPARSE and rocSPARSE singletons
+#include "KokkosKernels_tpl_handles_decl.hpp"
+#endif
+
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_LAPACK
+#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER
+#include "KokkosLapack_cusolver.hpp"
+#endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA
+#include "KokkosLapack_magma.hpp"
+#endif
+#endif
+
+namespace KokkosKernels {
+void eager_initialize() {
+  if (!Kokkos::is_initialized()) {
+    throw std::runtime_error("Kokkos::intialize must be called before KokkosKernels::eager_initialize");
+  }
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_BLAS
+#ifdef KOKKOSKERNELS_ENABLE_TPL_CUBLAS
+  (void)KokkosBlas::Impl::CudaBlasSingleton::singleton();
+#endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCBLAS
+  (void)KokkosBlas::Impl::RocBlasSingleton::singleton();
+#endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA
+  (void)KokkosBlas::Impl::MagmaSingleton::singleton();
+#endif
+#endif
+
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_SPARSE
+#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSPARSE
+  (void)KokkosKernels::Impl::CusparseSingleton::singleton();
+#endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE
+  (void)KokkosKernels::Impl::RocsparseSingleton::singleton();
+#endif
+#endif
+
+#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_LAPACK
+#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER
+  (void)KokkosLapack::Impl::CudaLapackSingleton::singleton();
+#endif
+#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA
+  (void)KokkosLapack::Impl::MagmaSingleton::singleton();
+#endif
+#endif
+}
+}  // namespace KokkosKernels

common/src/KokkosKernels_EagerInitialize.hpp

@@ -0,0 +1,38 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) 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.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//@HEADER
+
+#ifndef KOKKOKERNELS_EAGER_INITIALIZE_HPP
+#define KOKKOKERNELS_EAGER_INITIALIZE_HPP
+
+namespace KokkosKernels {
+// \brief Eagerly initialize handles for all enabled TPLs, as well
+// as any other globally shared resources that would otherwise be lazily initialized.
+//
+// Eagerly initializing a TPL means that it doesn't have to be
+// lazily initialized when first calling a kernel that uses it.
+// For example, \c eager_initialize() will call \c cusparseCreate() upfront
+// so that the first call to \c KokkosSparse::spmv doesn't have to.
+// This can add a significant amount of apparent runtime to that first kernel
+// call, even though the added time isn't really spent in the kernel.
+//
+// Calling this before using any kernels/TPLs is optional.
+// This function is idempotent (any calls after the first have no effect).
+//
+// \pre \c Kokkos::initialize() has been called.
+void eager_initialize();
+}  // namespace KokkosKernels
+
+#endif

common/unit_test/CMakeLists.txt

@@ -95,3 +95,10 @@ IF (KOKKOS_ENABLE_THREADS)
   )
 ENDIF ()
 
+# Add eager_initialize test, which is not backend-specific
+KOKKOSKERNELS_ADD_UNIT_TEST(
+  common_eager_initialize
+  SOURCES Test_Common_EagerInitialize.cpp
+  COMPONENTS common
+)
+