xinyu-intel commented on a change in pull request #16931: [MKL-DNN] Enable and
Optimization for s8 eltwise_add
URL: https://github.com/apache/incubator-mxnet/pull/16931#discussion_r352401970
##########
File path: src/operator/quantization/mkldnn/mkldnn_quantized_elemwise_add.cc
##########
@@ -39,6 +39,57 @@ static inline float GetScale(const NDArray& data, float
min, float max) {
return data_range / MaxAbs(min, max);
}
+class MKLDNNQuantizedElemwiseAddFwd {
+ public:
+ mkldnn::sum::primitive_desc fwd_pd;
+
+ MKLDNNQuantizedElemwiseAddFwd(
+ const mkldnn::memory::desc &output_desc,
+ const std::vector<float> &scales,
+ const std::vector<mkldnn::memory::desc> &data_md)
+ : fwd_pd(output_desc, scales, data_md, CpuEngine::Get()->get_engine()) {
+ fwd_ = std::make_shared<mkldnn::sum>(fwd_pd);
+ data_.resize(data_md.size());
+ }
+
+ const mkldnn::sum &GetFwd() const { return *fwd_; }
+
+ private:
+ std::shared_ptr<mkldnn::sum> fwd_;
+ std::vector<std::shared_ptr<mkldnn::memory>> data_;
+ std::shared_ptr<mkldnn::memory> out_;
+};
+
+static MKLDNNQuantizedElemwiseAddFwd &GetQuantizedElemwiseAddForward(
+ const mkldnn::memory::desc &output_desc, const std::vector<float> &scales,
+ const std::vector<NDArray> &in_data, const std::vector<NDArray> &out_data,
+ const std::vector<mkldnn::memory::desc> &data_md) {
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local std::unordered_map<OpSignature,
+ MKLDNNQuantizedElemwiseAddFwd, OpHash> fwds;
+#else
+ static MX_THREAD_LOCAL std::unordered_map<OpSignature,
+ MKLDNNQuantizedElemwiseAddFwd, OpHash> fwds;
+#endif
+ OpSignature key;
+ key.AddSign(in_data);
+
key.AddSign(in_data[quantized_elemwise_add_enum::kAMin].data().dptr<float>()[0]);
Review comment:
caching `min`, `max` will be more efficient compared with caching `scales`.
----------------------------------------------------------------
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
