bgawrych commented on a change in pull request #20759:
URL: https://github.com/apache/incubator-mxnet/pull/20759#discussion_r761043105
##########
File path: src/operator/nn/mkldnn/mkldnn_rnn.cc
##########
@@ -437,35 +472,27 @@ void MKLDNNRnnForward::SetNewDataMem(void* x,
}
}
-inline void MKLDNNMemoryReorder(const mkldnn::memory& src, const
mkldnn::memory& dst) {
-#if DMLC_CXX11_THREAD_LOCAL
- static thread_local std::unordered_map<OpSignature, mkldnn::reorder, OpHash>
reorderPrimitives;
-#else
- static MX_THREAD_LOCAL std::unordered_map<OpSignature, mkldnn::reorder,
OpHash> reorderPrimitives;
-#endif
- OpSignature key{};
- key.AddSign(src);
- key.AddSign(dst);
-
- auto it = reorderPrimitives.find(key);
- if (it == reorderPrimitives.end()) {
- auto reorder = mkldnn::reorder(src, dst);
- it = AddToCache(&reorderPrimitives, key, reorder);
- }
-
- mkldnn_args_map_t net_args;
- net_args.emplace(MKLDNN_ARG_SRC, src);
- net_args.emplace(MKLDNN_ARG_DST, dst);
- MKLDNNStream::Get()->RegisterPrimArgs(it->second, net_args);
-}
-
/*
* Reorder the concatenated weights memory to a efficient memory block
* with primitive-prefered format.
*/
void MKLDNNRnnForward::ReorderWeights() {
- MKLDNNMemoryReorder(*weights_layer_r_, *weights_layer_);
- MKLDNNMemoryReorder(*weights_iter_r_, *weights_iter_);
+if (param_.quantized) {
+ const mkldnn::primitive_attr& attr = this->fwd_inf_.GetPrimAttr();
+ auto ReorderWithAttr = [&](mkldnn::memory& src, mkldnn::memory& dst) {
+ auto reorder_pd = mkldnn::reorder::primitive_desc(src, dst, attr);
+ mkldnn_args_map_t net_args;
+ net_args[MKLDNN_ARG_SRC] = src;
+ net_args[MKLDNN_ARG_DST] = dst;
+ MKLDNNStream::Get()->RegisterPrimArgs(mkldnn::reorder(reorder_pd),
net_args);
+ };
+ ReorderWithAttr(*weights_layer_r_, *weights_layer_);
+ ReorderWithAttr(*weights_iter_r_, *weights_iter_);
+ } else {
+ MKLDNNMemoryReorder(*weights_layer_r_, *weights_layer_);
+ MKLDNNMemoryReorder(*weights_iter_r_, *weights_iter_);
+ // if (param_.proj_size > 0) MKLDNNMemoryReorder(*weights_proj_r_,
*weights_proj_);
Review comment:
Nope, projection was not added in 1.x - this can be addressed as other
issue
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
