1#pragma once

2template<typename T>

3struct MaxOp {

4__device__ inline T operator()(T const & x, T const & y) { return x > y ? x : y; }

5};

6

7template <>

8struct MaxOp<float> {

9// This is slightly faster

10__device__ inline float operator()(float const &x, float const &y) { return max(x, y); }

11};

12

13////////////////////////////////////////////////////////////////////////////////////////////////////

14

15template<typename T>

16struct SumOp {

17__device__ inline T operator()(T const & x, T const & y) { return x + y; }

18};

19

20////////////////////////////////////////////////////////////////////////////////////////////////////

21

22template<int THREADS>

23struct Allreduce {

24 static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);

25 template<typename T, typename Operator>

26 static __device__ inline T run(T x, Operator &op) {

27 constexpr int OFFSET = THREADS / 2;

28 x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));

29 return Allreduce<OFFSET>::run(x, op);

30 }

31};

32

33////////////////////////////////////////////////////////////////////////////////////////////////////

34

35template<>

36struct Allreduce<2> {

37template<typename T, typename Operator>

38static __device__ inline T run(T x, Operator &op) {

39 x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));

40 return x;

41}

42};

43

44template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>

45__device__ inline void thread_reduce_(Tensor<Engine0, Layout0> const &tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {

46 static_assert(Layout0::rank == 2, "Only support 2D Tensor");

47 static_assert(Layout1::rank == 1, "Only support 1D Tensor");

48 CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));

49 #pragma unroll

50 for (int mi = 0; mi < size<0>(tensor); mi++) {

51 summary(mi) = zero_init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));

52 #pragma unroll

53 for (int ni = 1; ni < size<1>(tensor); ni++) {

54 summary(mi) = op(summary(mi), tensor(mi, ni));

55 }

56 }

57}

58

59template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>

60__device__ inline void quad_allreduce_(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &src, Operator &op) {

61 CUTE_STATIC_ASSERT_V(size(dst) == size(src));

62 #pragma unroll

63 for (int i = 0; i < size(dst); i++){

64 // NOTE: 4表示4个线程, 因为在SM80_16x8x16_F32F16F16F32_TN中,

65 // 每组每行就是4个线程处理8个value的, 每个线程处理2个value

66 dst(i) = Allreduce<4>::run(src(i), op);

67 }

68}

69

70template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>

71__device__ inline void reduce_(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {

72 // NOTE: 遍历tensor每行, 记录到summary中

73 // reduce 当前thread的max

74 thread_reduce_<zero_init>(tensor, summary, op);

75 // NOTE: 二分法对summary[]进行reduce

76 // reduce thread间的max

77 quad_allreduce_(summary, summary, op);

78}

79

80

81template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>

82__device__ inline void reduce_max(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &max){

83 MaxOp<float> max_op;

84 reduce_<zero_init>(tensor, max, max_op);

85}

86

87template<typename Engine0, typename Layout0, typename Engine1, typename Layout1>

88__device__ inline void reduce_sum(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &sum){

89 SumOp<float> sum_op;

90 reduce_(tensor, sum, sum_op);

91}

92

93