| 297 | } |
| 298 | |
| 299 | Status NmsGpu(const float* d_sorted_boxes_float_ptr, const int num_boxes, |
| 300 | const float iou_threshold, int* d_selected_indices, int* h_nkeep, |
| 301 | OpKernelContext* context, const int max_boxes, bool flip_boxes) { |
| 302 | // Making sure we respect the __align(16)__ |
| 303 | // we promised to the compiler. |
| 304 | auto iptr = reinterpret_cast<std::uintptr_t>(d_sorted_boxes_float_ptr); |
| 305 | if ((iptr & 15) != 0) { |
| 306 | return errors::InvalidArgument("Boxes should be aligned to 16 Bytes."); |
| 307 | } |
| 308 | // allocate bitmask arrays on host and on device |
| 309 | Tensor h_num_selected, d_nms_mask; |
| 310 | const int bit_mask_len = |
| 311 | (num_boxes + kNmsBoxesPerThread - 1) / kNmsBoxesPerThread; |
| 312 | |
| 313 | int64 max_nms_mask_size = num_boxes * bit_mask_len; |
| 314 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 315 | DataType::DT_INT32, TensorShape({max_nms_mask_size}), &d_nms_mask)); |
| 316 | // reset data sensitive tensors |
| 317 | auto device = context->eigen_gpu_device(); |
| 318 | auto config = GetGpuLaunchConfig(d_nms_mask.NumElements(), device); |
| 319 | TF_CHECK_OK(GpuLaunchKernel(SetZero<int>, config.block_count, |
| 320 | config.thread_per_block, 0, device.stream(), |
| 321 | config.virtual_thread_count, |
| 322 | d_nms_mask.flat<int32>().data())); |
| 323 | |
| 324 | AllocatorAttributes alloc_attr; |
| 325 | alloc_attr.set_on_host(true); |
| 326 | alloc_attr.set_gpu_compatible(true); |
| 327 | // Size of this buffer can be reduced to kNmsChunkSize*bit_mask_len*2 and |
| 328 | // using it as a ring buffer. However savings should be a few MB . |
| 329 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 330 | DataType::DT_INT32, TensorShape({1}), &h_num_selected, alloc_attr)); |
| 331 | |
| 332 | int* d_delete_mask = d_nms_mask.flat<int>().data(); |
| 333 | int* h_selected_count = h_num_selected.flat<int>().data(); |
| 334 | const Box* d_sorted_boxes = |
| 335 | reinterpret_cast<const Box*>(d_sorted_boxes_float_ptr); |
| 336 | dim3 block_dim, thread_block; |
| 337 | int num_blocks = (num_boxes + kNmsBlockDim - 1) / kNmsBlockDim; |
| 338 | num_blocks = std::max(std::min(num_blocks, kNmsBlockDimMax), 1); |
| 339 | block_dim.x = num_blocks; |
| 340 | block_dim.y = num_blocks; |
| 341 | block_dim.z = 1; |
| 342 | thread_block.x = kNmsBlockDim; |
| 343 | thread_block.y = kNmsBlockDim; |
| 344 | thread_block.z = 1; |
| 345 | if (flip_boxes) { |
| 346 | TF_CHECK_OK(GpuLaunchKernel(NMSKernel<true>, block_dim, thread_block, 0, |
| 347 | device.stream(), d_sorted_boxes, num_boxes, |
| 348 | iou_threshold, bit_mask_len, d_delete_mask)); |
| 349 | } else { |
| 350 | TF_CHECK_OK(GpuLaunchKernel(NMSKernel<false>, block_dim, thread_block, 0, |
| 351 | device.stream(), d_sorted_boxes, num_boxes, |
| 352 | iou_threshold, bit_mask_len, d_delete_mask)); |
| 353 | } |
| 354 | TF_RETURN_IF_CUDA_ERROR(cudaGetLastError()); |
| 355 | // Overlapping CPU computes and D2H memcpy |
| 356 | // both take about the same time |
no test coverage detected