Merge pull request #13520 from brian-kelley/oneapi_mkl

Tribits: find SYCL libraries for TPL MKL
trilinos · Oct 11, 2024 · f0d5905 · f0d5905
2 parents 156678d + 0aee2ae
commit f0d5905
Show file tree

Hide file tree

Showing 12 changed files with 35 additions and 25 deletions.
diff --git a/cmake/TPLs/FindTPLMKL.cmake b/cmake/TPLs/FindTPLMKL.cmake
@@ -15,10 +15,32 @@
 # pseudorandom number generators.  That's why we require a header
 # file, to access the function declarations.
 
-TRIBITS_TPL_FIND_INCLUDE_DIRS_AND_LIBRARIES( MKL
-  REQUIRED_HEADERS mkl.h
-  REQUIRED_LIBS_NAMES mkl_rt
-  )
+IF(Kokkos_ENABLE_SYCL)
+  # For OneAPI MKL on GPU, use the CMake target
+  # Temporarily change CMAKE_CXX_COMPILER to icpx to convince MKL to add the DPCPP target
+  # If it sees that CMAKE_CXX_COMPILER is just ".../mpicxx", it won't do this.
+  set(CMAKE_CXX_COMPILER_PREVIOUS "${CMAKE_CXX_COMPILER}")
+  set(CMAKE_CXX_COMPILER "icpx")
+  # Use the BLAS95 and LAPACK95 interfaces (int32_t for dimensions and indices)
+  set(MKL_INTERFACE lp64)
+  find_package(MKL REQUIRED COMPONENTS MKL::MKL MKL::MKL_SYCL)
+  IF (NOT MKL_FOUND)
+    MESSAGE(FATAL_ERROR "MKL (as CMake package) was not found! This is required for SYCL+MKL")
+  ENDIF()
+  set(CMAKE_CXX_COMPILER "${CMAKE_CXX_COMPILER_PREVIOUS}")
+
+  tribits_extpkg_create_imported_all_libs_target_and_config_file( MKL
+    INNER_FIND_PACKAGE_NAME MKL
+    IMPORTED_TARGETS_FOR_ALL_LIBS MKL::MKL MKL::MKL_SYCL
+    )
+ELSE ()
+  # For host MKL, the single library libmkl_rt is sufficient.
+  # This works for older versions of MKL that don't provide MKLConfig.cmake.
+  TRIBITS_TPL_FIND_INCLUDE_DIRS_AND_LIBRARIES( MKL
+    REQUIRED_HEADERS mkl.h
+    REQUIRED_LIBS_NAMES mkl_rt
+    )
+ENDIF()
 
 # In the past, MKL users had to link with a long list of libraries.
 # The choice of libraries enables specific functionality.  Intel
@@ -69,7 +91,3 @@ TRIBITS_TPL_FIND_INCLUDE_DIRS_AND_LIBRARIES( MKL
 # where the MKLROOT environment variable points to my MKL install
 # directory.
 
-
-
-
-
diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm1_tpl_spec_avail.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm1_tpl_spec_avail.hpp
@@ -95,7 +95,7 @@ KOKKOSBLAS1_NRM1_TPL_SPEC_AVAIL_ROCBLAS(Kokkos::complex<float>, Kokkos::LayoutLe
 // oneMKL
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 
 #define KOKKOSBLAS1_NRM1_TPL_SPEC_AVAIL_MKL_SYCL(SCALAR, LAYOUT, MEMSPACE)                                             \
   template <class ExecSpace>                                                                                           \

diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm1_tpl_spec_decl.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm1_tpl_spec_decl.hpp
@@ -263,7 +263,7 @@ KOKKOSBLAS1_NRM1_TPL_SPEC_DECL_ROCBLAS(Kokkos::complex<double>, Kokkos::LayoutLe
 // oneMKL
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 
 #include <KokkosBlas_tpl_spec.hpp>
 #include <oneapi/mkl/blas.hpp>

diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm2_tpl_spec_avail.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm2_tpl_spec_avail.hpp
@@ -75,8 +75,7 @@ KOKKOSBLAS1_NRM2_TPL_SPEC_AVAIL(Kokkos::LayoutLeft, Kokkos::Cuda, Kokkos::CudaSp
 KOKKOSBLAS1_NRM2_TPL_SPEC_AVAIL(Kokkos::LayoutLeft, Kokkos::HIP, Kokkos::HIPSpace)
 #endif
 
-#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE) && \
-    defined(KOKKOS_ENABLE_SYCL)
+#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && defined(KOKKOS_ENABLE_SYCL)
 KOKKOSBLAS1_NRM2_TPL_SPEC_AVAIL(Kokkos::LayoutLeft, Kokkos::Experimental::SYCL,
                                 Kokkos::Experimental::SYCLDeviceUSMSpace)
 #endif

diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm2_tpl_spec_decl.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas1_nrm2_tpl_spec_decl.hpp
@@ -283,8 +283,7 @@ KOKKOSBLAS1_NRM2_TPL_SPEC_DECL_ROCBLAS_EXT(false)
 
 #endif
 
-#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE) && \
-    defined(KOKKOS_ENABLE_SYCL)
+#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && defined(KOKKOS_ENABLE_SYCL)
 #include <mkl.h>
 #include <oneapi/mkl/blas.hpp>
 #include <KokkosBlas_tpl_spec.hpp>

diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas2_gemv_tpl_spec_avail.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas2_gemv_tpl_spec_avail.hpp
@@ -129,7 +129,7 @@ KOKKOSBLAS2_GEMV_TPL_SPEC_AVAIL_ROCBLAS(Kokkos::complex<float>, Kokkos::LayoutRi
 
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MKL
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 
 #define KOKKOSBLAS2_GEMV_TPL_SPEC_AVAIL_ONEMKL(SCALAR, LAYOUT)                                           \
   template <class ExecSpace>                                                                             \

diff --git a/packages/kokkos-kernels/blas/tpls/KokkosBlas2_gemv_tpl_spec_decl.hpp b/packages/kokkos-kernels/blas/tpls/KokkosBlas2_gemv_tpl_spec_decl.hpp
@@ -594,8 +594,7 @@ KOKKOSBLAS2_CGEMV_ROCBLAS(Kokkos::LayoutRight, Kokkos::HIPSpace, false)
 #endif  // KOKKOSKERNELS_ENABLE_TPL_ROCBLAS
 
 // ONEMKL
-#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE) && \
-    defined(KOKKOS_ENABLE_SYCL)
+#if defined(KOKKOSKERNELS_ENABLE_TPL_MKL) && defined(KOKKOS_ENABLE_SYCL)
 #include <mkl.h>
 #include <oneapi/mkl/blas.hpp>
 #include <KokkosBlas_tpl_spec.hpp>

diff --git a/packages/kokkos-kernels/cmake/KokkosKernels_config.h.in b/packages/kokkos-kernels/cmake/KokkosKernels_config.h.in
@@ -29,7 +29,6 @@
    requires (a) header file(s) as well, and may use functions other
    than just BLAS and LAPACK functions.  */
 #cmakedefine HAVE_KOKKOSKERNELS_MKL
-#cmakedefine KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE
 
 #cmakedefine KOKKOSKERNELS_ENABLE_TESTS_AND_PERFSUITE
 #cmakedefine KOKKOSKERNELS_ENABLE_BENCHMARK

diff --git a/packages/kokkos-kernels/cmake/kokkoskernels_tpls.cmake b/packages/kokkos-kernels/cmake/kokkoskernels_tpls.cmake
@@ -31,10 +31,6 @@ MACRO(KOKKOSKERNELS_ADD_TPL_OPTION NAME DEFAULT_VALUE DOCSTRING)
   SET(ROOT_DEFAULT $ENV{${_NAME_ORIG}_ROOT})
   KOKKOSKERNELS_ADD_OPTION(${_NAME_ORIG}_ROOT "${ROOT_DEFAULT}" PATH "Location of ${_NAME} install root. Default: None or the value of the environment variable ${_NAME}_ROOT if set")
   IF (DEFINED TPL_ENABLE_${_NAME})
-    IF (${_NAME} STREQUAL MKL AND KOKKOSKERNELS_HAS_TRILINOS)
-      MESSAGE("Trilinos has enabled MKL and SYCL but it does not detect oneMKL correctly so we disable it!")
-      SET(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE ON)
-    ENDIF ()
     IF (TPL_ENABLE_${_NAME} AND NOT KOKKOSKERNELS_ENABLE_TPL_${_NAME})
       MESSAGE("Overriding KOKKOSKERNELS_ENABLE_TPL_${_NAME_ORIG}=OFF with TPL_ENABLE_${_NAME}=ON")
       SET(KOKKOSKERNELS_ENABLE_TPL_${_NAME_ORIG} ON)

diff --git a/packages/kokkos-kernels/sparse/src/KokkosSparse_spmv_handle.hpp b/packages/kokkos-kernels/sparse/src/KokkosSparse_spmv_handle.hpp
@@ -190,7 +190,7 @@ struct MKL_SpMV_Data : public TPL_SpMV_Data<ExecutionSpace> {
 };
 #endif
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 struct OneMKL_SpMV_Data : public TPL_SpMV_Data<Kokkos::Experimental::SYCL> {
   OneMKL_SpMV_Data(const Kokkos::Experimental::SYCL& exec_) : TPL_SpMV_Data(exec_) {}
   ~OneMKL_SpMV_Data() {

diff --git a/packages/kokkos-kernels/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp b/packages/kokkos-kernels/sparse/tpls/KokkosSparse_spmv_tpl_spec_avail.hpp
@@ -179,7 +179,7 @@ KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_MKL(Kokkos::complex<float>, Kokkos::OpenMP)
 KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_MKL(Kokkos::complex<double>, Kokkos::OpenMP)
 #endif
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 #define KOKKOSSPARSE_SPMV_TPL_SPEC_AVAIL_ONEMKL(SCALAR, ORDINAL, MEMSPACE)                                       \
   template <>                                                                                                    \
   struct spmv_tpl_spec_avail<                                                                                    \

diff --git a/packages/kokkos-kernels/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp b/packages/kokkos-kernels/sparse/tpls/KokkosSparse_spmv_tpl_spec_decl.hpp
@@ -548,7 +548,7 @@ KOKKOSSPARSE_SPMV_MKL(Kokkos::complex<double>, Kokkos::OpenMP)
 #undef KOKKOSSPARSE_SPMV_MKL
 #endif
 
-#if defined(KOKKOS_ENABLE_SYCL) && !defined(KOKKOSKERNELS_ENABLE_TPL_MKL_SYCL_OVERRIDE)
+#if defined(KOKKOS_ENABLE_SYCL)
 inline oneapi::mkl::transpose mode_kk_to_onemkl(char mode_kk) {
   switch (toupper(mode_kk)) {
     case 'N': return oneapi::mkl::transpose::nontrans;