| 402 | { |
| 403 | template<typename _Input, typename _Output, typename _Op, typename _Scalar, int BlockSize> |
| 404 | __global__ void ReduceBlockKernel(dim3 virtual_size, |
| 405 | _Input input, _Output output, _Op op, _Scalar initial, Index N) |
| 406 | { |
| 407 | const int part = threadIdx.x; |
| 408 | |
| 409 | typedef cub::BlockReduce<_Scalar, BlockSize> BlockReduceT; |
| 410 | __shared__ typename BlockReduceT::TempStorage temp_storage; |
| 411 | |
| 412 | for (Index i = blockIdx.x; i < virtual_size.x; ++i) { |
| 413 | const Index O = i * N; |
| 414 | //local reduce |
| 415 | _Scalar v = initial; |
| 416 | for (Index n = part; n < N; n += BlockSize) |
| 417 | v = op(v, input[O + n]); |
| 418 | //block reduce |
| 419 | v = BlockReduceT(temp_storage).Reduce(v, op); |
| 420 | if (part == 0) |
| 421 | output[i] = v; |
| 422 | } |
| 423 | } |
| 424 | } |
| 425 | template<typename _Input, typename _Output, int _Axis, typename _Op, typename _Scalar, int BlockSize> |
| 426 | struct ReductionEvaluator<_Input, _Output, _Axis, _Op, _Scalar, ReductionAlg::Block<BlockSize> > |