From 8e3c676993176d679cbbb9e69fb465d1608cc7b0 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 12:07:26 -0700 Subject: [PATCH 1/7] Make some MFEM changes in shaper support to permit computing volume fractions on larger meshes. --- RELEASE-NOTES.md | 3 + .../detail/shaping/shaping_helpers_mfem.cpp | 337 ++++++++++++------ 2 files changed, 234 insertions(+), 106 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index e9bc6f526c..29eafcdfdf 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -60,6 +60,9 @@ The Axom project release numbers follow [Semantic Versioning](http://semver.org/ ### Fixed - 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/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index ee95110f0c..e1fff90f1a 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -8,6 +8,9 @@ #if defined(AXOM_USE_MFEM) + #include + #include + #include #include #include @@ -559,143 +562,265 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, 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)) + const bool usesAnisotropicQuadrature = + usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType); + const std::int64_t elemTensorEntries = static_cast(dofs) * dofs; + const std::int64_t totalTensorEntries = elemTensorEntries * NE; + constexpr std::int64_t MAX_CACHED_MASS_BYTES = 1LL << 30; + const std::int64_t cachedMassBytes = totalTensorEntries * 3 * sizeof(double); + const bool useChunkedMassProcessing = + totalTensorEntries > std::numeric_limits::max() || + cachedMassBytes > MAX_CACHED_MASS_BYTES; + + if(useChunkedMassProcessing) { - mass_mat = mfemState.m_inoutTensors.Get(mass_matrix_name); + SLIC_INFO_ROOT( + axom::fmt::format(axom::utilities::locale(), + "Using chunked local mass assembly for '{}' (dofs={}, elements={:L}, " + "dense cache would require {:.2f} GiB)", + vf_name, + dofs, + NE, + static_cast(cachedMassBytes) / + (1024.0 * 1024.0 * 1024.0))); } - else + + axom::utilities::Timer timer(true); { - AXOM_ANNOTATE_SCOPE("mass integrator assemble"); + mfem::Vector b(fes->GetVSize()); + SLIC_ASSERT(b.Size() == dofs * NE); + { + AXOM_ANNOTATE_SCOPE("domain lf integrator assemble"); - mass_mat = new mfem::DenseTensor(dofs, dofs, NE); - mass_mat->HostWrite(); - (*mass_mat) = 0.; - mass_mat->ReadWrite(); + inout->ReadWrite(); + b.HostWrite(); + b = 0.; + b.ReadWrite(); - mfem::ConstantCoefficient one_coef(1.0); - mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); + assembleVolumeFractionRHS( + *fes, + *inout, + sampleIR, + usesAnisotropicQuadrature, + b); + } + inout->HostReadWrite(); - if(usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType)) + constexpr double minY = 0.; + constexpr double maxY = 1.; + if(useChunkedMassProcessing) { + AXOM_ANNOTATE_SCOPE("chunked mass solve"); + + mfem::ConstantCoefficient one_coef(1.0); + mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); mfem::DenseMatrix elemMat; - mass_mat->HostWrite(); - for(int elem = 0; elem < NE; ++elem) + + constexpr std::int64_t TARGET_CHUNK_BYTES = 256LL * 1024 * 1024; + const std::int64_t bytesPerElement = + (3 * elemTensorEntries * sizeof(double)) + (static_cast(dofs) * sizeof(int)); + const std::int64_t elemsPerChunk = + std::max(1, TARGET_CHUNK_BYTES / std::max(1, bytesPerElement)); + const int chunkSize = static_cast(std::min(NE, elemsPerChunk)); + + auto b_d = mfem::Reshape(b.HostReadWrite(), dofs, NE); + auto vf_d = mfem::Reshape(vf->HostReadWrite(), dofs, NE); + + for(int elemBegin = 0; elemBegin < NE; elemBegin += chunkSize) { - mass_integrator.AssembleElementMatrix(*fes->GetFE(elem), - *fes->GetElementTransformation(elem), - elemMat); - for(int j = 0; j < dofs; ++j) + const int chunkNE = std::min(chunkSize, NE - elemBegin); + + mfem::DenseTensor mass_mat(dofs, dofs, chunkNE); + mfem::DenseTensor mass_mat_inv(dofs, dofs, chunkNE); + mfem::DenseTensor shaping_scratch_buffer(dofs, dofs, chunkNE); + mfem::Array mass_mat_pivots(dofs * chunkNE); + mfem::Vector rhs_chunk(dofs * chunkNE); + mfem::Vector vf_chunk(dofs * chunkNE); + + mass_mat.HostWrite(); + mass_mat_inv.HostWrite(); + for(int elem = 0; elem < chunkNE; ++elem) + { + const int globalElem = elemBegin + elem; + mass_integrator.AssembleElementMatrix(*fes->GetFE(globalElem), + *fes->GetElementTransformation(globalElem), + elemMat); + for(int j = 0; j < dofs; ++j) + { + for(int i = 0; i < dofs; ++i) + { + const double value = elemMat(i, j); + mass_mat(i, j, elem) = value; + mass_mat_inv(i, j, elem) = value; + } + } + } + + auto rhs_chunk_d = mfem::Reshape(rhs_chunk.HostWrite(), dofs, chunkNE); + auto vf_chunk_d = mfem::Reshape(vf_chunk.HostWrite(), dofs, chunkNE); + for(int elem = 0; elem < chunkNE; ++elem) + { + const int globalElem = elemBegin + elem; + for(int i = 0; i < dofs; ++i) + { + const double value = b_d(i, globalElem); + rhs_chunk_d(i, elem) = value; + vf_chunk_d(i, elem) = value; + } + } + + mass_mat_inv.ReadWrite(); + mass_mat_pivots.Write(); + mfem::BatchLUFactor(mass_mat_inv, mass_mat_pivots); + + mass_mat_inv.Read(); + mass_mat_pivots.Read(); + vf_chunk.ReadWrite(); + mfem::BatchLUSolve(mass_mat_inv, mass_mat_pivots, vf_chunk); + + auto m_d = mfem::Reshape(mass_mat.HostReadWrite(), dofs, dofs, chunkNE); + auto fct_rhs_d = mfem::Reshape(rhs_chunk.HostReadWrite(), dofs, chunkNE); + auto fct_vf_d = mfem::Reshape(vf_chunk.HostReadWrite(), dofs, chunkNE); + auto fct_mat_d = + mfem::Reshape(shaping_scratch_buffer.HostWrite(), dofs, dofs, chunkNE); + + AXOM_ANNOTATE_BEGIN("fct project"); + axom::for_all(0, chunkNE, [=](int elem) { + FCT_correct(&m_d(0, 0, elem), + dofs, + &fct_rhs_d(0, elem), + minY, + maxY, + &fct_vf_d(0, elem), + &fct_mat_d(0, 0, elem)); + }); + AXOM_ANNOTATE_END("fct project"); + + for(int elem = 0; elem < chunkNE; ++elem) { + const int globalElem = elemBegin + elem; for(int i = 0; i < dofs; ++i) { - (*mass_mat)(i, j, elem) = elemMat(i, j); + vf_d(i, globalElem) = fct_vf_d(i, elem); } } } } 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()); - } + 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"); - mfemState.m_inoutTensors.Register(mass_matrix_name, mass_mat, true); - } + mass_mat = new mfem::DenseTensor(dofs, dofs, NE); + mass_mat->HostWrite(); + (*mass_mat) = 0.; + mass_mat->ReadWrite(); - 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"); + mfem::ConstantCoefficient one_coef(1.0); + mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); - mass_mat->ReadWrite(); - mass_mat_inv = new mfem::DenseTensor(*mass_mat); - mass_mat_pivots = new mfem::Array(dofs * NE); + if(usesAnisotropicQuadrature) + { + 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()); + } - mass_mat_inv->ReadWrite(); - mass_mat_pivots->Write(); - mfem::BatchLUFactor(*mass_mat_inv, *mass_mat_pivots); + mfemState.m_inoutTensors.Register(mass_matrix_name, mass_mat, true); + } - mfemState.m_inoutTensors.Register(minv_name, mass_mat_inv, true); - mfemState.m_inoutArrays.Register(pivots_name, mass_mat_pivots, 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"); - 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); - } + mass_mat->ReadWrite(); + mass_mat_inv = new mfem::DenseTensor(*mass_mat); + mass_mat_pivots = new mfem::Array(dofs * NE); - axom::utilities::Timer timer(true); - { - mfem::Vector b(fes->GetVSize()); - SLIC_ASSERT(b.Size() == dofs * NE); - { - AXOM_ANNOTATE_SCOPE("domain lf integrator assemble"); + mass_mat_inv->ReadWrite(); + mass_mat_pivots->Write(); + mfem::BatchLUFactor(*mass_mat_inv, *mass_mat_pivots); - inout->ReadWrite(); - b.HostWrite(); - b = 0.; - b.ReadWrite(); + mfemState.m_inoutTensors.Register(minv_name, mass_mat_inv, true); + mfemState.m_inoutArrays.Register(pivots_name, mass_mat_pivots, true); + } - assembleVolumeFractionRHS( - *fes, - *inout, - sampleIR, - usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType), - b); - } - inout->HostReadWrite(); + 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); + } - { - AXOM_ANNOTATE_SCOPE("batch lu solve"); + { + AXOM_ANNOTATE_SCOPE("batch lu solve"); - mass_mat_inv->Read(); - mass_mat_pivots->Read(); + mass_mat_inv->Read(); + mass_mat_pivots->Read(); - vf->HostReadWrite(); - (*vf) = b; - vf->ReadWrite(); - mfem::BatchLUSolve(*mass_mat_inv, *mass_mat_pivots, *vf); + vf->HostReadWrite(); + (*vf) = b; + vf->ReadWrite(); + mfem::BatchLUSolve(*mass_mat_inv, *mass_mat_pivots, *vf); + } + mass_mat_inv->HostReadWrite(); + mass_mat_pivots->HostReadWrite(); + + 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"); } - 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(); From 6ef51d080b85ddb831bcdff39f99367d175feb20 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 12:22:30 -0700 Subject: [PATCH 2/7] Refactor some MFEM support into smaller functions. --- .../detail/shaping/shaping_helpers_mfem.cpp | 692 +++++++++++------- 1 file changed, 418 insertions(+), 274 deletions(-) diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index e1fff90f1a..52c18f9888 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -36,6 +36,409 @@ 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 {}; +}; + +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)); +} + +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.ReadWrite(); + rhs.HostWrite(); + rhs = 0.; + rhs.ReadWrite(); + + assembleVolumeFractionRHS(fes, inout, sampleIR, usesAnisotropicQuadrature, rhs); + inout.HostReadWrite(); +} + +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.ReadWrite(); + 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, + mfem::Vector& rhs, + const int dofs, + const int numElements, + mfem::Vector& vfData, + mfem::DenseTensor& scratchBuffer) +{ + constexpr double minY = 0.; + constexpr double maxY = 1.; + + auto m_d = mfem::Reshape(massMat.HostReadWrite(), dofs, dofs, numElements); + auto b_d = mfem::Reshape(rhs.HostReadWrite(), dofs, numElements); + auto vf_d = mfem::Reshape(vfData.HostReadWrite(), dofs, numElements); + auto fct_mat_d = mfem::Reshape(scratchBuffer.HostReadWrite(), dofs, dofs, numElements); + + AXOM_ANNOTATE_BEGIN("fct project"); + 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)); + }); + AXOM_ANNOTATE_END("fct project"); +} + +void solveVolumeFractionsCached(MFEMState& mfemState, + const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config, + 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, rhs, config.dofs, config.numElements, vf, *scratchBuffer); +} + +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)); +} + +void assembleChunkMassMatrices(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; + } + } + } +} + +void copyChunkRHS(mfem::Vector& rhs, + const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::Vector& rhsChunk, + mfem::Vector& vfChunk) +{ + auto rhs_d = mfem::Reshape(rhs.HostReadWrite(), 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; + } + } +} + +void copyChunkResultToGridFunction(const VolumeFractionMassConfig& config, + int elemBegin, + int chunkNE, + mfem::Vector& vfChunk, + mfem::GridFunction& vf) +{ + auto vf_chunk_d = mfem::Reshape(vfChunk.HostReadWrite(), 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); + } + } +} + +void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, + const mfem::IntegrationRule& sampleIR, + const VolumeFractionMassConfig& config, + 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); + + assembleChunkMassMatrices( + fes, + massIntegrator, + config, + elemBegin, + chunkNE, + massMat, + massMatInv); + + copyChunkRHS(rhs, config, elemBegin, chunkNE, rhsChunk, vfChunk); + + massMatInv.ReadWrite(); + massMatPivots.Write(); + mfem::BatchLUFactor(massMatInv, massMatPivots); + + massMatInv.Read(); + massMatPivots.Read(); + vfChunk.ReadWrite(); + mfem::BatchLUSolve(massMatInv, massMatPivots, vfChunk); + + applyFCTProjection(massMat, rhsChunk, config.dofs, chunkNE, vfChunk, scratchBuffer); + copyChunkResultToGridFunction(config, elemBegin, chunkNE, vfChunk, vf); + } +} + } // namespace bool usesAnisotropicCustomTensorQuadrature(const mfem::Mesh& mesh, @@ -533,293 +936,34 @@ 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(); - const bool usesAnisotropicQuadrature = - usesAnisotropicCustomTensorQuadrature(*fes->GetMesh(), sampleResolution, quadratureType); - const std::int64_t elemTensorEntries = static_cast(dofs) * dofs; - const std::int64_t totalTensorEntries = elemTensorEntries * NE; - constexpr std::int64_t MAX_CACHED_MASS_BYTES = 1LL << 30; - const std::int64_t cachedMassBytes = totalTensorEntries * 3 * sizeof(double); - const bool useChunkedMassProcessing = - totalTensorEntries > std::numeric_limits::max() || - cachedMassBytes > MAX_CACHED_MASS_BYTES; - - if(useChunkedMassProcessing) - { - SLIC_INFO_ROOT( - axom::fmt::format(axom::utilities::locale(), - "Using chunked local mass assembly for '{}' (dofs={}, elements={:L}, " - "dense cache would require {:.2f} GiB)", - vf_name, - dofs, - NE, - static_cast(cachedMassBytes) / - (1024.0 * 1024.0 * 1024.0))); - } + const VolumeFractionMassConfig config = + makeVolumeFractionMassConfig(*fes, sampleResolution, quadratureType); + logChunkedMassProcessing(vf_name, config); 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, - usesAnisotropicQuadrature, - b); - } - inout->HostReadWrite(); - - constexpr double minY = 0.; - constexpr double maxY = 1.; - if(useChunkedMassProcessing) - { - AXOM_ANNOTATE_SCOPE("chunked mass solve"); - - mfem::ConstantCoefficient one_coef(1.0); - mfem::MassIntegrator mass_integrator(one_coef, &sampleIR); - mfem::DenseMatrix elemMat; - - constexpr std::int64_t TARGET_CHUNK_BYTES = 256LL * 1024 * 1024; - const std::int64_t bytesPerElement = - (3 * elemTensorEntries * sizeof(double)) + (static_cast(dofs) * sizeof(int)); - const std::int64_t elemsPerChunk = - std::max(1, TARGET_CHUNK_BYTES / std::max(1, bytesPerElement)); - const int chunkSize = static_cast(std::min(NE, elemsPerChunk)); - - auto b_d = mfem::Reshape(b.HostReadWrite(), dofs, NE); - auto vf_d = mfem::Reshape(vf->HostReadWrite(), dofs, NE); - - for(int elemBegin = 0; elemBegin < NE; elemBegin += chunkSize) - { - const int chunkNE = std::min(chunkSize, NE - elemBegin); - - mfem::DenseTensor mass_mat(dofs, dofs, chunkNE); - mfem::DenseTensor mass_mat_inv(dofs, dofs, chunkNE); - mfem::DenseTensor shaping_scratch_buffer(dofs, dofs, chunkNE); - mfem::Array mass_mat_pivots(dofs * chunkNE); - mfem::Vector rhs_chunk(dofs * chunkNE); - mfem::Vector vf_chunk(dofs * chunkNE); - - mass_mat.HostWrite(); - mass_mat_inv.HostWrite(); - for(int elem = 0; elem < chunkNE; ++elem) - { - const int globalElem = elemBegin + elem; - mass_integrator.AssembleElementMatrix(*fes->GetFE(globalElem), - *fes->GetElementTransformation(globalElem), - elemMat); - for(int j = 0; j < dofs; ++j) - { - for(int i = 0; i < dofs; ++i) - { - const double value = elemMat(i, j); - mass_mat(i, j, elem) = value; - mass_mat_inv(i, j, elem) = value; - } - } - } - - auto rhs_chunk_d = mfem::Reshape(rhs_chunk.HostWrite(), dofs, chunkNE); - auto vf_chunk_d = mfem::Reshape(vf_chunk.HostWrite(), dofs, chunkNE); - for(int elem = 0; elem < chunkNE; ++elem) - { - const int globalElem = elemBegin + elem; - for(int i = 0; i < dofs; ++i) - { - const double value = b_d(i, globalElem); - rhs_chunk_d(i, elem) = value; - vf_chunk_d(i, elem) = value; - } - } - - mass_mat_inv.ReadWrite(); - mass_mat_pivots.Write(); - mfem::BatchLUFactor(mass_mat_inv, mass_mat_pivots); - - mass_mat_inv.Read(); - mass_mat_pivots.Read(); - vf_chunk.ReadWrite(); - mfem::BatchLUSolve(mass_mat_inv, mass_mat_pivots, vf_chunk); - - auto m_d = mfem::Reshape(mass_mat.HostReadWrite(), dofs, dofs, chunkNE); - auto fct_rhs_d = mfem::Reshape(rhs_chunk.HostReadWrite(), dofs, chunkNE); - auto fct_vf_d = mfem::Reshape(vf_chunk.HostReadWrite(), dofs, chunkNE); - auto fct_mat_d = - mfem::Reshape(shaping_scratch_buffer.HostWrite(), dofs, dofs, chunkNE); - - AXOM_ANNOTATE_BEGIN("fct project"); - axom::for_all(0, chunkNE, [=](int elem) { - FCT_correct(&m_d(0, 0, elem), - dofs, - &fct_rhs_d(0, elem), - minY, - maxY, - &fct_vf_d(0, elem), - &fct_mat_d(0, 0, elem)); - }); - AXOM_ANNOTATE_END("fct project"); - - for(int elem = 0; elem < chunkNE; ++elem) - { - const int globalElem = elemBegin + elem; - for(int i = 0; i < dofs; ++i) - { - vf_d(i, globalElem) = fct_vf_d(i, elem); - } - } - } + solveVolumeFractionsChunked(*fes, sampleIR, config, b, *vf); } else { - 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(usesAnisotropicQuadrature) - { - 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); - } - - { - 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); - } - mass_mat_inv->HostReadWrite(); - mass_mat_pivots->HostReadWrite(); - - 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"); + solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, b, *vf); } } timer.stop(); From 225a7323fed7ff5b779195f0afeb322f0f622789 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 12:23:06 -0700 Subject: [PATCH 3/7] make style --- .../detail/shaping/shaping_helpers_mfem.cpp | 76 ++++++------------- 1 file changed, 24 insertions(+), 52 deletions(-) diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index 52c18f9888..be43197b7d 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -59,11 +59,10 @@ std::string formatSamplesPerDimension(axom::ArrayView sampleResolution, int case 2: return axom::fmt::format(" ({} * {})", sampleResolution[0], sampleResolution[1]); case 3: - return axom::fmt::format( - " ({} * {} * {})", - sampleResolution[0], - sampleResolution[1], - sampleResolution[2]); + return axom::fmt::format(" ({} * {} * {})", + sampleResolution[0], + sampleResolution[1], + sampleResolution[2]); default: return std::string(); } @@ -84,17 +83,15 @@ void logVolumeFractionInputs(int sampleNQ, sampleOrder, sampleSZ)); - SLIC_INFO_ROOT(axom::fmt::format( - axom::utilities::locale(), - "Mesh has dim {} and {:L} elements", - dim, - numElements)); + SLIC_INFO_ROOT(axom::fmt::format(axom::utilities::locale(), + "Mesh has dim {} and {:L} elements", + dim, + numElements)); } -VolumeFractionMassConfig makeVolumeFractionMassConfig( - const mfem::FiniteElementSpace& fes, - axom::ArrayView sampleResolution, - axom::numerics::QuadratureType quadratureType) +VolumeFractionMassConfig makeVolumeFractionMassConfig(const mfem::FiniteElementSpace& fes, + axom::ArrayView sampleResolution, + axom::numerics::QuadratureType quadratureType) { constexpr std::int64_t MAX_CACHED_MASS_BYTES = 1LL << 30; @@ -107,15 +104,13 @@ VolumeFractionMassConfig makeVolumeFractionMassConfig( const std::int64_t totalTensorEntries = config.elemTensorEntries * config.numElements; config.cachedMassBytes = totalTensorEntries * 3 * sizeof(double); - config.useChunkedMassProcessing = - totalTensorEntries > std::numeric_limits::max() || + config.useChunkedMassProcessing = totalTensorEntries > std::numeric_limits::max() || config.cachedMassBytes > MAX_CACHED_MASS_BYTES; return config; } -void logChunkedMassProcessing(const std::string& vfName, - const VolumeFractionMassConfig& config) +void logChunkedMassProcessing(const std::string& vfName, const VolumeFractionMassConfig& config) { if(!config.useChunkedMassProcessing) { @@ -129,8 +124,7 @@ void logChunkedMassProcessing(const std::string& vfName, vfName, config.dofs, config.numElements, - static_cast(config.cachedMassBytes) / - (1024.0 * 1024.0 * 1024.0))); + static_cast(config.cachedMassBytes) / (1024.0 * 1024.0 * 1024.0))); } void assembleVolumeFractionRHSVector(const mfem::FiniteElementSpace& fes, @@ -177,10 +171,9 @@ mfem::DenseTensor* getOrAssembleMassMatrix(MFEMState& mfemState, massMat->HostWrite(); for(int elem = 0; elem < config.numElements; ++elem) { - mass_integrator.AssembleElementMatrix( - *fes.GetFE(elem), - *fes.GetElementTransformation(elem), - elemMat); + 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) @@ -266,14 +259,7 @@ void applyFCTProjection(mfem::DenseTensor& massMat, AXOM_ANNOTATE_BEGIN("fct project"); 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)); + FCT_correct(&m_d(0, 0, elem), dofs, &b_d(0, elem), minY, maxY, &vf_d(0, elem), &fct_mat_d(0, 0, elem)); }); AXOM_ANNOTATE_END("fct project"); } @@ -309,8 +295,7 @@ void solveVolumeFractionsCached(MFEMState& mfemState, 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)) + + 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)); @@ -333,10 +318,9 @@ void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, for(int elem = 0; elem < chunkNE; ++elem) { const int globalElem = elemBegin + elem; - massIntegrator.AssembleElementMatrix( - *fes.GetFE(globalElem), - *fes.GetElementTransformation(globalElem), - elemMat); + 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) @@ -414,14 +398,7 @@ void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, mfem::Vector rhsChunk(config.dofs * chunkNE); mfem::Vector vfChunk(config.dofs * chunkNE); - assembleChunkMassMatrices( - fes, - massIntegrator, - config, - elemBegin, - chunkNE, - massMat, - massMatInv); + assembleChunkMassMatrices(fes, massIntegrator, config, elemBegin, chunkNE, massMat, massMatInv); copyChunkRHS(rhs, config, elemBegin, chunkNE, rhsChunk, vfChunk); @@ -950,12 +927,7 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, { mfem::Vector b(fes->GetVSize()); SLIC_ASSERT(b.Size() == config.dofs * config.numElements); - assembleVolumeFractionRHSVector( - *fes, - *inout, - sampleIR, - config.usesAnisotropicQuadrature, - b); + assembleVolumeFractionRHSVector(*fes, *inout, sampleIR, config.usesAnisotropicQuadrature, b); if(config.useChunkedMassProcessing) { From f55381dab7d0bab77b51d83ee0b82d330f25077d Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 16:37:30 -0700 Subject: [PATCH 4/7] Pass execution policy into shaping::computeVolumeFractionsForMaterial --- src/axom/quest/SamplingShaper.cpp | 3 +- .../detail/shaping/shaping_helpers_mfem.cpp | 215 +++++++++++++++--- .../detail/shaping/shaping_helpers_mfem.hpp | 3 +- 3 files changed, 193 insertions(+), 28 deletions(-) diff --git a/src/axom/quest/SamplingShaper.cpp b/src/axom/quest/SamplingShaper.cpp index f186d60ad5..e0b0eca6ef 100644 --- a/src/axom/quest/SamplingShaper.cpp +++ b/src/axom/quest/SamplingShaper.cpp @@ -429,7 +429,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 be43197b7d..1d51ee44d9 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -46,6 +46,39 @@ struct VolumeFractionMassConfig std::int64_t cachedMassBytes {}; }; +axom::runtime_policy::Policy selectFCTExecutionPolicy(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( + "FCT projection for '{}' uses MFEM 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( + "FCT projection 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, @@ -243,31 +276,67 @@ mfem::DenseTensor* getOrAllocateScratchBuffer(MFEMState& mfemState, } 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) + 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.HostReadWrite(), dofs, dofs, numElements); auto b_d = mfem::Reshape(rhs.HostReadWrite(), dofs, numElements); auto vf_d = mfem::Reshape(vfData.HostReadWrite(), dofs, numElements); auto fct_mat_d = mfem::Reshape(scratchBuffer.HostReadWrite(), dofs, dofs, numElements); - AXOM_ANNOTATE_BEGIN("fct project"); - axom::for_all(0, numElements, [=](int elem) { + 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)); }); - AXOM_ANNOTATE_END("fct project"); +} + +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; + } } 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) { @@ -289,7 +358,7 @@ void solveVolumeFractionsCached(MFEMState& mfemState, massMatInv->HostReadWrite(); massMatPivots->HostReadWrite(); - applyFCTProjection(*massMat, rhs, config.dofs, config.numElements, vf, *scratchBuffer); + applyFCTProjection(*massMat, execPolicy, rhs, config.dofs, config.numElements, vf, *scratchBuffer); } int computeChunkSize(const VolumeFractionMassConfig& config) @@ -303,13 +372,14 @@ int computeChunkSize(const VolumeFractionMassConfig& config) return static_cast(std::min(config.numElements, elemsPerChunk)); } -void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, - mfem::MassIntegrator& massIntegrator, - const VolumeFractionMassConfig& config, - int elemBegin, - int chunkNE, - mfem::DenseTensor& massMat, - mfem::DenseTensor& massMatInv) +/// Serial version that does not have to create MFEM objects each iteration. +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(); @@ -333,6 +403,73 @@ void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, } } +/// OpenMP parallelizable version that builds MFEM objects in the for_all. +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; + } + } + }); +} + +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, @@ -378,6 +515,7 @@ void copyChunkResultToGridFunction(const VolumeFractionMassConfig& config, void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, const mfem::IntegrationRule& sampleIR, const VolumeFractionMassConfig& config, + axom::runtime_policy::Policy execPolicy, mfem::Vector& rhs, mfem::GridFunction& vf) { @@ -398,21 +536,44 @@ void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, mfem::Vector rhsChunk(config.dofs * chunkNE); mfem::Vector vfChunk(config.dofs * chunkNE); - assembleChunkMassMatrices(fes, massIntegrator, config, elemBegin, chunkNE, massMat, massMatInv); + { + AXOM_ANNOTATE_SCOPE("chunk mass assembly"); + assembleChunkMassMatrices(fes, + sampleIR, + massIntegrator, + config, + execPolicy, + elemBegin, + chunkNE, + massMat, + massMatInv); + } - copyChunkRHS(rhs, config, elemBegin, chunkNE, rhsChunk, vfChunk); + { + AXOM_ANNOTATE_SCOPE("chunk rhs copy"); + copyChunkRHS(rhs, config, elemBegin, chunkNE, rhsChunk, vfChunk); + } - massMatInv.ReadWrite(); - massMatPivots.Write(); - mfem::BatchLUFactor(massMatInv, massMatPivots); + { + AXOM_ANNOTATE_SCOPE("chunk batch lu factor"); + massMatInv.ReadWrite(); + massMatPivots.Write(); + mfem::BatchLUFactor(massMatInv, massMatPivots); + } - massMatInv.Read(); - massMatPivots.Read(); - vfChunk.ReadWrite(); - mfem::BatchLUSolve(massMatInv, massMatPivots, vfChunk); + { + AXOM_ANNOTATE_SCOPE("chunk batch lu solve"); + massMatInv.Read(); + massMatPivots.Read(); + vfChunk.ReadWrite(); + mfem::BatchLUSolve(massMatInv, massMatPivots, vfChunk); + } - applyFCTProjection(massMat, rhsChunk, config.dofs, chunkNE, vfChunk, scratchBuffer); - copyChunkResultToGridFunction(config, elemBegin, chunkNE, vfChunk, vf); + applyFCTProjection(massMat, execPolicy, rhsChunk, config.dofs, chunkNE, vfChunk, scratchBuffer); + { + AXOM_ANNOTATE_SCOPE("chunk result copy"); + copyChunkResultToGridFunction(config, elemBegin, chunkNE, vfChunk, vf); + } } } @@ -893,7 +1054,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"); @@ -922,6 +1084,7 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, const VolumeFractionMassConfig config = makeVolumeFractionMassConfig(*fes, sampleResolution, quadratureType); logChunkedMassProcessing(vf_name, config); + const auto fctExecPolicy = selectFCTExecutionPolicy(execPolicy, vf_name); axom::utilities::Timer timer(true); { @@ -931,11 +1094,11 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, if(config.useChunkedMassProcessing) { - solveVolumeFractionsChunked(*fes, sampleIR, config, b, *vf); + solveVolumeFractionsChunked(*fes, sampleIR, config, fctExecPolicy, b, *vf); } else { - solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, b, *vf); + solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, fctExecPolicy, b, *vf); } } 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 6caff6dd88..b5d1625730 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.hpp @@ -227,7 +227,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); void computeVolumeFractionsIdentity(mfem::DataCollection* dc, mfem::QuadratureFunction* inout, From 3f2dcfaf6072530723de44a7f5d3d0c3700affb1 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 17:00:19 -0700 Subject: [PATCH 5/7] Speedup for LU factoring phase --- .../detail/shaping/shaping_helpers_mfem.cpp | 48 +++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index 1d51ee44d9..84499f1125 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -14,6 +14,8 @@ #include #include + #include "mfem/linalg/kernels.hpp" + namespace axom { namespace quest @@ -493,6 +495,48 @@ void copyChunkRHS(mfem::Vector& rhs, } } +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.Write(), config.dofs, chunkNE); + + axom::for_all(0, chunkNE, [=](int elem) { + mfem::kernels::LUFactor(&massMatInv_d(0, 0, elem), config.dofs, &massMatPivots_d(0, elem)); + }); +} + +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, @@ -556,9 +600,7 @@ void solveVolumeFractionsChunked(const mfem::FiniteElementSpace& fes, { AXOM_ANNOTATE_SCOPE("chunk batch lu factor"); - massMatInv.ReadWrite(); - massMatPivots.Write(); - mfem::BatchLUFactor(massMatInv, massMatPivots); + factorChunkMassMatrices(config, execPolicy, chunkNE, massMatInv, massMatPivots); } { From fb54fe2ee9010fc7d7b05304c79faf2d62c01cc5 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 17:21:17 -0700 Subject: [PATCH 6/7] Polishing changes. --- .../detail/shaping/shaping_helpers_mfem.cpp | 69 ++++++++++++++----- 1 file changed, 51 insertions(+), 18 deletions(-) diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index 84499f1125..0eeaeed01b 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -48,8 +48,12 @@ struct VolumeFractionMassConfig std::int64_t cachedMassBytes {}; }; -axom::runtime_policy::Policy selectFCTExecutionPolicy(axom::runtime_policy::Policy execPolicy, - const std::string& vfName) +/*! + * \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; @@ -69,13 +73,13 @@ axom::runtime_policy::Policy selectFCTExecutionPolicy(axom::runtime_policy::Poli #endif #if defined(AXOM_RUNTIME_POLICY_USE_CUDA) || defined(AXOM_RUNTIME_POLICY_USE_HIP) SLIC_WARNING_ROOT(axom::fmt::format( - "FCT projection for '{}' uses MFEM host data and currently falls back to sequential execution for device runtime policies.", + "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( - "FCT projection for '{}' falls back to sequential execution because the requested runtime policy is not available in this build.", + "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; } @@ -124,6 +128,9 @@ void logVolumeFractionInputs(int sampleNQ, 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) @@ -170,15 +177,17 @@ void assembleVolumeFractionRHSVector(const mfem::FiniteElementSpace& fes, { AXOM_ANNOTATE_SCOPE("domain lf integrator assemble"); - inout.ReadWrite(); + inout.Read(); rhs.HostWrite(); rhs = 0.; rhs.ReadWrite(); assembleVolumeFractionRHS(fes, inout, sampleIR, usesAnisotropicQuadrature, rhs); - inout.HostReadWrite(); } +/*! + * \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, @@ -246,7 +255,7 @@ std::pair*> getOrFactorMassMatrix( AXOM_ANNOTATE_SCOPE("batch lu factor"); - massMat.ReadWrite(); + massMat.Read(); auto* massMatInv = new mfem::DenseTensor(massMat); auto* massMatPivots = new mfem::Array(config.dofs * config.numElements); @@ -298,8 +307,8 @@ void applyFCTProjectionImpl(mfem::DenseTensor& massMat, AXOM_ANNOTATE_SCOPE("fct project"); - auto m_d = mfem::Reshape(massMat.HostReadWrite(), dofs, dofs, numElements); - auto b_d = mfem::Reshape(rhs.HostReadWrite(), dofs, numElements); + 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); @@ -334,6 +343,9 @@ void applyFCTProjection(mfem::DenseTensor& massMat, } } +/*! + * \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, @@ -363,6 +375,9 @@ void solveVolumeFractionsCached(MFEMState& mfemState, 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; @@ -374,7 +389,9 @@ int computeChunkSize(const VolumeFractionMassConfig& config) return static_cast(std::min(config.numElements, elemsPerChunk)); } -/// Serial version that does not have to create MFEM objects each iteration. +/*! + * \brief Assemble chunk mass matrices serially while reusing MFEM objects across elements. + */ void assembleChunkMassMatricesSequential(const mfem::FiniteElementSpace& fes, mfem::MassIntegrator& massIntegrator, const VolumeFractionMassConfig& config, @@ -405,7 +422,9 @@ void assembleChunkMassMatricesSequential(const mfem::FiniteElementSpace& fes, } } -/// OpenMP parallelizable version that builds MFEM objects in the for_all. +/*! + * \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, @@ -441,6 +460,9 @@ void assembleChunkMassMatricesImpl(const mfem::FiniteElementSpace& fes, }); } +/*! + * \brief Dispatch chunk mass assembly to the requested host execution policy. + */ void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, const mfem::IntegrationRule& sampleIR, mfem::MassIntegrator& massIntegrator, @@ -479,7 +501,7 @@ void copyChunkRHS(mfem::Vector& rhs, mfem::Vector& rhsChunk, mfem::Vector& vfChunk) { - auto rhs_d = mfem::Reshape(rhs.HostReadWrite(), config.dofs, config.numElements); + 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); @@ -495,6 +517,9 @@ void copyChunkRHS(mfem::Vector& rhs, } } +/*! + * \brief Factor each element mass matrix in a chunk for the OpenMP chunked solve path. + */ template void factorChunkMassMatricesImpl(const VolumeFractionMassConfig& config, int chunkNE, @@ -502,13 +527,16 @@ void factorChunkMassMatricesImpl(const VolumeFractionMassConfig& config, mfem::Array& massMatPivots) { auto massMatInv_d = mfem::Reshape(massMatInv.HostReadWrite(), config.dofs, config.dofs, chunkNE); - auto massMatPivots_d = mfem::Reshape(massMatPivots.Write(), 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, @@ -543,7 +571,7 @@ void copyChunkResultToGridFunction(const VolumeFractionMassConfig& config, mfem::Vector& vfChunk, mfem::GridFunction& vf) { - auto vf_chunk_d = mfem::Reshape(vfChunk.HostReadWrite(), config.dofs, chunkNE); + 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) @@ -556,6 +584,9 @@ void copyChunkResultToGridFunction(const VolumeFractionMassConfig& config, } } +/*! + * \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, @@ -711,7 +742,7 @@ mfem::GridFunction* getOrAllocateL2GridFunction(mfem::DataCollection* dc, } gf->MakeOwner(fec); - gf->HostReadWrite(); + gf->HostWrite(); *gf = 0.; dc->RegisterField(gf_name, gf); @@ -1126,7 +1157,8 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, const VolumeFractionMassConfig config = makeVolumeFractionMassConfig(*fes, sampleResolution, quadratureType); logChunkedMassProcessing(vf_name, config); - const auto fctExecPolicy = selectFCTExecutionPolicy(execPolicy, vf_name); + const auto volumeFractionExecPolicy = + selectVolumeFractionExecutionPolicy(execPolicy, vf_name); axom::utilities::Timer timer(true); { @@ -1136,11 +1168,12 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, if(config.useChunkedMassProcessing) { - solveVolumeFractionsChunked(*fes, sampleIR, config, fctExecPolicy, b, *vf); + solveVolumeFractionsChunked(*fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); } else { - solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, fctExecPolicy, b, *vf); + solveVolumeFractionsCached( + mfemState, *fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); } } timer.stop(); From 361040816af5c94632fd3e88964f86a0f794d406 Mon Sep 17 00:00:00 2001 From: Brad Whitlock Date: Thu, 25 Jun 2026 17:21:43 -0700 Subject: [PATCH 7/7] make style --- .../detail/shaping/shaping_helpers_mfem.cpp | 70 +++++++++++-------- 1 file changed, 39 insertions(+), 31 deletions(-) diff --git a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp index 0eeaeed01b..194c75388a 100644 --- a/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp +++ b/src/axom/quest/detail/shaping/shaping_helpers_mfem.cpp @@ -51,9 +51,8 @@ struct VolumeFractionMassConfig /*! * \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) +axom::runtime_policy::Policy selectVolumeFractionExecutionPolicy(axom::runtime_policy::Policy execPolicy, + const std::string& vfName) { using RuntimePolicy = axom::runtime_policy::Policy; @@ -61,25 +60,27 @@ axom::runtime_policy::Policy selectVolumeFractionExecutionPolicy( { case RuntimePolicy::seq: return RuntimePolicy::seq; -#if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) case RuntimePolicy::omp: return RuntimePolicy::omp; -#endif -#if defined(AXOM_RUNTIME_POLICY_USE_CUDA) + #endif + #if defined(AXOM_RUNTIME_POLICY_USE_CUDA) case RuntimePolicy::cuda: -#endif -#if defined(AXOM_RUNTIME_POLICY_USE_HIP) + #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)); + #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 + #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.", + "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; } @@ -332,11 +333,11 @@ void applyFCTProjection(mfem::DenseTensor& massMat, case RuntimePolicy::seq: applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); break; -#if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) case RuntimePolicy::omp: applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); break; -#endif + #endif default: applyFCTProjectionImpl(massMat, rhs, dofs, numElements, vfData, scratchBuffer); break; @@ -478,18 +479,27 @@ void assembleChunkMassMatrices(const mfem::FiniteElementSpace& fes, switch(execPolicy) { case RuntimePolicy::seq: - assembleChunkMassMatricesSequential( - fes, massIntegrator, config, elemBegin, chunkNE, massMat, massMatInv); + assembleChunkMassMatricesSequential(fes, + massIntegrator, + config, + elemBegin, + chunkNE, + massMat, + massMatInv); break; -#if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) case RuntimePolicy::omp: - assembleChunkMassMatricesImpl( - fes, sampleIR, config, elemBegin, chunkNE, massMat, massMatInv); + assembleChunkMassMatricesImpl(fes, sampleIR, config, elemBegin, chunkNE, massMat, massMatInv); break; -#endif + #endif default: - assembleChunkMassMatricesSequential( - fes, massIntegrator, config, elemBegin, chunkNE, massMat, massMatInv); + assembleChunkMassMatricesSequential(fes, + massIntegrator, + config, + elemBegin, + chunkNE, + massMat, + massMatInv); break; } } @@ -552,11 +562,11 @@ void factorChunkMassMatrices(const VolumeFractionMassConfig& config, massMatPivots.Write(); mfem::BatchLUFactor(massMatInv, massMatPivots); break; -#if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) + #if defined(AXOM_RUNTIME_POLICY_USE_OPENMP) case RuntimePolicy::omp: factorChunkMassMatricesImpl(config, chunkNE, massMatInv, massMatPivots); break; -#endif + #endif default: massMatInv.ReadWrite(); massMatPivots.Write(); @@ -1157,8 +1167,7 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, const VolumeFractionMassConfig config = makeVolumeFractionMassConfig(*fes, sampleResolution, quadratureType); logChunkedMassProcessing(vf_name, config); - const auto volumeFractionExecPolicy = - selectVolumeFractionExecutionPolicy(execPolicy, vf_name); + const auto volumeFractionExecPolicy = selectVolumeFractionExecutionPolicy(execPolicy, vf_name); axom::utilities::Timer timer(true); { @@ -1172,8 +1181,7 @@ void computeVolumeFractionsForMaterial(MFEMState& mfemState, } else { - solveVolumeFractionsCached( - mfemState, *fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); + solveVolumeFractionsCached(mfemState, *fes, sampleIR, config, volumeFractionExecPolicy, b, *vf); } } timer.stop();