| 483 | } |
| 484 | |
| 485 | Status DoNMS(OpKernelContext* context, const Tensor& boxes, |
| 486 | const Tensor& scores, const int64_t max_output_size, |
| 487 | const float iou_threshold_val, const float score_threshold) { |
| 488 | const int output_size = max_output_size; |
| 489 | int num_boxes = boxes.dim_size(0); |
| 490 | size_t cub_sort_temp_storage_bytes = 0; |
| 491 | auto cuda_stream = GetGpuStream(context); |
| 492 | auto device = context->eigen_gpu_device(); |
| 493 | // Calling cub with nullptrs as inputs will make it return |
| 494 | // workspace size needed for the operation instead of doing the operation. |
| 495 | // In this specific instance, cub_sort_temp_storage_bytes will contain the |
| 496 | // necessary workspace size for sorting after the call. |
| 497 | if (num_boxes == 0) { |
| 498 | Tensor* output_indices = nullptr; |
| 499 | TF_RETURN_IF_ERROR( |
| 500 | context->allocate_output(0, TensorShape({0}), &output_indices)); |
| 501 | return Status::OK(); |
| 502 | } |
| 503 | |
| 504 | cudaError_t cuda_ret = cub::DeviceRadixSort::SortPairsDescending( |
| 505 | nullptr, cub_sort_temp_storage_bytes, |
| 506 | static_cast<float*>(nullptr), // scores |
| 507 | static_cast<float*>(nullptr), // sorted scores |
| 508 | static_cast<int*>(nullptr), // input indices |
| 509 | static_cast<int*>(nullptr), // sorted indices |
| 510 | num_boxes, // num items |
| 511 | 0, 8 * sizeof(float), // sort all bits |
| 512 | cuda_stream); |
| 513 | TF_RETURN_IF_CUDA_ERROR(cuda_ret); |
| 514 | Tensor d_cub_sort_buffer; |
| 515 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 516 | DataType::DT_INT8, TensorShape({(int64)cub_sort_temp_storage_bytes}), |
| 517 | &d_cub_sort_buffer)); |
| 518 | Tensor d_indices; |
| 519 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 520 | DataType::DT_INT32, TensorShape({num_boxes}), &d_indices)); |
| 521 | Tensor d_sorted_indices; |
| 522 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 523 | DataType::DT_INT32, TensorShape({num_boxes}), &d_sorted_indices)); |
| 524 | Tensor d_selected_indices; |
| 525 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 526 | DataType::DT_INT32, TensorShape({num_boxes}), &d_selected_indices)); |
| 527 | Tensor d_sorted_scores; |
| 528 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 529 | DataType::DT_FLOAT, TensorShape({num_boxes}), &d_sorted_scores)); |
| 530 | Tensor d_sorted_boxes; |
| 531 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 532 | DataType::DT_FLOAT, TensorShape({num_boxes, 4}), &d_sorted_boxes)); |
| 533 | |
| 534 | // this will return sorted scores and their indices |
| 535 | auto config = GetGpuLaunchConfig(num_boxes, device); |
| 536 | // initialize box and score indices |
| 537 | TF_CHECK_OK(GpuLaunchKernel(Iota<int>, config.block_count, |
| 538 | config.thread_per_block, 0, device.stream(), |
| 539 | config.virtual_thread_count, 0, |
| 540 | d_indices.flat<int>().data(), 1)); |
| 541 | TF_RETURN_IF_CUDA_ERROR(cudaGetLastError()); |
| 542 | cuda_ret = cub::DeviceRadixSort::SortPairsDescending( |
no test coverage detected