| 1 | #pragma once |
| 2 | |
| 3 | #ifndef USE_ROCM |
| 4 | #include <cub/cub.cuh> |
| 5 | #if CUB_VERSION >= 200800 |
| 6 | #include <cuda/std/functional> |
| 7 | using CubAddOp = cuda::std::plus<>; |
| 8 | using CubMaxOp = cuda::maximum<>; |
| 9 | #else // if CUB_VERSION < 200800 |
| 10 | using CubAddOp = cub::Sum; |
| 11 | using CubMaxOp = cub::Max; |
| 12 | #endif // CUB_VERSION |
| 13 | #else |
| 14 | #include <hipcub/hipcub.hpp> |
| 15 | namespace cub = hipcub; |
| 16 | using CubAddOp = hipcub::Sum; |
| 17 | using CubMaxOp = hipcub::Max; |
| 18 | #endif // USE_ROCM |