Add DynamicQuantizeLinear to DNNL EP (#9620)

implement dynamicquantizelinear in DNNL EP

add debug log in EP for operator coverage

block gpu elementwise op with 5 dims or more

Signed-off-by: Wang <zhaoyang.wang@intel.com>

onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc

@@ -58,6 +58,10 @@ DNNLExecutionProvider::~DNNLExecutionProvider() {
 std::vector<std::vector<NodeIndex>> DNNLExecutionProvider::GetSupportedNodes(const GraphViewer& graph_viewer) const {
   std::vector<std::vector<size_t>> supported_node_vecs;
   std::vector<size_t> supported_node_vec;
+  
+  std::unordered_map<std::string,int> all_nodes_count;
+  std::unordered_map<std::string,int> supported_nodes_count;
+
   const auto& node_indices = graph_viewer.GetNodesInTopologicalOrder();
   for (size_t i = 0; i < node_indices.size(); i++) {
     auto node_idx = node_indices[i];
@@ -65,7 +69,14 @@ std::vector<std::vector<NodeIndex>> DNNLExecutionProvider::GetSupportedNodes(con
 
     bool supported = opManager_.IsNodeSupported(node, graph_viewer);
 
-    if (debug_log_) {
+    //update count
+    if(debug_log_){
+      auto node_optype_ver = node->OpType() + "_" + std::to_string(node->SinceVersion());
+      all_nodes_count[node_optype_ver]++;
+      if (supported) {
+        supported_nodes_count[node_optype_ver]++;
+      }
+
       LOGS_DEFAULT(ERROR) << "Operator type: [" << node->OpType()
                           << "] index: [" << node_idx
                           << "] name: [" << node->Name()
@@ -87,6 +98,23 @@ std::vector<std::vector<NodeIndex>> DNNLExecutionProvider::GetSupportedNodes(con
     supported_node_vecs.push_back(supported_node_vec);
   }
 
+  //collect statistics and report
+  if (debug_log_) {
+    int all_counts = 0;
+    int support_counts = 0;
+    for(auto e: all_nodes_count){
+      auto optype_ver = e.first;
+      auto all_count = e.second;
+      auto support_count = supported_nodes_count[optype_ver];
+      all_counts += all_count;
+      support_counts += support_count;
+      LOGS_DEFAULT(ERROR) << "Operator type: [" << optype_ver << "] coverage: " << support_count << ":" << all_count << " percentage: " << (float)support_count / (float)all_count;
+    }
+    LOGS_DEFAULT(ERROR) << "Total coverge: " << support_counts << ":" << all_counts
+                          << " percentage: " << (float)support_counts / (float)all_counts;
+  }
+  
+
   return supported_node_vecs;
 }
 
@@ -333,6 +361,11 @@ Status DNNLExecutionProvider::Compile(const std::vector<Node*>& fused_nodes,
           auto output_name = subgraph_primitive->GetOrderedOutputs()[i];
           auto output_md = subgraph_primitive->GetOutputInfo(output_name);
           auto output_shape = output_md.dims();
+          //if an output is a scaler, onednn internally uses tensor representation (eg, (1,1,...))
+          //but allocating an output with no shape instead of the equivalent tensorshape to avoid shape mismatch
+          if (subgraph_primitive->IsScalarOutput(output_name)) {
+            output_shape.clear();
+          }
           auto* output_tensor =
               ort.KernelContext_GetOutput(context, i, output_shape.data(), output_shape.size());
           auto* tensor_info = ort.GetTensorTypeAndShape(output_tensor);

onnxruntime/core/providers/dnnl/dnnl_node_capability.cc

@@ -426,6 +426,13 @@ bool DnnlElementwiseCapability::IsDimensionSupported(const Node* node) const {
     return true;
   }
 
+  //reject gpu elmentwise op with 5 dims or more
+  if (dnnl_engine_get_count(dnnl_engine_kind_t::dnnl_gpu)) {
+    if(node_inputs[0]->Shape()->dim_size() > 5 ){
+      return false;
+    } 
+  }
+
   // OneDNN will silently convert scaler values to a {1} tensor which causes issues for
   // for Onnruntime when it expects an empty tensor i.e. {}
   // TODO convert {1} outputs back to scaler {} once that is done DnnlElementwiseCapability
@@ -502,4 +509,14 @@ bool DnnlReshapeNodeCapability::IsDimensionSupported(const Node* node) const {
 
   return true;
 }
+
+// DnnlDynamicQuantizeLinearNodeCapability class
+// reserve for future capability change
+//-------------------------------------
+bool DnnlDynamicQuantizeLinearNodeCapability::Supported(const Node* node, const GraphViewer& graph_viewer) const {
+  ORT_UNUSED_PARAMETER(graph_viewer);
+  if (!IsTypeSupported(node)) return false;
+  return true;
+}
+
 }  // namespace onnxruntime

onnxruntime/core/providers/dnnl/dnnl_node_capability.h

@@ -289,4 +289,16 @@ class DnnlReshapeNodeCapability : public DnnlDefaultNodeCapability {
   bool IsDimensionSupported(const Node* node) const;
 };
 
+/**
+ * Decide if a DynamicQuantizeLinear op is supported by DnnlExecutionProvider
+ */
+class DnnlDynamicQuantizeLinearNodeCapability : public DnnlDefaultNodeCapability {
+ public:
+  DnnlDynamicQuantizeLinearNodeCapability() : DnnlDefaultNodeCapability({type_float32}) {}
+
+  bool Supported(const Node* node, const GraphViewer& graph_viewer) const override;
+
+ private:
+};
+
 }  // namespace onnxruntime

onnxruntime/core/providers/dnnl/dnnl_op_manager.cc

@@ -12,6 +12,7 @@ DnnlOpManager::DnnlOpManager() {
   dnnl_ops_map_.emplace(std::make_pair("BatchNormalization", std::unique_ptr<DnnlNodeCapability>(new DnnlBatchNormalizationNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Conv", std::unique_ptr<DnnlNodeCapability>(new DnnlDefaultNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Div", std::unique_ptr<DnnlNodeCapability>(new DnnlBinaryNodeCapability())));
+  dnnl_ops_map_.emplace(std::make_pair("DynamicQuantizeLinear", std::unique_ptr<DnnlNodeCapability>(new DnnlDynamicQuantizeLinearNodeCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Elu", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Exp", std::unique_ptr<DnnlNodeCapability>(new DnnlElementwiseCapability())));
   dnnl_ops_map_.emplace(std::make_pair("Gemm", std::unique_ptr<DnnlNodeCapability>(new DnnlGemmNodeCapability())));

onnxruntime/core/providers/dnnl/subgraph/dnnl_dynamicquantizelinear.cc

@@ -0,0 +1,178 @@
+// Copyright(C) 2021 Intel Corporation
+// Licensed under the MIT License
+
+#include "dnnl_dynamicquantizelinear.h"
+#include "dnnl_subgraph.h"
+#include "dnnl_subgraph_primitive.h"
+
+namespace onnxruntime {
+namespace ort_dnnl {
+
+
+/*
+x_min = np.minimum(0, np.min(X))
+x_max = np.maximum(0, np.max(X))
+Y_Scale = np.float32((x_max - x_min) / (255 - 0))  # uint8 -> [0, 255]
+Y_ZeroPoint = np.clip(round((0 - x_min) / Y_Scale), 0, 255).astype(np.uint8)
+Y = np.clip(np.round(X / Y_Scale) + Y_ZeroPoint, 0, 255).astype(np.uint8)
+*/
+void DnnlDynamicQuantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node) {
+  auto eng = sp.GetEngine();
+  auto x_memory = sp.GetMemory(node.Input(IN_X).Name());
+  x_memory = sp.GetMemoryAndReshape(node.Input(IN_X), x_memory.get_desc(), eng);
+  auto x_memory_desc = x_memory.get_desc();
+  auto x_memory_dims = x_memory_desc.dims();
+  auto x_memory_dt = x_memory_desc.data_type();
+
+  //dims of all ones
+  dnnl::memory::dims min_max_dst_dims(x_memory_dims.size(), 1);
+
+  auto min_max_dst_mem_desc = dnnl::memory::desc(min_max_dst_dims, x_memory_dt, sp.GetDnnlFormat(x_memory_dims.size()));
+
+  //max_reduction responsible for producing scale
+  auto max_reduction_d = dnnl::reduction::desc(
+      dnnl::algorithm::reduction_max, x_memory_desc, min_max_dst_mem_desc, 0.f, 0.f);
+  auto min_reduction_d = dnnl::reduction::desc(
+      dnnl::algorithm::reduction_min, x_memory_desc, min_max_dst_mem_desc, 0.f, 0.f);
+
+  //prepare a zero memory, used for adding 0 to data range for min max operation
+  auto zero_mem = dnnl::memory(min_max_dst_mem_desc, eng);
+  WriteZeroToMem(zero_mem);
+
+  //max(x) with 0 added to range -> sub min(x) -> div 255
+  dnnl::primitive_attr max_reduction_attr;
+  {
+    dnnl::post_ops sub_min_div_255;
+    //max(0,reduce_max(x))
+    sub_min_div_255.append_binary(dnnl::algorithm::binary_max, zero_mem.get_desc());
+    //max - min
+    sub_min_div_255.append_binary(dnnl::algorithm::binary_sub, min_max_dst_mem_desc);
+    // /255
+    sub_min_div_255.append_eltwise(1.0f, dnnl::algorithm::eltwise_linear, 1.0f / 255.0f, 0.0f);
+    max_reduction_attr.set_post_ops(sub_min_div_255);
+  }
+
+  //add 0 to reduce min range
+  dnnl::primitive_attr min_reduction_attr;
+  {
+    dnnl::post_ops add_0_to_range;
+    add_0_to_range.append_binary(dnnl::algorithm::binary_min, zero_mem.get_desc());
+    min_reduction_attr.set_post_ops(add_0_to_range);
+  }
+
+  auto max_reduction_pd = dnnl::reduction::primitive_desc(max_reduction_d, max_reduction_attr, eng);
+  auto min_reduction_pd = dnnl::reduction::primitive_desc(min_reduction_d, min_reduction_attr, eng);
+
+  auto max_reduction_prim = dnnl::reduction(max_reduction_pd);
+  auto min_reduction_prim = dnnl::reduction(min_reduction_pd);
+
+  auto y_scale_mem = dnnl::memory(min_max_dst_mem_desc, eng);
+  auto min_reduction_dst_mem = dnnl::memory(min_max_dst_mem_desc, eng);
+
+  std::unordered_map<int, dnnl::memory> min_reduction_args = {{DNNL_ARG_SRC, x_memory}, {DNNL_ARG_DST, min_reduction_dst_mem}};
+  min_reduction_args[DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1] = zero_mem;
+
+  std::unordered_map<int, dnnl::memory> max_reduction_args = {{DNNL_ARG_SRC, x_memory}, {DNNL_ARG_DST, y_scale_mem}};
+  max_reduction_args[DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1] = zero_mem;
+  max_reduction_args[DNNL_ARG_ATTR_MULTIPLE_POST_OP(1) | DNNL_ARG_SRC_1] = min_reduction_dst_mem;
+
+  //compute min first since max_reduction needs min dst as post op arg
+  sp.AddPrimitive(min_reduction_prim, min_reduction_args);
+  sp.AddPrimitive(max_reduction_prim, max_reduction_args);
+
+
+  //prepare y zero point kernel
+  auto y_zero_point_d = dnnl::binary::desc(dnnl::algorithm::binary_div, min_reduction_dst_mem.get_desc(), y_scale_mem.get_desc(), min_reduction_dst_mem.get_desc());
+
+  dnnl::primitive_attr y_zero_point_attr;
+  {
+    y_zero_point_attr.set_scales(DNNL_ARG_SRC_0, 0, {-1.0f});
+    dnnl::post_ops div_saturate_round;
+    div_saturate_round.append_eltwise(1.0f, dnnl::algorithm::eltwise_round, 0.0f, 0.0f);
+    //clip might not be necessary as reorder cast will saturate on lower precision
+    //might still need it as compute y needs saturated zero point already
+    div_saturate_round.append_eltwise(1.0f, dnnl::algorithm::eltwise_clip_v2, 0.0f, 255.0f);
+    y_zero_point_attr.set_post_ops(div_saturate_round);
+  }
+  auto y_zero_point_pd = dnnl::binary::primitive_desc(y_zero_point_d, y_zero_point_attr, eng);
+  auto y_zero_point_prim = dnnl::binary(y_zero_point_pd);
+
+  auto y_zero_point_dst_mem = dnnl::memory(y_zero_point_pd.dst_desc(), eng);
+  std::unordered_map<int, dnnl::memory> y_zero_point_args = {{DNNL_ARG_SRC_0, min_reduction_dst_mem}, {DNNL_ARG_SRC_1, y_scale_mem}, {DNNL_ARG_DST, y_zero_point_dst_mem}};
+
+  sp.AddPrimitive(y_zero_point_prim, y_zero_point_args);
+
+
+  //prepare y kernel
+  //x/y -> round() -> + y_zp -> clip 0,255
+  auto y_d = dnnl::binary::desc(dnnl::algorithm::binary_div, x_memory.get_desc(), y_scale_mem.get_desc(), x_memory.get_desc());
+  dnnl::primitive_attr y_attr;
+  {
+    dnnl::post_ops round_zp_saturate;
+    round_zp_saturate.append_eltwise(1.0f, dnnl::algorithm::eltwise_round, 0.0f, 0.0f);
+    round_zp_saturate.append_binary(dnnl::algorithm::binary_add, y_zero_point_dst_mem.get_desc());
+    //clip might not be necessary as reorder cast will saturate on lower precision
+    round_zp_saturate.append_eltwise(1.0f, dnnl::algorithm::eltwise_clip_v2, 0.0f, 255.0f);
+    y_attr.set_post_ops(round_zp_saturate);
+  }
+  auto y_pd = dnnl::binary::primitive_desc(y_d, y_attr, eng);
+  auto y_prim = dnnl::binary(y_pd);
+
+  auto y_mem = dnnl::memory(y_pd.dst_desc(), eng);
+  std::unordered_map<int, dnnl::memory> y_args = {{DNNL_ARG_SRC_0, x_memory}, {DNNL_ARG_SRC_1, y_scale_mem}, {DNNL_ARG_DST, y_mem}};
+  y_args[DNNL_ARG_ATTR_MULTIPLE_POST_OP(1) | DNNL_ARG_SRC_1] = y_zero_point_dst_mem;
+  sp.AddPrimitive(y_prim, y_args);
+
+
+  //set output y scale
+  sp.SetMemory(node.Output(OUT_Y_SCALE), y_scale_mem, false, true);
+
+  //data type change for y_zp and set memory
+  auto y_zero_point_dst_md_uint8 = ChangeMemoryDescDataType(y_zero_point_dst_mem.get_desc(), dnnl::memory::data_type::u8);
+  auto y_zero_point_dst_mem_uint8 = dnnl::memory(y_zero_point_dst_md_uint8, eng);
+  sp.AddPrimitive(dnnl::reorder(y_zero_point_dst_mem, y_zero_point_dst_mem_uint8), {{DNNL_ARG_FROM, y_zero_point_dst_mem}, {DNNL_ARG_TO, y_zero_point_dst_mem_uint8}});
+  sp.SetMemory(node.Output(OUT_Y_ZP), y_zero_point_dst_mem_uint8, false, true);
+
+  //data type change for y and set memory
+  auto y_md_uint8 = ChangeMemoryDescDataType(y_mem.get_desc(), dnnl::memory::data_type::u8);
+  auto y_mem_uint8 = dnnl::memory(y_md_uint8, eng);
+  sp.AddPrimitive(dnnl::reorder(y_mem, y_mem_uint8), {{DNNL_ARG_FROM, y_mem}, {DNNL_ARG_TO, y_mem_uint8}});
+  sp.SetMemory(node.Output(OUT_Y), y_mem_uint8);
+
+}
+
+//change md to targeted data type of cast op dst
+dnnl::memory::desc DnnlDynamicQuantizeLinear::ChangeMemoryDescDataType(dnnl::memory::desc md, dnnl::memory::data_type dt) {
+  auto dims = md.dims();
+  auto strides = md.data.format_desc.blocking.strides;
+  dnnl::memory::dims strides_vec;
+  for (size_t i = 0; i < dims.size(); i++) {
+    strides_vec.push_back(strides[i]);
+  }
+  auto result = dnnl::memory::desc(dims, dt, strides_vec);
+  return result;
+}
+
+//write zero to memory
+void DnnlDynamicQuantizeLinear::WriteZeroToMem(dnnl::memory& mem) {
+  bool on_gpu = false;
+  if (mem.get_engine().get_kind() == dnnl::engine::kind::gpu) {
+    on_gpu = true;
+  }
+  if (!on_gpu) {
+    auto dst = mem.get_data_handle();
+    size_t size = mem.get_desc().get_size();
+    memset(dst, 0, size);
+  } else {
+    //create a memory on cpu and do a reorder to gpu
+    auto cpu_engine = dnnl::engine(dnnl::engine::kind::cpu, 0);
+    auto cpu_memory = dnnl::memory(mem.get_desc(),cpu_engine);
+    memset(cpu_memory.get_data_handle(),0,cpu_memory.get_desc().get_size());
+    dnnl::stream s{mem.get_engine()};
+    //mem now contains all zero
+    dnnl::reorder(cpu_memory, mem).execute(s, cpu_memory, mem);
+  }
+}
+
+}  // namespace ort_dnnl
+}  // namespace onnxruntime

onnxruntime/core/providers/dnnl/subgraph/dnnl_dynamicquantizelinear.h

@@ -0,0 +1,32 @@
+// Copyright(C) 2021 Intel Corporation
+// Licensed under the MIT License
+
+#pragma once
+#include "dnnl_subgraph.h"
+#include "dnnl_subgraph_primitive.h"
+
+namespace onnxruntime {
+namespace ort_dnnl {
+
+class DnnlDynamicQuantizeLinear {
+ public:
+  enum InputTensors : int {
+    IN_X = 0,  // Input tensor float32
+  };
+
+  enum OutputTensors : int {
+    OUT_Y = 0,        // Quantized output tensor, tensor uint8
+    OUT_Y_SCALE = 1,  // Output scale. It's a scalar, which means a per-tensor/layer quantization, tensor float32
+    OUT_Y_ZP = 2,     // Output zero point. It's a scalar, which means a per-tensor/layer quantization, tensor uint8
+  };
+
+  DnnlDynamicQuantizeLinear() = default;
+  void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node);
+
+ private:
+  void WriteZeroToMem(dnnl::memory& mem);
+  dnnl::memory::desc ChangeMemoryDescDataType(dnnl::memory::desc md, dnnl::memory::data_type dt);
+};
+
+}  // namespace ort_dnnl
+}  // namespace onnxruntime

onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.cc

@@ -6,6 +6,7 @@
 #include "dnnl_batchnorm.h"
 #include "dnnl_binary.h"
 #include "dnnl_conv.h"
+#include "dnnl_dynamicquantizelinear.h"
 #include "dnnl_elementwise.h"
 #include "dnnl_gemm.h"
 #include "dnnl_lrn.h"
@@ -52,6 +53,8 @@ void DnnlSubgraphPrimitive::AddKernels() {
       DnnlBinary().CreatePrimitive(*this, node);
     } else if (node.OpType() == "Conv") {
       DnnlConv().CreatePrimitive(*this, node);
+    } else if (node.OpType() == "DynamicQuantizeLinear") {
+      DnnlDynamicQuantizeLinear().CreatePrimitive(*this, node);
     } else if (elementwise_ops.count(node.OpType())) {
       DnnlElementwise().CreatePrimitive(*this, node);
     } else if (node.OpType() == "Gemm") {
@@ -144,6 +147,7 @@ void DnnlSubgraphPrimitive::Compile(const std::unordered_map<std::string, OnnxTe
   net_.clear();
   net_args_.clear();
   reshapes_.clear();
+  scalar_outputs_.clear();
   //initializer should not be cleared upon recompile
   //initializers_.clear();
 
@@ -297,10 +301,13 @@ bool DnnlSubgraphPrimitive::HasMemory(std::string memory_name, dnnl::memory::des
   return false;
 }
 
-void DnnlSubgraphPrimitive::SetMemory(DnnlTensor tensor, dnnl::memory mem, bool always_copy_output) {
+void DnnlSubgraphPrimitive::SetMemory(DnnlTensor tensor, dnnl::memory mem, bool always_copy_output, bool is_scalar) {
   if (always_copy_output) {
     outputs_are_always_copied_.insert(tensor.Name());
   }
+  if (is_scalar) {
+    scalar_outputs_.insert(tensor.Name());
+  }
   SetMemory(tensor.Name(), mem);
 }
 
@@ -513,6 +520,10 @@ onnxruntime::common::Status DnnlSubgraphPrimitive::Predict(const std::unordered_
   return Status::OK();
 }
 
+bool DnnlSubgraphPrimitive::IsScalarOutput(const std::string& name) {
+  return Contains(scalar_outputs_,name);
+}
+
 dnnl::memory::desc DnnlSubgraphPrimitive::GetOutputInfo(std::string name) {
   if (Contains(outputs_, name)) {
     return outputs_.at(name).get_desc();

onnxruntime/core/providers/dnnl/subgraph/dnnl_subgraph_primitive.h

@@ -55,11 +55,13 @@ class DnnlSubgraphPrimitive {
   dnnl::memory GetMemory(const DnnlTensor& tensor);
   dnnl::memory GetMemory(const DnnlTensor& tensor, dnnl::memory::desc mem_desc, dnnl::engine eng);
   //set memory to a tensor (output)
-  // if always_copy_output is true a copy of the memory will be made when the output is leaving the subgraph.
-  void SetMemory(DnnlTensor tensor, dnnl::memory mem, bool always_copy_output = false);
+  //if always_copy_output is true a copy of the memory will be made when the output is leaving the subgraph.
+  //is_scalar is true to indicate a scalar output in order to allocate the correct onnxruntime output buffer
+  void SetMemory(DnnlTensor tensor, dnnl::memory mem, bool always_copy_output = false, bool is_scalar = false);
   void SetMemory(std::string memory_name, dnnl::memory mem);
   void SetInitializer(std::string memory_name, dnnl::memory mem);
   dnnl::memory::desc GetOutputInfo(std::string name);
+  bool IsScalarOutput(const std::string& name);
   bool IsDynamic();
   OrtMutex& GetMutex() { return mutex_; }
 
@@ -87,6 +89,7 @@ class DnnlSubgraphPrimitive {
   std::vector<std::unordered_map<int, dnnl::memory>> net_args_;
 
   std::vector<std::pair<dnnl::memory, dnnl::memory>> reshapes_;
+  std::unordered_set<std::string> scalar_outputs_;
 
   ort_dnnl::DnnlSubgraph* subgraph_;