提交 caf95767 编写于 作者: X xiaolil1

modify mkldnn_helpler for int8

上级 3e5c1915
...@@ -193,7 +193,10 @@ class MKLDNNHandler { ...@@ -193,7 +193,10 @@ class MKLDNNHandler {
const std::shared_ptr<mkldnn::memory> user_memory_p, const std::shared_ptr<mkldnn::memory> user_memory_p,
const std::string& suffix, const std::string& suffix,
std::vector<mkldnn::primitive>& pipeline, // NOLINT std::vector<mkldnn::primitive>& pipeline, // NOLINT
bool is_persistent = false) { bool is_persistent = false,
bool is_INT8 = false,
std::vector<float> scale_data = {1.0f},
int mask = 0) {
// create reorder primitive if the input format is not the preferred one // create reorder primitive if the input format is not the preferred one
auto local_key = key_ + suffix; auto local_key = key_ + suffix;
auto key_reorder_p = key_ + suffix + "reorder_p"; auto key_reorder_p = key_ + suffix + "reorder_p";
...@@ -207,9 +210,18 @@ class MKLDNNHandler { ...@@ -207,9 +210,18 @@ class MKLDNNHandler {
std::shared_ptr<mkldnn::primitive> reorder_p; std::shared_ptr<mkldnn::primitive> reorder_p;
if (mpd != user_mpd) { if (mpd != user_mpd) {
target_memory_p = std::make_shared<mkldnn::memory>(mpd); target_memory_p = std::make_shared<mkldnn::memory>(mpd);
if(is_INT8){
mkldnn::primitive_attr attri;
attri.set_output_scales(mask, scale_data);
auto reorder_pd = std::shared_ptr<mkldnn::reorder::primitive_desc>(
new mkldnn::reorder::primitive_desc(mpd, user_mpd, attri));
auto reorder_p =
std::shared_ptr<mkldnn::reorder>(new mkldnn::reorder(*reorder_pd, *user_memory_p, *target_memory_p));
} else{
auto reorder_p = auto reorder_p =
std::make_shared<mkldnn::reorder>(*user_memory_p, *target_memory_p); std::make_shared<mkldnn::reorder>(*user_memory_p, *target_memory_p);
}
dev_ctx_.SetBlob(key_reorder_p, reorder_p); dev_ctx_.SetBlob(key_reorder_p, reorder_p);
pipeline.push_back(*reorder_p); pipeline.push_back(*reorder_p);
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册