| 213 | |
| 214 | template<bool nonmax> |
| 215 | __global__ void non_max_counts(unsigned *d_counts, unsigned *d_offsets, |
| 216 | unsigned *d_total, float *flags, |
| 217 | const float *score, const unsigned idim0, |
| 218 | const unsigned idim1, const unsigned edge) { |
| 219 | const int xid = blockIdx.x * blockDim.x * 2 + threadIdx.x; |
| 220 | const int yid = blockIdx.y * blockDim.y * 8 + threadIdx.y; |
| 221 | const int tid = blockDim.x * threadIdx.y + threadIdx.x; |
| 222 | |
| 223 | const int xoff = blockDim.x; |
| 224 | const int yoff = blockDim.y; |
| 225 | |
| 226 | const int xend = (blockIdx.x + 1) * blockDim.x * 2; |
| 227 | const int yend = (blockIdx.y + 1) * blockDim.y * 8; |
| 228 | |
| 229 | const int bid = blockIdx.y * gridDim.x + blockIdx.x; |
| 230 | using BlockReduce = |
| 231 | cub::BlockReduce<unsigned, 32, cub::BLOCK_REDUCE_WARP_REDUCTIONS, 8>; |
| 232 | |
| 233 | __shared__ typename BlockReduce::TempStorage temp_storage; |
| 234 | |
| 235 | unsigned count = 0; |
| 236 | for (int y = yid; y < yend; y += yoff) { |
| 237 | if (y >= idim1 - edge - 1 || y <= edge + 1) continue; |
| 238 | for (int x = xid; x < xend; x += xoff) { |
| 239 | if (x >= idim0 - edge - 1 || x <= edge + 1) continue; |
| 240 | |
| 241 | float v = score[y * idim0 + x]; |
| 242 | if (v == 0) { |
| 243 | if (nonmax) flags[y * idim0 + x] = 0; |
| 244 | continue; |
| 245 | } |
| 246 | |
| 247 | if (nonmax) { |
| 248 | float max_v = v; |
| 249 | max_v = max_val(score[x - 1 + idim0 * (y - 1)], |
| 250 | score[x - 1 + idim0 * y]); |
| 251 | max_v = max_val(max_v, score[x - 1 + idim0 * (y + 1)]); |
| 252 | max_v = max_val(max_v, score[x + idim0 * (y - 1)]); |
| 253 | max_v = max_val(max_v, score[x + idim0 * (y + 1)]); |
| 254 | max_v = max_val(max_v, score[x + 1 + idim0 * (y - 1)]); |
| 255 | max_v = max_val(max_v, score[x + 1 + idim0 * (y)]); |
| 256 | max_v = max_val(max_v, score[x + 1 + idim0 * (y + 1)]); |
| 257 | |
| 258 | v = (v > max_v) ? v : 0; |
| 259 | flags[y * idim0 + x] = v; |
| 260 | if (v == 0) continue; |
| 261 | } |
| 262 | |
| 263 | count++; |
| 264 | } |
| 265 | } |
| 266 | |
| 267 | int sum = BlockReduce(temp_storage).Sum(count); |
| 268 | |
| 269 | if (tid == 0) { |
| 270 | unsigned total = sum ? atomicAdd(d_total, sum) : 0; |
| 271 | d_counts[bid] = sum; |
| 272 | d_offsets[bid] = total; |