Kokkos Kernels: Fix for batched QR

Signed-off-by: Christian Glusa <caglusa@sandia.gov>
trilinos · Oct 7, 2024 · 060a997 · 060a997
1 parent df226ba
commit 060a997
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/packages/kokkos-kernels/batched/dense/impl/KokkosBatched_QR_Serial_Internal.hpp b/packages/kokkos-kernels/batched/dense/impl/KokkosBatched_QR_Serial_Internal.hpp
@@ -53,7 +53,7 @@ struct SerialQR_Internal {
     A_part2x2.partWithATL(A, m, n, 0, 0);
     t_part2x1.partWithAT(t, m, 0);
 
-    for (int m_atl = 0; m_atl < m; ++m_atl) {
+    for (int m_atl = 0; m_atl < Kokkos::min(m, n); ++m_atl) {
       // part 2x2 into 3x3
       A_part3x3.partWithABR(A_part2x2, 1, 1);
       const int m_A22 = m - m_atl - 1;

diff --git a/packages/kokkos-kernels/batched/dense/impl/KokkosBatched_QR_TeamVector_Internal.hpp b/packages/kokkos-kernels/batched/dense/impl/KokkosBatched_QR_TeamVector_Internal.hpp
@@ -54,7 +54,7 @@ struct TeamVectorQR_Internal {
     A_part2x2.partWithATL(A, m, n, 0, 0);
     t_part2x1.partWithAT(t, m, 0);
 
-    for (int m_atl = 0; m_atl < m; ++m_atl) {
+    for (int m_atl = 0; m_atl < Kokkos::min(m, n); ++m_atl) {
       // part 2x2 into 3x3
       A_part3x3.partWithABR(A_part2x2, 1, 1);
       const int m_A22 = m - m_atl - 1;