Make cpr weight calculation thread parallel

opm/simulators/linalg/ISTLSolver.hpp

@@ -30,6 +30,8 @@
 #include <opm/common/Exceptions.hpp>
 #include <opm/common/TimingMacros.hpp>
 
+#include <opm/grid/utility/ElementChunks.hpp>
+
 #include <opm/models/discretization/common/fvbaseproperties.hh>
 #include <opm/models/common/multiphasebaseproperties.hh>
 #include <opm/models/utils/parametersystem.hpp>
@@ -161,6 +163,8 @@ std::unique_ptr<Matrix> blockJacobiAdjacency(const Grid& grid,
         using AbstractPreconditionerType = Dune::PreconditionerWithUpdate<Vector, Vector>;
         using WellModelOperator = WellModelAsLinearOperator<WellModel, Vector, Vector>;
         using ElementMapper = GetPropType<TypeTag, Properties::ElementMapper>;
+        using ElementChunksType = ElementChunks<GridView, Dune::Partitions::All>;
+
         constexpr static std::size_t pressureIndex = GetPropType<TypeTag, Properties::Indices>::pressureSwitchIdx;
 
         enum { enablePolymerMolarWeight = getPropValue<TypeTag, Properties::EnablePolymerMW>() };
@@ -304,6 +308,8 @@ std::unique_ptr<Matrix> blockJacobiAdjacency(const Grid& grid,
 
             // Print parameters to PRT/DBG logs.
             detail::printLinearSolverParameters(parameters_, activeSolverNum_, prm_,  simulator_.gridView().comm());
+
+            element_chunks_ = std::make_unique<ElementChunksType>(simulator_.vanguard().gridView(), Dune::Partitions::all, ThreadManager::maxThreads());
         }
 
         // nothing to clean here
@@ -590,10 +596,12 @@ std::unique_ptr<Matrix> blockJacobiAdjacency(const Grid& grid,
                         {
                             Vector weights(rhs_->size());
                             ElementContext elemCtx(simulator_);
-                            Amg::getTrueImpesWeights(pressIndex, weights,
-                                                             simulator_.vanguard().gridView(),
-                                                             elemCtx, simulator_.model(),
-                                                             ThreadManager::threadId());
+                            Amg::getTrueImpesWeights(pressIndex,
+                                                     weights,
+                                                     elemCtx,
+                                                     simulator_.model(),
+                                                     *element_chunks_
+                            );
                             return weights;
                         };
                 } else if  (weightsType == "trueimpesanalytic" ) {
@@ -602,10 +610,12 @@ std::unique_ptr<Matrix> blockJacobiAdjacency(const Grid& grid,
                         {
                             Vector weights(rhs_->size());
                             ElementContext elemCtx(simulator_);
-                            Amg::getTrueImpesWeightsAnalytic(pressIndex, weights,
-                                                             simulator_.vanguard().gridView(),
-                                                             elemCtx, simulator_.model(),
-                                                             ThreadManager::threadId());
+                            Amg::getTrueImpesWeightsAnalytic(pressIndex,
+                                                             weights,
+                                                             elemCtx,
+                                                             simulator_.model(),
+                                                             *element_chunks_
+                            );
                             return weights;
                         };
                 } else {
@@ -652,6 +662,7 @@ std::unique_ptr<Matrix> blockJacobiAdjacency(const Grid& grid,
         std::vector<PropertyTree> prm_;
 
         std::shared_ptr< CommunicationType > comm_;
+        std::unique_ptr<ElementChunksType> element_chunks_;
     }; // end ISTLSolver
 
 } // namespace Opm

opm/simulators/linalg/getQuasiImpesWeights.hpp

@@ -22,8 +22,10 @@
 
 #include <dune/common/fvector.hh>
 
+#include <opm/grid/utility/ElementChunks.hpp>
 #include <opm/simulators/utils/DeferredLoggingErrorHelpers.hpp>
 #include <opm/material/common/MathToolbox.hpp>
+#include <opm/models/parallel/threadmanager.hpp>
 #include <algorithm>
 #include <cmath>
 
@@ -52,18 +54,25 @@ namespace Amg
         using VectorBlockType = typename Vector::block_type;
         using MatrixBlockType = typename Matrix::block_type;
         const Matrix& A = matrix;
+
         VectorBlockType rhs(0.0);
         rhs[pressureVarIndex] = 1.0;
-        const auto endi = A.end();
-        for (auto i = A.begin(); i != endi; ++i) {
-            const auto endj = (*i).end();
+        // Use OpenMP to parallelize over matrix rows
+        #ifdef _OPENMP
+        #pragma omp parallel for
+        #endif
+        for (int row_idx = 0; row_idx < static_cast<int>(A.N()); ++row_idx) {
             MatrixBlockType diag_block(0.0);
-            for (auto j = (*i).begin(); j != endj; ++j) {
-                if (i.index() == j.index()) {
+            // Find diagonal block for this row
+            const auto row_it = A.begin() + row_idx;
+            const auto endj = (*row_it).end();
+            for (auto j = (*row_it).begin(); j != endj; ++j) {
+                if (row_it.index() == j.index()) {
                     diag_block = (*j);
                     break;
                 }
             }
+
             VectorBlockType bweights;
             if (transpose) {
                 diag_block.solve(bweights, rhs);
@@ -71,12 +80,12 @@ namespace Amg
                 auto diag_block_transpose = Details::transposeDenseMatrix(diag_block);
                 diag_block_transpose.solve(bweights, rhs);
             }
+
             double abs_max = *std::max_element(
                 bweights.begin(), bweights.end(), [](double a, double b) { return std::fabs(a) < std::fabs(b); });
             bweights /= std::fabs(abs_max);
-            weights[i.index()] = bweights;
+            weights[row_idx] = bweights;
         }
-        // return weights;
     }
 
     template <class Matrix, class Vector>
@@ -87,59 +96,75 @@ namespace Amg
         return weights;
     }
 
-    template<class Vector, class GridView, class ElementContext, class Model>
-    void getTrueImpesWeights(int pressureVarIndex, Vector& weights, const GridView& gridView,
-                             ElementContext& elemCtx, const Model& model, std::size_t threadId)
+    template<class Vector, class ElementContext, class Model, class ElementChunksType>
+    void getTrueImpesWeights(int pressureVarIndex, Vector& weights,
+                             const ElementContext& elemCtx,
+                             const Model& model,
+                             const ElementChunksType& element_chunks)
     {
         using VectorBlockType = typename Vector::block_type;
         using Matrix = typename std::decay_t<decltype(model.linearizer().jacobian())>;
         using MatrixBlockType = typename Matrix::MatrixBlock;
         constexpr int numEq = VectorBlockType::size();
-        using Evaluation = typename std::decay_t<decltype(model.localLinearizer(threadId).localResidual().residual(0))>
+        using Evaluation = typename std::decay_t<decltype(model.localLinearizer(ThreadManager::threadId()).localResidual().residual(0))>
             ::block_type;
-        VectorBlockType rhs(0.0);
-        rhs[pressureVarIndex] = 1.0;
-        int index = 0;
+
         OPM_BEGIN_PARALLEL_TRY_CATCH();
-        for (const auto& elem : elements(gridView)) {
-            elemCtx.updatePrimaryStencil(elem);
-            elemCtx.updatePrimaryIntensiveQuantities(/*timeIdx=*/0);
+        #ifdef _OPENMP
+        #pragma omp parallel for
+        #endif
+        for (const auto& chunk : element_chunks) {
+            std::size_t thread_id = ThreadManager::threadId();
+            VectorBlockType rhs(0.0);
+            rhs[pressureVarIndex] = 1.0;
             Dune::FieldVector<Evaluation, numEq> storage;
-            model.localLinearizer(threadId).localResidual().computeStorage(storage,elemCtx,/*spaceIdx=*/0, /*timeIdx=*/0);
-            auto extrusionFactor = elemCtx.intensiveQuantities(0, /*timeIdx=*/0).extrusionFactor();
-            auto scvVolume = elemCtx.stencil(/*timeIdx=*/0).subControlVolume(0).volume() * extrusionFactor;
-            auto storage_scale = scvVolume / elemCtx.simulator().timeStepSize();
             MatrixBlockType block;
-            double pressure_scale = 50e5;
-            for (int ii = 0; ii < numEq; ++ii) {
-                for (int jj = 0; jj < numEq; ++jj) {
-                    block[ii][jj] = storage[ii].derivative(jj)/storage_scale;
-                    if (jj == pressureVarIndex) {
-                        block[ii][jj] *= pressure_scale;
+            VectorBlockType bweights;
+            MatrixBlockType block_transpose;
+
+            ElementContext localElemCtx(elemCtx.simulator());
+
+            for (const auto& elem : chunk) {
+                localElemCtx.updatePrimaryStencil(elem);
+                localElemCtx.updatePrimaryIntensiveQuantities(/*timeIdx=*/0);
+
+                model.localLinearizer(thread_id).localResidual().computeStorage(storage, localElemCtx, /*spaceIdx=*/0, /*timeIdx=*/0);
+
+                auto extrusionFactor = localElemCtx.intensiveQuantities(0, /*timeIdx=*/0).extrusionFactor();
+                auto scvVolume = localElemCtx.stencil(/*timeIdx=*/0).subControlVolume(0).volume() * extrusionFactor;
+                auto storage_scale = scvVolume / localElemCtx.simulator().timeStepSize();
+
+                const double pressure_scale = 50e5;
+                for (int ii = 0; ii < numEq; ++ii) {
+                    for (int jj = 0; jj < numEq; ++jj) {
+                        block[ii][jj] = storage[ii].derivative(jj)/storage_scale;
+                        if (jj == pressureVarIndex) {
+                            block[ii][jj] *= pressure_scale;
+                        }
                     }
                 }
-            }
-            VectorBlockType bweights;
-            MatrixBlockType block_transpose = Details::transposeDenseMatrix(block);
-            block_transpose.solve(bweights, rhs);
-            double abs_max = *std::max_element(
-                bweights.begin(), bweights.end(), [](double a, double b) { return std::fabs(a) < std::fabs(b); });
-            // probably a scaling which could give approximately total compressibility would be better
-            bweights /=  std::fabs(abs_max); // given normal densities this scales weights to about 1.
 
-            weights[index] = bweights;
-            ++index;
+                block_transpose = Details::transposeDenseMatrix(block);
+                block_transpose.solve(bweights, rhs);
+
+                double abs_max = *std::max_element(
+                    bweights.begin(), bweights.end(), [](double a, double b) { return std::fabs(a) < std::fabs(b); });
+                // probably a scaling which could give approximately total compressibility would be better
+                bweights /=  std::fabs(abs_max); // given normal densities this scales weights to about 1.
+
+                const auto index = localElemCtx.globalSpaceIndex(/*spaceIdx=*/0, /*timeIdx=*/0);
+                weights[index] = bweights;
+            }
         }
         OPM_END_PARALLEL_TRY_CATCH("getTrueImpesWeights() failed: ", elemCtx.simulator().vanguard().grid().comm());
     }
 
-    template <class Vector, class GridView, class ElementContext, class Model>
+    template <class Vector, class ElementContext, class Model, class ElementChunksType>
     void getTrueImpesWeightsAnalytic(int /*pressureVarIndex*/,
                                      Vector& weights,
-                                     const GridView& gridView,
-                                     ElementContext& elemCtx,
+                                     const ElementContext& elemCtx,
                                      const Model& model,
-                                     std::size_t threadId)
+                                     const ElementChunksType& element_chunks)
     {
         // The sequential residual is a linear combination of the
         // mass balance residuals, with coefficients equal to (for
@@ -156,56 +181,70 @@ namespace Amg
         using PrimaryVariables = typename Model::PrimaryVariables;
         using VectorBlockType = typename Vector::block_type;
         using Evaluation =
-            typename std::decay_t<decltype(model.localLinearizer(threadId).localResidual().residual(0))>::block_type;
+            typename std::decay_t<decltype(model.localLinearizer(ThreadManager::threadId()).localResidual().residual(0))>::block_type;
         using Toolbox = MathToolbox<Evaluation>;
-        VectorBlockType rhs(0.0);
+
         const auto& solution = model.solution(/*timeIdx*/ 0);
+
+        // Use OpenMP to parallelize over element chunks
         OPM_BEGIN_PARALLEL_TRY_CATCH();
-        for (const auto& elem : elements(gridView)) {
-            elemCtx.updatePrimaryStencil(elem);
-            elemCtx.updatePrimaryIntensiveQuantities(/*timeIdx=*/0);
-            const auto& index = elemCtx.globalSpaceIndex(/*spaceIdx=*/0, /*timeIdx=*/0);
-            const auto& intQuants = elemCtx.intensiveQuantities(/*spaceIdx=*/0, /*timeIdx=*/0);
-            const auto& fs = intQuants.fluidState();
+        #ifdef _OPENMP
+        #pragma omp parallel for
+        #endif
+        for (const auto& chunk : element_chunks) {
+            // Thread-local variables
             VectorBlockType bweights;
 
-            if (FluidSystem::phaseIsActive(FluidSystem::waterPhaseIdx)) {
-                unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
-                    FluidSystem::solventComponentIndex(FluidSystem::waterPhaseIdx));
-                bweights[activeCompIdx]
-                    = Toolbox::template decay<LhsEval>(1 / fs.invB(FluidSystem::waterPhaseIdx));
-            }
+            // Each thread gets a unique copy of elemCtx
+            ElementContext localElemCtx(elemCtx.simulator());
 
-            double denominator = 1.0;
-            double rs = Toolbox::template decay<double>(fs.Rs());
-            double rv = Toolbox::template decay<double>(fs.Rv());
-            const auto& priVars = solution[index];
-            if (priVars.primaryVarsMeaningGas() == PrimaryVariables::GasMeaning::Rv) {
-                rs = 0.0;
-            }
-            if (priVars.primaryVarsMeaningGas() == PrimaryVariables::GasMeaning::Rs) {
-                rv = 0.0;
-            }
-            if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)
-                && FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
-                denominator = Toolbox::template decay<LhsEval>(1 - rs * rv);
-            }
+            for (const auto& elem : chunk) {
+                localElemCtx.updatePrimaryStencil(elem);
+                localElemCtx.updatePrimaryIntensiveQuantities(/*timeIdx=*/0);
 
-            if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)) {
-                unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
-                    FluidSystem::solventComponentIndex(FluidSystem::oilPhaseIdx));
-                bweights[activeCompIdx] = Toolbox::template decay<LhsEval>(
-                    (1 / fs.invB(FluidSystem::oilPhaseIdx) - rs / fs.invB(FluidSystem::gasPhaseIdx))
-                    / denominator);
+                const auto index = localElemCtx.globalSpaceIndex(/*spaceIdx=*/0, /*timeIdx=*/0);
+                const auto& intQuants = localElemCtx.intensiveQuantities(/*spaceIdx=*/0, /*timeIdx=*/0);
+                const auto& fs = intQuants.fluidState();
+
+                if (FluidSystem::phaseIsActive(FluidSystem::waterPhaseIdx)) {
+                    unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
+                        FluidSystem::solventComponentIndex(FluidSystem::waterPhaseIdx));
+                    bweights[activeCompIdx]
+                        = Toolbox::template decay<LhsEval>(1 / fs.invB(FluidSystem::waterPhaseIdx));
+                }
+
+                double denominator = 1.0;
+                double rs = Toolbox::template decay<double>(fs.Rs());
+                double rv = Toolbox::template decay<double>(fs.Rv());
+                const auto& priVars = solution[index];
+                if (priVars.primaryVarsMeaningGas() == PrimaryVariables::GasMeaning::Rv) {
+                    rs = 0.0;
+                }
+                if (priVars.primaryVarsMeaningGas() == PrimaryVariables::GasMeaning::Rs) {
+                    rv = 0.0;
+                }
+                if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)
+                    && FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
+                    denominator = Toolbox::template decay<LhsEval>(1 - rs * rv);
+                }
+
+                if (FluidSystem::phaseIsActive(FluidSystem::oilPhaseIdx)) {
+                    unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
+                        FluidSystem::solventComponentIndex(FluidSystem::oilPhaseIdx));
+                    bweights[activeCompIdx] = Toolbox::template decay<LhsEval>(
+                        (1 / fs.invB(FluidSystem::oilPhaseIdx) - rs / fs.invB(FluidSystem::gasPhaseIdx))
+                        / denominator);
+                }
+                if (FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
+                    unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
+                        FluidSystem::solventComponentIndex(FluidSystem::gasPhaseIdx));
+                    bweights[activeCompIdx] = Toolbox::template decay<LhsEval>(
+                        (1 / fs.invB(FluidSystem::gasPhaseIdx) - rv / fs.invB(FluidSystem::oilPhaseIdx))
+                        / denominator);
+                }
+
+                weights[index] = bweights;
             }
-            if (FluidSystem::phaseIsActive(FluidSystem::gasPhaseIdx)) {
-                unsigned activeCompIdx = Indices::canonicalToActiveComponentIndex(
-                    FluidSystem::solventComponentIndex(FluidSystem::gasPhaseIdx));
-                bweights[activeCompIdx] = Toolbox::template decay<LhsEval>(
-                    (1 / fs.invB(FluidSystem::gasPhaseIdx) - rv / fs.invB(FluidSystem::oilPhaseIdx))
-                    / denominator);
-            }
-            weights[index] = bweights;
         }
         OPM_END_PARALLEL_TRY_CATCH("getTrueImpesAnalyticWeights() failed: ", elemCtx.simulator().vanguard().grid().comm());
     }