| 354 | { |
| 355 | template<typename _Input, typename _Output, typename _Op, typename _Scalar> |
| 356 | __global__ void ReduceWarpKernel(dim3 virtual_size, |
| 357 | _Input input, _Output output, _Op op, _Scalar initial, Index N) |
| 358 | { |
| 359 | CUMAT_KERNEL_1D_LOOP(i_, virtual_size) |
| 360 | const Index i = i_ / 32;//i_ >> 5; |
| 361 | const int warp = i_ % 32; // i_ & 32; |
| 362 | const Index O = i * N; |
| 363 | //local reduce |
| 364 | _Scalar v = initial; |
| 365 | for (Index n = warp; n < N; n += 32) |
| 366 | v = op(v, input[O + n]); |
| 367 | //final warp reduce |
| 368 | #pragma unroll |
| 369 | for (int offset = 16; offset > 0; offset /= 2) |
| 370 | v = op(v, __shfl_down_sync(0xffffffff, v, offset)); |
| 371 | //write output |
| 372 | if (warp == 0) output[i] = v; |
| 373 | CUMAT_KERNEL_1D_LOOP_END |
| 374 | } |
| 375 | } |
| 376 | template<typename _Input, typename _Output, int _Axis, typename _Op, typename _Scalar> |
| 377 | struct ReductionEvaluator<_Input, _Output, _Axis, _Op, _Scalar, ReductionAlg::Warp> |