Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 7230bb9

Browse files
huangzhiyuaneric-haibin-lin
huangzhiyuan
authored and
eric-haibin-lin
committed
MKLDNN Forward FullyConnected op cache (#11611)
* Enable primitive allocation cache for FullyConnected * Enable primitive allocation cache for FullyConnected * fix indent and pass in_data as last argument for CreateMKLDNNMem * fix indent and pass in_data as last argument for CreateMKLDNNMem
1 parent 54d5777 commit 7230bb9

2 files changed

Lines changed: 123 additions & 12 deletions

File tree

‎src/operator/nn/fully_connected-inl.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,11 @@ struct FullyConnectedParam : public dmlc::Parameter<FullyConnectedParam> {
6161
DMLC_DECLARE_FIELD(flatten).set_default(true)
6262
.describe("Whether to collapse all but the first axis of the input data tensor.");
6363
}
64+
bool operator==(const FullyConnectedParam& other) const {
65+
return this->num_hidden == other.num_hidden &&
66+
this->no_bias == other.no_bias &&
67+
this->flatten == other.flatten;
68+
}
6469
};
6570

6671
template<typename xpu, typename DType>
@@ -228,4 +233,16 @@ void FullyConnectedGradCompute(const nnvm::NodeAttrs& attrs,
228233

229234
} // namespace op
230235
} // namespace mxnet
236+
namespace std {
237+
template<>
238+
struct hash<mxnet::op::FullyConnectedParam> {
239+
size_t operator()(const mxnet::op::FullyConnectedParam& val) {
240+
size_t ret = 0;
241+
ret = dmlc::HashCombine(ret, val.num_hidden);
242+
ret = dmlc::HashCombine(ret, val.no_bias);
243+
ret = dmlc::HashCombine(ret, val.flatten);
244+
return ret;
245+
}
246+
};
247+
} // namespace std
231248
#endif // MXNET_OPERATOR_NN_FULLY_CONNECTED_INL_H_

‎src/operator/nn/mkldnn/mkldnn_fully_connected.cc‎

Lines changed: 106 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,100 @@ inline static mkldnn::inner_product_backward_weights::primitive_desc GetIPBwdWei
8282
}
8383
}
8484

85+
class MKLDNNFullyConnectForward {
86+
std::shared_ptr<mkldnn::memory> data;
87+
std::shared_ptr<mkldnn::memory> weight;
88+
std::shared_ptr<mkldnn::memory> out;
89+
std::shared_ptr<mkldnn::memory> bias;
90+
std::shared_ptr<mkldnn::inner_product_forward> ipFwd;
91+
92+
public:
93+
mkldnn::inner_product_forward::primitive_desc ipFwd_pd;
94+
95+
MKLDNNFullyConnectForward(const FullyConnectedParam &param, bool is_train,
96+
const NDArray &data, const NDArray &weight,
97+
const NDArray *bias,
98+
const mkldnn::memory::desc &output)
99+
: ipFwd_pd(GetIPFwd(data, weight, bias, output, is_train)) {}
100+
101+
void SetNewMem(const mkldnn::memory &data, const mkldnn::memory &weight,
102+
const mkldnn::memory *bias, const mkldnn::memory &output) {
103+
if (this->data == nullptr)
104+
this->data = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
105+
ipFwd_pd.src_primitive_desc(), data.get_data_handle()));
106+
else
107+
this->data->set_data_handle(data.get_data_handle());
108+
109+
if (this->weight == nullptr)
110+
this->weight = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
111+
ipFwd_pd.weights_primitive_desc(), weight.get_data_handle()));
112+
else
113+
this->weight->set_data_handle(weight.get_data_handle());
114+
115+
if (this->out == nullptr)
116+
this->out = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
117+
ipFwd_pd.dst_primitive_desc(), output.get_data_handle()));
118+
else
119+
this->out->set_data_handle(output.get_data_handle());
120+
121+
if (bias != nullptr) {
122+
if (this->bias == nullptr)
123+
this->bias = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
124+
ipFwd_pd.bias_primitive_desc(), bias->get_data_handle()));
125+
else
126+
this->bias->set_data_handle(bias->get_data_handle());
127+
if (this->ipFwd == nullptr)
128+
this->ipFwd = std::shared_ptr<mkldnn::inner_product_forward>(
129+
new mkldnn::inner_product_forward(
130+
ipFwd_pd, mkldnn::primitive::at(*this->data),
131+
mkldnn::primitive::at(*this->weight),
132+
mkldnn::primitive::at(*this->bias), *this->out));
133+
} else if (this->ipFwd == nullptr) {
134+
this->ipFwd = std::shared_ptr<mkldnn::inner_product_forward>(
135+
new mkldnn::inner_product_forward(
136+
ipFwd_pd, mkldnn::primitive::at(*this->data),
137+
mkldnn::primitive::at(*this->weight), *this->out));
138+
}
139+
}
140+
const mkldnn::inner_product_forward &GetIpFwd() const {
141+
return *ipFwd;
142+
}
143+
};
144+
145+
typedef ParamOpSign<FullyConnectedParam> MKLDNNFullyconSignature;
146+
147+
static inline MKLDNNFullyConnectForward &GetFCFwd(
148+
const nnvm::NodeAttrs &attrs, const NDArray &data, const NDArray &weight,
149+
const NDArray *bias, const mkldnn::memory::desc &output,
150+
const bool is_train) {
151+
#if DMLC_CXX11_THREAD_LOCAL
152+
static thread_local std::unordered_map<MKLDNNFullyconSignature,
153+
MKLDNNFullyConnectForward, OpHash> fcFwds;
154+
#else
155+
static MX_THREAD_LOCAL std::unordered_map<MKLDNNFullyconSignature,
156+
MKLDNNFullyConnectForward, OpHash> fcFwds;
157+
#endif
158+
const FullyConnectedParam& param = nnvm::get<FullyConnectedParam>(attrs.parsed);
159+
MKLDNNFullyconSignature key(param);
160+
key.AddSign(data);
161+
key.AddSign(weight);
162+
key.AddSign(is_train);
163+
164+
if (bias)
165+
key.AddSign(*bias);
166+
167+
auto it = fcFwds.find(key);
168+
if (it == fcFwds.end()) {
169+
MKLDNNFullyConnectForward fcFwd(param, is_train, data, weight, bias,
170+
output);
171+
auto ins_ret = fcFwds.insert(
172+
std::pair<MKLDNNFullyconSignature, MKLDNNFullyConnectForward>(key, fcFwd));
173+
CHECK(ins_ret.second);
174+
it = ins_ret.first;
175+
}
176+
return it->second;
177+
}
178+
85179
void MKLDNNFCForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
86180
const std::vector<NDArray> &in_data,
87181
const std::vector<OpReqType> &req,
@@ -112,21 +206,21 @@ void MKLDNNFCForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
112206
out_md = mkldnn::memory::desc(out_dims, get_mkldnn_type(out_data[fullc::kOut].dtype()),
113207
mkldnn::memory::format::any);
114208
}
115-
116-
mkldnn::inner_product_forward::primitive_desc ipFwd_pd = GetIPFwd(data, weight,
117-
param.no_bias ? nullptr : &in_data[fullc::kBias], out_md, ctx.is_train);
118-
auto data_mem = data.GetMKLDNNDataReorder(ipFwd_pd.src_primitive_desc());
119-
auto weight_mem = weight.GetMKLDNNDataReorder(ipFwd_pd.weights_primitive_desc());
209+
MKLDNNFullyConnectForward &FCFwd =
210+
GetFCFwd(attrs, data, weight, param.no_bias ? nullptr : &in_data[fullc::kBias],
211+
out_md, ctx.is_train);
212+
auto data_mem = data.GetMKLDNNDataReorder(FCFwd.ipFwd_pd.src_primitive_desc());
213+
auto weight_mem = weight.GetMKLDNNDataReorder(FCFwd.ipFwd_pd.weights_primitive_desc());
120214
auto out_mem = CreateMKLDNNMem(out_data[fullc::kOut],
121-
ipFwd_pd.dst_primitive_desc(), req[fullc::kOut]);
122-
if (param.no_bias) {
123-
MKLDNNStream::Get()->RegisterPrim(mkldnn::inner_product_forward(
124-
ipFwd_pd, *data_mem, *weight_mem, *out_mem.second));
215+
FCFwd.ipFwd_pd.dst_primitive_desc(), req[fullc::kOut], &data);
216+
if (!param.no_bias) {
217+
auto bias_mem = in_data[fullc::kBias].GetMKLDNNDataReorder(
218+
FCFwd.ipFwd_pd.bias_primitive_desc());
219+
FCFwd.SetNewMem(*data_mem, *weight_mem, bias_mem, *out_mem.second);
125220
} else {
126-
auto bias_mem = in_data[fullc::kBias].GetMKLDNNDataReorder(ipFwd_pd.bias_primitive_desc());
127-
MKLDNNStream::Get()->RegisterPrim(mkldnn::inner_product_forward(ipFwd_pd,
128-
*data_mem, *weight_mem, *bias_mem, *out_mem.second));
221+
FCFwd.SetNewMem(*data_mem, *weight_mem, nullptr, *out_mem.second);
129222
}
223+
MKLDNNStream::Get()->RegisterPrim(FCFwd.GetIpFwd());
130224
CommitOutput(out_data[fullc::kOut], out_mem);
131225
MKLDNNStream::Get()->Submit();
132226
}

0 commit comments

Comments
 (0)