| 220 | template <typename InputIteratorT, typename OutputIteratorT, |
| 221 | typename OffsetIteratorT, typename ReduceOp, typename T> |
| 222 | Status GpuSegmentedReduce( |
| 223 | OpKernelContext* context, int num_segments, ReduceOp reduce_op, |
| 224 | const T& initial_value, |
| 225 | InputIteratorT input, // [any] |
| 226 | OffsetIteratorT segment_offsets, // [num_segments + 1] |
| 227 | OutputIteratorT output) { // [num_segments] |
| 228 | if (num_segments == 0) return Status::OK(); |
| 229 | const auto& cu_stream = GetGpuStream(context); |
| 230 | size_t temp_storage_bytes; |
| 231 | auto err = gpuprim::DeviceSegmentedReduce::Reduce( |
| 232 | nullptr, temp_storage_bytes, input, output, num_segments, segment_offsets, |
| 233 | segment_offsets + 1, reduce_op, initial_value, cu_stream); |
| 234 | if (err != 0) { |
| 235 | return errors::Internal( |
| 236 | "Failed to launch gpuprim::DeviceSegmentedReduce::Reduce to calculate " |
| 237 | "temp_storage_bytes, status: ", |
| 238 | cudaGetErrorString(err)); |
| 239 | } |
| 240 | Tensor temp_storage; |
| 241 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 242 | DT_INT8, TensorShape({static_cast<int64>(temp_storage_bytes)}), |
| 243 | &temp_storage)); |
| 244 | err = gpuprim::DeviceSegmentedReduce::Reduce( |
| 245 | temp_storage.flat<int8>().data(), temp_storage_bytes, input, output, |
| 246 | num_segments, segment_offsets, segment_offsets + 1, reduce_op, |
| 247 | initial_value, cu_stream); |
| 248 | if (err != 0) { |
| 249 | return errors::Internal( |
| 250 | "Failed to launch gpuprim::DeviceSegmentedReduce::Reduce" |
| 251 | ", temp_storage_bytes: ", |
| 252 | temp_storage_bytes, ", status: ", cudaGetErrorString(err)); |
| 253 | } |
| 254 | return Status::OK(); |
| 255 | } |
| 256 | |
| 257 | template <typename InputIteratorT, typename FlagIteratorT, |
| 258 | typename OutputIteratorT, typename NumSelectedT = int> |