Batched - QR: Implementing a version correct for complex scalars

batched/dense/impl/KokkosBatched_ApplyHouseholder_Serial_Impl.hpp

@@ -31,9 +31,15 @@ template <>
 template <typename uViewType, typename tauViewType, typename AViewType, typename wViewType>
 KOKKOS_INLINE_FUNCTION int SerialApplyHouseholder<Side::Left>::invoke(const uViewType &u2, const tauViewType &tau,
                                                                       const AViewType &A, const wViewType &w) {
-  return SerialApplyLeftHouseholderInternal::invoke(A.extent(0) - 1, A.extent(1), tau.data(), u2.data(), u2.stride(0),
-                                                    A.data(), A.stride(1), A.data() + A.stride(0), A.stride(0),
-                                                    A.stride(1), w.data());
+  if constexpr (AViewType::rank() == 1) {
+    return SerialApplyLeftHouseholderInternal::invoke(A.extent(0) - 1, 1, tau.data(), u2.data(), u2.stride(0),
+						      A.data(), 1, A.data() + A.stride(0), A.stride(0),
+						      1, w.data());
+  } else {
+    return SerialApplyLeftHouseholderInternal::invoke(A.extent(0) - 1, A.extent(1), tau.data(), u2.data(), u2.stride(0),
+						      A.data(), A.stride(1), A.data() + A.stride(0), A.stride(0),
+						      A.stride(1), w.data());
+  }
 }
 
 template <>

batched/dense/impl/KokkosBatched_Householder_Serial_Internal.hpp

@@ -34,16 +34,18 @@ struct SerialLeftHouseholderInternal {
                                            /* */ ValueType* chi1,
                                            /* */ ValueType* x2, const int x2s,
                                            /* */ ValueType* tau) {
-    typedef ValueType value_type;
-    typedef typename Kokkos::ArithTraits<ValueType>::mag_type mag_type;
+    using value_type = ValueType;
+    using KAT        = Kokkos::ArithTraits<value_type>;
+    using mag_type   = typename KAT::mag_type;
+    using KAT_mag    = Kokkos::ArithTraits<mag_type>;
 
-    const mag_type zero(0);
-    const mag_type half(0.5);
-    const mag_type one(1);
-    const mag_type minus_one(-1);
+    const mag_type zero = KAT_mag::zero();
+    const mag_type one = KAT_mag::one();
+    const mag_type half = one / (one + one);
+    const mag_type minus_one = - one;
 
     /// compute the 2norm of x2
-    mag_type norm_x2_square(0);
+    mag_type norm_x2_square = zero;
     for (int i = 0; i < m_x2; ++i) {
       const auto x2_at_i = x2[i * x2s];
       norm_x2_square += Kokkos::abs(x2_at_i) * Kokkos::abs(x2_at_i);
@@ -52,7 +54,7 @@ struct SerialLeftHouseholderInternal {
     /// if norm_x2 is zero, return with trivial values
     if (norm_x2_square == zero) {
       *chi1 = -(*chi1);
-      *tau  = half;
+      *tau  = half * KAT::one();
 
       return 0;
     }
@@ -64,7 +66,7 @@ struct SerialLeftHouseholderInternal {
     const mag_type norm_x = Kokkos::ArithTraits<mag_type>::sqrt(norm_x2_square + norm_chi1 * norm_chi1);
 
     /// compute alpha
-    const mag_type alpha = (*chi1 < Kokkos::ArithTraits<value_type>::zero() ? one : minus_one) * norm_x;
+    const mag_type alpha = (Kokkos::real(*chi1) < Kokkos::ArithTraits<mag_type>::zero() ? one : minus_one) * norm_x;
 
     /// overwrite x2 with u2
     const value_type chi1_minus_alpha     = *chi1 - alpha;
@@ -76,7 +78,11 @@ struct SerialLeftHouseholderInternal {
 
     /// compute tau
     const mag_type chi1_minus_alpha_square = Kokkos::abs(chi1_minus_alpha) * Kokkos::abs(chi1_minus_alpha);
-    *tau                                   = half + half * (norm_x2_square / chi1_minus_alpha_square);
+    if constexpr(KAT::is_complex) {
+      *tau = alpha / (alpha - Kokkos::conj(*chi1));
+    } else {
+      *tau = half + half * (norm_x2_square / chi1_minus_alpha_square);
+    }
 
     /// overwrite chi1 with alpha
     *chi1 = alpha;

batched/dense/impl/KokkosBatched_Householder_TeamVector_Impl.hpp

@@ -27,8 +27,8 @@ namespace KokkosBatched {
 /// TeamVector Impl
 /// ===============
 
-template <typename MemberType>
-template <typename aViewType, typename tauViewType>
+template <>
+template <typename MemberType, typename aViewType, typename tauViewType>
 KOKKOS_INLINE_FUNCTION int TeamVectorHouseholder<Side::Left>::invoke(const MemberType &member, const aViewType &a,
                                                                      const tauViewType &tau) {
   return TeamVectorLeftHouseholderInternal::invoke(member, a.extent(0) - 1, a.data(), a.data() + a.stride(0),

batched/dense/src/KokkosBatched_ApplyHouseholder_Decl.hpp

@@ -31,15 +31,15 @@ template <typename ArgSide>
 struct SerialApplyHouseholder {
   template <typename uViewType, typename tauViewType, typename AViewType, typename wViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(const uViewType &u2, const tauViewType &tau,
-                                           const AViewType const wViewType &w);
+                                           const AViewType &A, const wViewType &w);
 };
 
 // level 1 operation
 template <typename MemberType, typename ArgSide>
 struct TeamVectorApplyHouseholder {
   template <typename uViewType, typename tauViewType, typename AViewType, typename wViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const uViewType &u2, const tauViewType &tau,
-                                           const AViewType const wViewType &w);
+                                           const AViewType &A, const wViewType &w);
 };
 
 }  // namespace KokkosBatched

batched/dense/src/KokkosBatched_Householder_Decl.hpp

@@ -38,7 +38,7 @@ struct SerialHouseholder {
 ///
 
 // level 1 operation
-template <typename MemberType, typename ArgSide>
+template <typename ArgSide>
 struct TeamVectorHouseholder {
   template <typename MemberType, typename aViewType, typename tauViewType>
   KOKKOS_INLINE_FUNCTION static int invoke(const MemberType &member, const aViewType &a, const tauViewType &tau);

batched/dense/unit_test/Test_Batched_Dense.hpp

@@ -63,6 +63,7 @@
 #include "Test_Batched_SerialLacgv.hpp"
 #include "Test_Batched_SerialGbtrf.hpp"
 #include "Test_Batched_SerialGbtrs.hpp"
+#include "Test_Batched_SerialHouseholder.hpp"
 
 // Team Kernels
 #include "Test_Batched_TeamAxpy.hpp"

batched/dense/unit_test/Test_Batched_SerialHouseholder.hpp

@@ -0,0 +1,133 @@
+//@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
+/// \author Luc Berger-Vergiat ([email protected])
+/// \author Cameron Smith ([email protected])
+
+#include "gtest/gtest.h"
+#include "KokkosBatched_Householder_Decl.hpp"
+#include "KokkosBatched_ApplyHouseholder_Decl.hpp"
+
+namespace Test {
+
+template <class Device, class Scalar>
+void test_Householder_analytic_real() {
+  using ExecutionSpace = typename Device::execution_space;
+  using vec_type = Kokkos::View<Scalar*, ExecutionSpace>;
+
+  const Scalar tol = 20 * Kokkos::ArithTraits<Scalar>::eps();
+
+  vec_type vec("vector to reflect", 4), reflector("reflector", 4), tau("tau", 1);
+  auto reflector_h = Kokkos::create_mirror_view(reflector);
+  reflector_h(0) = 4; reflector_h(1) = 2; reflector_h(2) = 2; reflector_h(3) = 1;
+  Kokkos::deep_copy(reflector, reflector_h);
+  Kokkos::deep_copy(vec, reflector_h);
+
+  auto tau_h = Kokkos::create_mirror_view(tau);
+  SerialHouseholder<KokkosBatched::Side::Left>::invoke(reflector_h, tau_h);
+
+  // x = [4, 2, 2, 1]
+  // output = [-5, 2/9, 2/9, 1/9]
+  Test::EXPECT_NEAR_KK_REL(reflector_h(0), -5.0, tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(1),  2.0 /  9.0, tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(2),  2.0 /  9.0, tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(3),  1.0 /  9.0, tol);
+  Test::EXPECT_NEAR_KK_REL(tau_h(0), 90.0 / 162.0, tol);
+
+  vec_type workspace("workspace", 1);
+  auto u = Kokkos::subview(reflector, Kokkos::pair<int, int>(1, 4));
+  KokkosBatched::SerialApplyHouseholder<KokkosBatched::Side::Left>::invoke(u, tau, vec, workspace);
+
+  auto vec_h = Kokkos::create_mirror_view(vec);
+  Kokkos::deep_copy(vec_h, vec);
+  Test::EXPECT_NEAR_KK_REL(vec_h(0), reflector_h(0), tol);
+  Test::EXPECT_NEAR_KK(vec_h(1),  0.0, tol);
+  Test::EXPECT_NEAR_KK(vec_h(2),  0.0, tol);
+  Test::EXPECT_NEAR_KK(vec_h(3),  0.0, tol);
+}
+
+template <class Device, class Scalar>
+void test_Householder_analytic_cplx() {
+  using ExecutionSpace = typename Device::execution_space;
+  using vec_type = Kokkos::View<Scalar*, ExecutionSpace>;
+
+  // const Scalar zero = Kokkos::ArithTraits<Scalar>::zero();
+  const Scalar tol = 20 * Kokkos::ArithTraits<Scalar>::eps();
+
+  vec_type vec("vector to reflect", 4), reflector("reflector", 4), tau("tau", 1);
+  auto reflector_h = Kokkos::create_mirror(vec);
+  reflector_h(0) = Kokkos::complex(4., 3.);
+  reflector_h(1) = Kokkos::complex(2., 5.);
+  reflector_h(2) = Kokkos::complex(2., 5.);
+  reflector_h(3) = Kokkos::complex(1., 4.);
+  Kokkos::deep_copy(reflector, reflector_h);
+  Kokkos::deep_copy(vec, reflector_h);
+
+  auto tau_h = Kokkos::create_mirror_view(tau);
+  SerialHouseholder<KokkosBatched::Side::Left>::invoke(reflector_h, tau_h);
+
+  // x = [4+3i, 2+5i, 2+5i, 1+4i]
+  // 1 / (x(0) + ||x||) -> (14-3i) / (14*14 + 3*3) = (14-3i) / 205
+  // output = [-10, (43 + 64i) / 205, (43 + 64i) / 205, (26 + 53i) / 205]
+  Test::EXPECT_NEAR_KK_REL(reflector_h(0), Scalar(-10.0,  0.0), tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(1), Scalar( 43.0, 64.0) / 205.0, tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(2), Scalar( 43.0, 64.0) / 205.0, tol);
+  Test::EXPECT_NEAR_KK_REL(reflector_h(3), Scalar( 26.0, 53.0) / 205.0, tol);
+  Test::EXPECT_NEAR_KK_REL(tau_h(0), Scalar( 1.4,  -0.3) / 2.05, tol);
+
+  std::cout << "Reflector vector: {" << reflector_h(0) << ",  " << reflector_h(1)
+	    << ",  " << reflector_h(2) << ",  " << reflector_h(3) << "}" << std::endl;
+  std::cout << "tau=" << tau_h(0) << std::endl;
+
+  vec_type workspace("workspace", 1);
+  auto u = Kokkos::subview(reflector, Kokkos::pair<int, int>(1, 4));
+  KokkosBatched::SerialApplyHouseholder<KokkosBatched::Side::Left>::invoke(u, tau, vec, workspace);
+
+  auto vec_h = Kokkos::create_mirror_view(vec);
+  Kokkos::deep_copy(vec_h, vec);
+  // Test::EXPECT_NEAR_KK_REL(vec_h(0), reflector_h(0), tol);
+  // Test::EXPECT_NEAR_KK(vec_h(1),  zero, tol);
+  // Test::EXPECT_NEAR_KK(vec_h(2),  zero, tol);
+  // Test::EXPECT_NEAR_KK(vec_h(3),  zero, tol);
+
+  std::cout << "Reflected vector: {" << vec_h(0) << ",  " << vec_h(1) << ",  " << vec_h(2) << ",  " << vec_h(3) << "}" << std::endl;
+}
+
+}  // namespace Test
+
+
+#if defined(KOKKOSKERNELS_INST_FLOAT)
+TEST_F(TestCategory, serial_householder_float) {
+  ::Test::test_Householder_analytic_real<TestDevice, float>();
+}
+#endif
+
+#if defined(KOKKOSKERNELS_INST_DOUBLE)
+TEST_F(TestCategory, serial_householder_double) {
+  ::Test::test_Householder_analytic_real<TestDevice, double>();
+}
+#endif
+
+#if defined(KOKKOSKERNELS_INST_COMPLEX_FLOAT)
+TEST_F(TestCategory, serial_householder_cplx_float) {
+  ::Test::test_Householder_analytic_cplx<TestDevice, Kokkos::complex<float>>();
+}
+#endif
+
+#if defined(KOKKOSKERNELS_INST_COMPLEX_DOUBLE)
+TEST_F(TestCategory, serial_householder_cplx_double) {
+  ::Test::test_Householder_analytic_cplx<TestDevice, Kokkos::complex<double>>();
+}
+#endif

batched/dense/unit_test/Test_Batched_SerialQR.hpp

@@ -472,10 +472,11 @@ void test_QR_batch(const int numMat, const int numRows, const int numCols) {
   // Check that Q*R = A
 
   using ExecutionSpace = typename Device::execution_space;
+  using size_type = typename Kokkos::View<Scalar*, ExecutionSpace>::size_type;
 
   {
     Kokkos::View<Scalar**, ExecutionSpace> tau("tau", numMat, numCols);
-    Kokkos::View<Scalar*, ExecutionSpace> tmp("work buffer", numMat * Kokkos::max(numRows, numCols));
+    Kokkos::View<Scalar*, ExecutionSpace> tmp("work buffer", static_cast<size_type>(numMat) * Kokkos::max(numRows, numCols));
     Kokkos::View<Scalar***, ExecutionSpace> As("A matrices", numMat, numRows, numCols);
     Kokkos::View<Scalar***, ExecutionSpace> Bs("B matrices", numMat, numRows, numCols);
     Kokkos::View<Scalar***, ExecutionSpace> Qs("Q matrices", numMat, numRows, numRows);

benchmarks/blas/KokkosBlas1_team_dot_benchmark.cpp

@@ -51,8 +51,7 @@ struct teamDotFunctor {
 
 template <class ExecSpace>
 static void run(benchmark::State& state) {
-  const auto m = state.range(0);
-  // const auto repeat = state.range(1);
+  const auto m      = state.range(0);
   // Declare type aliases
   using Scalar   = double;
   using MemSpace = typename ExecSpace::memory_space;