Add WarpX_UNITY_BUILD (#5702)

Unity builds combine all `.cpp` files into a single one (a single
translation unit, TU). This is useful for performance optimization,
computer science experiments, or differentiable programming.

FYI: A CPU unity build of ImpactX is around 18sec on my laptop
(BLAST-ImpactX/impactx#861), for WarpX around
3:30min. In both cases, this refers to the `lib*` of ImpactX/WarpX part.

---------

Co-authored-by: Edoardo Zoni <[email protected]>

Author: ax3l

CMakeLists.txt

@@ -166,6 +166,10 @@ option(WarpX_PYTHON_IPO
     ${_WarpX_PYTHON_IPO_DEFAULT}
 )
 
+# Unity builds combine all .cpp files into a single one (a single translation
+# unit, TU).
+option(WarpX_UNITY_BUILD          "WarpX library as unity build" OFF)
+
 set(pyWarpX_VERSION_INFO "" CACHE STRING
     "PEP-440 conformant version (set by setup.py)")
 
@@ -277,6 +281,17 @@ foreach(D IN LISTS WarpX_DIMS)
             POSITION_INDEPENDENT_CODE ON
             WINDOWS_EXPORT_ALL_SYMBOLS ON
         )
+
+        # Optional: build only a single TU
+        if(WarpX_UNITY_BUILD)
+            set_target_properties(lib_${SD} PROPERTIES
+                UNITY_BUILD ON
+                UNITY_BUILD_MODE BATCH
+                UNITY_BUILD_UNIQUE_ID "WARPX_UNITY_ID"
+                # Number must be more than the number of .cpp files in WarpX
+                UNITY_BUILD_BATCH_SIZE 10000
+            )
+        endif()
     endif()
 
     # executable application

Docs/source/install/cmake.rst

@@ -121,6 +121,7 @@ CMake Option                  Default & Values                               Des
 ============================= ============================================== ===========================================================
 ``BUILD_SHARED_LIBS``         ON/**OFF**                                     `Build shared libraries for dependencies <https://cmake.org/cmake/help/latest/variable/BUILD_SHARED_LIBS.html>`__
 ``WarpX_CCACHE``              **ON**/OFF                                     Search and use CCache to speed up rebuilds.
+``WarpX_UNITY_BUILD``         ON/**OFF**                                     WarpX library as unity build (single TU)
 ``AMReX_CUDA_PTX_VERBOSE``    ON/**OFF**                                     Print CUDA code generation statistics from ``ptxas``.
 ``WarpX_amrex_src``           *None*                                         Path to AMReX source directory (preferred if set)
 ``WarpX_amrex_repo``          ``https://github.com/AMReX-Codes/amrex.git``   Repository URI to pull and build AMReX from

Source/BoundaryConditions/PML_RZ.cpp

@@ -43,6 +43,8 @@ PML_RZ::PML_RZ (int lev, amrex::BoxArray const& grid_ba, amrex::DistributionMapp
       m_do_pml_in_domain(do_pml_in_domain),
       m_geom(geom)
 {
+    using ablastr::fields::Direction;
+
     auto & warpx = WarpX::GetInstance();
 
     bool const remake = false;
@@ -73,6 +75,7 @@ PML_RZ::ApplyDamping (amrex::MultiFab* Et_fp, amrex::MultiFab* Ez_fp,
                       amrex::MultiFab* Bt_fp, amrex::MultiFab* Bz_fp,
                       amrex::Real dt, ablastr::fields::MultiFabRegister& fields)
 {
+    using ablastr::fields::Direction;
 
     amrex::Real const dr = m_geom->CellSize(0);
     amrex::Real const cdt_over_dr = PhysConst::c*dt/dr;
@@ -132,6 +135,8 @@ PML_RZ::ApplyDamping (amrex::MultiFab* Et_fp, amrex::MultiFab* Ez_fp,
 void
 PML_RZ::FillBoundaryE (ablastr::fields::MultiFabRegister& fields, PatchType patch_type, std::optional<bool> nodal_sync)
 {
+    using ablastr::fields::Direction;
+
     amrex::MultiFab * pml_Er = fields.get(FieldType::pml_E_fp, Direction{0}, 0);
     amrex::MultiFab * pml_Et = fields.get(FieldType::pml_E_fp, Direction{1}, 0);
 
@@ -148,6 +153,8 @@ PML_RZ::FillBoundaryB (ablastr::fields::MultiFabRegister& fields, PatchType patc
 {
     if (patch_type == PatchType::fine)
     {
+        using ablastr::fields::Direction;
+
         amrex::MultiFab * pml_Br = fields.get(FieldType::pml_B_fp, Direction{0}, 0);
         amrex::MultiFab * pml_Bt = fields.get(FieldType::pml_B_fp, Direction{1}, 0);
 
@@ -160,6 +167,8 @@ PML_RZ::FillBoundaryB (ablastr::fields::MultiFabRegister& fields, PatchType patc
 void
 PML_RZ::CheckPoint (ablastr::fields::MultiFabRegister& fields, std::string const& dir) const
 {
+    using ablastr::fields::Direction;
+
     if (fields.has(FieldType::pml_E_fp, Direction{0}, 0)) {
         amrex::VisMF::AsyncWrite(*fields.get(FieldType::pml_E_fp, Direction{0}, 0), dir+"_Er_fp");
         amrex::VisMF::AsyncWrite(*fields.get(FieldType::pml_E_fp, Direction{1}, 0), dir+"_Et_fp");
@@ -171,6 +180,8 @@ PML_RZ::CheckPoint (ablastr::fields::MultiFabRegister& fields, std::string const
 void
 PML_RZ::Restart (ablastr::fields::MultiFabRegister& fields, std::string const& dir)
 {
+    using ablastr::fields::Direction;
+
     if (fields.has(FieldType::pml_E_fp, Direction{0}, 0)) {
         amrex::VisMF::Read(*fields.get(FieldType::pml_E_fp, Direction{0}, 0), dir+"_Er_fp");
         amrex::VisMF::Read(*fields.get(FieldType::pml_E_fp, Direction{1}, 0), dir+"_Et_fp");
@@ -195,6 +206,8 @@ PML_RZ::PushPMLPSATDSinglePatchRZ (
     SpectralSolverRZ& solver,
     ablastr::fields::MultiFabRegister& fields)
 {
+    using ablastr::fields::Direction;
+
     SpectralFieldIndex const& Idx = solver.m_spectral_index;
     amrex::MultiFab * pml_Er = fields.get(FieldType::pml_E_fp, Direction{0}, 0);
     amrex::MultiFab * pml_Et = fields.get(FieldType::pml_E_fp, Direction{1}, 0);

Source/Diagnostics/FlushFormats/FlushFormatCheckpoint.cpp

@@ -22,13 +22,20 @@
 #include <AMReX_Utility.H>
 #include <AMReX_VisMF.H>
 
+#ifndef WARPX_UNITY_ID
+#define WARPX_UNITY_ID
+#endif
+
 using namespace amrex;
 using warpx::fields::FieldType;
 
 namespace
+{
+namespace WARPX_UNITY_ID
 {
     const std::string default_level_prefix {"Level_"};
 }
+}
 
 void
 FlushFormatCheckpoint::WriteToFile (
@@ -48,6 +55,7 @@ FlushFormatCheckpoint::WriteToFile (
         bool /*isLastBTDFlush*/) const
 {
     using ablastr::fields::Direction;
+    using WARPX_UNITY_ID::default_level_prefix;
 
     WARPX_PROFILE("FlushFormatCheckpoint::WriteToFile()");
 

Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp

@@ -49,13 +49,20 @@
 #include <utility>
 #include <vector>
 
+#ifndef WARPX_UNITY_ID
+#define WARPX_UNITY_ID
+#endif
+
 using namespace amrex;
 using warpx::fields::FieldType;
 
 namespace
+{
+namespace WARPX_UNITY_ID
 {
     const std::string default_level_prefix {"Level_"};
 }
+}
 
 void
 FlushFormatPlotfile::WriteToFile (
@@ -570,6 +577,7 @@ FlushFormatPlotfile::WriteAllRawFields(
     const bool plot_raw_fields_guards) const
 {
     using ablastr::fields::Direction;
+    using WARPX_UNITY_ID::default_level_prefix;
 
     if (!plot_raw_fields) { return; }
     auto & warpx = WarpX::GetInstance();

Source/Diagnostics/ReducedDiags/DifferentialLuminosity.cpp

@@ -54,17 +54,12 @@
 #include <memory>
 #include <vector>
 
-using ParticleType = WarpXParticleContainer::ParticleType;
-using ParticleTileType = WarpXParticleContainer::ParticleTileType;
-using ParticleTileDataType = ParticleTileType::ParticleTileDataType;
-using ParticleBins = amrex::DenseBins<ParticleTileDataType>;
-using index_type = ParticleBins::index_type;
-
-using namespace amrex;
 
 DifferentialLuminosity::DifferentialLuminosity (const std::string& rd_name)
 : ReducedDiags{rd_name}
 {
+    using namespace amrex::literals;
+
     // read colliding species names - must be 2
     const amrex::ParmParse pp_rd_name(m_rd_name);
     pp_rd_name.getarr("species", m_beam_name);
@@ -112,7 +107,7 @@ DifferentialLuminosity::DifferentialLuminosity (const std::string& rd_name)
             {
                 ofs << m_sep;
                 ofs << "[" << off++ << "]";
-                const Real b = m_bin_min + m_bin_size*(Real(i)+0.5_rt);
+                const amrex::Real b = m_bin_min + m_bin_size*(amrex::Real(i)+0.5_rt);
                 ofs << "bin" << 1+i << "=" << b << "(eV)";
             }
             ofs << "\n";
@@ -129,6 +124,12 @@ void DifferentialLuminosity::ComputeDiags (int step)
 #else
     WARPX_PROFILE("DifferentialLuminosity::ComputeDiags");
 
+    using namespace amrex;
+    using ParticleTileType = WarpXParticleContainer::ParticleTileType;
+    using ParticleTileDataType = ParticleTileType::ParticleTileDataType;
+    using ParticleBins = amrex::DenseBins<ParticleTileDataType>;
+    using index_type = ParticleBins::index_type;
+
     // Since this diagnostic *accumulates* the luminosity in the
     // array d_data, we add contributions at *each timestep*, but
     // we only write the data to file at intervals specified by the user.

Source/Diagnostics/ReducedDiags/DifferentialLuminosity2D.cpp

@@ -60,17 +60,10 @@
 #include <memory>
 #include <vector>
 
-using ParticleType = WarpXParticleContainer::ParticleType;
-using ParticleTileType = WarpXParticleContainer::ParticleTileType;
-using ParticleTileDataType = ParticleTileType::ParticleTileDataType;
-using ParticleBins = amrex::DenseBins<ParticleTileDataType>;
-using index_type = ParticleBins::index_type;
-
 #ifdef WARPX_USE_OPENPMD
 namespace io = openPMD;
 #endif
 
-using namespace amrex;
 
 DifferentialLuminosity2D::DifferentialLuminosity2D (const std::string& rd_name)
 : ReducedDiags{rd_name}
@@ -80,6 +73,7 @@ DifferentialLuminosity2D::DifferentialLuminosity2D (const std::string& rd_name)
     WARPX_ABORT_WITH_MESSAGE(
         "DifferentialLuminosity2D diagnostics does not work in RZ geometry.");
 #endif
+    using namespace amrex::literals;
 
     // read colliding species names - must be 2
     amrex::ParmParse pp_rd_name(m_rd_name);
@@ -120,9 +114,9 @@ DifferentialLuminosity2D::DifferentialLuminosity2D (const std::string& rd_name)
     m_bin_size_2 = (bin_max_2 - bin_min_2) / bin_num_2;
 
     // resize data array on the host
-    Array<int,2> tlo{0,0}; // lower bounds
-    Array<int,2> thi{m_bin_num_1-1, m_bin_num_2-1}; // inclusive upper bounds
-    m_h_data_2D.resize(tlo, thi, The_Pinned_Arena());
+    amrex::Array<int,2> tlo{0,0}; // lower bounds
+    amrex::Array<int,2> thi{m_bin_num_1-1, m_bin_num_2-1}; // inclusive upper bounds
+    m_h_data_2D.resize(tlo, thi, amrex::The_Pinned_Arena());
 
     auto const& h_table_data = m_h_data_2D.table();
     // initialize data on the host
@@ -136,17 +130,22 @@ DifferentialLuminosity2D::DifferentialLuminosity2D (const std::string& rd_name)
     m_d_data_2D.resize(tlo, thi);
     // copy data from host to device
     m_d_data_2D.copy(m_h_data_2D);
-    Gpu::streamSynchronize();
+    amrex::Gpu::streamSynchronize();
 } // end constructor
 
 void DifferentialLuminosity2D::ComputeDiags (int step)
 {
 #if defined(WARPX_DIM_RZ)
     amrex::ignore_unused(step);
 #else
-
     WARPX_PROFILE("DifferentialLuminosity2D::ComputeDiags");
 
+    using namespace amrex;
+    using ParticleTileType = WarpXParticleContainer::ParticleTileType;
+    using ParticleTileDataType = ParticleTileType::ParticleTileDataType;
+    using ParticleBins = amrex::DenseBins<ParticleTileDataType>;
+    using index_type = ParticleBins::index_type;
+
     // Since this diagnostic *accumulates* the luminosity in the
     // table m_d_data_2D, we add contributions at *each timestep*, but
     // we only write the data to file at intervals specified by the user.
@@ -336,7 +335,7 @@ void DifferentialLuminosity2D::WriteToFile (int step) const
 
 #ifdef WARPX_USE_OPENPMD
     // only IO processor writes
-    if ( !ParallelDescriptor::IOProcessor() ) { return; }
+    if ( !amrex::ParallelDescriptor::IOProcessor() ) { return; }
 
     // TODO: support different filename templates
     std::string filename = "openpmd";