reduce.h 1.99 KB
Newer Older
1
2
3
4
5
6
7
8
9
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
#pragma once

#include "common.h"

namespace tl {

struct SumOp {
10
  template <typename T> TL_DEVICE T operator()(T const &x, T const &y) {
11
12
13
14
15
    return x + y;
  }
};

struct MaxOp {
16
  template <typename T> TL_DEVICE T operator()(T const &x, T const &y) {
17
18
19
20
21
    return cutlass::fast_max(x, y);
  }
};

struct MinOp {
22
  template <typename T> TL_DEVICE T operator()(T const &x, T const &y) {
23
24
25
26
    return cutlass::fast_min(x, y);
  }
};

27
28
template <class Reducer, int threads, int scale, int all_threads = threads>
struct AllReduce {
29
30
31
  static_assert(threads == 1024 or threads == 512 or threads == 256 or
                threads == 128 or threads == 64 or threads == 32 or
                threads == 16 or threads == 8 or threads == 4 or threads == 2);
32
  static_assert(threads % scale == 0);
33
  template <typename T> static TL_DEVICE T run(T x, T *red_buf = nullptr) {
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    constexpr int offset = threads / 2;
    if constexpr (offset >= 32) {
      __syncthreads();
      red_buf[threadIdx.x] = x;
      __syncthreads();
      x = Reducer()(x, red_buf[threadIdx.x ^ offset]);
    } else {
      x = Reducer()(x, T(__shfl_xor_sync(uint32_t(-1), x, offset)));
    }
    if constexpr (offset == scale) {
      return x;
    } else {
      return AllReduce<Reducer, offset, scale>::run(x, red_buf);
    }
  }
49
50
51
52
53

  template <typename T>
  static TL_DEVICE T run_hopper(T x, T *red_buf = nullptr) {
    constexpr int offset = threads / 2;
    if constexpr (offset >= 32) {
54
      asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(all_threads));
55
      red_buf[threadIdx.x] = x;
56
      asm volatile("bar.sync %0, %1;" : : "r"(2), "r"(all_threads));
57
58
59
60
61
62
63
      x = Reducer()(x, red_buf[threadIdx.x ^ offset]);
    } else {
      x = Reducer()(x, T(__shfl_xor_sync(uint32_t(-1), x, offset)));
    }
    if constexpr (offset == scale) {
      return x;
    } else {
64
65
      return AllReduce<Reducer, offset, scale, all_threads>::run_hopper(
          x, red_buf);
66
67
    }
  }
68
69
};

70
} // namespace tl