From: Peter Munch Date: Wed, 24 Feb 2021 20:54:28 +0000 (+0100) Subject: Modularized MatrixFreeTools::compute_diagonal() X-Git-Tag: v9.3.0-rc1~407^2 X-Git-Url: https://gitweb.dealii.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=18bde0713cf2917d74d79ca3186ef3fdcd67c765;p=dealii.git Modularized MatrixFreeTools::compute_diagonal() --- diff --git a/include/deal.II/matrix_free/tools.h b/include/deal.II/matrix_free/tools.h index 55376a701b..d3498fba29 100644 --- a/include/deal.II/matrix_free/tools.h +++ b/include/deal.II/matrix_free/tools.h @@ -247,110 +247,130 @@ namespace MatrixFreeTools std::vector col; std::vector val; }; - } // namespace internal - template - void - compute_diagonal( - const MatrixFree &matrix_free, - LinearAlgebra::distributed::Vector & diagonal_global, - const std::function &)> &local_vmult, - const unsigned int dof_no, - const unsigned int quad_no, - const unsigned int first_selected_component) - { - using VectorType = LinearAlgebra::distributed::Vector; + template + class ComputeDiagonalHelper + { + public: + static const unsigned int n_lanes = VectorizedArrayType::size(); + + ComputeDiagonalHelper(FEEvaluation &phi) + : phi(phi) + {} - // initialize vector - matrix_free.initialize_dof_vector(diagonal_global); + void + reinit(const unsigned int cell) + { + this->phi.reinit(cell); + // STEP 1: get relevant information from FEEvaluation + const unsigned int first_selected_component = + phi.get_first_selected_component(); + const auto & dof_info = phi.get_dof_info(); + const unsigned int n_fe_components = dof_info.start_components.back(); + const unsigned int dofs_per_component = phi.dofs_per_component; + const auto & matrix_free = phi.get_matrix_free(); + + const unsigned int n_lanes_filled = + matrix_free.n_active_entries_per_cell_batch(cell); + + std::array dof_indices{}; + { + for (unsigned int v = 0; v < n_lanes_filled; ++v) + dof_indices[v] = + dof_info.dof_indices.data() + + dof_info + .row_starts[(cell * n_lanes + v) * n_fe_components + + first_selected_component] + .first; + } - int dummy; + // STEP 2: setup CSR storage of transposed locally-relevant + // constraint matrix + c_pools = std::array, n_lanes>(); - matrix_free.template cell_loop( - [&](const MatrixFree &matrix_free, - LinearAlgebra::distributed::Vector & diagonal_global, - const int &, - const std::pair &range) mutable { - FEEvaluation - phi(matrix_free, range, dof_no, quad_no, first_selected_component); + for (unsigned int v = 0; v < n_lanes_filled; ++v) + { + unsigned int index_indicators, next_index_indicators; + + index_indicators = + dof_info + .row_starts[(cell * n_lanes + v) * n_fe_components + + first_selected_component] + .second; + next_index_indicators = + dof_info + .row_starts[(cell * n_lanes + v) * n_fe_components + + first_selected_component + 1] + .second; + + // STEP 2a: setup locally-relevant constraint matrix in a + // coordinate list (COO) + std::vector> + locally_relevant_constrains; // (constrained local index, + // global index of dof which + // constrains, weight) + + if (n_components == 1 || n_fe_components == 1) + { + AssertDimension(n_components, + 1); // TODO: currently no block vector supported - const unsigned int n_lanes = VectorizedArrayType::size(); + unsigned int ind_local = 0; + for (; index_indicators != next_index_indicators; + ++index_indicators, ++ind_local) + { + const std::pair indicator = + dof_info.constraint_indicator[index_indicators]; - // local storage: buffer so that we access the global vector once - // note: may be larger then dofs_per_cell in the presence of - // constraints! - std::array, n_lanes> diagonals_local_constrained; + for (unsigned int j = 0; j < indicator.first; + ++j, ++ind_local) + locally_relevant_constrains.emplace_back( + ind_local, dof_indices[v][j], 1.0); - for (unsigned int cell = range.first; cell < range.second; ++cell) - { - // STEP 1: get relevant information from FEEvaluation - const auto & dof_info = matrix_free.get_dof_info(dof_no); - const unsigned int n_fe_components = - dof_info.start_components.back(); - const unsigned int dofs_per_component = phi.dofs_per_component; + dof_indices[v] += indicator.first; - const unsigned int n_lanes_filled = - matrix_free.n_active_entries_per_cell_batch(cell); + const Number *data_val = + matrix_free.constraint_pool_begin(indicator.second); + const Number *end_pool = + matrix_free.constraint_pool_end(indicator.second); - std::array dof_indices{}; - { - for (unsigned int v = 0; v < n_lanes_filled; ++v) - dof_indices[v] = - dof_info.dof_indices.data() + - dof_info - .row_starts[(cell * n_lanes + v) * n_fe_components + - first_selected_component] - .first; - } + for (; data_val != end_pool; ++data_val, ++dof_indices[v]) + locally_relevant_constrains.emplace_back(ind_local, + *dof_indices[v], + *data_val); + } - // STEP 2: setup CSR storage of transposed locally-relevant - // constraint matrix - std::array, n_lanes> c_pools; + AssertIndexRange(ind_local, dofs_per_component + 1); - for (unsigned int v = 0; v < n_lanes_filled; ++v) + for (; ind_local < dofs_per_component; + ++dof_indices[v], ++ind_local) + locally_relevant_constrains.emplace_back(ind_local, + *dof_indices[v], + 1.0); + } + else { - unsigned int index_indicators, next_index_indicators; - - index_indicators = - dof_info - .row_starts[(cell * n_lanes + v) * n_fe_components + - first_selected_component] - .second; - next_index_indicators = - dof_info - .row_starts[(cell * n_lanes + v) * n_fe_components + - first_selected_component + 1] - .second; - - // STEP 2a: setup locally-relevant constraint matrix in a - // coordinate list (COO) - std::vector> - locally_relevant_constrains; // (constrained local index, - // global index of dof which - // constrains, weight) - - if (n_components == 1 || n_fe_components == 1) + // case with vector-valued finite elements where all + // components are included in one single vector. Assumption: + // first come all entries to the first component, then all + // entries to the second one, and so on. This is ensured by + // the way MatrixFree reads out the indices. + for (unsigned int comp = 0; comp < n_components; ++comp) { - AssertDimension( - n_components, - 1); // TODO: currently no block vector supported - unsigned int ind_local = 0; + + // check whether there is any constraint on the current + // cell for (; index_indicators != next_index_indicators; ++index_indicators, ++ind_local) { @@ -358,11 +378,13 @@ namespace MatrixFreeTools indicator = dof_info.constraint_indicator[index_indicators]; + // run through values up to next constraint for (unsigned int j = 0; j < indicator.first; ++j, ++ind_local) locally_relevant_constrains.emplace_back( - ind_local, dof_indices[v][j], 1.0); - + comp * dofs_per_component + ind_local, + dof_indices[v][j], + 1.0); dof_indices[v] += indicator.first; const Number *data_val = @@ -373,218 +395,257 @@ namespace MatrixFreeTools for (; data_val != end_pool; ++data_val, ++dof_indices[v]) locally_relevant_constrains.emplace_back( - ind_local, *dof_indices[v], *data_val); + comp * dofs_per_component + ind_local, + *dof_indices[v], + *data_val); } AssertIndexRange(ind_local, dofs_per_component + 1); + // get the dof values past the last constraint for (; ind_local < dofs_per_component; ++dof_indices[v], ++ind_local) - locally_relevant_constrains.emplace_back(ind_local, - *dof_indices[v], - 1.0); - } - else - { - // case with vector-valued finite elements where all - // components are included in one single vector. Assumption: - // first come all entries to the first component, then all - // entries to the second one, and so on. This is ensured by - // the way MatrixFree reads out the indices. - for (unsigned int comp = 0; comp < n_components; ++comp) - { - unsigned int ind_local = 0; - - // check whether there is any constraint on the current - // cell - for (; index_indicators != next_index_indicators; - ++index_indicators, ++ind_local) - { - const std::pair - indicator = - dof_info.constraint_indicator[index_indicators]; - - // run through values up to next constraint - for (unsigned int j = 0; j < indicator.first; - ++j, ++ind_local) - locally_relevant_constrains.emplace_back( - comp * dofs_per_component + ind_local, - dof_indices[v][j], - 1.0); - dof_indices[v] += indicator.first; - - const Number *data_val = - matrix_free.constraint_pool_begin( - indicator.second); - const Number *end_pool = - matrix_free.constraint_pool_end(indicator.second); - - for (; data_val != end_pool; - ++data_val, ++dof_indices[v]) - locally_relevant_constrains.emplace_back( - comp * dofs_per_component + ind_local, - *dof_indices[v], - *data_val); - } - - AssertIndexRange(ind_local, dofs_per_component + 1); - - // get the dof values past the last constraint - for (; ind_local < dofs_per_component; - ++dof_indices[v], ++ind_local) - locally_relevant_constrains.emplace_back( - comp * dofs_per_component + ind_local, - *dof_indices[v], - 1.0); + locally_relevant_constrains.emplace_back( + comp * dofs_per_component + ind_local, + *dof_indices[v], + 1.0); - if (comp + 1 < n_components) - { - next_index_indicators = - dof_info - .row_starts[(cell * n_lanes + v) * - n_fe_components + - first_selected_component + comp + 2] - .second; - } + if (comp + 1 < n_components) + { + next_index_indicators = + dof_info + .row_starts[(cell * n_lanes + v) * n_fe_components + + first_selected_component + comp + 2] + .second; } } + } - // STEP 2b: transpose COO - - // presort vector for transposed access - std::sort(locally_relevant_constrains.begin(), - locally_relevant_constrains.end(), - [](const auto &a, const auto &b) { - if (std::get<1>(a) < std::get<1>(b)) - return true; - return (std::get<1>(a) == std::get<1>(b)) && - (std::get<0>(a) < std::get<0>(b)); - }); - - // make sure that all entries are unique - locally_relevant_constrains.erase( - unique(locally_relevant_constrains.begin(), - locally_relevant_constrains.end(), - [](const auto &a, const auto &b) { - return (std::get<1>(a) == std::get<1>(b)) && - (std::get<0>(a) == std::get<0>(b)); - }), - locally_relevant_constrains.end()); - - // STEP 2c: translate COO to CRS - auto &c_pool = c_pools[v]; + // STEP 2b: transpose COO + + // presort vector for transposed access + std::sort(locally_relevant_constrains.begin(), + locally_relevant_constrains.end(), + [](const auto &a, const auto &b) { + if (std::get<1>(a) < std::get<1>(b)) + return true; + return (std::get<1>(a) == std::get<1>(b)) && + (std::get<0>(a) < std::get<0>(b)); + }); + + // make sure that all entries are unique + locally_relevant_constrains.erase( + unique(locally_relevant_constrains.begin(), + locally_relevant_constrains.end(), + [](const auto &a, const auto &b) { + return (std::get<1>(a) == std::get<1>(b)) && + (std::get<0>(a) == std::get<0>(b)); + }), + locally_relevant_constrains.end()); + + // STEP 2c: translate COO to CRS + auto &c_pool = c_pools[v]; + { + if (locally_relevant_constrains.size() > 0) + c_pool.row_lid_to_gid.emplace_back( + std::get<1>(locally_relevant_constrains.front())); + for (const auto &j : locally_relevant_constrains) { - if (locally_relevant_constrains.size() > 0) - c_pool.row_lid_to_gid.emplace_back( - std::get<1>(locally_relevant_constrains.front())); - for (const auto &j : locally_relevant_constrains) + if (c_pool.row_lid_to_gid.back() != std::get<1>(j)) { - if (c_pool.row_lid_to_gid.back() != std::get<1>(j)) - { - c_pool.row_lid_to_gid.push_back(std::get<1>(j)); - c_pool.row.push_back(c_pool.val.size()); - } - - c_pool.col.emplace_back(std::get<0>(j)); - c_pool.val.emplace_back(std::get<2>(j)); + c_pool.row_lid_to_gid.push_back(std::get<1>(j)); + c_pool.row.push_back(c_pool.val.size()); } - if (c_pool.val.size() > 0) - c_pool.row.push_back(c_pool.val.size()); + c_pool.col.emplace_back(std::get<0>(j)); + c_pool.val.emplace_back(std::get<2>(j)); } - } - // STEP 3: compute element matrix A_e, apply - // locally-relevant constraints C_e^T * A_e * C_e, and get the - // the diagonal entry - // (C_e^T * A_e * C_e)(i,i) - // or - // C_e^T(i,:) * A_e * C_e(:,i). - // - // Since, we compute the element matrix column-by-column and as a - // result never actually have the full element matrix, we actually - // perform following steps: - // 1) loop over all columns of the element matrix - // a) compute column i - // b) compute for each j (rows of C_e^T): - // (C_e^T(j,:) * A_e(:,i)) * C_e(i,j) - // or - // (C_e^T(j,:) * A_e(:,i)) * C_e^T(j,i) - // This gives a contribution the the j-th entry of the - // locally-relevant diagonal and comprises the multiplication - // by the locally-relevant constraint matrix from the left and - // the right. There is no contribution to the j-th vector - // entry if the j-th row of C_e^T is empty or C_e^T(j,i) is - // zero. - - // set size locally-relevant diagonal - for (unsigned int v = 0; v < n_lanes_filled; ++v) - diagonals_local_constrained[v].assign( - c_pools[v].row_lid_to_gid.size(), Number(0.0)); - - phi.reinit(cell); - - // loop over all columns of element stiffness matrix - for (unsigned int i = 0; i < phi.dofs_per_cell; ++i) + if (c_pool.val.size() > 0) + c_pool.row.push_back(c_pool.val.size()); + } + } + // STEP 3: compute element matrix A_e, apply + // locally-relevant constraints C_e^T * A_e * C_e, and get the + // the diagonal entry + // (C_e^T * A_e * C_e)(i,i) + // or + // C_e^T(i,:) * A_e * C_e(:,i). + // + // Since, we compute the element matrix column-by-column and as a + // result never actually have the full element matrix, we actually + // perform following steps: + // 1) loop over all columns of the element matrix + // a) compute column i + // b) compute for each j (rows of C_e^T): + // (C_e^T(j,:) * A_e(:,i)) * C_e(i,j) + // or + // (C_e^T(j,:) * A_e(:,i)) * C_e^T(j,i) + // This gives a contribution the the j-th entry of the + // locally-relevant diagonal and comprises the multiplication + // by the locally-relevant constraint matrix from the left and + // the right. There is no contribution to the j-th vector + // entry if the j-th row of C_e^T is empty or C_e^T(j,i) is + // zero. + + // set size locally-relevant diagonal + for (unsigned int v = 0; v < n_lanes_filled; ++v) + diagonals_local_constrained[v].assign( + c_pools[v].row_lid_to_gid.size(), Number(0.0)); + } + + void + prepare_basis_vector(const unsigned int i) + { + this->i = i; + + // compute i-th column of element stiffness matrix: + // this could be simply performed as done at the moment with + // matrix-free operator evaluation applied to a ith-basis vector + for (unsigned int j = 0; j < phi.dofs_per_cell; ++j) + phi.begin_dof_values()[j] = static_cast(i == j); + } + + void + submit() + { + const auto ith_column = phi.begin_dof_values(); + + // apply local constraint matrix from left and from right: + // loop over all rows of transposed constrained matrix + for (unsigned int v = 0; + v < phi.get_matrix_free().n_active_entries_per_cell_batch( + phi.get_current_cell_index()); + ++v) + { + const auto &c_pool = c_pools[v]; + + for (unsigned int j = 0; j < c_pool.row.size() - 1; ++j) { - // compute i-th column of element stiffness matrix: - // this could be simply performed as done at the moment with - // matrix-free operator evaluation applied to a ith-basis vector + // check if the result will be zero, so that we can skip + // the following computations -> binary search + const auto scale_iterator = + std::lower_bound(c_pool.col.begin() + c_pool.row[j], + c_pool.col.begin() + c_pool.row[j + 1], + i); + + // explanation: j-th row of C_e^T is empty (see above) + if (scale_iterator == c_pool.col.begin() + c_pool.row[j + 1]) + continue; + + // explanation: C_e^T(j,i) is zero (see above) + if (*scale_iterator != i) + continue; + + // apply constraint matrix from the left + Number temp = 0.0; + for (unsigned int k = c_pool.row[j]; k < c_pool.row[j + 1]; ++k) + temp += c_pool.val[k] * ith_column[c_pool.col[k]][v]; + + // apply constraint matrix from the right + diagonals_local_constrained[v][j] += + temp * + c_pool.val[std::distance(c_pool.col.begin(), scale_iterator)]; + } + } + } - for (unsigned int j = 0; j < phi.dofs_per_cell; ++j) - phi.begin_dof_values()[j] = static_cast(i == j); + void + distribute_local_to_global( + LinearAlgebra::distributed::Vector &diagonal_global) + { + // STEP 4: assembly results: add into global vector + for (unsigned int v = 0; + v < phi.get_matrix_free().n_active_entries_per_cell_batch( + phi.get_current_cell_index()); + ++v) + for (unsigned int j = 0; j < c_pools[v].row.size() - 1; ++j) + ::dealii::internal::vector_access_add( + diagonal_global, + c_pools[v].row_lid_to_gid[j], + diagonals_local_constrained[v][j]); + } - local_vmult(phi); + private: + FEEvaluation φ - const auto ith_column = phi.begin_dof_values(); + unsigned int i; - // apply local constraint matrix from left and from right: - // loop over all rows of transposed constrained matrix - for (unsigned int v = 0; v < n_lanes_filled; ++v) - { - const auto &c_pool = c_pools[v]; + std::array, n_lanes> c_pools; - for (unsigned int j = 0; j < c_pool.row.size() - 1; ++j) - { - // check if the result will be zero, so that we can skip - // the following computations -> binary search - const auto scale_iterator = - std::lower_bound(c_pool.col.begin() + c_pool.row[j], - c_pool.col.begin() + - c_pool.row[j + 1], - i); - - // explanation: j-th row of C_e^T is empty (see above) - if (scale_iterator == - c_pool.col.begin() + c_pool.row[j + 1]) - continue; - - // explanation: C_e^T(j,i) is zero (see above) - if (*scale_iterator != i) - continue; - - // apply constraint matrix from the left - Number temp = 0.0; - for (unsigned int k = c_pool.row[j]; - k < c_pool.row[j + 1]; - ++k) - temp += c_pool.val[k] * ith_column[c_pool.col[k]][v]; - - // apply constraint matrix from the right - diagonals_local_constrained[v][j] += - temp * c_pool.val[std::distance(c_pool.col.begin(), - scale_iterator)]; - } - } + // local storage: buffer so that we access the global vector once + // note: may be larger then dofs_per_cell in the presence of + // constraints! + std::array, n_lanes> diagonals_local_constrained; + }; + + } // namespace internal + + template + void + compute_diagonal( + const MatrixFree &matrix_free, + LinearAlgebra::distributed::Vector & diagonal_global, + const std::function &)> &local_vmult, + const unsigned int dof_no, + const unsigned int quad_no, + const unsigned int first_selected_component) + { + using VectorType = LinearAlgebra::distributed::Vector; + + // initialize vector + matrix_free.initialize_dof_vector(diagonal_global, dof_no); + + int dummy; + + matrix_free.template cell_loop( + [&](const MatrixFree &matrix_free, + LinearAlgebra::distributed::Vector & diagonal_global, + const int &, + const std::pair &range) mutable { + FEEvaluation + phi(matrix_free, range, dof_no, quad_no, first_selected_component); + + internal::ComputeDiagonalHelper + helper(phi); + + for (unsigned int cell = range.first; cell < range.second; ++cell) + { + helper.reinit(cell); + + for (unsigned int i = 0; i < phi.dofs_per_cell; ++i) + { + helper.prepare_basis_vector(i); + local_vmult(phi); + helper.submit(); } - // STEP 4: assembly results: add into global vector - for (unsigned int v = 0; v < n_lanes_filled; ++v) - for (unsigned int j = 0; j < c_pools[v].row.size() - 1; ++j) - ::dealii::internal::vector_access_add( - diagonal_global, - c_pools[v].row_lid_to_gid[j], - diagonals_local_constrained[v][j]); + helper.distribute_local_to_global(diagonal_global); } }, diagonal_global,