Kokkos Kernels: Fix for batched QR

Signed-off-by: Christian Glusa <caglusa@sandia.gov>
trilinos · Sep 18, 2024 · e10c935 · e10c935
1 parent 273b9af
commit e10c935
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;