Adding support to memory collection of Particles and MultiFabs within a grid #4629

jessdagostini · 2025-08-27T18:07:05Z

Summary

This PR adds functions to return memory usage from Particles and FabArray/MultiFabs within a grid/box. The main goal is to use this information for load balancing. We aim to consider the memory allocated to a worker as a factor in the load-balancing algorithms for AMReX, so there is a need to capture specific grid/box memory information.

Additional background

Checklist

The proposed changes:

fix a bug or incorrect behavior in AMReX
add new capabilities to AMReX
changes answers in the test suite to more than roundoff level
are likely to significantly affect the results of downstream AMReX users
include documentation in the code and/or rst files, if appropriate

… box

WeiqunZhang · 2025-08-28T23:37:29Z

Src/Base/AMReX_MultiFabUtil.H

+// We have tried to work with more templated version, but it didn't work. What is the most appropriate way to do it?
+LayoutData<std::size_t> CapacityOfFabs (BoxArray const& ba,
+                 DistributionMapping const& dm,
+                 std::vector<MultiFab*> const& mfs);


Could we remove BoxArray and DistributionMapping from the function arguments? The assumption here is all multifabs must have the same boxarray (ignoring the index type) and distributionmapping.

Could make argument be Vector<MultiFab const*> const& mfs. Vector provides bound check if assertion is on. We are not changing the data in mulitfabs. Thus we should pass them in as const*.

If we want it be consistent with how it's done in particle container, we could make this a member function of FabArray. That might be better. That will also solve the template "issue".

I believe it's possible to remove the BoxArray and the DistributionMapping if we have a way to access this information from MultiFabUtil. In the implementation at the .cpp file, we need those to create our LayoutData object that will hold the memory information for each box (7a29e92#diff-e104e0f29e7a0e37ad2ee839c42e2c317091f995556dcc35c18b47d66506761fR67)

Sounds good about the Vector! We weren't sure if our approach was the best :)

We have tried to implement this function over FabArrayUtil, and we had a hard time calling it from an "user-end" code. If you have a suggestion on how to make it work in the FabArray I would appreciate!
For context, I have tested both implementations using examples from the AMReX Tutorials.

More specifically, I have used this example for the Particles https://github.com/AMReX-Codes/amrex-tutorials/blob/main/ExampleCodes/Particles/NeighborList/main.cpp and this for the MultiFabs https://github.com/AMReX-Codes/amrex-tutorials/blob/main/ExampleCodes/Basic/HeatEquation_EX0_C/Source/main.cpp
On the MultiFabs, I used the following code to get the memory information

std::vector<amrex::MultiFab*> multifabs_to_check; multifabs_to_check.push_back(&phi_old); multifabs_to_check.push_back(&phi_new); amrex::LayoutData<std::size_t> mem_usage = amrex::CapacityOfFabs(ba, dm, multifabs_to_check); long total_bytes = 0; for (amrex::MFIter mfi(mem_usage); mfi.isValid(); ++mfi) { amrex::Print() << "Box " << mfi.index() << " uses " << mem_usage[mfi] << " bytes\n"; }

Please, let me know if I am missing something or if there's another approach for this!

MultiFab is derived from FabArray, which has member functions that return BoxArray and DistributionMapping used to build the MultiFab. Something like

auto const& ba = mf_ptr->boxArray(); auto const& dm = mf_ptr->DistributionMap();

Thank you! I have pushed an updated version with adjustments. I have tested with the same code example as shared above and got the same results, so I believe it is producing correct data.
About the point on

Vector provides bound check if assertion is on.

I am not doing any assertions at this point. Should I?

I agree with @WeiqunZhang's suggestion to make CapacityofFabs a member function of FabArray. Other than that, looks good to me.

Nice, do you also want to change the particle version to use the same interface? I think this will be helpful when there is more than on ParticleContainer in an application (as there is in WarpX).

Yes, I think it makes sense to change ParticleContainer to this format! The only concern is that doing so will result in this function having a slightly different signature than other ParticleContainer functions that perform similar tasks, but are not directly memory-related (like https://github.com/AMReX-Codes/amrex/blob/development/Src/Particle/AMReX_ParticleContainerI.H#L496). Do you think this is ok?

I'm okay with that. That function has a slightly different use case. I believe it was added for IO, where we write separate files for each species anyway, so we don't need to aggregate the totals over multiple species, as you do here.

Ok! Working on it!

Done! 7fb20d4

… discussion

Src/Particle/AMReX_ParticleContainer.H

atmyers · 2025-09-03T23:20:25Z

Src/Base/AMReX_FabArray.H

+template <class F, std::enable_if_t<IsBaseFab<F>::value,int> FOO>
+void
+FabArray<FAB>::capacityOfFabs(LayoutData<std::size_t>& mem) {
+    // Should we have assertions here? What should be asserted?


It would be good to assert that mem and *this use the same BoxArray and DistributionMap. You can use the CellEqual method for BoxArray and the == operator for the DistributionMap.

Got it! Just pushed the assertion on 7a29e92

Src/Base/AMReX_MultiFabUtil.H

…bArray

Src/Base/AMReX_FabArray.H

Co-authored-by: Andrew Myers <atmyers@lbl.gov>

Src/Base/AMReX_FabArray.H

jessdagostini · 2025-09-07T16:03:06Z

I have committed fixes for errors pointed out by two pipelines. There were two other pipelines "oneAPI" that failed because of a GPG key error. Not sure if I can fix this on my side

WeiqunZhang

The function is straightforward enough that I don't think we need to add tests for it.

Src/Base/AMReX_FabArray.H

Src/Particle/AMReX_ParticleContainerI.H

Src/Particle/AMReX_ParticleContainer.H

WeiqunZhang · 2025-09-08T04:30:15Z

Src/Base/AMReX_FabArray.H

    void setDomainBndry (value_type val, int strt_comp, int ncomp, const Geometry& geom);

+    // Get capacity of the FabArray
+    template <class F=FAB, std::enable_if_t<IsBaseFab<F>::value,int> = 0>


For flexibility, we could also make the integer type a template typename here and for in particle containers. For example,

template <typename I, class F=FAB, std::enable_if_t<IsBaseFab<F>::value && std::is_integral_v<I> && sizeof(I) >= sizeof(Long), int> = 0> void capacityOfFabs (LayoutData<I>& mem) const;

And in the implementation, we also want to do explicit static_cast before adding the values.

Ok! Working on the update for this!

Actually, do you think it's worth doing it? Because I know that it is a best practice in C++ to use std::size_t for memory-related math (size_t was created for holding sizeof returns https://en.cppreference.com/w/cpp/types/size_t.html)

Actually the best practice is proably to avoid unsigned integers in almost all cases.

Even Stroustrup does not like unsigned integers. https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1428r0.pdf Stroustrup said,

Unfortunately, sizeof yields an unsigned (and it would be hard to change that), but we don’t have to
follow that for all types with something to do with sizes

Java does not even have build-in unsigned integer types. https://www.artima.com/articles/james-gosling-on-java-may-2001 Goshling said,

In programming language design, one of the standard problems is that the language grows so complex that nobody can understand it. One of the little experiments I tried was asking people about the rules for unsigned arithmetic in C. It turns out nobody understands how unsigned arithmetic in C works. There are a few obvious things that people understand, but many people don't understand it.

Hmm, got it! Interesting that the unsigned is such unknown territory. Will start working on including the template you suggested in both functions

Just update both function signatures with the new template 55d4a34

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

atmyers · 2025-09-11T17:43:32Z

The function is straightforward enough that I don't think we need to add tests for it.

I asked Jessica to add those. Maybe these new functions could just be called somewhere in an existing test, instead creating new ones?

jessdagostini · 2025-09-12T17:28:04Z

The function is straightforward enough that I don't think we need to add tests for it.

I asked Jessica to add those. Maybe these new functions could just be called somewhere in an existing test, instead creating new ones?

I could include the Particles call in the NeighborList test. Not sure if it's the best option. I am not sure where to put the FabArray one, though.

Src/Particle/AMReX_ParticleContainerI.H

+    AMREX_ASSERT(BoxArray::SameRefs(mem.boxArray(), ParticleBoxArray(lev)) &&
+                 mem.DistributionMap() == ParticleDistributionMap(lev));
+
+    [[maybe_unused]] Gpu::NoSyncRegion no_sync{};


atmyers · 2025-09-16T17:44:00Z

The function is straightforward enough that I don't think we need to add tests for it.

I asked Jessica to add those. Maybe these new functions could just be called somewhere in an existing test, instead creating new ones?

I could include the Particles call in the NeighborList test. Not sure if it's the best option. I am not sure where to put the FabArray one, though.

I think that if you 1) Remove the new particle test added in this PR, in lieu of the modification to the existing NeighborList one, and 2) move/rename the MultiFab one to amrex/Tests/HeatEquation, then this should be good to go.

jessdagostini · 2025-09-16T20:53:52Z

The function is straightforward enough that I don't think we need to add tests for it.

I asked Jessica to add those. Maybe these new functions could just be called somewhere in an existing test, instead creating new ones?

I could include the Particles call in the NeighborList test. Not sure if it's the best option. I am not sure where to put the FabArray one, though.

I think that if you 1) Remove the new particle test added in this PR, in lieu of the modification to the existing NeighborList one, and 2) move/rename the MultiFab one to amrex/Tests/HeatEquation, then this should be good to go.

Sounds good! Did that on 64e1e84

WeiqunZhang · 2025-09-16T21:54:49Z

LGTM. Hopefully, one last thing. For the test to run, we need to add it to Line 128 of amrex/Tests/CMakeLists.txt.

jessdagostini · 2025-09-16T22:33:59Z

LGTM. Hopefully, one last thing. For the test to run, we need to add it to Line 128 of amrex/Tests/CMakeLists.txt.

Cool! Pushed here 2b01097

Tests/CMakeLists.txt

Pull in AMReX-Codes/amrex#4629 for #6166. Signed-off-by: Axel Huebl <axel.huebl@plasma.ninja>

jessdagostini added 2 commits August 27, 2025 10:58

Adding func to return mem capacity/usage of fabs within a box

7a29e92

Adding functions to return mem capacity/usage from particles within a…

20d8658

… box

WeiqunZhang requested changes Aug 28, 2025

View reviewed changes

Update: update CapacityOfFabs parameters based on AMReX-Codes#4629 PR…

b615a4e

… discussion

atmyers reviewed Sep 3, 2025

View reviewed changes

Src/Particle/AMReX_ParticleContainer.H Outdated Show resolved Hide resolved

Moving capacityOfFabs function from MultiFab to FabArray

db786f3

atmyers reviewed Sep 3, 2025

View reviewed changes

jessdagostini added 2 commits September 4, 2025 11:44

Adding assertion on capacityOfFabs

1bbf3ca

Fixing reference

5f0f788

atmyers reviewed Sep 4, 2025

View reviewed changes

Src/Base/AMReX_MultiFabUtil.H Outdated Show resolved Hide resolved

jessdagostini added 2 commits September 4, 2025 12:45

Removing reference to capacityOfFabs function on MultiFabUtil

a390353

Change capacity function on particles to have similar signature as Fa…

7fb20d4

…bArray

atmyers reviewed Sep 5, 2025

View reviewed changes

Src/Base/AMReX_FabArray.H Outdated Show resolved Hide resolved

Update Src/Base/AMReX_FabArray.H

fb94457

Co-authored-by: Andrew Myers <atmyers@lbl.gov>

atmyers reviewed Sep 5, 2025

View reviewed changes

Src/Base/AMReX_FabArray.H Outdated Show resolved Hide resolved

jessdagostini added 3 commits September 5, 2025 15:07

Fixing reference to CellEqual func

6c0e6e2

Pushing tests for memory capturing (based on amrex-tutorials)

6b231b7

Fixing pipeline issues

bad61d0

WeiqunZhang requested changes Sep 8, 2025

View reviewed changes

WeiqunZhang reviewed Sep 8, 2025

View reviewed changes

ax3l added enhancement load balancing and removed enhancement labels Sep 11, 2025

jessdagostini and others added 6 commits September 11, 2025 10:36

Update Src/Base/AMReX_FabArray.H

8034c87

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Update Src/Base/AMReX_FabArray.H

2d2ddd0

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Update Src/Base/AMReX_FabArray.H

6cf35fa

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Update Src/Particle/AMReX_ParticleContainerI.H

327a241

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Update Src/Particle/AMReX_ParticleContainerI.H

6c0d68e

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Update Src/Particle/AMReX_ParticleContainer.H

58114a0

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>

Updating template signature for FabArray capacity func

84bbf37

Updating Particles function signature and testing

55d4a34

github-advanced-security bot found potential problems Sep 12, 2025

View reviewed changes

Finishing reorganization of tests

64e1e84

atmyers approved these changes Sep 16, 2025

View reviewed changes

Including HeatEquation test on CMakeLists

2b01097

Merge branch 'development' into development

c399f60

WeiqunZhang reviewed Sep 17, 2025

View reviewed changes

Tests/CMakeLists.txt Outdated Show resolved Hide resolved

WeiqunZhang and others added 4 commits September 16, 2025 19:21

Update Tests/CMakeLists.txt

63b21a5

Merge branch 'development' into jessdagostini/development

11dd878

cat doxygen.out

30ec777

workaround doxygen issue

d648a26

WeiqunZhang approved these changes Sep 17, 2025

View reviewed changes

WeiqunZhang merged commit f4295e3 into AMReX-Codes:development Sep 17, 2025
73 of 75 checks passed

ax3l mentioned this pull request Sep 17, 2025

AMReX/pyAMReX/PICSAR: Weekly Update BLAST-WarpX/warpx#6182

Merged

EZoni pushed a commit to BLAST-WarpX/warpx that referenced this pull request Sep 17, 2025

AMReX/pyAMReX/PICSAR: Weekly Update (#6182)

856017d

Pull in AMReX-Codes/amrex#4629 for #6166. Signed-off-by: Axel Huebl <axel.huebl@plasma.ninja>

Adding support to memory collection of Particles and MultiFabs within a grid #4629

Adding support to memory collection of Particles and MultiFabs within a grid #4629

Uh oh!

Conversation

jessdagostini commented Aug 27, 2025

Summary

Additional background

Checklist

Uh oh!

WeiqunZhang Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessdagostini Aug 29, 2025 • edited by ax3l Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

WeiqunZhang Aug 29, 2025 • edited by ax3l Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jessdagostini Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessdagostini commented Sep 7, 2025

Uh oh!

WeiqunZhang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

WeiqunZhang Sep 8, 2025 • edited by ax3l Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

WeiqunZhang Sep 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

atmyers commented Sep 11, 2025

Uh oh!

jessdagostini commented Sep 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Check notice

atmyers commented Sep 16, 2025

Uh oh!

WeiqunZhang Aug 28, 2025 •

edited

Loading

jessdagostini Aug 29, 2025 •

edited by ax3l

Loading

WeiqunZhang Aug 29, 2025 •

edited by ax3l

Loading

jessdagostini Aug 29, 2025 •

edited

Loading

WeiqunZhang Sep 8, 2025 •

edited by ax3l

Loading

WeiqunZhang Sep 11, 2025 •

edited

Loading

jessdagostini commented Sep 12, 2025 •

edited

Loading