diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index ead5bd3ad9..69110d19d7 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -108,6 +108,9 @@ The Axom project release numbers follow [Semantic Versioning](http://semver.org/ value into device kernels. - Primal: Fixes signs of `compute_moments` to match orientation convention in `primal::evaluate_area_integral` - Quest: Improves error handling/reporting when loading an invalid c2c contour +- Quest: Sampling-based MFEM shaping now falls back to chunked local mass assembly/solve for large + meshes instead of overflowing `mfem::DenseTensor` size calculations when caching full + `dofs x dofs x NE` tensors. - Primal: Improves reproducibility of 3D GWN methods by removing some sources of randomness - Core: ArrayView assigments/copies now copy the stride - Core: Array construction from strided ArrayView now correctly copies the strided elements diff --git a/src/axom/quest/SamplingShaper.cpp b/src/axom/quest/SamplingShaper.cpp index 183f1a9c44..ba5fe13375 100644 --- a/src/axom/quest/SamplingShaper.cpp +++ b/src/axom/quest/SamplingShaper.cpp @@ -411,7 +411,8 @@ void SamplingShaper::computeVolumeFractionsForMaterial(const std::string& matFie matField, m_volfracOrder, m_samplingResolution, - m_quadratureType); + m_quadratureType, + m_execPolicy); return; } #endif diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index 734f8e6bbb..05b0e38b92 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -8,9 +8,14 @@ #if defined(AXOM_USE_MFEM) + #include + #include + #include #include #include + #include "mfem/linalg/kernels.hpp" + namespace axom { namespace quest @@ -33,6 +38,628 @@ class OwnedQuadratureSpace : public mfem::QuadratureSpace std::unique_ptr m_ir; }; +struct VolumeFractionMassConfig +{ + int dofs {}; + int numElements {}; + bool usesAnisotropicQuadrature {false}; + bool useChunkedMassProcessing {false}; + std::int64_t elemTensorEntries {}; + std::int64_t cachedMassBytes {}; +}; + +/*! + * \brief Select an execution policy supported by the host-side MFEM volume fraction path. + */ +axom::runtime_policy::Policy selectVolumeFractionExecutionPolicy(axom::runtime_policy::Policy execPolicy, + const std::string& vfName) +{ + using RuntimePolicy = axom::runtime_policy::Policy; + + switch(execPolicy) + { + case RuntimePolicy::seq: + return RuntimePolicy::seq; + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + case RuntimePolicy::omp: + return RuntimePolicy::omp; + #endif + #if defined(AXOM_RUNTIME_POLICY_USE_CUDA) + case RuntimePolicy::cuda: + #endif + #if defined(AXOM_RUNTIME_POLICY_USE_HIP) + case RuntimePolicy::hip: + #endif + #if defined(AXOM_RUNTIME_POLICY_USE_CUDA) || defined(AXOM_RUNTIME_POLICY_USE_HIP) + SLIC_WARNING_ROOT( + axom::fmt::format("MFEM volume fraction processing for '{}' uses host data and currently " + "falls back to sequential execution for device runtime policies.", + vfName)); + return RuntimePolicy::seq; + #endif + default: + SLIC_WARNING_ROOT(axom::fmt::format( + "MFEM volume fraction processing for '{}' falls back to sequential execution because the " + "requested runtime policy is not available in this build.", + vfName)); + return RuntimePolicy::seq; + } +} + +void assembleVolumeFractionRHS(const mfem::FiniteElementSpace& fes, + mfem::QuadratureFunction& inout, + const mfem::IntegrationRule& sampleIR, + bool usesAnisotropicQuadrature, + mfem::Vector& rhs); + +std::string formatSamplesPerDimension(axom::ArrayView sampleResolution, int dim) +{ + switch(dim) + { + case 2: + return axom::fmt::format(" ({} * {})", sampleResolution[0], sampleResolution[1]); + case 3: + return axom::fmt::format(" ({} * {} * {})", + sampleResolution[0], + sampleResolution[1], + sampleResolution[2]); + default: + return std::string(); + } +} + +void logVolumeFractionInputs(int sampleNQ, + int sampleOrder, + int sampleSZ, + int dim, + int numElements, + axom::ArrayView sampleResolution) +{ + SLIC_INFO_ROOT(axom::fmt::format(axom::utilities::locale(), + "In computeVolumeFractions(): num samples per element {}{} | " + "sample polynomial order {} | total samples {:L}", + sampleNQ, + formatSamplesPerDimension(sampleResolution, dim), + sampleOrder, + sampleSZ)); + + SLIC_INFO_ROOT(axom::fmt::format(axom::utilities::locale(), + "Mesh has dim {} and {:L} elements", + dim, + numElements)); +} + +/*! + * \brief Compute the cached and chunked mass-solve settings for this volume fraction field. + */ +VolumeFractionMassConfig makeVolumeFractionMassConfig(const mfem::FiniteElementSpace& fes, + axom::ArrayView sampleResolution, + axom::numerics::QuadratureType quadratureType) +{ + constexpr std::int64_t MAX_CACHED_MASS_BYTES = 1LL << 30; + + VolumeFractionMassConfig config; + config.dofs = fes.GetTypicalFE()->GetDof(); + config.numElements = fes.GetMesh()->GetNE(); + config.usesAnisotropicQuadrature = + usesAnisotropicCustomTensorQuadrature(*fes.GetMesh(), sampleResolution, quadratureType); + config.elemTensorEntries = static_cast(config.dofs) * config.dofs; + + const std::int64_t totalTensorEntries = config.elemTensorEntries * config.numElements; + config.cachedMassBytes = totalTensorEntries * 3 * sizeof(double); + config.useChunkedMassProcessing = totalTensorEntries > std::numeric_limits::max() || + config.cachedMassBytes > MAX_CACHED_MASS_BYTES; + + return config; +} + +void logChunkedMassProcessing(const std::string& vfName, const VolumeFractionMassConfig& config) +{ + if(!config.useChunkedMassProcessing) + { + return; + } + + SLIC_INFO_ROOT( + axom::fmt::format(axom::utilities::locale(), + "Using chunked local mass assembly for '{}' (dofs={}, elements={:L}, " + "dense cache would require {:.2f} GiB)", + vfName, + config.dofs, + config.numElements, + static_cast(config.cachedMassBytes) / (1024.0 * 1024.0 * 1024.0))); +} + +void assembleVolumeFractionRHSVector(const mfem::FiniteElementSpace& fes, + mfem::QuadratureFunction& inout, + const mfem::IntegrationRule& sampleIR, + bool usesAnisotropicQuadrature, + mfem::Vector& rhs) +{ + AXOM_ANNOTATE_SCOPE("domain lf integrator assemble"); + + inout.Read(); + rhs.HostWrite(); + rhs = 0.; + rhs.ReadWrite(); + + assembleVolumeFractionRHS(fes, inout, sampleIR, usesAnisotropicQuadrature, rhs); +} + +/*! + * \brief Assemble or reuse the whole-mesh local mass matrices for the cached solve path. + */ +mfem::DenseTensor* getOrAssembleMassMatrix(MFEMState& mfemState, + const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config) +{ + const std::string massMatrixName = "shaping_mass_matrix"; + if(mfemState.m_inoutTensors.Has(massMatrixName)) + { + return mfemState.m_inoutTensors.Get(massMatrixName); + } + + AXOM_ANNOTATE_SCOPE("mass integrator assemble"); + + auto* massMat = new mfem::DenseTensor(config.dofs, config.dofs, config.numElements); + massMat->HostWrite(); + (*massMat) = 0.; + massMat->ReadWrite(); + + mfem::ConstantCoefficient one_coef(1.0); + mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); + + if(config.usesAnisotropicQuadrature) + { + mfem::DenseMatrix elemMat; + massMat->HostWrite(); + for(int elem = 0; elem < config.numElements; ++elem) + { + mass_integrator.AssembleElementMatrix(*fes.GetFE(elem), + *fes.GetElementTransformation(elem), + elemMat); + for(int j = 0; j < config.dofs; ++j) + { + for(int i = 0; i < config.dofs; ++i) + { + (*massMat)(i, j, elem) = elemMat(i, j); + } + } + } + } + else + { + const int sz = massMat->TotalSize(); + mfem::Vector mass_vec; + mfem::Swap(massMat->GetMemory(), mass_vec.GetMemory()); + mass_vec.SetSize(sz); + mass_integrator.AssembleEA(fes, mass_vec, false); + mfem::Swap(massMat->GetMemory(), mass_vec.GetMemory()); + } + + mfemState.m_inoutTensors.Register(massMatrixName, massMat, true); + return massMat; +} + +std::pair*> getOrFactorMassMatrix( + MFEMState& mfemState, + mfem::DenseTensor& massMat, + const VolumeFractionMassConfig& config) +{ + const std::string minvName = "shaping_mass_matrix_inv"; + const std::string pivotsName = "shaping_mass_matrix_pivots"; + if(mfemState.m_inoutTensors.Has(minvName) && mfemState.m_inoutArrays.Has(pivotsName)) + { + return {mfemState.m_inoutTensors.Get(minvName), mfemState.m_inoutArrays.Get(pivotsName)}; + } + + AXOM_ANNOTATE_SCOPE("batch lu factor"); + + massMat.Read(); + auto* massMatInv = new mfem::DenseTensor(massMat); + auto* massMatPivots = new mfem::Array(config.dofs * config.numElements); + + massMatInv->ReadWrite(); + massMatPivots->Write(); + mfem::BatchLUFactor(*massMatInv, *massMatPivots); + + mfemState.m_inoutTensors.Register(minvName, massMatInv, true); + mfemState.m_inoutArrays.Register(pivotsName, massMatPivots, true); + + return {massMatInv, massMatPivots}; +} + +mfem::DenseTensor* getOrAllocateScratchBuffer(MFEMState& mfemState, + const VolumeFractionMassConfig& config) +{ + const std::string scratchBufferName = "shaping_scratch_buffer"; + if(mfemState.m_inoutTensors.Has(scratchBufferName)) + { + return mfemState.m_inoutTensors.Get(scratchBufferName); + } + + auto* scratchBuffer = new mfem::DenseTensor(config.dofs, config.dofs, config.numElements); + scratchBuffer->HostWrite(); + (*scratchBuffer) = 0.; + mfemState.m_inoutTensors.Register(scratchBufferName, scratchBuffer, true); + + return scratchBuffer; +} + +void applyFCTProjection(mfem::DenseTensor& massMat, + axom::runtime_policy::Policy execPolicy, + mfem::Vector& rhs, + const int dofs, + const int numElements, + mfem::Vector& vfData, + mfem::DenseTensor& scratchBuffer); + +template +void applyFCTProjectionImpl(mfem::DenseTensor& massMat, + mfem::Vector& rhs, + const int dofs, + const int numElements, + mfem::Vector& vfData, + mfem::DenseTensor& scratchBuffer) +{ + constexpr double minY = 0.; + constexpr double maxY = 1.; + + AXOM_ANNOTATE_SCOPE("fct project"); + + auto m_d = mfem::Reshape(massMat.HostRead(), dofs, dofs, numElements); + auto b_d = mfem::Reshape(rhs.HostRead(), dofs, numElements); + auto vf_d = mfem::Reshape(vfData.HostReadWrite(), dofs, numElements); + auto fct_mat_d = mfem::Reshape(scratchBuffer.HostReadWrite(), dofs, dofs, numElements); + + axom::for_all(0, numElements, [=](int elem) { + FCT_correct(&m_d(0, 0, elem), dofs, &b_d(0, elem), minY, maxY, &vf_d(0, elem), &fct_mat_d(0, 0, elem)); + }); +} + +void applyFCTProjection(mfem::DenseTensor& massMat, + axom::runtime_policy::Policy execPolicy, + mfem::Vector& rhs, + const int dofs, + const int numElements, + mfem::Vector& vfData, + mfem::DenseTensor& scratchBuffer) +{ + using RuntimePolicy = axom::runtime_policy::Policy; + + switch(execPolicy) + { + case RuntimePolicy::seq: + applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); + break; + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + case RuntimePolicy::omp: + applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); + break; + #endif + default: + applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); + break; + } +} + +/*! + * \brief Solve the cached whole-mesh mass systems and apply FCT to the result. + */ +void solveVolumeFractionsCached(MFEMState& mfemState, + const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config, + axom::runtime_policy::Policy execPolicy, + mfem::Vector& rhs, + mfem::GridFunction& vf) +{ + auto* massMat = getOrAssembleMassMatrix(mfemState, fes, sampleIR, config); + auto [massMatInv, massMatPivots] = getOrFactorMassMatrix(mfemState, *massMat, config); + auto* scratchBuffer = getOrAllocateScratchBuffer(mfemState, config); + + { + AXOM_ANNOTATE_SCOPE("batch lu solve"); + + massMatInv->Read(); + massMatPivots->Read(); + + vf.HostReadWrite(); + vf = rhs; + vf.ReadWrite(); + mfem::BatchLUSolve(*massMatInv, *massMatPivots, vf); + } + massMatInv->HostReadWrite(); + massMatPivots->HostReadWrite(); + + applyFCTProjection(*massMat, execPolicy, rhs, config.dofs, config.numElements, vf, *scratchBuffer); +} + +/*! + * \brief Choose a chunk size that bounds the temporary chunk workspace near the target size. + */ +int computeChunkSize(const VolumeFractionMassConfig& config) +{ + constexpr std::int64_t TARGET_CHUNK_BYTES = 256LL * 1024 * 1024; + const std::int64_t bytesPerElement = (3 * config.elemTensorEntries * sizeof(double)) + + (static_cast(config.dofs) * sizeof(int)); + const std::int64_t elemsPerChunk = + std::max(1, TARGET_CHUNK_BYTES / std::max(1, bytesPerElement)); + + return static_cast(std::min(config.numElements, elemsPerChunk)); +} + +/*! + * \brief Assemble chunk mass matrices serially while reusing MFEM objects across elements. + */ +void assembleChunkMassMatricesSequential(const mfem::FiniteElementSpace& fes, + mfem::MassIntegrator& massIntegrator, + const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::DenseTensor& massMat, + mfem::DenseTensor& massMatInv) +{ + mfem::DenseMatrix elemMat; + massMat.HostWrite(); + massMatInv.HostWrite(); + + for(int elem = 0; elem < chunkNE; ++elem) + { + const int globalElem = elemBegin + elem; + massIntegrator.AssembleElementMatrix(*fes.GetFE(globalElem), + *fes.GetElementTransformation(globalElem), + elemMat); + for(int j = 0; j < config.dofs; ++j) + { + for(int i = 0; i < config.dofs; ++i) + { + const double value = elemMat(i, j); + massMat(i, j, elem) = value; + massMatInv(i, j, elem) = value; + } + } + } +} + +/*! + * \brief Assemble chunk mass matrices with per-element MFEM state for host parallel execution. + */ +template +void assembleChunkMassMatricesImpl(const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::DenseTensor& massMat, + mfem::DenseTensor& massMatInv) +{ + auto massMat_d = mfem::Reshape(massMat.HostWrite(), config.dofs, config.dofs, chunkNE); + auto massMatInv_d = mfem::Reshape(massMatInv.HostWrite(), config.dofs, config.dofs, chunkNE); + auto* mesh = fes.GetMesh(); + + axom::for_all(0, chunkNE, [=, &fes, &sampleIR](int elem) { + mfem::ConstantCoefficient oneCoef(1.0); + mfem::MassIntegrator massIntegrator(oneCoef, &sampleIR); + mfem::DenseMatrix elemMat; + mfem::IsoparametricTransformation tr; + + const int globalElem = elemBegin + elem; + mesh->GetElementTransformation(globalElem, &tr); + massIntegrator.AssembleElementMatrix(*fes.GetFE(globalElem), tr, elemMat); + + for(int j = 0; j < config.dofs; ++j) + { + for(int i = 0; i < config.dofs; ++i) + { + const double value = elemMat(i, j); + massMat_d(i, j, elem) = value; + massMatInv_d(i, j, elem) = value; + } + } + }); +} + +/*! + * \brief Dispatch chunk mass assembly to the requested host execution policy. + */ +void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + mfem::MassIntegrator& massIntegrator, + const VolumeFractionMassConfig& config, + axom::runtime_policy::Policy execPolicy, + int elemBegin, + int chunkNE, + mfem::DenseTensor& massMat, + mfem::DenseTensor& massMatInv) +{ + using RuntimePolicy = axom::runtime_policy::Policy; + + switch(execPolicy) + { + case RuntimePolicy::seq: + assembleChunkMassMatricesSequential(fes, + massIntegrator, + config, + elemBegin, + chunkNE, + massMat, + massMatInv); + break; + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + case RuntimePolicy::omp: + assembleChunkMassMatricesImpl(fes, sampleIR, config, elemBegin, chunkNE, massMat, massMatInv); + break; + #endif + default: + assembleChunkMassMatricesSequential(fes, + massIntegrator, + config, + elemBegin, + chunkNE, + massMat, + massMatInv); + break; + } +} + +void copyChunkRHS(mfem::Vector& rhs, + const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::Vector& rhsChunk, + mfem::Vector& vfChunk) +{ + auto rhs_d = mfem::Reshape(rhs.HostRead(), config.dofs, config.numElements); + auto rhs_chunk_d = mfem::Reshape(rhsChunk.HostWrite(), config.dofs, chunkNE); + auto vf_chunk_d = mfem::Reshape(vfChunk.HostWrite(), config.dofs, chunkNE); + + for(int elem = 0; elem < chunkNE; ++elem) + { + const int globalElem = elemBegin + elem; + for(int i = 0; i < config.dofs; ++i) + { + const double value = rhs_d(i, globalElem); + rhs_chunk_d(i, elem) = value; + vf_chunk_d(i, elem) = value; + } + } +} + +/*! + * \brief Factor each element mass matrix in a chunk for the OpenMP chunked solve path. + */ +template +void factorChunkMassMatricesImpl(const VolumeFractionMassConfig& config, + int chunkNE, + mfem::DenseTensor& massMatInv, + mfem::Array& massMatPivots) +{ + auto massMatInv_d = mfem::Reshape(massMatInv.HostReadWrite(), config.dofs, config.dofs, chunkNE); + auto massMatPivots_d = mfem::Reshape(massMatPivots.HostWrite(), config.dofs, chunkNE); + + axom::for_all(0, chunkNE, [=](int elem) { + mfem::kernels::LUFactor(&massMatInv_d(0, 0, elem), config.dofs, &massMatPivots_d(0, elem)); + }); +} + +/*! + * \brief Dispatch chunk LU factorization to the requested host execution policy. + */ +void factorChunkMassMatrices(const VolumeFractionMassConfig& config, + axom::runtime_policy::Policy execPolicy, + int chunkNE, + mfem::DenseTensor& massMatInv, + mfem::Array& massMatPivots) +{ + using RuntimePolicy = axom::runtime_policy::Policy; + + switch(execPolicy) + { + case RuntimePolicy::seq: + massMatInv.ReadWrite(); + massMatPivots.Write(); + mfem::BatchLUFactor(massMatInv, massMatPivots); + break; + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + case RuntimePolicy::omp: + factorChunkMassMatricesImpl(config, chunkNE, massMatInv, massMatPivots); + break; + #endif + default: + massMatInv.ReadWrite(); + massMatPivots.Write(); + mfem::BatchLUFactor(massMatInv, massMatPivots); + break; + } +} + +void copyChunkResultToGridFunction(const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::Vector& vfChunk, + mfem::GridFunction& vf) +{ + auto vf_chunk_d = mfem::Reshape(vfChunk.HostRead(), config.dofs, chunkNE); + auto vf_d = mfem::Reshape(vf.HostReadWrite(), config.dofs, config.numElements); + + for(int elem = 0; elem < chunkNE; ++elem) + { + const int globalElem = elemBegin + elem; + for(int i = 0; i < config.dofs; ++i) + { + vf_d(i, globalElem) = vf_chunk_d(i, elem); + } + } +} + +/*! + * \brief Solve the chunked mass systems and apply FCT when the cached path is too large. + */ +void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config, + axom::runtime_policy::Policy execPolicy, + mfem::Vector& rhs, + mfem::GridFunction& vf) +{ + AXOM_ANNOTATE_SCOPE("chunked mass solve"); + + mfem::ConstantCoefficient one_coef(1.0); + mfem::MassIntegrator massIntegrator(one_coef, &sampleIR); + const int chunkSize = computeChunkSize(config); + + for(int elemBegin = 0; elemBegin < config.numElements; elemBegin += chunkSize) + { + const int chunkNE = std::min(chunkSize, config.numElements - elemBegin); + + mfem::DenseTensor massMat(config.dofs, config.dofs, chunkNE); + mfem::DenseTensor massMatInv(config.dofs, config.dofs, chunkNE); + mfem::DenseTensor scratchBuffer(config.dofs, config.dofs, chunkNE); + mfem::Array massMatPivots(config.dofs * chunkNE); + mfem::Vector rhsChunk(config.dofs * chunkNE); + mfem::Vector vfChunk(config.dofs * chunkNE); + + { + AXOM_ANNOTATE_SCOPE("chunk mass assembly"); + assembleChunkMassMatrices(fes, + sampleIR, + massIntegrator, + config, + execPolicy, + elemBegin, + chunkNE, + massMat, + massMatInv); + } + + { + AXOM_ANNOTATE_SCOPE("chunk rhs copy"); + copyChunkRHS(rhs, config, elemBegin, chunkNE, rhsChunk, vfChunk); + } + + { + AXOM_ANNOTATE_SCOPE("chunk batch lu factor"); + factorChunkMassMatrices(config, execPolicy, chunkNE, massMatInv, massMatPivots); + } + + { + AXOM_ANNOTATE_SCOPE("chunk batch lu solve"); + massMatInv.Read(); + massMatPivots.Read(); + vfChunk.ReadWrite(); + mfem::BatchLUSolve(massMatInv, massMatPivots, vfChunk); + } + + applyFCTProjection(massMat, execPolicy, rhsChunk, config.dofs, chunkNE, vfChunk, scratchBuffer); + { + AXOM_ANNOTATE_SCOPE("chunk result copy"); + copyChunkResultToGridFunction(config, elemBegin, chunkNE, vfChunk, vf); + } + } +} + } // namespace bool usesAnisotropicCustomTensorQuadrature(const mfem::Mesh& mesh, @@ -125,7 +752,7 @@ mfem::GridFunction* getOrAllocateL2GridFunction(mfem::DataCollection* dc, } gf->MakeOwner(fec); - gf->HostReadWrite(); + gf->HostWrite(); *gf = 0.; dc->RegisterField(gf_name, gf); @@ -510,7 +1137,8 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, const std::string& matField, int volfracOrder, axom::ArrayView sampleResolution, - axom::numerics::QuadratureType quadratureType) + axom::numerics::QuadratureType quadratureType, + axom::runtime_policy::Policy execPolicy) { AXOM_ANNOTATE_SCOPE("computeVolumeFractionsForMaterial"); @@ -530,172 +1158,31 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, mfem::Mesh* mesh = dc->GetMesh(); const int dim = mesh->Dimension(); const int NE = mesh->GetNE(); - - auto samples_per_dim = [=](auto sampleRes, int dimValue) -> std::string { - switch(dimValue) - { - case 2: - return axom::fmt::format(" ({} * {})", sampleRes[0], sampleRes[1]); - case 3: - return axom::fmt::format(" ({} * {} * {})", sampleRes[0], sampleRes[1], sampleRes[2]); - default: - return std::string(); - } - }; - - SLIC_INFO_ROOT(axom::fmt::format(axom::utilities::locale(), - "In computeVolumeFractions(): num samples per element {}{} | " - "sample polynomial order {} | total samples {:L}", - sampleNQ, - samples_per_dim(sampleResolution, dim), - sampleOrder, - sampleSZ)); - - SLIC_INFO_ROOT( - axom::fmt::format(axom::utilities::locale(), "Mesh has dim {} and {:L} elements", dim, NE)); + logVolumeFractionInputs(sampleNQ, sampleOrder, sampleSZ, dim, NE, sampleResolution); const auto vf_name = shaping::volumeFractionFieldName(materialName); mfem::GridFunction* vf = getOrAllocateL2GridFunction(dc, vf_name, volfracOrder, dim, mfem::BasisType::Positive); const mfem::FiniteElementSpace* fes = vf->FESpace(); - const int dofs = fes->GetTypicalFE()->GetDof(); - - mfem::DenseTensor* mass_mat {nullptr}; - const std::string mass_matrix_name = "shaping_mass_matrix"; - if(mfemState.m_inoutTensors.Has(mass_matrix_name)) - { - mass_mat = mfemState.m_inoutTensors.Get(mass_matrix_name); - } - else - { - AXOM_ANNOTATE_SCOPE("mass integrator assemble"); - - mass_mat = new mfem::DenseTensor(dofs, dofs, NE); - mass_mat->HostWrite(); - (*mass_mat) = 0.; - mass_mat->ReadWrite(); - - mfem::ConstantCoefficient one_coef(1.0); - mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); - - if(usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType)) - { - mfem::DenseMatrix elemMat; - mass_mat->HostWrite(); - for(int elem = 0; elem < NE; ++elem) - { - mass_integrator.AssembleElementMatrix(*fes->GetFE(elem), - *fes->GetElementTransformation(elem), - elemMat); - for(int j = 0; j < dofs; ++j) - { - for(int i = 0; i < dofs; ++i) - { - (*mass_mat)(i, j, elem) = elemMat(i, j); - } - } - } - } - else - { - const int sz = mass_mat->TotalSize(); - mfem::Vector mass_vec; - mfem::Swap(mass_mat->GetMemory(), mass_vec.GetMemory()); - mass_vec.SetSize(sz); - mass_integrator.AssembleEA(*fes, mass_vec, false); - mfem::Swap(mass_mat->GetMemory(), mass_vec.GetMemory()); - } - - mfemState.m_inoutTensors.Register(mass_matrix_name, mass_mat, true); - } - - mfem::DenseTensor* mass_mat_inv {nullptr}; - mfem::Array* mass_mat_pivots {nullptr}; - const std::string minv_name = "shaping_mass_matrix_inv"; - const std::string pivots_name = "shaping_mass_matrix_pivots"; - if(mfemState.m_inoutTensors.Has(minv_name) && mfemState.m_inoutArrays.Has(pivots_name)) - { - mass_mat_inv = mfemState.m_inoutTensors.Get(minv_name); - mass_mat_pivots = mfemState.m_inoutArrays.Get(pivots_name); - } - else - { - AXOM_ANNOTATE_SCOPE("batch lu factor"); - - mass_mat->ReadWrite(); - mass_mat_inv = new mfem::DenseTensor(*mass_mat); - mass_mat_pivots = new mfem::Array(dofs * NE); - - mass_mat_inv->ReadWrite(); - mass_mat_pivots->Write(); - mfem::BatchLUFactor(*mass_mat_inv, *mass_mat_pivots); - - mfemState.m_inoutTensors.Register(minv_name, mass_mat_inv, true); - mfemState.m_inoutArrays.Register(pivots_name, mass_mat_pivots, true); - } - - mfem::DenseTensor* shaping_scratch_buffer {nullptr}; - const std::string scratch_buffer_name = "shaping_scratch_buffer"; - if(mfemState.m_inoutTensors.Has(scratch_buffer_name)) - { - shaping_scratch_buffer = mfemState.m_inoutTensors.Get(scratch_buffer_name); - } - else - { - shaping_scratch_buffer = new mfem::DenseTensor(dofs, dofs, NE); - shaping_scratch_buffer->HostWrite(); - (*shaping_scratch_buffer) = 0.; - mfemState.m_inoutTensors.Register(scratch_buffer_name, shaping_scratch_buffer, true); - } + const VolumeFractionMassConfig config = + makeVolumeFractionMassConfig(*fes, sampleResolution, quadratureType); + logChunkedMassProcessing(vf_name, config); + const auto volumeFractionExecPolicy = selectVolumeFractionExecutionPolicy(execPolicy, vf_name); axom::utilities::Timer timer(true); { mfem::Vector b(fes->GetVSize()); - SLIC_ASSERT(b.Size() == dofs * NE); + SLIC_ASSERT(b.Size() == config.dofs * config.numElements); + assembleVolumeFractionRHSVector(*fes, *inout, sampleIR, config.usesAnisotropicQuadrature, b); + + if(config.useChunkedMassProcessing) { - AXOM_ANNOTATE_SCOPE("domain lf integrator assemble"); - - inout->ReadWrite(); - b.HostWrite(); - b = 0.; - b.ReadWrite(); - - assembleVolumeFractionRHS( - *fes, - *inout, - sampleIR, - usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType), - b); + solveVolumeFractionsChunked(*fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); } - inout->HostReadWrite(); - + else { - AXOM_ANNOTATE_SCOPE("batch lu solve"); - - mass_mat_inv->Read(); - mass_mat_pivots->Read(); - - vf->HostReadWrite(); - (*vf) = b; - vf->ReadWrite(); - mfem::BatchLUSolve(*mass_mat_inv, *mass_mat_pivots, *vf); + solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); } - mass_mat_inv->HostReadWrite(); - mass_mat_pivots->HostReadWrite(); - - constexpr double minY = 0.; - constexpr double maxY = 1.; - - auto m_d = mfem::Reshape(mass_mat->HostReadWrite(), dofs, dofs, NE); - auto b_d = mfem::Reshape(b.HostReadWrite(), dofs, NE); - auto vf_d = mfem::Reshape(vf->HostReadWrite(), dofs, NE); - auto fct_mat_d = mfem::Reshape(shaping_scratch_buffer->HostReadWrite(), dofs, dofs, NE); - - AXOM_ANNOTATE_BEGIN("fct project"); - axom::for_all(0, NE, [=](int i) { - FCT_correct(&m_d(0, 0, i), dofs, &b_d(0, i), minY, maxY, &vf_d(0, i), &fct_mat_d(0, 0, i)); - }); - AXOM_ANNOTATE_END("fct project"); } timer.stop(); diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp index c40e4cca44..f872587eed 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp @@ -230,7 +230,8 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, const std::string& matField, int volfracOrder, axom::ArrayView sampleResolution, - axom::numerics::QuadratureType quadratureType); + axom::numerics::QuadratureType quadratureType, + axom::runtime_policy::Policy execPolicy); /*! * \brief Creates a new GridFunction based on \a inout and registers it with the \a dc DataCollection.