From 6ea703c18d103bebf057889d88dbde746bf63d10 Mon Sep 17 00:00:00 2001 From: Brian Kelley Date: Fri, 27 Sep 2024 18:16:50 -0600 Subject: [PATCH] TPL singletons: allow query of whether initialized And test KokkosKernels::eager_initialize() using this --- blas/tpls/KokkosBlas_Cuda_tpl.hpp | 18 ++- blas/tpls/KokkosBlas_Magma_tpl.hpp | 18 ++- blas/tpls/KokkosBlas_Rocm_tpl.hpp | 20 +++- blas/tpls/KokkosBlas_magma.hpp | 4 + blas/tpls/KokkosBlas_tpl_spec.hpp | 9 ++ common/src/KokkosKernels_EagerInitialize.cpp | 1 - common/unit_test/CMakeLists.txt | 7 ++ common/unit_test/Test_Common.hpp | 1 - .../unit_test/Test_Common_EagerInitialize.cpp | 113 ++++++++++++++++++ .../unit_test/Test_Common_EagerInitialize.hpp | 27 ----- lapack/tpls/KokkosLapack_Cuda_tpl.hpp | 18 ++- lapack/tpls/KokkosLapack_Magma_tpl.hpp | 18 ++- lapack/tpls/KokkosLapack_cusolver.hpp | 5 + lapack/tpls/KokkosLapack_magma.hpp | 5 + .../tpls/KokkosKernels_tpl_handles_decl.hpp | 8 ++ sparse/tpls/KokkosKernels_tpl_handles_def.hpp | 40 +++++-- 16 files changed, 256 insertions(+), 56 deletions(-) create mode 100644 common/unit_test/Test_Common_EagerInitialize.cpp delete mode 100644 common/unit_test/Test_Common_EagerInitialize.hpp diff --git a/blas/tpls/KokkosBlas_Cuda_tpl.hpp b/blas/tpls/KokkosBlas_Cuda_tpl.hpp index d80e3a23d8..fa2749c980 100644 --- a/blas/tpls/KokkosBlas_Cuda_tpl.hpp +++ b/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& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + cublasDestroy(instance->handle); + instance.reset(); + }); + } + return *instance; +} + +bool CudaBlasSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& CudaBlasSingleton::get_instance() { + static std::unique_ptr s; return s; } diff --git a/blas/tpls/KokkosBlas_Magma_tpl.hpp b/blas/tpls/KokkosBlas_Magma_tpl.hpp index f149a790df..bce5d4057a 100644 --- a/blas/tpls/KokkosBlas_Magma_tpl.hpp +++ b/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& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + magma_finalize(); + instance.reset(); + }); + } + return *instance; +} + +bool MagmaSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& MagmaSingleton::get_instance() { + static std::unique_ptr s; return s; } diff --git a/blas/tpls/KokkosBlas_Rocm_tpl.hpp b/blas/tpls/KokkosBlas_Rocm_tpl.hpp index b5a7dabf6f..699c0b5db5 100644 --- a/blas/tpls/KokkosBlas_Rocm_tpl.hpp +++ b/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& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + 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::get_instance() { + static std::unique_ptr s; return s; } diff --git a/blas/tpls/KokkosBlas_magma.hpp b/blas/tpls/KokkosBlas_magma.hpp index 5f5fcfe4e1..e86d2ee2cd 100644 --- a/blas/tpls/KokkosBlas_magma.hpp +++ b/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& get_instance(); }; } // namespace Impl diff --git a/blas/tpls/KokkosBlas_tpl_spec.hpp b/blas/tpls/KokkosBlas_tpl_spec.hpp index 7f40edf435..8ad70595fc 100644 --- a/blas/tpls/KokkosBlas_tpl_spec.hpp +++ b/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& 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& get_instance(); }; inline void rocblas_internal_error_throw(rocblas_status rocblasState, const char* name, const char* file, diff --git a/common/src/KokkosKernels_EagerInitialize.cpp b/common/src/KokkosKernels_EagerInitialize.cpp index ec4c6e1225..a8c8a39f48 100644 --- a/common/src/KokkosKernels_EagerInitialize.cpp +++ b/common/src/KokkosKernels_EagerInitialize.cpp @@ -53,7 +53,6 @@ void eager_initialize() { (void)KokkosBlas::Impl::RocBlasSingleton::singleton(); #endif #ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA -#include "KokkosBlas_Magma_tpl.hpp" (void)KokkosBlas::Impl::MagmaSingleton::singleton(); #endif #endif diff --git a/common/unit_test/CMakeLists.txt b/common/unit_test/CMakeLists.txt index c0d8fc116f..c963e908e5 100644 --- a/common/unit_test/CMakeLists.txt +++ b/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 +) + diff --git a/common/unit_test/Test_Common.hpp b/common/unit_test/Test_Common.hpp index eb18918c5c..fb93a494d6 100644 --- a/common/unit_test/Test_Common.hpp +++ b/common/unit_test/Test_Common.hpp @@ -28,6 +28,5 @@ #include #include #include -#include #endif // TEST_COMMON_HPP diff --git a/common/unit_test/Test_Common_EagerInitialize.cpp b/common/unit_test/Test_Common_EagerInitialize.cpp new file mode 100644 index 0000000000..20910069fc --- /dev/null +++ b/common/unit_test/Test_Common_EagerInitialize.cpp @@ -0,0 +1,113 @@ +//@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 KK_EAGERINIT_TEST_HPP +#define KK_EAGERINIT_TEST_HPP + +#include +#include "Kokkos_Core.hpp" +#include "KokkosKernels_config.h" +#include "KokkosKernels_EagerInitialize.hpp" + +#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_BLAS +#include "KokkosBlas_tpl_spec.hpp" //cuBLAS, rocBLAS +#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA +#include "KokkosBlas_Magma_tpl.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 + +// Count the number of singletons which are currently initialized, +// and the numInitialized number of singleton classes that are currently enabled +// (based on which TPLs and components were enabled at configure-time) +void countSingletons(int& numInitialized, int& numEnabled) { + numInitialized = 0; + numEnabled = 0; +#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_BLAS +#ifdef KOKKOSKERNELS_ENABLE_TPL_CUBLAS + numEnabled++; + if (KokkosBlas::Impl::CudaBlasSingleton::is_initialized()) numInitialized++; +#endif +#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCBLAS + numEnabled++; + if (KokkosBlas::Impl::RocBlasSingleton::is_initialized()) numInitialized++; +#endif +#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA + numEnabled++; + if (KokkosBlas::Impl::MagmaSingleton::is_initialized()) numInitialized++; +#endif +#endif + +#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_SPARSE +#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSPARSE + numEnabled++; + if (KokkosKernels::Impl::CusparseSingleton::is_initialized()) numInitialized++; +#endif +#ifdef KOKKOSKERNELS_ENABLE_TPL_ROCSPARSE + numEnabled++; + if (KokkosKernels::Impl::RocsparseSingleton::is_initialized()) numInitialized++; +#endif +#endif + +#ifdef KOKKOSKERNELS_ENABLE_COMPONENT_LAPACK +#ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER + numEnabled++; + if (KokkosLapack::Impl::CudaLapackSingleton::is_initialized()) numInitialized++; +#endif +#ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA + numEnabled++; + if (KokkosLapack::Impl::MagmaSingleton::is_initialized()) numInitialized++; +#endif +#endif +} + +int main() { + int numInitialized, numEnabled; + Kokkos::initialize(); + { + // Check that no singletons are already initialized. + countSingletons(numInitialized, numEnabled); + if (numInitialized != 0) + throw std::runtime_error("At least one singleton was initialized before it should have been"); + KokkosKernels::eager_initialize(); + // Check that all singletons are now initialized. + countSingletons(numInitialized, numEnabled); + std::cout << "Kokkos::eager_initialize() set up " << numInitialized << " of " << numEnabled << " TPL singletons.\n"; + if (numInitialized != numEnabled) + throw std::runtime_error("At least one singleton was not initialized by eager_initialize()"); + } + Kokkos::finalize(); + // Finally, make sure that all singletons were finalized during Kokkos::finalize(). + countSingletons(numInitialized, numEnabled); + if (numInitialized != 0) + throw std::runtime_error("At least one singleton was not correctly finalized by Kokkos::finalize()"); + return 0; +} + +#endif diff --git a/common/unit_test/Test_Common_EagerInitialize.hpp b/common/unit_test/Test_Common_EagerInitialize.hpp deleted file mode 100644 index a6f55b7022..0000000000 --- a/common/unit_test/Test_Common_EagerInitialize.hpp +++ /dev/null @@ -1,27 +0,0 @@ -//@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 KK_EAGERINIT_TEST_HPP -#define KK_EAGERINIT_TEST_HPP - -#include "KokkosKernels_EagerInitialize.hpp" - -TEST_F(TestCategory, common_eager_initialize) { - KokkosKernels::eager_initialize(); - KokkosKernels::eager_initialize(); -}; - -#endif diff --git a/lapack/tpls/KokkosLapack_Cuda_tpl.hpp b/lapack/tpls/KokkosLapack_Cuda_tpl.hpp index 3ead12d5f4..e3191ea93b 100644 --- a/lapack/tpls/KokkosLapack_Cuda_tpl.hpp +++ b/lapack/tpls/KokkosLapack_Cuda_tpl.hpp @@ -25,12 +25,24 @@ namespace Impl { CudaLapackSingleton::CudaLapackSingleton() { cusolverStatus_t stat = cusolverDnCreate(&handle); if (stat != CUSOLVER_STATUS_SUCCESS) Kokkos::abort("CUSOLVER initialization failed\n"); - - Kokkos::push_finalize_hook([&]() { cusolverDnDestroy(handle); }); } CudaLapackSingleton& CudaLapackSingleton::singleton() { - static CudaLapackSingleton s; + std::unique_ptr& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + cusolverDnDestroy(instance->handle); + instance.reset(); + }); + } + return *instance; +} + +bool CudaLapackSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& CudaLapackSingleton::get_instance() { + static std::unique_ptr s; return s; } diff --git a/lapack/tpls/KokkosLapack_Magma_tpl.hpp b/lapack/tpls/KokkosLapack_Magma_tpl.hpp index 636c40735d..542f681281 100644 --- a/lapack/tpls/KokkosLapack_Magma_tpl.hpp +++ b/lapack/tpls/KokkosLapack_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& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + magma_finalize(); + instance.reset(); + }); + } + return *instance; +} + +bool MagmaSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& MagmaSingleton::get_instance() { + static std::unique_ptr s; return s; } diff --git a/lapack/tpls/KokkosLapack_cusolver.hpp b/lapack/tpls/KokkosLapack_cusolver.hpp index 272fb8b3b8..15a88714af 100644 --- a/lapack/tpls/KokkosLapack_cusolver.hpp +++ b/lapack/tpls/KokkosLapack_cusolver.hpp @@ -32,6 +32,11 @@ struct CudaLapackSingleton { CudaLapackSingleton(); static CudaLapackSingleton& singleton(); + + static bool is_initialized(); + + private: + static std::unique_ptr& get_instance(); }; inline void cusolver_internal_error_throw(cusolverStatus_t cusolverStatus, const char* name, const char* file, diff --git a/lapack/tpls/KokkosLapack_magma.hpp b/lapack/tpls/KokkosLapack_magma.hpp index dfde113fa6..b1b7bb1ab6 100644 --- a/lapack/tpls/KokkosLapack_magma.hpp +++ b/lapack/tpls/KokkosLapack_magma.hpp @@ -30,6 +30,11 @@ struct MagmaSingleton { MagmaSingleton(); static MagmaSingleton& singleton(); + + static bool is_initialized(); + + private: + static std::unique_ptr& get_instance(); }; } // namespace Impl diff --git a/sparse/tpls/KokkosKernels_tpl_handles_decl.hpp b/sparse/tpls/KokkosKernels_tpl_handles_decl.hpp index a1cd3c97f5..4a14d43df0 100644 --- a/sparse/tpls/KokkosKernels_tpl_handles_decl.hpp +++ b/sparse/tpls/KokkosKernels_tpl_handles_decl.hpp @@ -30,7 +30,11 @@ struct CusparseSingleton { CusparseSingleton(); + static bool is_initialized(); static CusparseSingleton& singleton(); + + private: + static std::unique_ptr& get_instance(); }; } // namespace Impl @@ -48,7 +52,11 @@ struct RocsparseSingleton { RocsparseSingleton(); + static bool is_initialized(); static RocsparseSingleton& singleton(); + + private: + static std::unique_ptr& get_instance(); }; } // namespace Impl diff --git a/sparse/tpls/KokkosKernels_tpl_handles_def.hpp b/sparse/tpls/KokkosKernels_tpl_handles_def.hpp index a88ad12130..68757841dc 100644 --- a/sparse/tpls/KokkosKernels_tpl_handles_def.hpp +++ b/sparse/tpls/KokkosKernels_tpl_handles_def.hpp @@ -25,14 +25,24 @@ namespace KokkosKernels { namespace Impl { -CusparseSingleton::CusparseSingleton() { - KOKKOS_CUSPARSE_SAFE_CALL(cusparseCreate(&cusparseHandle)); +CusparseSingleton::CusparseSingleton() { KOKKOS_CUSPARSE_SAFE_CALL(cusparseCreate(&cusparseHandle)); } - Kokkos::push_finalize_hook([&]() { cusparseDestroy(cusparseHandle); }); +CusparseSingleton& CusparseSingleton::singleton() { + std::unique_ptr& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + KOKKOS_CUSPARSE_SAFE_CALL(cusparseDestroy(instance->cusparseHandle)); + instance.reset(); + }); + } + return *instance; } -CusparseSingleton& CusparseSingleton::singleton() { - static CusparseSingleton s; +bool CusparseSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& CusparseSingleton::get_instance() { + static std::unique_ptr s; return s; } @@ -46,14 +56,24 @@ CusparseSingleton& CusparseSingleton::singleton() { namespace KokkosKernels { namespace Impl { -RocsparseSingleton::RocsparseSingleton() { - KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_handle(&rocsparseHandle)); +RocsparseSingleton::RocsparseSingleton() { KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_create_handle(&rocsparseHandle)); } - Kokkos::push_finalize_hook([&]() { KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_destroy_handle(rocsparseHandle)); }); +RocsparseSingleton& RocsparseSingleton::singleton() { + std::unique_ptr& instance = get_instance(); + if (!instance) { + instance = std::make_unique(); + Kokkos::push_finalize_hook([&]() { + KOKKOS_ROCSPARSE_SAFE_CALL_IMPL(rocsparse_destroy_handle(instance->rocsparseHandle)); + instance.reset(); + }); + } + return *instance; } -RocsparseSingleton& RocsparseSingleton::singleton() { - static RocsparseSingleton s; +bool RocsparseSingleton::is_initialized() { return get_instance() != nullptr; } + +std::unique_ptr& RocsparseSingleton::get_instance() { + static std::unique_ptr s; return s; }