WIP: waiting for GPU implementation of permuteJustLocal

eth-cscs · Nov 30, 2023 · d8ea3dd · d8ea3dd
1 parent 6c0ed42
commit d8ea3dd
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/include/dlaf/eigensolver/tridiag_solver/merge.h b/include/dlaf/eigensolver/tridiag_solver/merge.h
@@ -1682,10 +1682,21 @@ void mergeDistSubproblems(comm::CommunicatorGrid grid,
   auto k = ex::split(stablePartitionIndexForDeflation(dist_evecs, i_begin, i_end, ws_h.c, ws_h.d0,
                                                       ws_hm.i2, ws_h.i3, ws_hm.i5));
 
-  copy(idx_begin_tiles_vec, sz_tiles_vec, ws_hm.i5, ws.i5);
-  dlaf::permutations::internal::permuteJustLocal<B, D, T, Coord::Col>(i_begin, i_end, ws.i5, ws.e0,
-                                                                      ws.e1);
+  // Reorder Eigenvectors
+  if constexpr (Backend::MC == B) {
+    copy(idx_begin_tiles_vec, sz_tiles_vec, ws_hm.i5, ws.i5);
+    dlaf::permutations::internal::permuteJustLocal<B, D, T, Coord::Col>(i_begin, i_end, ws.i5, ws.e0,
+                                                                        ws.e1);
+  }
+  else {
+    // TODO remove this branch. It exists just because GPU permuteJustLocal is not implemented yet
+    copy(idx_loc_begin, sz_loc_tiles, ws.e0, ws_hm.e0);
+    dlaf::permutations::internal::permuteJustLocal<Backend::MC, Device::CPU, T, Coord::Col>(
+        i_begin, i_end, ws_hm.i5, ws_hm.e0, ws_hm.e2);
+    copy(idx_loc_begin, sz_loc_tiles, ws_hm.e2, ws.e1);
+  }
 
+  // Reorder Eigenvalues
   applyIndex(i_begin, i_end, ws_h.i3, ws_h.d0, ws_hm.d1);
   applyIndex(i_begin, i_end, ws_h.i3, ws_hm.z0, ws_hm.z1);
   copy(idx_begin_tiles_vec, sz_tiles_vec, ws_hm.d1, ws_h.d0);