File tree Expand file tree Collapse file tree 1 file changed +11
-6
lines changed Expand file tree Collapse file tree 1 file changed +11
-6
lines changed Original file line number Diff line number Diff line change @@ -195,14 +195,16 @@ __global__ void compute_descriptor_se_r (VALUETYPE* descript,
195195 const VALUETYPE* coord,
196196 const VALUETYPE rmin,
197197 const VALUETYPE rmax,
198- compute_t * sel_diff_dev)
198+ compute_t * sel_diff_dev,
199+ const int sec_size)
199200{
200201 // <<<nloc, sec.back()>>>
201- const unsigned int idx = blockIdx .x ;
202- const unsigned int idy = threadIdx .x ;
202+ const unsigned int idx = blockIdx .y ;
203+ const unsigned int idy = blockIdx . x * blockDim . x + threadIdx .x ;
203204 const int idx_deriv = idy * 3 ; // 1 components time 3 directions
204205 const int idx_value = idy; // 1 components
205-
206+ if (idy >= sec_size) {return ;}
207+
206208 // else {return;}
207209 VALUETYPE * row_descript = descript + idx * ndescrpt;
208210 VALUETYPE * row_descript_deriv = descript_deriv + idx * descript_deriv_size;
@@ -310,7 +312,9 @@ void DescrptSeRLauncher(const VALUETYPE* coord,
310312 nei_iter
311313 );
312314 }
313- compute_descriptor_se_r<<<nloc, sec.back()>>> (
315+ const int nblock_ = (sec.back () + LEN -1 ) / LEN;
316+ dim3 block_grid (nblock_, nloc);
317+ compute_descriptor_se_r<<<block_grid, LEN>>> (
314318 descript,
315319 ndescrpt,
316320 descript_deriv,
@@ -325,6 +329,7 @@ void DescrptSeRLauncher(const VALUETYPE* coord,
325329 coord,
326330 rcut_smth,
327331 rcut,
328- sel_diff
332+ sel_diff,
333+ sec.back ()
329334 );
330335}
You can’t perform that action at this time.
0 commit comments