Removed shmem to fix bug in HIP vvar_grad kernel

Author: ryanstocks00

src/xc_integrator/local_work_driver/device/hip/kernels/uvvars.hip

@@ -501,8 +501,6 @@ __global__ void eval_vvar_grad_kern( size_t        ntasks,
   double* den_y_eval_device = nullptr;
   double* den_z_eval_device = nullptr;
 
-  constexpr auto warp_size = hip::warp_size;
-
   if constexpr (den_select == DEN_S) {
     den_eval_device   = task.den_s;
     den_x_eval_device = task.dden_sx;
@@ -534,62 +532,61 @@ __global__ void eval_vvar_grad_kern( size_t        ntasks,
   const auto* dbasis_z_eval_device = task.dbfz;
 
   const auto* den_basis_prod_device = task.zmat;
-  
-  __shared__ double den_shared[4][warp_size][VVAR_KERNEL_SM_BLOCK+1];
 
-  for ( int bid_x = blockIdx.x * blockDim.x; 
-        bid_x < nbf;
-        bid_x += blockDim.x * gridDim.x ) {
-    
-    for ( int bid_y = blockIdx.y * VVAR_KERNEL_SM_BLOCK; 
-          bid_y < npts;
-          bid_y += VVAR_KERNEL_SM_BLOCK * gridDim.y ) {
-        
-      for (int sm_y = threadIdx.y; sm_y < VVAR_KERNEL_SM_BLOCK; sm_y += blockDim.y) {
-        den_shared[0][threadIdx.x][sm_y] = 0.;
-        den_shared[1][threadIdx.x][sm_y] = 0.;
-        den_shared[2][threadIdx.x][sm_y] = 0.;
-        den_shared[3][threadIdx.x][sm_y] = 0.;
+  // We always launch enough blocks to cover npts, so blocks aren't doing multiple results
+  double den_reg = 0.;
+  double dx_reg = 0.;
+  double dy_reg = 0.;
+  double dz_reg = 0.;
+
+  // Have each thread accumulate its own reduction result into a register.
+  // There's no real _need_ for LDS because the reductions are small and
+  // therefore can be done without sharing.
+  for( int ibf = 0; ibf < nbf; ibf++ ) {
+
+    for( int  ipt = blockIdx.x * blockDim.x + threadIdx.x; ipt < npts; ipt += blockDim.x * gridDim.x ) {
+
+      const double* bf_col   = basis_eval_device     + ibf*npts;
+      const double* bf_x_col = dbasis_x_eval_device  + ibf*npts;
+      const double* bf_y_col = dbasis_y_eval_device  + ibf*npts;
+      const double* bf_z_col = dbasis_z_eval_device  + ibf*npts;
+      const double* db_col   = den_basis_prod_device + ibf*npts;
+
+      den_reg += bf_col[ ipt ]   * db_col[ ipt ];
+      dx_reg += 2 * bf_x_col[ ipt ] * db_col[ ipt ];
+      dy_reg += 2 * bf_y_col[ ipt ] * db_col[ ipt ];
+      dz_reg += 2 * bf_z_col[ ipt ] * db_col[ ipt ];
+    }
+  }
 
-        if (bid_y + threadIdx.x < npts and bid_x + sm_y < nbf) { 
-          const double* db_col   = den_basis_prod_device + (bid_x + sm_y)*npts;
-          const double* bf_col   = basis_eval_device     + (bid_x + sm_y)*npts;
-          const double* bf_x_col = dbasis_x_eval_device  + (bid_x + sm_y)*npts;
-          const double* bf_y_col = dbasis_y_eval_device  + (bid_x + sm_y)*npts;
-          const double* bf_z_col = dbasis_z_eval_device  + (bid_x + sm_y)*npts;
 
-          den_shared[0][threadIdx.x][sm_y] = bf_col  [ bid_y + threadIdx.x ] * db_col[ bid_y + threadIdx.x ];
-          den_shared[1][threadIdx.x][sm_y] = bf_x_col[ bid_y + threadIdx.x ] * db_col[ bid_y + threadIdx.x ];
-          den_shared[2][threadIdx.x][sm_y] = bf_y_col[ bid_y + threadIdx.x ] * db_col[ bid_y + threadIdx.x ];
-          den_shared[3][threadIdx.x][sm_y] = bf_z_col[ bid_y + threadIdx.x ] * db_col[ bid_y + threadIdx.x ];
-        }
-      }
-      __syncthreads();
+  for( int  ipt = blockIdx.x * blockDim.x + threadIdx.x; ipt < npts; ipt += blockDim.x * gridDim.x ) {
+    den_eval_device   [ipt] = den_reg;
+    den_x_eval_device [ipt] = dx_reg ;
+    den_y_eval_device [ipt] = dy_reg ;
+    den_z_eval_device [ipt] = dz_reg ;
+  }
 
+}
 
-      for (int sm_y = threadIdx.y; sm_y < VVAR_KERNEL_SM_BLOCK; sm_y += blockDim.y) {
-        const int tid_y = bid_y + sm_y;
-         double den_reg = den_shared[0][sm_y][threadIdx.x];
-        double dx_reg  = den_shared[1][sm_y][threadIdx.x];
-        double dy_reg  = den_shared[2][sm_y][threadIdx.x];
-        double dz_reg  = den_shared[3][sm_y][threadIdx.x];
 
-        // Warp blocks are stored col major
-        den_reg =     hip::warp_reduce_sum<warp_size>( den_reg );
-        dx_reg  = 2. * hip::warp_reduce_sum<warp_size>( dx_reg );
-        dy_reg  = 2. * hip::warp_reduce_sum<warp_size>( dy_reg );
-        dz_reg  = 2. * hip::warp_reduce_sum<warp_size>( dz_reg );
+__global__ void eval_vvars_gga_kernel( 
+  size_t        npts,
+  const double* den_x_eval_device,
+  const double* den_y_eval_device,
+  const double* den_z_eval_device,
+        double* gamma_eval_device
+) {
+
+  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  if( tid < npts ) {
 
+    const double dx = den_x_eval_device[ tid ];
+    const double dy = den_y_eval_device[ tid ];
+    const double dz = den_z_eval_device[ tid ];
+
+    gamma_eval_device[tid] = dx*dx + dy*dy + dz*dz;
 
-        if( threadIdx.x == 0 and tid_y < npts ) {
-          atomicAdd( den_eval_device   + tid_y, den_reg );
-          atomicAdd( den_x_eval_device + tid_y, dx_reg  );
-          atomicAdd( den_y_eval_device + tid_y, dy_reg  );
-          atomicAdd( den_z_eval_device + tid_y, dz_reg  );
-        }
-      }
-      __syncthreads();
-    }
   }
 
 }
@@ -656,10 +653,9 @@ void eval_vvar( size_t ntasks, int32_t nbf_max, int32_t npts_max, bool do_grad,
   dim3 threads;
   dim3 blocks;
   if( do_grad ) {
-    threads = dim3( hip::warp_size, hip::max_warps_per_thread_block / 2, 1 );
-    blocks = dim3( std::min(uint64_t(4), util::div_ceil( nbf_max, 4 )),
-            std::min(uint64_t(16), util::div_ceil( nbf_max, 16 )),
-            ntasks );
+    threads = dim3(hip::max_warps_per_thread_block, 1, 1);
+    blocks = dim3( util::div_ceil( npts_max, threads.x),
+           1, 1);
   } else {
     threads = dim3( hip::warp_size, hip::max_warps_per_thread_block, 1 );
     blocks = dim3( util::div_ceil( nbf_max,  threads.x ),