1
/*
2
Copyright 2020 The OneFlow Authors. All rights reserved.
3

4
Licensed under the Apache License, Version 2.0 (the "License");
5
you may not use this file except in compliance with the License.
6
You may obtain a copy of the License at
7

8
    http://www.apache.org/licenses/LICENSE-2.0
9

10
Unless required by applicable law or agreed to in writing, software
11
distributed under the License is distributed on an "AS IS" BASIS,
12
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
See the License for the specific language governing permissions and
14
limitations under the License.
15
*/
16

17

18

19
#ifndef ONEFLOW_CORE_CUDA_ELEMENTWISE_H_
20
#define ONEFLOW_CORE_CUDA_ELEMENTWISE_H_
21

22
#include <cuda_runtime.h>
23
#include <cstdint>
24
#include <algorithm>
25
#include <type_traits>
26

27
namespace oneflow {
28

29
namespace cuda {
30

31
namespace elementwise {
32

33
constexpr int kBlockSize = 256;
34
constexpr int kNumWaves = 32;
35

36
inline cudaError_t GetNumBlocks(int64_t n, int* num_blocks) {
37
  int dev;
38
  {
39
    cudaError_t err = cudaGetDevice(&dev);
40
    if (err != cudaSuccess) { return err; }
41
  }
42
  int sm_count;
43
  {
44
    cudaError_t err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
45
    if (err != cudaSuccess) { return err; }
46
  }
47
  int tpm;
48
  {
49
    cudaError_t err = cudaDeviceGetAttribute(&tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev);
50
    if (err != cudaSuccess) { return err; }
51
  }
52
  *num_blocks = std::max<int>(1, std::min<int64_t>((n + kBlockSize - 1) / kBlockSize,
53
                                                   sm_count * tpm / kBlockSize * kNumWaves));
54
  return cudaSuccess;
55
}
56

57
template<typename T, int pack_size>
58
struct GetPackType {
59
  using type = typename std::aligned_storage<pack_size * sizeof(T), pack_size * sizeof(T)>::type;
60
};
61

62
template<typename T, int pack_size>
63
using PackType = typename GetPackType<T, pack_size>::type;
64

65
template<typename T, int pack_size>
66
union Pack {
67
  static_assert(sizeof(PackType<T, pack_size>) == sizeof(T) * pack_size, "");
68
  __device__ Pack() {
69
    // do nothing
70
  }
71
  PackType<T, pack_size> storage;
72
  T elem[pack_size];
73
};
74

75
template<typename T, int pack_size>
76
struct alignas(sizeof(T) * pack_size) Packed {
77
  __device__ Packed() {
78
    // do nothing
79
  }
80
  union {
81
    T elem[pack_size];
82
  };
83
};
84

85
constexpr int kMaxPackBytes = 128 / 8;
86
constexpr int kMaxPackSize = 8;
87

88
constexpr int Min(int a, int b) { return a < b ? a : b; }
89

90
template<typename T>
91
constexpr int PackSize() {
92
  return Min(kMaxPackBytes / sizeof(T), kMaxPackSize);
93
}
94

95
template<typename T, typename U, typename... Args>
96
constexpr int PackSize() {
97
  return Min(PackSize<T>(), PackSize<U, Args...>());
98
}
99

100
template<typename T>
101
class HasApply2 {
102
  typedef char one;
103
  struct two {
104
    char x[2];
105
  };
106

107
  template<typename C>
108
  static one test(decltype(&C::Apply2));
109
  template<typename C>
110
  static two test(...);
111

112
 public:
113
  enum { value = sizeof(test<T>(0)) == sizeof(char) };
114
};
115

116
template<int pack_size, typename FunctorT, typename R, typename... IN>
117
__device__ typename std::enable_if<HasApply2<FunctorT>::value == true && pack_size % 2 == 0,
118
                                   Packed<R, pack_size>>::type
119
ApplyPack(const FunctorT& functor, const Packed<IN, pack_size>... in) {
120
  Packed<R, pack_size> ret;
121
#pragma unroll
122
  for (int j = 0; j < pack_size; j += 2) { functor.Apply2(ret.elem + j, (in.elem + j)...); }
123
  return ret;
124
}
125

126
template<int pack_size, typename FunctorT, typename R, typename... IN>
127
__device__ typename std::enable_if<HasApply2<FunctorT>::value == false || pack_size % 2 != 0,
128
                                   Packed<R, pack_size>>::type
129
ApplyPack(const FunctorT& functor, const Packed<IN, pack_size>... in) {
130
  Packed<R, pack_size> ret;
131
#pragma unroll
132
  for (int j = 0; j < pack_size; ++j) { ret.elem[j] = functor((in.elem[j])...); }
133
  return ret;
134
}
135

136
template<int pack_size, typename FactoryT, typename R, typename... IN>
137
__global__ void __launch_bounds__(kBlockSize)
138
    ApplyGeneric(FactoryT factory, int64_t n_pack, Packed<R, pack_size>* pack_r,
139
                 const Packed<IN, pack_size>*... pack_in, int64_t n_tail, R* tail_r,
140
                 const IN*... tail_in) {
141
  auto functor = factory();
142
  const int global_tid = blockIdx.x * kBlockSize + threadIdx.x;
143
  for (int64_t i = global_tid; i < n_pack; i += blockDim.x * gridDim.x) {
144
    pack_r[i] = ApplyPack<pack_size, decltype(functor), R, IN...>(functor, (pack_in[i])...);
145
  }
146
  if (global_tid < n_tail) { tail_r[global_tid] = functor((tail_in[global_tid])...); }
147
}
148

149
template<typename FunctorT>
150
struct SimpleFactory {
151
  explicit SimpleFactory(FunctorT functor) : tpl(functor) {}
152
  __device__ FunctorT operator()() const { return tpl; }
153

154
 private:
155
  FunctorT tpl;
156
};
157

158
template<size_t pack_size>
159
bool IsAlignedForPack() {
160
  return true;
161
}
162

163
template<size_t pack_size, typename T, typename... Args>
164
bool IsAlignedForPack(const T* ptr, const Args*... others) {
165
  return reinterpret_cast<uintptr_t>(ptr) % sizeof(Pack<T, pack_size>) == 0
166
         && IsAlignedForPack<pack_size, Args...>(others...);
167
}
168

169
template<size_t pack_size, typename FactoryT, typename R, typename... IN>
170
cudaError_t LaunchKernel(FactoryT factory, int64_t n, R* r, const IN*... in, cudaStream_t stream) {
171
  const int64_t n_pack = n / pack_size;
172
  const int64_t tail_offset = n_pack * pack_size;
173
  const int64_t n_tail = n - tail_offset;
174
  int num_blocks;
175
  {
176
    cudaError_t err = GetNumBlocks(n_pack, &num_blocks);
177
    if (err != cudaSuccess) { return err; }
178
  }
179
  ApplyGeneric<pack_size, FactoryT, R, IN...><<<num_blocks, kBlockSize, 0, stream>>>(
180
      factory, n_pack, reinterpret_cast<Packed<R, pack_size>*>(r),
181
      (reinterpret_cast<const Packed<IN, pack_size>*>(in))..., n_tail, r + tail_offset,
182
      (in + tail_offset)...);
183
  return cudaPeekAtLastError();
184
}
185

186
template<typename FactoryT, typename R, typename... IN>
187
struct GenericLauncher {
188
  static cudaError_t Launch(FactoryT factory, int64_t n, R* r, const IN*... in,
189
                            cudaStream_t stream) {
190
    constexpr int max_pack_size = PackSize<R, IN...>();
191
    if (IsAlignedForPack<max_pack_size, R, IN...>(r, in...)) {
192
      return LaunchKernel<max_pack_size, FactoryT, R, IN...>(factory, n, r, in..., stream);
193
    } else {
194
      return LaunchKernel<1, FactoryT, R, IN...>(factory, n, r, in..., stream);
195
    }
196
  }
197
};
198

199
template<typename FactoryT, typename R, typename A>
200
inline cudaError_t UnaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a,
201
                                    cudaStream_t stream) {
202
  return GenericLauncher<FactoryT, R, A>::Launch(factory, n, r, a, stream);
203
}
204

205
template<typename FunctorT, typename R, typename A>
206
inline cudaError_t Unary(FunctorT functor, int64_t n, R* r, const A* a, cudaStream_t stream) {
207
  return UnaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, stream);
208
}
209

210
template<typename FactoryT, typename R, typename A, typename B>
211
inline cudaError_t BinaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a, const B* b,
212
                                     cudaStream_t stream) {
213
  return GenericLauncher<FactoryT, R, A, B>::Launch(factory, n, r, a, b, stream);
214
}
215

216
template<typename FunctorT, typename R, typename A, typename B>
217
inline cudaError_t Binary(FunctorT functor, int64_t n, R* r, const A* a, const B* b,
218
                          cudaStream_t stream) {
219
  return BinaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, b, stream);
220
}
221

222
template<typename FactoryT, typename R, typename A, typename B, typename C>
223
inline cudaError_t TernaryWithFactory(FactoryT factory, int64_t n, R* r, const A* a, const B* b,
224
                                      const C* c, cudaStream_t stream) {
225
  return GenericLauncher<FactoryT, R, A, B, C>::Launch(factory, n, r, a, b, c, stream);
226
}
227

228
template<typename FunctorT, typename R, typename A, typename B, typename C>
229
inline cudaError_t Ternary(FunctorT functor, int64_t n, R* r, const A* a, const B* b, const C* c,
230
                           cudaStream_t stream) {
231
  return TernaryWithFactory(SimpleFactory<FunctorT>(functor), n, r, a, b, c, stream);
232
}
233

234
}  // namespace elementwise
235

236
}  // namespace cuda
237

238
}  // namespace oneflow
239

240
#endif  // ONEFLOW_CORE_CUDA_ELEMENTWISE_H_