Skip to content

Commit 80ed55d

Browse files
ZzSeanZzSean
committed
fix
1 parent aa181f0 commit 80ed55d

3 files changed

Lines changed: 75 additions & 82 deletions

File tree

‎paddle/fluid/operators/elementwise/elementwise_add_op.cu‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,16 +24,29 @@ namespace plat = paddle::platform;
2424
namespace paddle {
2525
namespace operators {
2626

27+
/*
28+
input: an array;
29+
return: the result of the math functor
30+
1. For Unary Op, the length of input array is 1,
31+
e.g. Relu: return args[0] > 0 ? args[0] : 0;
32+
2. For Binary Op, the length of input array is 2,
33+
e.g. Add: return args[0] + args[1];
34+
*/
35+
template <typename T>
36+
struct CudaAddFunctor {
37+
inline HOSTDEVICE T operator()(T args[]) const { return args[0] + args[1]; }
38+
};
39+
2740
template <typename T>
2841
struct SameDimsElemwiseAdd<platform::CUDADeviceContext, T> {
2942
void operator()(const framework::ExecutionContext& ctx,
3043
const framework::Tensor* x, const framework::Tensor* y,
3144
framework::Tensor* z) {
32-
auto size = x->numel();
33-
std::vector<const T*> ins = {x->data<T>(), y->data<T>()};
34-
std::vector<T*> outs = {z->data<T>()};
35-
LaunchElementwiseCudaKernel<ElementwiseType::kBinary>(ctx, ins, outs, size,
36-
CudaAddFunctor<T>());
45+
std::vector<const framework::Tensor*> ins = {x, y};
46+
std::vector<framework::Tensor*> outs = {z};
47+
LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T>(
48+
ctx.template device_context<platform::CUDADeviceContext>(), ins, &outs,
49+
CudaAddFunctor<T>());
3750
}
3851
};
3952

‎paddle/fluid/operators/elementwise/elementwise_add_op.h‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ limitations under the License. */
1818
#include "paddle/fluid/operators/elementwise/elementwise_op.h"
1919
#include "paddle/fluid/operators/elementwise/elementwise_op_function.cu.h"
2020
#include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
21-
#include "paddle/fluid/operators/elementwise/elementwise_op_impl.cu.h"
2221
#include "paddle/fluid/operators/math/blas.h"
2322
#include "paddle/fluid/operators/math/math_function.h"
2423
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
@@ -38,11 +37,6 @@ namespace cub = hipcub;
3837
namespace paddle {
3938
namespace operators {
4039

41-
template <typename T, class Enable = void>
42-
struct CudaAddFunctor {
43-
inline HOSTDEVICE T operator()(T args[]) const { return args[0] + args[1]; }
44-
};
45-
4640
template <typename DeviceContext, typename T>
4741
void default_elementwise_add(const framework::ExecutionContext &ctx,
4842
const framework::Tensor *x,
@@ -141,6 +135,11 @@ elementwise_add_grad(const framework::ExecutionContext &ctx,
141135
#ifdef PADDLE_WITH_CUDA
142136
#ifdef __NVCC__
143137

138+
template <typename T, int Size>
139+
struct alignas(sizeof(T) * Size) AlignedVector {
140+
T val[Size];
141+
};
142+
144143
template <typename T>
145144
inline int VectorizedSize(const T *pointer) {
146145
uint64_t address = reinterpret_cast<uint64_t>(pointer);

‎paddle/fluid/operators/elementwise/elementwise_op_impl.cu.h‎

Lines changed: 52 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -13,41 +13,23 @@ See the License for the specific language governing permissions and
1313
limitations under the License. */
1414
#pragma once
1515

16-
#include "paddle/fluid/operators/elementwise/elementwise_op.h"
17-
#include "paddle/fluid/operators/elementwise/elementwise_op_function.cu.h"
18-
#include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
19-
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
20-
#ifdef __NVCC__
21-
#include <cuda.h>
22-
#include <cuda_fp16.h>
23-
#include "cub/cub.cuh"
24-
#endif
25-
#ifdef __HIPCC__
26-
#include <hip/hip_fp16.h>
27-
#include <hip/hip_runtime.h>
28-
#include <hipcub/hipcub.hpp>
29-
namespace cub = hipcub;
30-
#endif
31-
#endif
32-
3316
namespace paddle {
3417
namespace operators {
3518

36-
#ifdef PADDLE_WITH_CUDA
37-
#ifdef __NVCC__
38-
3919
enum ElementwiseType { kUnary = 1, kBinary = 2 };
4020

4121
template <typename T, int Size>
42-
struct alignas(sizeof(T) * Size) AlignedVector {
22+
struct alignas(sizeof(T) * Size) CudaAlignedVector {
4323
T val[Size];
4424
};
4525

4626
template <typename T>
4727
int GetVectorizedSizeImpl(const T *pointer) {
4828
uint64_t address = reinterpret_cast<uint64_t>(pointer);
49-
constexpr int vec4 = std::alignment_of<AlignedVector<T, 4>>::value; // NOLINT
50-
constexpr int vec2 = std::alignment_of<AlignedVector<T, 2>>::value; // NOLINT
29+
constexpr int vec4 =
30+
std::alignment_of<CudaAlignedVector<T, 4>>::value; // NOLINT
31+
constexpr int vec2 =
32+
std::alignment_of<CudaAlignedVector<T, 2>>::value; // NOLINT
5133
if (address % vec4 == 0) {
5234
return 4;
5335
} else if (address % vec2 == 0) {
@@ -57,19 +39,21 @@ int GetVectorizedSizeImpl(const T *pointer) {
5739
}
5840

5941
template <typename T>
60-
int GetVectorizedSize(const std::vector<const T *> ins,
61-
const std::vector<T *> outs) {
42+
int GetVectorizedSize(const std::vector<const framework::Tensor *> &ins,
43+
const std::vector<framework::Tensor *> &outs) {
6244
int vec_size = 4;
6345
for (auto iter = ins.begin(); iter != ins.end(); ++iter) {
64-
vec_size = std::min<int>(vec_size, GetVectorizedSizeImpl(*iter));
46+
vec_size =
47+
std::min<int>(vec_size, GetVectorizedSizeImpl((*iter)->data<T>()));
6548
}
6649
for (auto iter = outs.begin(); iter != outs.end(); ++iter) {
67-
vec_size = std::min<int>(vec_size, GetVectorizedSizeImpl(*iter));
50+
vec_size =
51+
std::min<int>(vec_size, GetVectorizedSizeImpl((*iter)->data<T>()));
6852
}
6953
return vec_size;
7054
}
7155

72-
template <ElementwiseType N, int VecSize, typename T>
56+
template <ElementwiseType ET, int VecSize, typename T>
7357
struct ElementwiseDataWrapper {
7458
T *out;
7559
const T *in0;
@@ -78,20 +62,20 @@ struct ElementwiseDataWrapper {
7862
const T *in1 = nullptr)
7963
: out(out), in0(in0), in1(in1) {}
8064

81-
using VecType = AlignedVector<T, VecSize>;
65+
using VecType = CudaAlignedVector<T, VecSize>;
8266

8367
inline __device__ void load_vector(VecType args[], int idx) {
8468
const VecType *x_vec = reinterpret_cast<const VecType *>(in0);
8569
args[0] = x_vec[idx];
86-
if (N == ElementwiseType::kBinary) {
70+
if (ET == ElementwiseType::kBinary) {
8771
const VecType *y_vec = reinterpret_cast<const VecType *>(in1);
8872
args[1] = y_vec[idx];
8973
}
9074
}
9175

9276
inline __device__ void load_scalar(T args[], int idx) {
9377
args[0] = in0[idx];
94-
if (N == ElementwiseType::kBinary) {
78+
if (ET == ElementwiseType::kBinary) {
9579
args[1] = in1[idx];
9680
}
9781
}
@@ -104,16 +88,17 @@ struct ElementwiseDataWrapper {
10488
inline __device__ void store_scalar(T res, int idx) { out[idx] = res; }
10589
};
10690

107-
template <ElementwiseType N, int VecSize, typename T, typename Functor>
108-
__device__ void VectorizedKernelImpl(ElementwiseDataWrapper<N, VecSize, T> data,
109-
int size, Functor func, int tid) {
110-
using VecType = AlignedVector<T, VecSize>;
111-
VecType ins_vec[N];
91+
template <ElementwiseType ET, int VecSize, typename T, typename Functor>
92+
__device__ void VectorizedKernelImpl(
93+
ElementwiseDataWrapper<ET, VecSize, T> data, int size, Functor func,
94+
int tid) {
95+
using VecType = CudaAlignedVector<T, VecSize>;
96+
VecType ins_vec[ET];
11297
VecType out_vec;
113-
T *ins_ptr[N];
98+
T *ins_ptr[ET];
11499
T *out_ptr;
115100
#pragma unroll
116-
for (int i = 0; i < N; ++i) {
101+
for (int i = 0; i < ET; ++i) {
117102
ins_ptr[i] = reinterpret_cast<T *>(&(ins_vec[i]));
118103
}
119104
out_ptr = reinterpret_cast<T *>(&out_vec);
@@ -124,9 +109,9 @@ __device__ void VectorizedKernelImpl(ElementwiseDataWrapper<N, VecSize, T> data,
124109
// compute
125110
#pragma unroll
126111
for (int i = 0; i < VecSize; ++i) {
127-
T ins[N];
112+
T ins[ET];
128113
#pragma unroll
129-
for (int j = 0; j < N; ++j) {
114+
for (int j = 0; j < ET; ++j) {
130115
ins[j] = ins_ptr[j][i];
131116
}
132117
out_ptr[i] = func(ins);
@@ -136,10 +121,11 @@ __device__ void VectorizedKernelImpl(ElementwiseDataWrapper<N, VecSize, T> data,
136121
data.store_vector(out_vec, tid);
137122
}
138123

139-
template <ElementwiseType N, typename T, typename Functor>
140-
__device__ void ScalarKernelImpl(ElementwiseDataWrapper<N, 1, T> data, int size,
141-
Functor func, int start, int remain) {
142-
T ins[N];
124+
template <ElementwiseType ET, typename T, typename Functor>
125+
__device__ void ScalarKernelImpl(ElementwiseDataWrapper<ET, 1, T> data,
126+
int size, Functor func, int start,
127+
int remain) {
128+
T ins[ET];
143129
T out;
144130

145131
for (int i = 0; i < remain; ++i) {
@@ -153,72 +139,67 @@ __device__ void ScalarKernelImpl(ElementwiseDataWrapper<N, 1, T> data, int size,
153139
}
154140
}
155141

156-
template <ElementwiseType N, int VecSize, typename T, typename Functor>
142+
template <ElementwiseType ET, int VecSize, typename T, typename Functor>
157143
__global__ void VectorizedKernel(const T *__restrict__ in0,
158144
const T *__restrict__ in1, T *out, int size,
159145
Functor func) {
160146
int tid = blockIdx.x * blockDim.x + threadIdx.x;
161147
int remain = size - VecSize * tid;
162148
remain = remain > 0 ? remain : 0;
163149
if (remain >= VecSize) {
164-
auto data = ElementwiseDataWrapper<N, VecSize, T>(out, in0, in1);
150+
auto data = ElementwiseDataWrapper<ET, VecSize, T>(out, in0, in1);
165151
VectorizedKernelImpl(data, size, func, tid);
166152
} else {
167-
auto data = ElementwiseDataWrapper<N, 1, T>(out, in0, in1);
153+
auto data = ElementwiseDataWrapper<ET, 1, T>(out, in0, in1);
168154
ScalarKernelImpl(data, size, func, tid * VecSize, remain);
169155
}
170156
}
171157

172-
template <ElementwiseType N, typename T, typename Functor>
158+
template <ElementwiseType ET, typename T, typename Functor>
173159
__global__ void ScalarKernel(const T *__restrict__ in0,
174160
const T *__restrict__ in1, T *out, int size,
175161
Functor func) {
176-
auto data = ElementwiseDataWrapper<N, 1, T>(out, in0, in1);
162+
auto data = ElementwiseDataWrapper<ET, 1, T>(out, in0, in1);
177163
int tid = blockIdx.x * blockDim.x + threadIdx.x;
178164
int remain = tid < size ? 1 : 0;
179165
ScalarKernelImpl(data, size, func, tid, remain);
180166
}
181167

182-
template <ElementwiseType N, typename T, typename Functor>
183-
void LaunchElementwiseCudaKernel(const framework::ExecutionContext &ctx,
184-
const std::vector<const T *> &ins,
185-
std::vector<T *> outs, int size,
186-
Functor func) {
168+
template <ElementwiseType ET, typename T, typename Functor>
169+
void LaunchElementwiseCudaKernel(
170+
const platform::CUDADeviceContext &ctx,
171+
const std::vector<const framework::Tensor *> &ins,
172+
std::vector<framework::Tensor *> *outs, Functor func) {
187173
// calculate the max vec_size for all ins and outs
188-
int vec_size = GetVectorizedSize(ins, outs);
174+
auto size = ins[0]->numel();
175+
int vec_size = GetVectorizedSize<T>(ins, *outs);
189176
int block_size = PADDLE_CUDA_THREAD_SIZE;
190177
int grid_size =
191178
((size + vec_size - 1) / vec_size + block_size - 1) / block_size;
192-
const T *in0 = ins[0];
193-
const T *in1 = nullptr;
194-
if (N == ElementwiseType::kBinary) {
195-
in1 = ins[1];
196-
}
197-
T *out = outs[0];
179+
const T *in0 = ins[0]->data<T>();
180+
const T *in1 = (ET == ElementwiseType::kBinary) ? ins[1]->data<T>() : nullptr;
181+
T *out = (*outs)[0]->data<T>();
198182
// cuda kernel
199-
auto stream =
200-
ctx.template device_context<platform::CUDADeviceContext>().stream();
183+
auto stream = ctx.stream();
201184
switch (vec_size) {
202185
case 4:
203-
VectorizedKernel<N, 4><<<grid_size, block_size, 0, stream>>>(
186+
VectorizedKernel<ET, 4><<<grid_size, block_size, 0, stream>>>(
204187
in0, in1, out, size, func);
205188
break;
206189
case 2:
207-
VectorizedKernel<N, 2><<<grid_size, block_size, 0, stream>>>(
190+
VectorizedKernel<ET, 2><<<grid_size, block_size, 0, stream>>>(
208191
in0, in1, out, size, func);
209192
break;
210193
case 1:
211-
ScalarKernel<N><<<grid_size, block_size, 0, stream>>>(in0, in1, out, size,
212-
func);
194+
ScalarKernel<ET><<<grid_size, block_size, 0, stream>>>(in0, in1, out,
195+
size, func);
213196
break;
214197
default:
215-
PADDLE_THROW(
216-
platform::errors::Unimplemented("Unsupported vectorized size!"));
198+
PADDLE_THROW(platform::errors::Unimplemented(
199+
"Unsupported vectorized size: %d !", vec_size));
217200
break;
218201
}
219202
}
220203

221-
#endif
222-
#endif
223204
} // namespace operators
224205
} // namespace paddle

0 commit comments

Comments
 (0)