OpenMP, how does it work for PyAMReX?

[Simon067]

How can OpenMP be used to run parallel over parts of the grid? As far as I understand, when you make a view (like a not-copied numpy array) for data in the MultiFabs you no longer use OpenMP to parallelize actions over it? But can you even do calculations on the data without using views? I took HeatEquation.py as an example of how to do calculations on fields, but are there other ways (that are OpenMP parallelizable)?

edit: it appears my mistake it that it doesn't do any tiling within boxes when iterating in a typical

for mfi in mfab:
    do_something()

loop? do you need to globally turn on static tiling?

[ax3l]

Hi @Simon067 ,

If you compile (py)AMReX with OpenMP support, then intrinsic functions called on the MultiFabs will make use of OpenMP threads. Examples: mfab.saxpy(...), mfab.norm2(0), etc.

We have not yet exposed detailed OpenMP support on the MFiter level, but I am happy to explore it. In C++ this looks like this:
https://amrex-cC++odes.github.io/amrex/docs_html/Basics.html#sec-basics-mfiter-tiling

  #ifdef AMREX_USE_OMP
  #pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
  #endif
  for (MFIter mfi(divEfield, TilingIfNotGPU()); mfi.isValid(); ++mfi) {
      Array4<Real> const& divE = divEfield.array(mfi);  //  global-indexed
      Box const& tdive = mfi.tilebox(divEfield.ixType().toIntVect());
      amrex::ParallelFor(tdive, [=] AMREX_GPU_DEVICE (int i, int j, int k) {
          divE(i,j,k) = ... Ex(i-1,j,k) ...;                             // neighbor access
      });
  }

Note that tiling is only a technique to manufacture work units for a thread pool (e.g., using ThreadPoolExecutor in Python) when you have fewer boxes than threads on a process (MPI rank). On itself it does only add more runs through the mfi loop that can be parallelized over threads.

If we expose tiling by implementing proper amr.MFIter.__iter__, so one can add amr.MFItInfo with tiling info, we will need to expose also some helper for slicing the mfi.tilebox(). Currently, mfab.array(mfi) exposes the whole fab box.

An alternative you can explore in the meantime to manufacture more parallel work on CPUs: over-decompose the grid and use more MPI ranks.

[ax3l]

Started a draft on this in #610