Apply suggestions from code review

Co-authored-by: Alexander Sinn <[email protected]>

Author: WeiqunZhang

Src/Base/AMReX_GpuLaunch.H

@@ -180,7 +180,7 @@ namespace Gpu {
 
     struct ExecConfig
     {
-        Long ntotalthreads;
+        Long start_idx;
         int nblocks;
     };
 
@@ -195,6 +195,7 @@ namespace Gpu {
         // loops inside GPU kernels.
         auto nlaunches = int((N+nmax-1)/nmax);
         Vector<ExecConfig> r(nlaunches);
+        Long ndone = 0;
         for (int i = 0; i < nlaunches; ++i) {
             int nblocks;
             if (N > nmax) {
@@ -203,8 +204,9 @@ namespace Gpu {
             } else {
                 nblocks = int((N+MT-1)/MT);
             }
-            // Total # of threads in this launch
-            r[i].ntotalthreads = Long(nblocks) * MT;
+            // At which element ID the kernel should start
+            r[i].start_idx = ndone;
+            ndone += Long(nblocks) * MT;
             // # of blocks in this launch
             r[i].nblocks = nblocks;
         }

Src/Base/AMReX_GpuLaunchFunctsG.H

@@ -800,22 +800,19 @@ ParallelFor (Gpu::KernelInfo const&, T n, L const& f) noexcept
     static_assert(sizeof(T) >= 2);
     if (amrex::isEmpty(n)) { return; }
     const auto& nec = Gpu::makeNExecutionConfigs<MT>(n);
-    T ndone = 0;
     for (auto const& ec : nec) {
-        T nleft = n - ndone;
+        const T start_idx = T(ec.start_idx);
+        const T nleft = n - start_idx;
         AMREX_LAUNCH_KERNEL(MT, ec.nblocks, MT, 0, Gpu::gpuStream(),
         [=] AMREX_GPU_DEVICE () noexcept {
             // This will not overflow, even though nblocks*MT might.
             auto tid = T(MT)*T(blockIdx.x)+T(threadIdx.x);
             if (tid < nleft) {
-                detail::call_f_scalar_handler(f, tid+ndone,
+                detail::call_f_scalar_handler(f, tid+start_idx,
                     Gpu::Handler(amrex::min((std::uint64_t(nleft-tid)+(std::uint64_t)threadIdx.x),
                     (std::uint64_t)blockDim.x)));
             }
         });
-        if (Long(nleft) > ec.ntotalthreads) {
-            ndone += ec.ntotalthreads;
-        }
     }
     AMREX_GPU_ERROR_CHECK();
 }