Skip to content

Commit 0d0d19a

Browse files
dcslindcslin
committed
- Integrated MKL dnn for operations below fwd/bwd w/ tests
- Conv2D - BatchNorm 4D/2D - Pooling Max/Avg - Updated CMake for flag: USE_MKLDNN
1 parent eec0d52 commit 0d0d19a

16 files changed

Lines changed: 1203 additions & 7 deletions

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ OPTION(USE_OPENCL "Use OpenCL" OFF)
6666
OPTION(ENABLE_DIST "Enable distributed training" OFF)
6767
OPTION(DISABLE_WARNINGS "Disable warnings under windows" ON)
6868
OPTION(USE_MODULES "Compile dependent libs as submodules together with singa" OFF)
69+
OPTION(USE_MKLDNN "Use mkl-dnn libs" ON)
6970

7071

7172
# TODO: remove all USE_CBLAS in codes

‎cmake/Dependencies.cmake‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,3 +141,7 @@ IF(USE_JAVA)
141141
FIND_PACKAGE(JNI REQUIRED)
142142
FIND_PACKAGE(SWIG 3.0 REQUIRED)
143143
ENDIF()
144+
145+
IF(USE_MKLDNN)
146+
LIST(APPEND SINGA_LINKER_LIBS mkldnn)
147+
ENDIF()

‎cmake/Templates/singa_config.h.in‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,5 @@
5151
// #cmakedefine CUDNN_MINOR_VERSION @CUDNN_MINOR_VERSION@
5252
// #cmakedefine CUDNN_PATCH_VERSION @CUDNN_PATCH_VERSION@
5353
// #cmakedefine CUDNN_VERSION @CUDNN_VERSION@
54+
55+
#cmakedefine USE_MKLDNN

‎include/singa/core/common.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@
3939
#include "singa/utils/opencl_utils.h"
4040
#endif // USE_OPENCL
4141

42+
#ifdef USE_MKLDNN
43+
#include <mkldnn.hpp>
44+
#endif // USE_MKLDNN
45+
4246
using std::atomic;
4347

4448
namespace singa {
@@ -111,6 +115,10 @@ typedef struct _Context {
111115
long vcl_ctx_id;
112116
#endif
113117

118+
#ifdef USE_MKLDNN
119+
mkldnn::engine *engine;
120+
#endif // USE_MKLDNN
121+
114122
} Context;
115123

116124
} // namespace singa

‎include/singa/core/device.h‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@
3939
#include "singa/utils/opencl_utils.h"
4040
#endif // USE_OPENCL
4141

42+
#ifdef USE_MKLDNN
43+
#include "singa/utils/mkldnn_utils.h"
44+
#endif // USE_MKLDNN
45+
4246
using std::vector;
4347
using std::string;
4448
using std::function;
@@ -141,7 +145,7 @@ extern std::shared_ptr<Device> defaultDevice;
141145
/// It runs cpp code.
142146
class CppCPU : public Device {
143147
public:
144-
~CppCPU() {};
148+
~CppCPU();
145149
CppCPU();
146150

147151
std::shared_ptr<Device> host() const override { return defaultDevice;}

‎include/singa/utils/mkldnn_utils.h‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
#ifndef SINGA_UTILS_MKLDNN_UTILS_H_
2+
#define SINGA_UTILS_MKLDNN_UTILS_H_
3+
4+
#include <mkldnn.hpp>
5+
6+
namespace singa {
7+
/*
8+
supported data type by mkldnn
9+
mkldnn_f32 - 32-bit/single-precision floating point.
10+
mkldnn_s32 - 32-bit signed integer.
11+
mkldnn_s16 - 16-bit signed integer.
12+
mkldnn_s8 - 8-bit signed integer.
13+
mkldnn_u8 - 8-bit unsigned integer.
14+
*/
15+
inline mkldnn::memory::data_type GetMKLDNNDataType(DataType dtype) {
16+
mkldnn::memory::data_type ret = mkldnn::memory::data_type::f32;
17+
switch (dtype) {
18+
case kFloat32:
19+
ret = mkldnn::memory::data_type::f32;
20+
break;
21+
case kDouble:
22+
LOG(FATAL) << "The data type " << DataType_Name(dtype)
23+
<< " is not support by mkldnn";
24+
break;
25+
case kFloat16:
26+
LOG(FATAL) << "The data type " << DataType_Name(dtype)
27+
<< " is not support by mkldnn";
28+
break;
29+
default:
30+
LOG(FATAL) << "The data type " << DataType_Name(dtype)
31+
<< " is not support by mkldnn";
32+
}
33+
return ret;
34+
}
35+
}
36+
#endif // SINGA_UTILS_MKLDNN_UTILS_H_

‎src/core/device/cpp_cpu.cc‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,18 @@ std::shared_ptr<Device> defaultDevice=std::make_shared<CppCPU>();
2424

2525
CppCPU::CppCPU() : Device(-1, 1) {
2626
lang_ = kCpp;
27+
#ifdef USE_MKLDNN
28+
ctx_.engine = new mkldnn::engine(mkldnn::engine::cpu, 0);
29+
#endif //USE_MKLDNN
2730
//host_ = nullptr;
2831
}
2932

33+
CppCPU::~CppCPU() {
34+
#ifdef USE_MKLDNN
35+
delete(ctx_.engine);
36+
#endif //USE_MKLDNN
37+
38+
};
3039

3140
void CppCPU::SetRandSeed(unsigned seed) {
3241
ctx_.random_generator.seed(seed);

‎src/model/operation/batchnorm.cc‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,187 @@ BatchNormHandle::BatchNormHandle(const float momentum, const Tensor& input) {
1717
} else {
1818
LOG(FATAL) << "The dimension of input should either be 4D or 2D.";
1919
}
20+
21+
22+
#ifdef USE_MKLDNN
23+
dtype = GetMKLDNNDataType(input.data_type());
24+
epsilon =1e-5f;
25+
data_memory_format = is_2d ? mkldnn::memory::format::nc : mkldnn::memory::format::nchw;
26+
if (is_2d) {
27+
x_dims = {(int)batchsize, (int)channels};
28+
y_dims = {(int)batchsize, (int)channels};
29+
} else {
30+
x_dims = {(int)batchsize, (int)channels, (int)height, (int)width};
31+
y_dims = {(int)batchsize, (int)channels, (int)height, (int)width};
32+
}
33+
34+
auto eng = *input.device()->context(0)->engine;
35+
x_md = new mkldnn::memory::desc(x_dims, dtype, data_memory_format);
36+
dx_md = new mkldnn::memory::desc(x_dims, dtype, data_memory_format);
37+
bn_fwd_d = new mkldnn::batch_normalization_forward::desc(mkldnn::forward_training, *x_md, epsilon,
38+
mkldnn::use_scale_shift);
39+
bn_fwd_pd = new mkldnn::batch_normalization_forward::primitive_desc(*bn_fwd_d, eng);
40+
#endif // USE_MKLDNN
41+
2042
};
2143

44+
45+
BatchNormHandle::~BatchNormHandle() {
46+
#ifdef USE_MKLDNN
47+
delete (x_md);
48+
delete (dx_md);
49+
delete (bn_fwd_d);
50+
delete (bn_fwd_pd);
51+
#endif // USE_MKLDNN
52+
}
53+
54+
#ifdef USE_MKLDNN
55+
56+
Tensor CpuBatchNormForwardInference(const BatchNormHandle &bnh, const Tensor& x, const Tensor& bnScale, const Tensor& bnBias,
57+
Tensor& running_mean, Tensor& running_var){
58+
59+
CHECK_EQ(x.device()->lang(), kCpp);
60+
Tensor y;
61+
y.ResetLike(x);
62+
63+
64+
Tensor w = get_bn_weight_from(bnScale, bnBias);
65+
66+
y.device()->Exec([&y, &x, &w, &bnh](Context *ctx) {
67+
try {
68+
auto eng = *ctx->engine;
69+
using namespace mkldnn;
70+
auto x_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng}, x.block()->mutable_data());
71+
auto y_mem = memory({{{bnh.y_dims}, bnh.dtype, bnh.data_memory_format}, eng}, y.block()->mutable_data());
72+
73+
auto bn_fwd_d = batch_normalization_forward::desc(forward_inference, *bnh.x_md, bnh.epsilon, use_scale_shift);
74+
auto bn_fwd_pd = batch_normalization_forward::primitive_desc(bn_fwd_d, eng);
75+
76+
auto w_mem = memory(bn_fwd_pd.weights_primitive_desc(), w.block()->mutable_data());
77+
78+
auto bn = batch_normalization_forward(bn_fwd_pd, x_mem, w_mem, y_mem);
79+
80+
stream(stream::kind::eager).submit({bn}).wait();
81+
}
82+
catch (mkldnn::error &e) {
83+
InitLogging("");
84+
LOG(FATAL) << "MKLDNN Batch Norm" << "Status: " << e.status << " Message: " << e.message;
85+
}
86+
87+
}, {y.block(), x.block(), w.block()}, {y.block()});
88+
89+
return y;
90+
91+
}
92+
93+
const std::vector<Tensor>
94+
CpuBatchNormForwardTraining(const BatchNormHandle &bnh, const Tensor &x, const Tensor &bnScale, const Tensor &bnBias,
95+
Tensor &running_mean, Tensor &running_var) {
96+
97+
Tensor y;
98+
y.ResetLike(x);
99+
100+
// mean and var for local batch
101+
Tensor mean;
102+
mean.ResetLike(running_mean);
103+
Tensor var;
104+
var.ResetLike(running_var);
105+
106+
// combine scale and bias to construct weight tensor in required format for backward
107+
Tensor w = get_bn_weight_from(bnScale, bnBias);
108+
109+
y.device()->Exec([&x, &y, &mean, &var, &w, &bnh](Context *ctx) {
110+
try {
111+
auto eng = *ctx->engine;
112+
using namespace mkldnn;
113+
114+
auto x_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng},
115+
x.block()->mutable_data());
116+
auto y_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng},
117+
y.block()->mutable_data());
118+
auto m_mem = memory(bnh.bn_fwd_pd->mean_primitive_desc(), mean.block()->mutable_data());
119+
120+
auto v_mem = memory(bnh.bn_fwd_pd->variance_primitive_desc(), var.block()->mutable_data());
121+
122+
auto w_mem = memory(bnh.bn_fwd_pd->weights_primitive_desc(),w.block()->mutable_data());
123+
124+
auto bn_fwd = batch_normalization_forward(*bnh.bn_fwd_pd, x_mem, w_mem, y_mem, m_mem, v_mem);
125+
126+
stream(stream::kind::eager).submit({bn_fwd}).wait();
127+
}
128+
catch (mkldnn::error &e) {
129+
singa::InitLogging("");
130+
LOG(FATAL) << "MKLDNN Batch Norm Backward" << "Status: " << e.status << " Message: " << e.message;
131+
}
132+
}, {x.block(), w.block()}, {y.block(), mean.block(), var.block()});
133+
134+
135+
// local implemented running mean as mkldnn does not support it yet:
136+
// https://github.com/intel/mkl-dnn/issues/371
137+
running_mean = running_mean*bnh.factor + mean*(1-bnh.factor);
138+
running_var = running_var*bnh.factor + var*(1-bnh.factor);
139+
140+
141+
return {y, running_mean, running_var};
142+
143+
}
144+
145+
const std::vector<Tensor> CpuBatchNormBackwardx(const BatchNormHandle &bnh,
146+
const Tensor &y, const Tensor &dy,
147+
const Tensor &x,
148+
const Tensor &bnScale, const Tensor &bnBias,
149+
const Tensor &mean, const Tensor &var){
150+
Tensor dx;
151+
dx.ResetLike(dy);
152+
153+
// combine scale and bias to construct weight tensor in required format for backward
154+
Tensor w = get_bn_weight_from(bnScale, bnBias);
155+
156+
Tensor dw(Shape{bnScale.Size(),bnBias.Size()});
157+
158+
dx.device()->Exec([&dw, &x, &dx, &y, &dy, &w, &mean, &var, &bnh](Context *ctx) {
159+
160+
try {
161+
auto eng = *ctx->engine;
162+
using namespace mkldnn;
163+
164+
auto x_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng}, x.block()->mutable_data());
165+
auto dx_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng}, dx.block()->mutable_data());
166+
auto y_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng}, y.block()->mutable_data());
167+
auto dy_mem = memory({{{bnh.x_dims}, bnh.dtype, bnh.data_memory_format}, eng}, dy.block()->mutable_data());
168+
169+
auto m_mem = memory(bnh.bn_fwd_pd->mean_primitive_desc(), mean.block()->mutable_data());
170+
auto v_mem = memory(bnh.bn_fwd_pd->variance_primitive_desc(), var.block()->mutable_data());
171+
auto w_mem = memory(bnh.bn_fwd_pd->weights_primitive_desc(), w.block()->mutable_data());
172+
173+
174+
auto bn_bwd_d = batch_normalization_backward::desc(backward, *bnh.dx_md, *bnh.x_md, bnh.epsilon, use_scale_shift);
175+
auto bn_bwd_pd = batch_normalization_backward::primitive_desc(bn_bwd_d, eng, *bnh.bn_fwd_pd);
176+
177+
178+
auto dw_mem = memory(bn_bwd_pd.diff_weights_primitive_desc(), dw.block()->mutable_data());
179+
180+
auto bn_bwd = batch_normalization_backward(bn_bwd_pd, x_mem, m_mem, v_mem, dy_mem, w_mem, dx_mem, dw_mem);
181+
182+
stream(stream::kind::eager).submit({bn_bwd}).wait();
183+
}
184+
catch (mkldnn::error &e) {
185+
singa::InitLogging("");
186+
LOG(FATAL) << "MKLDNN Batch Norm Backward" << "Status: " << e.status << " Message: " << e.message;
187+
}
188+
189+
}, {x.block(), dy.block(), mean.block(), var.block()},
190+
{dx.block(), dw.block()});
191+
192+
Tensor dbnScale = CopyRows(dw, 0, bnScale.Size());
193+
Tensor dbnBias = CopyRows(dw, 1, bnBias.Size());
194+
195+
return {dx, dbnScale, dbnBias};
196+
}
197+
198+
199+
#endif // USE_MKLDNN
200+
22201
#ifdef USE_CUDNN
23202
CudnnBatchNormHandle::CudnnBatchNormHandle(const float momentum,
24203
const Tensor& input): BatchNormHandle(momentum, input) {

‎src/model/operation/batchnorm.h‎

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,28 @@
77
#ifdef USE_CUDNN
88
#include <cudnn.h>
99
#include "../layer/cudnn_utils.h" // check_cudnn
10-
#endif // USE_CUDNN
10+
#endif // USE_CUDNN
11+
12+
#ifdef USE_MKLDNN
13+
#include <mkldnn.hpp>
14+
15+
// combine scale and bias into weight format recognised by mkldnn api
16+
static inline singa::Tensor get_bn_weight_from(const singa::Tensor &s, const singa::Tensor &b) {
17+
singa::Tensor w(singa::Shape{s.Size(),b.Size()});
18+
CopyDataToFrom(&w, s,s.Size(),0,0);
19+
CopyDataToFrom(&w, b,b.Size(),s.Size(),0);
20+
return w;
21+
}
22+
23+
24+
#endif // USE_MKLDNN
1125

1226
namespace singa {
1327

1428
class BatchNormHandle {
1529
public:
1630
BatchNormHandle(const float momentum, const Tensor& input);
31+
~BatchNormHandle();
1732

1833
float factor;
1934

@@ -23,13 +38,37 @@ class BatchNormHandle {
2338
size_t width;
2439
bool is_2d;
2540
//bool train = true;
41+
#ifdef USE_MKLDNN
42+
mkldnn::memory::data_type dtype;
43+
mkldnn::memory::dims x_dims;
44+
mkldnn::memory::dims y_dims;
45+
mkldnn::memory::desc *x_md;
46+
mkldnn::memory::desc *dx_md;
47+
mkldnn::batch_normalization_forward::desc *bn_fwd_d;
48+
mkldnn::batch_normalization_forward::primitive_desc *bn_fwd_pd;
49+
float epsilon;
50+
mkldnn::memory::format data_memory_format;
51+
#endif //USE_MKLDNN
2652
};
2753

28-
//Tensor CpuBatchNormForwardTraining();
2954

30-
//Tensor CpuBatchNormForwardInference();
55+
#ifdef USE_MKLDNN
56+
57+
Tensor
58+
CpuBatchNormForwardInference(const BatchNormHandle &bnh, const Tensor &x, const Tensor &bnScale, const Tensor &bnBias,
59+
Tensor &running_mean, Tensor &running_var);
60+
61+
const std::vector<Tensor>
62+
CpuBatchNormForwardTraining(const BatchNormHandle &bnh, const Tensor &x, const Tensor &bnScale, const Tensor &bnBias,
63+
Tensor &running_mean, Tensor &running_var);
64+
65+
const std::vector<Tensor> CpuBatchNormBackwardx(const BatchNormHandle &bnh,
66+
const Tensor &y, const Tensor &dy,
67+
const Tensor &x,
68+
const Tensor &bnScale, const Tensor &bnBias,
69+
const Tensor &mean, const Tensor &var);
3170

32-
//Tensor CpuBatchNormBackwardx();
71+
#endif // USE_MKLDNN
3372

3473

3574
#ifdef USE_CUDNN

0 commit comments

Comments
 (0)