Function CountIf

tensorflow/core/kernels/non_max_suppression_op.cu.cc:423–453

// TODO(sami) Not really a good way. Perhaps consider using thrust?
template <typename Op>
Status CountIf(OpKernelContext* context, const float* dev_array, const Op& op,
               int num_elements, int* result) {
  Tensor scratch_output;
  Tensor workspace;
  Tensor element_count;
  size_t workspace_size = 0;
  auto cuda_stream = tensorflow::GetGpuStream(context);
  auto device = context->eigen_gpu_device();
  cub::DeviceSelect::If(nullptr, workspace_size, static_cast<float*>(nullptr),
                        static_cast<float*>(nullptr),
                        static_cast<int*>(nullptr), num_elements, op);

  TF_RETURN_IF_ERROR(context->allocate_temp(
      DataType::DT_FLOAT, TensorShape({num_elements}), &scratch_output));
  TF_RETURN_IF_ERROR(context->allocate_temp(
      DataType::DT_INT8, TensorShape({(int64)workspace_size}), &workspace));
  TF_RETURN_IF_ERROR(context->allocate_temp(DataType::DT_INT32,
                                            TensorShape({1}), &element_count));
  cudaEvent_t copy_done;
  TF_RETURN_IF_CUDA_ERROR(
      cudaEventCreateWithFlags(&copy_done, cudaEventDisableTiming));
  TF_RETURN_IF_CUDA_ERROR(cub::DeviceSelect::If(
      workspace.flat<int8>().data(), workspace_size, dev_array,
      scratch_output.flat<float>().data(), element_count.flat<int32>().data(),
      num_elements, op, cuda_stream));
  device.memcpyDeviceToHost(result, element_count.flat<int32>().data(),
                            sizeof(int));
  TF_RETURN_IF_CUDA_ERROR(cudaEventRecord(copy_done, device.stream()));
  TF_RETURN_IF_CUDA_ERROR(cudaEventSynchronize(copy_done));
  return Status::OK();
}

template <typename Op>
Status CountIf(OpKernelContext* context, const float* dev_array, const Op& op,