ciyongch commented on a change in pull request #17959: [MKLDNN] Add LSTMP to
v1.6.x
URL: https://github.com/apache/incubator-mxnet/pull/17959#discussion_r402685654
##########
File path: src/operator/nn/mkldnn/mkldnn_rnn.cc
##########
@@ -470,66 +488,92 @@ inline void EmplaceNetArgs(mkldnn_args_map_t* net_args,
const int arg_name,
* memory with preferred format_tag. Finally, naive bias is fused to MKLDNN
* bias memory.
*/
-void MKLDNNRnnForward::SetWeightsMem(MKLDNNRnnMemMgr* mgr, void *w_ptr, void
*b_ptr,
+void MKLDNNRnnForward::SetWeightsMem(void *w_ptr, void *b_ptr,
const bool is_train, const int dtype) {
using format_tag = mkldnn::memory::format_tag;
auto mkldnn_dtype = get_mkldnn_type(dtype);
+ const size_t dtype_bytes = mshadow::mshadow_sizeof(dtype);
+
+ const size_t buffer_bytes = this->GetSize() // byte number of the buffer
+ + (param_.workspace_size + param_.reserve_size) * dtype_bytes
+ + kMKLDNNAlign * 7; // Add margin for alignment of seven times
allocation for the
Review comment:
Please remove the white-space to fix lint error.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
