Skip to content

Commit fb79cbe

Browse files
LuLu
authored andcommitted
fix bug of 0 output when sec_a.back() is lager than 1024
1 parent b7940fc commit fb79cbe

File tree

1 file changed

+11
-6
lines changed

1 file changed

+11
-6
lines changed

source/op/cuda/descrpt_se_r.cu

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff 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
}

0 commit comments

Comments
 (0)