Merge pull request #221 from Exawind/reduce_scratch_memory_size

Reduced the amount of scratch memory used for assembling the system.
Exawind · Jul 30, 2024 · 47744e8 · 47744e8
2 parents e68529d + 3f187f1
commit 47744e8
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/restruct_poc/solver/assemble_system.hpp b/src/restruct_poc/solver/assemble_system.hpp
@@ -35,10 +35,9 @@ void AssembleSystem(Solver& solver, Beams& beams, Subview_N R_system) {
     );
 
     auto num_rows = solver.num_system_dofs;
-    auto num_columns = solver.num_system_dofs;
 
-    auto row_data_size = Kokkos::View<double*>::shmem_size(num_columns);
-    auto col_idx_size = Kokkos::View<int*>::shmem_size(num_columns);
+    auto row_data_size = Kokkos::View<double*>::shmem_size(solver.matrix_terms.extent(2));
+    auto col_idx_size = Kokkos::View<int*>::shmem_size(solver.matrix_terms.extent(2));
     auto sparse_matrix_policy = Kokkos::TeamPolicy<>(num_rows, Kokkos::AUTO());
     sparse_matrix_policy.set_scratch_size(1, Kokkos::PerTeam(row_data_size + col_idx_size));