Lapack: fixing issue with Magma TPL in gesv, trtri, etc...

blas/src/KokkosBlas1_swap.hpp

@@ -26,12 +26,12 @@ namespace KokkosBlas {
 /// \brief Swaps the entries of vectors x and y.
 ///
 /// \tparam execution_space an execution space to perform parallel work
-/// \tparam XVector Type of the first vector x; a 1-D Kokkos::View.
-/// \tparam YVector Type of the first vector y; a 1-D Kokkos::View.
+/// \tparam XVector Type of the first vector x; a rank 1 Kokkos::View.
+/// \tparam YVector Type of the first vector y; a rank 1 Kokkos::View.
 ///
 /// \param space [in] execution space passed to execution policies
-/// \param x [in/out] 1-D View.
-/// \param y [in/out] 1-D View.
+/// \param x [in/out] rank 1 View.
+/// \param y [in/out] rank 1 View.
 ///
 /// Swaps x and y. Note that this is akin to performing a deep_copy, swapping
 /// pointers inside view can only be performed if no aliasing, subviews, etc...
@@ -100,11 +100,11 @@ void swap(execution_space const& space, XVector const& x, YVector const& y) {
 
 /// \brief Swaps the entries of vectors x and y.
 ///
-/// \tparam XVector Type of the first vector x; a 1-D Kokkos::View.
-/// \tparam YVector Type of the first vector y; a 1-D Kokkos::View.
+/// \tparam XVector Type of the first vector x; a rank 1 Kokkos::View.
+/// \tparam YVector Type of the first vector y; a rank 1 Kokkos::View.
 ///
-/// \param x [in/out] 1-D View.
-/// \param y [in/out] 1-D View.
+/// \param x [in/out] rank 1 View.
+/// \param y [in/out] rank 1 View.
 ///
 /// This function is non-blocking unless the underlying TPL requested
 /// at compile time is itself blocking. Note that the kernel will be

blas/unit_test/Test_Blas1_swap.hpp

@@ -3,11 +3,12 @@
 namespace Test {
 namespace Impl {
 
-template <class VectorType>
+template <class ScalarType, class DeviceType>
 void test_swap(int const vector_length) {
-  using vector_type     = VectorType;
-  using execution_space = typename vector_type::execution_space;
-  using scalar_type     = typename VectorType::non_const_value_type;
+  using execution_space = typename DeviceType::execution_space;
+  using memory_space    = typename DeviceType::memory_space;
+  using vector_type     = Kokkos::View<ScalarType*, memory_space>;
+  using scalar_type     = typename vector_type::non_const_value_type;
   using mag_type        = typename Kokkos::ArithTraits<scalar_type>::mag_type;
 
   // Note that Xref and Yref need to always be copies of X and Y
@@ -43,14 +44,12 @@ void test_swap(int const vector_length) {
 }  // namespace Impl
 }  // namespace Test
 
-template <class scalar_type, class execution_space>
+template <class ScalarType, class DeviceType>
 int test_swap() {
-  using Vector = Kokkos::View<scalar_type*, execution_space>;
-
-  Test::Impl::test_swap<Vector>(0);
-  Test::Impl::test_swap<Vector>(10);
-  Test::Impl::test_swap<Vector>(256);
-  Test::Impl::test_swap<Vector>(1024);
+  Test::Impl::test_swap<ScalarType, DeviceType>(0);
+  Test::Impl::test_swap<ScalarType, DeviceType>(10);
+  Test::Impl::test_swap<ScalarType, DeviceType>(256);
+  Test::Impl::test_swap<ScalarType, DeviceType>(1024);
 
   return 0;
 }

lapack/src/KokkosLapack_gesv.hpp

@@ -40,6 +40,8 @@ namespace KokkosLapack {
 /// 1-D or 2-D Kokkos::View.
 /// \tparam IPIVV Output pivot indices, as a 1-D Kokkos::View
 ///
+/// \param space [in] execution space instance used to specified how to execute
+///   the gesv kernels.
 /// \param A [in,out] On entry, the N-by-N matrix to be solved. On exit, the
 /// factors L and U from
 ///   the factorization A = P*L*U; the unit diagonal elements of L are not
@@ -166,6 +168,29 @@ void gesv(const ExecutionSpace& space, const AMatrix& A, const BXMV& B,
   }
 }
 
+/// \brief Solve the dense linear equation system A*X = B.
+///
+/// \tparam AMatrix Input matrix/Output LU, as a 2-D Kokkos::View.
+/// \tparam BXMV Input (right-hand side)/Output (solution) (multi)vector, as a
+/// 1-D or 2-D Kokkos::View.
+/// \tparam IPIVV Output pivot indices, as a 1-D Kokkos::View
+///
+/// \param A [in,out] On entry, the N-by-N matrix to be solved. On exit, the
+/// factors L and U from
+///   the factorization A = P*L*U; the unit diagonal elements of L are not
+///   stored.
+/// \param B [in,out] On entry, the right hand side (multi)vector B. On exit,
+/// the solution (multi)vector X.
+/// \param IPIV [out] On exit, the pivot indices (for partial pivoting).
+/// If the View extents are zero and its data pointer is NULL, pivoting is not
+/// used.
+///
+template <class AMatrix, class BXMV, class IPIVV>
+void gesv(const AMatrix& A, const BXMV& B, const IPIVV& IPIV) {
+  typename AMatrix::execution_space space{};
+  gesv(space, A, B, IPIV);
+}
+
 }  // namespace KokkosLapack
 
 #endif  // KOKKOSLAPACK_GESV_HPP_

lapack/tpls/KokkosLapack_Cuda_tpl.hpp

@@ -40,7 +40,7 @@ CudaLapackSingleton& CudaLapackSingleton::singleton() {
 #endif  // defined (KOKKOSKERNELS_ENABLE_TPL_CUSOLVER)
 
 #if defined(KOKKOSKERNELS_ENABLE_TPL_MAGMA)
-#include <KokkosLapack_tpl_spec.hpp>
+#include <KokkosLapack_magma.hpp>
 
 namespace KokkosLapack {
 namespace Impl {

lapack/tpls/KokkosLapack_gesv_tpl_spec_avail.hpp

@@ -50,10 +50,15 @@ KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_LAPACK(Kokkos::complex<double>,
 KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_LAPACK(Kokkos::complex<float>,
                                         Kokkos::LayoutLeft, Kokkos::HostSpace)
 #endif
+}  // namespace Impl
+}  // namespace KokkosLapack
 
 // MAGMA
 #ifdef KOKKOSKERNELS_ENABLE_TPL_MAGMA
+#include "magma_v2.h"
 
+namespace KokkosLapack {
+namespace Impl {
 #define KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(SCALAR, LAYOUT, MEMSPACE)     \
   template <>                                                                \
   struct gesv_tpl_spec_avail<                                                \
@@ -62,7 +67,9 @@ KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_LAPACK(Kokkos::complex<float>,
                    Kokkos::MemoryTraits<Kokkos::Unmanaged> >,                \
       Kokkos::View<SCALAR**, LAYOUT, Kokkos::Device<Kokkos::Cuda, MEMSPACE>, \
                    Kokkos::MemoryTraits<Kokkos::Unmanaged> >,                \
-      Kokkos::View<int*, LAYOUT, Kokkos::HostSpace,                          \
+      Kokkos::View<magma_int_t*, LAYOUT,                                     \
+                   Kokkos::Device<Kokkos::DefaultHostExecutionSpace,         \
+                                  Kokkos::HostSpace>,                        \
                    Kokkos::MemoryTraits<Kokkos::Unmanaged> > > {             \
     enum : bool { value = true };                                            \
   };
@@ -75,9 +82,9 @@ KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<double>,
                                        Kokkos::LayoutLeft, Kokkos::CudaSpace)
 KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_MAGMA(Kokkos::complex<float>,
                                        Kokkos::LayoutLeft, Kokkos::CudaSpace)
-#endif
 }  // namespace Impl
 }  // namespace KokkosLapack
+#endif  // KOKKOSKERNELS_ENABLE_TPL_MAGMA
 
 // CUSOLVER
 #ifdef KOKKOSKERNELS_ENABLE_TPL_CUSOLVER
@@ -106,6 +113,19 @@ KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(Kokkos::complex<double>,
 KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(Kokkos::complex<float>,
                                           Kokkos::LayoutLeft, Kokkos::CudaSpace)
 
+#if defined(KOKKOSKERNELS_INST_MEMSPACE_CUDAUVMSPACE)
+KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(double, Kokkos::LayoutLeft,
+                                          Kokkos::CudaUVMSpace)
+KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(float, Kokkos::LayoutLeft,
+                                          Kokkos::CudaUVMSpace)
+KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(Kokkos::complex<double>,
+                                          Kokkos::LayoutLeft,
+                                          Kokkos::CudaUVMSpace)
+KOKKOSLAPACK_GESV_TPL_SPEC_AVAIL_CUSOLVER(Kokkos::complex<float>,
+                                          Kokkos::LayoutLeft,
+                                          Kokkos::CudaUVMSpace)
+#endif
+
 }  // namespace Impl
 }  // namespace KokkosLapack
 #endif  // CUSOLVER