From 2d5f8b52ac1789da316baf3b4dd326360d624bed Mon Sep 17 00:00:00 2001
From: Luca Fedeli <luca.fedeli@cea.fr>
Date: Thu, 19 Sep 2024 14:36:54 +0200
Subject: [PATCH] try to improve performances

---
 Source/Particles/Radiation/RadiationHandler.H |  3 ++-
 .../Particles/Radiation/RadiationHandler.cpp  | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/Source/Particles/Radiation/RadiationHandler.H b/Source/Particles/Radiation/RadiationHandler.H
index 670745dacde..010c75e52a0 100644
--- a/Source/Particles/Radiation/RadiationHandler.H
+++ b/Source/Particles/Radiation/RadiationHandler.H
@@ -19,6 +19,7 @@
 
 #include <AMReX_GpuContainers.H>
 #include <AMReX_REAL.H>
+#include <AMReX_Scan.H>
 #include <AMReX_Vector.H>
 
 #include <memory>
@@ -97,8 +98,8 @@ private:
     std::optional<int> m_step_skip = std::nullopt;
 
     std::optional<std::array<amrex::Real,2>> m_gamma_range = std::nullopt;
-    amrex::Gpu::DeviceVector<int> m_offset;
     amrex::Gpu::DeviceVector<int> m_mask;
+    amrex::Gpu::DeviceVector<int> m_offset;
     amrex::Gpu::DeviceVector<int> m_idx;
 };
 #endif // WARPX_PARTICLES_RADIATION_H
diff --git a/Source/Particles/Radiation/RadiationHandler.cpp b/Source/Particles/Radiation/RadiationHandler.cpp
index 6e10462974b..08bf8aacac9 100644
--- a/Source/Particles/Radiation/RadiationHandler.cpp
+++ b/Source/Particles/Radiation/RadiationHandler.cpp
@@ -564,9 +564,10 @@ void RadiationHandler::add_radiation_contribution(
                     else{ //m_gamma_range.has_value()
 
 
-                        m_offset.resize(np);
                         m_mask.resize(np);
+                        m_offset.resize(np);
                         int* const AMREX_RESTRICT p_mask = m_mask.dataPtr();
+                        int* const AMREX_RESTRICT p_offset = m_offset.dataPtr();
 
                         const auto gamma_min = m_gamma_range.value()[0];
                         const auto gamma_max = m_gamma_range.value()[1];
@@ -587,6 +588,18 @@ void RadiationHandler::add_radiation_contribution(
                                 p_mask[ip] =  is_in;
                         });
 
+                        const auto nrad = amrex::Scan::ExclusiveSum(np, p_mask, p_offset);
+
+                        m_idx.resize(nrad);
+                        int* const AMREX_RESTRICT p_idx = m_idx.dataPtr();
+
+                        amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(int ip){
+                            if (p_mask[ip]){
+                                p_idx[p_offset[ip]] = ip;
+                            }
+                        });
+
+
 #if defined(WARPX_DIM_3D)
                         amrex::ParallelFor(
                             np_omegas_detpos, [=] AMREX_GPU_DEVICE(int, int i_om, int i_det){
@@ -607,9 +620,9 @@ void RadiationHandler::add_radiation_contribution(
                             auto sum_cy = Complex{0.0_prt, 0.0_prt};
                             auto sum_cz = Complex{0.0_prt, 0.0_prt};
 
-                            for (int ip =  0; ip < np; ++ip){
+                            for (int irad =  0; irad < nrad; ++irad){
 
-                                if (!p_mask[ip]) {continue;}
+                                const auto ip = p_idx[irad];
 
                                 amrex::ParticleReal xp, yp, zp;
                                 GetPosition.AsStored(ip, xp, yp, zp);