Use llama for particle frame and shared memory DataBox layout

Also support LLAMA frames in the IO.
ComputationalRadiationPhysics · Dec 3, 2023 · d008726 · d008726
1 parent 5da4754
commit d008726
Show file tree

Hide file tree

Showing 54 changed files with 852 additions and 260 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "thirdParty/llama"]
+	path = thirdParty/llama
+	url = https://github.com/alpaka-group/llama
diff --git a/include/picongpu/algorithms/Set.hpp b/include/picongpu/algorithms/Set.hpp
@@ -33,9 +33,9 @@ namespace picongpu
         }
 
         template<typename Dst, typename T_Worker>
-        HDINLINE void operator()(T_Worker const&, Dst& dst) const
+        HDINLINE void operator()(T_Worker const&, Dst&& dst) const
         {
-            dst = value;
+            std::forward<Dst>(dst) = value;
         }
 
     private:

diff --git a/include/picongpu/algorithms/Velocity.hpp b/include/picongpu/algorithms/Velocity.hpp
@@ -28,7 +28,7 @@ namespace picongpu
     struct Velocity
     {
         template<typename MomType, typename MassType>
-        HDINLINE MomType operator()(const MomType mom, const MassType mass0)
+        HDINLINE auto operator()(const MomType mom, const MassType mass0)
         {
             const float_X rc2 = MUE0_EPS0;
             const float_X m0_2 = mass0 * mass0;

diff --git a/include/picongpu/fields/FieldJ.kernel b/include/picongpu/fields/FieldJ.kernel
@@ -131,7 +131,8 @@ namespace picongpu
                 // The rest uses normal weighting
                 const float_X weighting = particle[weighting_];
                 Velocity velocity;
-                const float3_X vel = velocity(particle[momentum_], attribute::getMass(weighting, particle));
+                const float3_X vel
+                    = velocity(static_cast<float3_X>(particle[momentum_]), attribute::getMass(weighting, particle));
                 auto fieldJShiftToParticle = jBox.shift(localCell);
                 ParticleAlgo perParticle;
                 perParticle(worker, fieldJShiftToParticle, pos, vel, charge, m_deltaTime);

diff --git a/include/picongpu/fields/FieldTmp.kernel b/include/picongpu/fields/FieldTmp.kernel
@@ -94,7 +94,9 @@ namespace picongpu
             if(!forEachParticle.hasParticles())
                 return;
 
-            auto cachedVal = CachedBox::create<0, typename T_TmpBox::ValueType>(worker, T_BlockDescription{});
+            auto cachedVal = CachedBox::create<0, SharedDataBoxMemoryLayout, typename T_TmpBox::ValueType>(
+                worker,
+                T_BlockDescription{});
             Set<typename T_TmpBox::ValueType> set(float_X(0.0));
 
             auto collective = makeThreadCollective<T_BlockDescription>();

diff --git a/include/picongpu/fields/MaxwellSolver/AddCurrentDensity.kernel b/include/picongpu/fields/MaxwellSolver/AddCurrentDensity.kernel
@@ -22,6 +22,8 @@
 
 #include "picongpu/simulation_defines.hpp"
 
+#include "picongpu/param/memory.param"
+
 #include <pmacc/dimensions/SuperCellDescription.hpp>
 #include <pmacc/lockstep.hpp>
 #include <pmacc/mappings/threads/ThreadCollective.hpp>
@@ -68,7 +70,9 @@ namespace picongpu::fields::maxwellSolver
 
             constexpr uint32_t cellsPerSuperCell = pmacc::math::CT::volume<SuperCellSize>::type::value;
 
-            auto cachedJ = CachedBox::create<0, typename FieldJ::DataBoxType::ValueType>(worker, BlockArea());
+            auto cachedJ = CachedBox::create<0, SharedDataBoxMemoryLayout, typename FieldJ::DataBoxType::ValueType>(
+                worker,
+                BlockArea());
 
             pmacc::math::operation::Assign assign;
             DataSpace<simDim> const block(

diff --git a/include/picongpu/fields/MaxwellSolver/FDTD/FDTDBase.kernel b/include/picongpu/fields/MaxwellSolver/FDTD/FDTDBase.kernel
@@ -179,7 +179,9 @@ namespace picongpu
                         auto srcFieldBlock = srcField.shift(beginCellIdx);
                         auto cacheStencilArea = makeThreadCollective<StencilCfg>();
                         auto cachedSrcField
-                            = CachedBox::create<0u, typename T_SrcBox::ValueType>(worker, StencilCfg{});
+                            = CachedBox::create<0u, SharedDataBoxMemoryLayout, typename T_SrcBox::ValueType>(
+                                worker,
+                                StencilCfg{});
                         cacheStencilArea(worker, assign, cachedSrcField, srcFieldBlock);
 
                         worker.sync();

diff --git a/include/picongpu/fields/currentDeposition/Cache.hpp b/include/picongpu/fields/currentDeposition/Cache.hpp
@@ -48,15 +48,11 @@ namespace picongpu
                  */
                 template<typename T_BlockDescription, typename T_Worker, typename T_FieldBox>
                 DINLINE static auto create(T_Worker const& worker, T_FieldBox const& fieldBox)
-#if(!BOOST_COMP_CLANG)
-                    -> decltype(CachedBox::create<0u, typename T_FieldBox::ValueType>(
-                        worker,
-                        std::declval<T_BlockDescription>()))
-#endif
                 {
                     using ValueType = typename T_FieldBox::ValueType;
                     /* this memory is used by all virtual blocks */
-                    auto cache = CachedBox::create<0u, ValueType>(worker, T_BlockDescription{});
+                    auto cache
+                        = CachedBox::create<0u, SharedDataBoxMemoryLayout, ValueType>(worker, T_BlockDescription{});
 
                     Set<ValueType> set(ValueType::create(0.0_X));
                     auto collectiveFill = makeThreadCollective<T_BlockDescription>();
@@ -90,9 +86,6 @@ namespace picongpu
                  */
                 template<typename T_BlockDescription, typename T_Worker, typename T_FieldBox>
                 DINLINE static auto create([[maybe_unused]] T_Worker const& worker, T_FieldBox const& fieldBox)
-#if(!BOOST_COMP_CLANG)
-                    -> T_FieldBox
-#endif
                 {
                     return fieldBox;
                 }

diff --git a/include/picongpu/fields/incidentField/Solver.hpp b/include/picongpu/fields/incidentField/Solver.hpp
@@ -283,8 +283,9 @@ namespace picongpu
                     using IntVector = pmacc::math::Vector<int, simDim>;
                     auto const beginLocalUserIdx
                         = Index{math::max(IntVector{beginUserIdx - totalCellOffset}, IntVector::create(0))};
-                    auto const endLocalUserIdx
-                        = Index{math::min(IntVector{endUserIdx - totalCellOffset}, IntVector{localDomain.size})};
+                    auto const endLocalUserIdx = Index{math::min(
+                        IntVector{endUserIdx - totalCellOffset},
+                        static_cast<const IntVector&>(localDomain.size))};
 
                     // Check if we have any active cells in the local domain
                     bool areAnyCellsInLocalDomain = true;

diff --git a/include/picongpu/param/memory.param b/include/picongpu/param/memory.param
@@ -114,4 +114,19 @@ namespace picongpu
      */
     constexpr bool fieldTmpSupportGatherCommunication = true;
 
+    struct ParticleFrameMemoryLayout
+        : llama::mapping::BindSoA<llama::mapping::Blobs::Single, llama::mapping::SubArrayAlignment::Align>
+    {
+        inline static constexpr bool splitVector = false;
+    };
+
+    struct ParticleFrameMemoryLayoutOpenPMD : llama::mapping::BindSoA<llama::mapping::Blobs::OnePerField>
+    {
+        inline static constexpr bool splitVector = false;
+    };
+
+    struct SharedDataBoxMemoryLayout : llama::mapping::BindAoS<>
+    {
+        inline static constexpr bool splitVector = false;
+    };
 } // namespace picongpu
diff --git a/include/picongpu/particles/Particles.hpp b/include/picongpu/particles/Particles.hpp
@@ -22,6 +22,7 @@
 
 #include "picongpu/fields/Fields.def"
 #include "picongpu/fields/Fields.hpp"
+#include "picongpu/param/memory.param"
 #include "picongpu/particles/boundary/Description.hpp"
 #include "picongpu/particles/boundary/Utility.hpp"
 #include "picongpu/particles/manipulators/manipulators.def"
@@ -89,6 +90,7 @@ namespace picongpu
                       pmacc::HandleGuardRegion<
                           pmacc::particles::policies::ExchangeParticles,
                           pmacc::particles::policies::DoNothing>>>,
+              ParticleFrameMemoryLayout,
               MappingDesc,
               DeviceHeap>
         , public ISimulationData
@@ -108,7 +110,8 @@ namespace picongpu
                 pmacc::HandleGuardRegion<
                     pmacc::particles::policies::ExchangeParticles,
                     pmacc::particles::policies::DoNothing>>>;
-        using ParticlesBaseType = ParticlesBase<SpeciesParticleDescription, picongpu::MappingDesc, DeviceHeap>;
+        using ParticlesBaseType
+            = ParticlesBase<SpeciesParticleDescription, ParticleFrameMemoryLayout, picongpu::MappingDesc, DeviceHeap>;
         using FrameType = typename ParticlesBaseType::FrameType;
         using FrameTypeBorder = typename ParticlesBaseType::FrameTypeBorder;
         using ParticlesBoxType = typename ParticlesBaseType::ParticlesBoxType;

diff --git a/include/picongpu/particles/Particles.kernel b/include/picongpu/particles/Particles.kernel
@@ -222,8 +222,10 @@ namespace picongpu
 
             onlyMaster([&]() { mustShiftSupercell = 0; });
 
-            auto cachedB = CachedBox::create<0, typename T_BBox::ValueType>(worker, T_DataDomain());
-            auto cachedE = CachedBox::create<1, typename T_EBox::ValueType>(worker, T_DataDomain());
+            auto cachedB
+                = CachedBox::create<0, SharedDataBoxMemoryLayout, typename T_BBox::ValueType>(worker, T_DataDomain());
+            auto cachedE
+                = CachedBox::create<1, SharedDataBoxMemoryLayout, typename T_EBox::ValueType>(worker, T_DataDomain());
 
             worker.sync();
 

diff --git a/include/picongpu/particles/Particles.tpp b/include/picongpu/particles/Particles.tpp
@@ -43,6 +43,7 @@
 #include <pmacc/traits/Resolve.hpp>
 
 #include <algorithm>
+#include <fstream>
 #include <iostream>
 #include <limits>
 #include <memory>
@@ -197,7 +198,9 @@ namespace picongpu
         const std::shared_ptr<DeviceHeap>& heap,
         picongpu::MappingDesc cellDescription,
         SimulationDataId datasetID)
-        : ParticlesBase<SpeciesParticleDescription, picongpu::MappingDesc, DeviceHeap>(heap, cellDescription)
+        : ParticlesBase<SpeciesParticleDescription, ParticleFrameMemoryLayout, picongpu::MappingDesc, DeviceHeap>(
+            heap,
+            cellDescription)
         , m_datasetID(datasetID)
     {
         constexpr bool particleHasShape = pmacc::traits::HasIdentifier<FrameType, shape<>>::type::value;
@@ -212,6 +215,21 @@ namespace picongpu
 
         size_t sizeOfExchanges = 0u;
 
+#if __has_include(<fmt/format.h>)
+        // dump the data layout of the particle frames
+        if constexpr(PIConGPUVerbose::log_level & picLog::MEMORY::lvl)
+        {
+            log<picLog::MEMORY>(
+                "Dumping LLAMA memory layout for frame and border into llama_frame.* and llama_border_fream.*");
+            auto fm = typename decltype(FrameType::view)::Mapping{};
+            std::ofstream{"llama_frame.html"} << llama::toHtml(fm);
+            std::ofstream{"llama_frame.svg"} << llama::toSvg(fm);
+            auto bfm = typename decltype(FrameTypeBorder::view)::Mapping{};
+            std::ofstream{"llama_border_frame.html"} << llama::toHtml(bfm);
+            std::ofstream{"llama_border_frame.svg"} << llama::toSvg(bfm);
+        }
+#endif
+
         const uint32_t commTag = pmacc::traits::getUniqueId();
         log<picLog::MEMORY>("communication tag for species %1%: %2%") % FrameType::getName() % commTag;
 

diff --git a/include/picongpu/particles/ionization/byCollision/ThomasFermi/ThomasFermi_Impl.hpp b/include/picongpu/particles/ionization/byCollision/ThomasFermi/ThomasFermi_Impl.hpp
@@ -23,6 +23,7 @@
 
 #include "picongpu/fields/CellType.hpp"
 #include "picongpu/fields/FieldTmp.hpp"
+#include "picongpu/param/memory.param"
 #include "picongpu/particles/atomicPhysics/SetChargeState.hpp"
 #include "picongpu/particles/ionization/byCollision/ThomasFermi/AlgorithmThomasFermi.hpp"
 #include "picongpu/particles/ionization/byCollision/ThomasFermi/ThomasFermi.def"
@@ -104,8 +105,20 @@ namespace picongpu
                 PMACC_ALIGN(eneBox, FieldTmp::DataBoxType);
 
                 /* shared memory EM-field device databoxes */
-                PMACC_ALIGN(cachedRho, DataBox<SharedBox<ValueType_Rho, typename BlockArea::FullSuperCellSize, 0>>);
-                PMACC_ALIGN(cachedEne, DataBox<SharedBox<ValueType_Ene, typename BlockArea::FullSuperCellSize, 1>>);
+                PMACC_ALIGN(
+                    cachedRho,
+                    DataBox<SharedBox<
+                        ValueType_Rho,
+                        typename BlockArea::FullSuperCellSize,
+                        0,
+                        SharedDataBoxMemoryLayout>>);
+                PMACC_ALIGN(
+                    cachedEne,
+                    DataBox<SharedBox<
+                        ValueType_Ene,
+                        typename BlockArea::FullSuperCellSize,
+                        1,
+                        SharedDataBoxMemoryLayout>>);
 
             public:
                 /* host constructor initializing member : random number generator */
@@ -185,8 +198,8 @@ namespace picongpu
                 DINLINE void collectiveInit(const T_Worker& worker, const DataSpace<simDim>& blockCell)
                 {
                     /* caching of density and "temperature" fields */
-                    cachedRho = CachedBox::create<0, ValueType_Rho>(worker, BlockArea());
-                    cachedEne = CachedBox::create<1, ValueType_Ene>(worker, BlockArea());
+                    cachedRho = CachedBox::create<0, SharedDataBoxMemoryLayout, ValueType_Rho>(worker, BlockArea());
+                    cachedEne = CachedBox::create<1, SharedDataBoxMemoryLayout, ValueType_Ene>(worker, BlockArea());
 
                     /* instance of nvidia assignment operator */
                     pmacc::math::operation::Assign assign;

diff --git a/include/picongpu/particles/ionization/byField/ADK/ADK_Impl.hpp b/include/picongpu/particles/ionization/byField/ADK/ADK_Impl.hpp
@@ -102,8 +102,14 @@ namespace picongpu
                 PMACC_ALIGN(bBox, FieldB::DataBoxType);
                 PMACC_ALIGN(jBox, FieldJ::DataBoxType);
                 /* shared memory EM-field device databoxes */
-                PMACC_ALIGN(cachedE, DataBox<SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1>>);
-                PMACC_ALIGN(cachedB, DataBox<SharedBox<ValueType_B, typename BlockArea::FullSuperCellSize, 0>>);
+                PMACC_ALIGN(
+                    cachedE,
+                    DataBox<
+                        SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1, SharedDataBoxMemoryLayout>>);
+                PMACC_ALIGN(
+                    cachedB,
+                    DataBox<
+                        SharedBox<ValueType_B, typename BlockArea::FullSuperCellSize, 0, SharedDataBoxMemoryLayout>>);
 
             public:
                 /* host constructor initializing member : random number generator */
@@ -137,8 +143,8 @@ namespace picongpu
                     jBox = jBox.shift(blockCell);
 
                     /* caching of E and B fields */
-                    cachedB = CachedBox::create<0, ValueType_B>(worker, BlockArea());
-                    cachedE = CachedBox::create<1, ValueType_E>(worker, BlockArea());
+                    cachedB = CachedBox::create<0, SharedDataBoxMemoryLayout, ValueType_B>(worker, BlockArea());
+                    cachedE = CachedBox::create<1, SharedDataBoxMemoryLayout, ValueType_E>(worker, BlockArea());
 
                     /* instance of nvidia assignment operator */
                     pmacc::math::operation::Assign assign;

diff --git a/include/picongpu/particles/ionization/byField/BSI/BSI_Impl.hpp b/include/picongpu/particles/ionization/byField/BSI/BSI_Impl.hpp
@@ -25,6 +25,7 @@
 #include "picongpu/fields/FieldB.hpp"
 #include "picongpu/fields/FieldE.hpp"
 #include "picongpu/fields/FieldJ.hpp"
+#include "picongpu/param/memory.param"
 #include "picongpu/particles/ParticlesFunctors.hpp"
 #include "picongpu/particles/atomicPhysics/SetChargeState.hpp"
 #include "picongpu/particles/ionization/byField/BSI/AlgorithmBSI.hpp"
@@ -93,7 +94,10 @@ namespace picongpu
                 FieldE::DataBoxType eBox;
                 FieldJ::DataBoxType jBox;
                 /* shared memory EM-field device databoxes */
-                PMACC_ALIGN(cachedE, DataBox<SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1>>);
+                PMACC_ALIGN(
+                    cachedE,
+                    DataBox<
+                        SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1, SharedDataBoxMemoryLayout>>);
 
             public:
                 /* host constructor */
@@ -125,7 +129,7 @@ namespace picongpu
                     jBox = jBox.shift(blockCell);
 
                     /* caching of E field */
-                    cachedE = CachedBox::create<1, ValueType_E>(worker, BlockArea());
+                    cachedE = CachedBox::create<1, SharedDataBoxMemoryLayout, ValueType_E>(worker, BlockArea());
 
                     /* instance of nvidia assignment operator */
                     pmacc::math::operation::Assign assign;

diff --git a/include/picongpu/particles/ionization/byField/Keldysh/Keldysh_Impl.hpp b/include/picongpu/particles/ionization/byField/Keldysh/Keldysh_Impl.hpp
@@ -102,8 +102,14 @@ namespace picongpu
                 PMACC_ALIGN(bBox, FieldB::DataBoxType);
                 PMACC_ALIGN(jBox, FieldJ::DataBoxType);
                 /* shared memory EM-field device databoxes */
-                PMACC_ALIGN(cachedE, DataBox<SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1>>);
-                PMACC_ALIGN(cachedB, DataBox<SharedBox<ValueType_B, typename BlockArea::FullSuperCellSize, 0>>);
+                PMACC_ALIGN(
+                    cachedE,
+                    DataBox<
+                        SharedBox<ValueType_E, typename BlockArea::FullSuperCellSize, 1, SharedDataBoxMemoryLayout>>);
+                PMACC_ALIGN(
+                    cachedB,
+                    DataBox<
+                        SharedBox<ValueType_B, typename BlockArea::FullSuperCellSize, 0, SharedDataBoxMemoryLayout>>);
 
             public:
                 /* host constructor initializing member : random number generator */
@@ -137,8 +143,8 @@ namespace picongpu
                     jBox = jBox.shift(blockCell);
 
                     /* caching of E and B fields */
-                    cachedB = CachedBox::create<0, ValueType_B>(worker, BlockArea());
-                    cachedE = CachedBox::create<1, ValueType_E>(worker, BlockArea());
+                    cachedB = CachedBox::create<0, SharedDataBoxMemoryLayout, ValueType_B>(worker, BlockArea());
+                    cachedE = CachedBox::create<1, SharedDataBoxMemoryLayout, ValueType_E>(worker, BlockArea());
 
                     /* instance of nvidia assignment operator */
                     pmacc::math::operation::Assign assign;

diff --git a/include/picongpu/plugins/PhaseSpace/PhaseSpaceFunctors.hpp b/include/picongpu/plugins/PhaseSpace/PhaseSpaceFunctors.hpp
@@ -160,7 +160,7 @@ namespace picongpu
             /* create shared mem */
             constexpr int blockCellsInDir = SuperCellSize::template at<r_dir>::type::value;
             using SharedMemSize = SuperCellDescription<pmacc::math::CT::Int<num_pbins, blockCellsInDir>>;
-            auto sharedMemHist = CachedBox::create<0u, float_PS>(worker, SharedMemSize{});
+            auto sharedMemHist = CachedBox::create<0u, SharedDataBoxMemoryLayout, float_PS>(worker, SharedMemSize{});
 
             Set<float_PS> set(float_PS{0.0});
             auto collectiveOnSharedHistogram = makeThreadCollective<SharedMemSize>();