From 80bc9a2c4ed071936897423b0bccd2c636ba142c Mon Sep 17 00:00:00 2001 From: Bruno Turcksin Date: Thu, 16 Mar 2023 17:01:54 -0400 Subject: [PATCH] Move inv_jacobian to Kokkos::View --- .../deal.II/matrix_free/cuda_fe_evaluation.h | 31 +++++++++----- .../deal.II/matrix_free/cuda_matrix_free.h | 15 ++++--- .../matrix_free/cuda_matrix_free.templates.h | 40 ++++++++----------- 3 files changed, 48 insertions(+), 38 deletions(-) diff --git a/include/deal.II/matrix_free/cuda_fe_evaluation.h b/include/deal.II/matrix_free/cuda_fe_evaluation.h index 541c0ab852..db1b3e0638 100644 --- a/include/deal.II/matrix_free/cuda_fe_evaluation.h +++ b/include/deal.II/matrix_free/cuda_fe_evaluation.h @@ -250,8 +250,20 @@ namespace CUDAWrappers const bool use_coloring; - Number *inv_jac; - // We would like to use + // FIXME We would like to use + // Kokkos::Subview, int, decltype(Kokkos::ALL), + // decltype(Kokkos::ALL), decltype(Kokkos::ALL)> but we get error: + // incomplete type is not allowed. I cannot reproduce outside of deal.II. + // Need to investigate more. + Kokkos::Subview< + Kokkos::View, + int, + Kokkos::pair, + Kokkos::pair, + Kokkos::pair> + inv_jac; + // FIXME We would like to use // Kokkos::Subview, int, decltype(Kokkos::ALL)> // but we get error: incomplete type is not allowed. I cannot reproduce @@ -283,6 +295,12 @@ namespace CUDAWrappers , mf_object_id(data->id) , constraint_mask(data->constraint_mask[cell_id]) , use_coloring(data->use_coloring) + , inv_jac(Kokkos::subview( + data->inv_jacobian, + cell_id, + Kokkos::pair(0, Utilities::pow(n_q_points_1d, dim)), + Kokkos::pair(0, dim), + Kokkos::pair(0, dim))) , JxW(Kokkos::subview( data->JxW, cell_id, @@ -290,7 +308,6 @@ namespace CUDAWrappers , values(shdata->values) { local_to_global = data->local_to_global + padding_length * cell_id; - inv_jac = data->inv_jacobian + padding_length * cell_id; for (unsigned int i = 0; i < dim; ++i) gradients[i] = shdata->gradients[i]; @@ -510,14 +527,12 @@ namespace CUDAWrappers // TODO optimize if the mesh is uniform const unsigned int q_point = internal::compute_index(); - const Number * inv_jacobian = &inv_jac[q_point]; gradient_type grad; for (unsigned int d_1 = 0; d_1 < dim; ++d_1) { Number tmp = 0.; for (unsigned int d_2 = 0; d_2 < dim; ++d_2) - tmp += inv_jacobian[padding_length * n_cells * (dim * d_2 + d_1)] * - gradients[d_2][q_point]; + tmp += inv_jac(q_point, d_2, d_1) * gradients[d_2][q_point]; grad[d_1] = tmp; } @@ -537,13 +552,11 @@ namespace CUDAWrappers { // TODO optimize if the mesh is uniform const unsigned int q_point = internal::compute_index(); - const Number * inv_jacobian = &inv_jac[q_point]; for (unsigned int d_1 = 0; d_1 < dim; ++d_1) { Number tmp = 0.; for (unsigned int d_2 = 0; d_2 < dim; ++d_2) - tmp += inv_jacobian[n_cells * padding_length * (dim * d_1 + d_2)] * - grad_in[d_2]; + tmp += inv_jac(q_point, d_1, d_2) * grad_in[d_2]; gradients[d_1][q_point] = tmp * JxW[q_point]; } } diff --git a/include/deal.II/matrix_free/cuda_matrix_free.h b/include/deal.II/matrix_free/cuda_matrix_free.h index 74d6955cd1..069e98e74a 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.h @@ -185,9 +185,10 @@ namespace CUDAWrappers types::global_dof_index *local_to_global; /** - * Pointer to the inverse Jacobian. + * Kokkos::View of the inverse Jacobian. */ - Number *inv_jacobian; + Kokkos::View + inv_jacobian; /** * Kokkos::View of the Jacobian times the weights. @@ -514,10 +515,12 @@ namespace CUDAWrappers std::vector local_to_global; /** - * Vector of pointer to the inverse Jacobian associated to the cells of each - * color. + * Vector of Kokkos::View of the inverse Jacobian associated to the cells of + * each color. */ - std::vector inv_jacobian; + std::vector< + Kokkos::View> + inv_jacobian; /** * Vector of Kokkos::View to the Jacobian times the weights associated to @@ -813,7 +816,7 @@ namespace CUDAWrappers if (update_flags & update_gradients) { data_host.inv_jacobian.resize(n_elements * dim * dim); - Utilities::CUDA::copy_to_host(data.inv_jacobian, + Utilities::CUDA::copy_to_host(data.inv_jacobian.data(), data_host.inv_jacobian); } diff --git a/include/deal.II/matrix_free/cuda_matrix_free.templates.h b/include/deal.II/matrix_free/cuda_matrix_free.templates.h index 69c9f885b6..c50828c180 100644 --- a/include/deal.II/matrix_free/cuda_matrix_free.templates.h +++ b/include/deal.II/matrix_free/cuda_matrix_free.templates.h @@ -236,7 +236,8 @@ namespace CUDAWrappers std::vector local_to_global_host; std::vector> q_points_host; Kokkos::View JxW; - std::vector inv_jacobian_host; + Kokkos::View + inv_jacobian; std::vector constraint_mask_host; // Local buffer @@ -359,7 +360,12 @@ namespace CUDAWrappers dofs_per_cell); if (update_flags & update_gradients) - inv_jacobian_host.resize(n_cells * padding_length * dim * dim); + inv_jacobian = + Kokkos::View( + Kokkos::view_alloc("inv_jacobian_" + std::to_string(color), + Kokkos::WithoutInitializing), + n_cells, + dofs_per_cell); constraint_mask_host.resize(n_cells); } @@ -423,13 +429,16 @@ namespace CUDAWrappers if (update_flags & update_gradients) { + // FIXME too many deep_copy + auto inv_jacobian_host = Kokkos::create_mirror_view_and_copy( + MemorySpace::Host::kokkos_space{}, inv_jacobian); const std::vector> &inv_jacobians = fe_values.get_inverse_jacobians(); - std::copy(&inv_jacobians[0][0][0], - &inv_jacobians[0][0][0] + - q_points_per_cell * sizeof(DerivativeForm<1, dim, dim>) / - sizeof(double), - &inv_jacobian_host[cell_id * padding_length * dim * dim]); + for (unsigned int i = 0; i < q_points_per_cell; ++i) + for (unsigned int j = 0; j < dim; ++j) + for (unsigned int k = 0; k < dim; ++k) + inv_jacobian_host(cell_id, i, j, k) = inv_jacobians[i][j][k]; + Kokkos::deep_copy(inv_jacobian, inv_jacobian_host); } } @@ -466,18 +475,7 @@ namespace CUDAWrappers // Inverse jacobians if (update_flags & update_gradients) { - // Reorder so that all J_11 elements are together, all J_12 elements - // are together, etc., i.e., reorder indices from - // cell_id*q_points_per_cell*dim*dim + q*dim*dim +i to - // i*q_points_per_cell*n_cells + cell_id*q_points_per_cell+q - transpose_in_place(inv_jacobian_host, - padding_length * n_cells, - dim * dim); - - alloc_and_copy(&data->inv_jacobian[color], - ArrayView(inv_jacobian_host.data(), - inv_jacobian_host.size()), - n_cells * dim * dim * padding_length); + data->inv_jacobian[color] = inv_jacobian; } alloc_and_copy( @@ -716,10 +714,6 @@ namespace CUDAWrappers Utilities::CUDA::free(local_to_global_color_ptr); local_to_global.clear(); - for (auto &inv_jacobian_color_ptr : inv_jacobian) - Utilities::CUDA::free(inv_jacobian_color_ptr); - inv_jacobian.clear(); - for (auto &constraint_mask_color_ptr : constraint_mask) Utilities::CUDA::free(constraint_mask_color_ptr); constraint_mask.clear(); -- 2.39.5