azai91 commented on a change in pull request #11611: MKLDNN Forward FullyConnected op cache URL: https://github.com/apache/incubator-mxnet/pull/11611#discussion_r211482743
########## File path: src/operator/nn/mkldnn/mkldnn_fully_connected.cc ########## @@ -82,6 +82,101 @@ inline static mkldnn::inner_product_backward_weights::primitive_desc GetIPBwdWei } } +class MKLDNNFullyConnectForward { + std::shared_ptr<mkldnn::memory> data; + std::shared_ptr<mkldnn::memory> weight; + std::shared_ptr<mkldnn::memory> out; + std::shared_ptr<mkldnn::memory> bias; + std::shared_ptr<mkldnn::inner_product_forward> ipFwd; + + public: + mkldnn::inner_product_forward::primitive_desc ipFwd_pd; + + MKLDNNFullyConnectForward(const FullyConnectedParam ¶m, bool is_train, + const NDArray &data, const NDArray &weight, + const NDArray *bias, + const mkldnn::memory::desc &output) + : ipFwd_pd(GetIPFwd(data, weight, bias, output, is_train)) {} + + void SetNewMem(const mkldnn::memory &data, const mkldnn::memory &weight, + const mkldnn::memory *bias, const mkldnn::memory &output) { + if (this->data == nullptr) + this->data = std::shared_ptr<mkldnn::memory>(new mkldnn::memory( + ipFwd_pd.src_primitive_desc(), data.get_data_handle())); + else + this->data->set_data_handle(data.get_data_handle()); + + if (this->weight == nullptr) + this->weight = std::shared_ptr<mkldnn::memory>(new mkldnn::memory( + ipFwd_pd.weights_primitive_desc(), weight.get_data_handle())); + else + this->weight->set_data_handle(weight.get_data_handle()); + + if (this->out == nullptr) + this->out = std::shared_ptr<mkldnn::memory>(new mkldnn::memory( + ipFwd_pd.dst_primitive_desc(), output.get_data_handle())); + else + this->out->set_data_handle(output.get_data_handle()); + + if (bias != nullptr) { + if (this->bias == nullptr) + this->bias = std::shared_ptr<mkldnn::memory>(new mkldnn::memory( + ipFwd_pd.bias_primitive_desc(), bias->get_data_handle())); + else + this->bias->set_data_handle(bias->get_data_handle()); + if (this->ipFwd == nullptr) Review comment: fix indent here? ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services