[oneDNN] Improve DequantizeLinear operator performance. (#12611)

* Detect when ZeroPoint = 0 and avoid sub op.

* Added tests to verify constant initializer behaviour.

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

@@ -25,6 +25,18 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
 
   // Check if scale and zp are scalars
   bool isScalar = sp.IsScalar(node.Input(IN_X_SCALE));
+  // Check if zp is needed
+  bool isZeroPointUseful = false;
+  if (node.Input(IN_X_ZERO_POINT).Exists()) {
+    // If zp exists then it's needed
+    isZeroPointUseful = true;
+    // If it's constant then we can evaluate if zp == 0
+    if (node.Input(IN_X_ZERO_POINT).IsConstant()) {
+      // if zp == 0 then isZeroPointUseful = false; else isZeroPointUseful = true
+      auto mem = sp.GetMemory(node.Input(IN_X_ZERO_POINT));
+      isZeroPointUseful = isZeroPointNonZero(&mem);
+    }
+  }
 
   // Get the x and scale mem
   auto x_mem = sp.GetMemory(node.Input(IN_X));
@@ -43,9 +55,13 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
   if (axis < 0) {
     axis += x_dims;
   }
-  // If scale is a vector, add padding for broadcasting
-  if (!isScalar) {
-    Padd(&x_scale_md, static_cast<uint64_t>(axis + 1), x_dims);
+  // Prepare the scale to prevent broacasting errors
+  if (isScalar) {
+    // For scalar scale
+    Padd(&x_scale_md, x_dims, false);
+  } else {
+    // For N-D scale
+    Padd(&x_scale_md, static_cast<uint64_t>(axis) + 1, x_dims);
   }
 
   // Create dst mem
@@ -53,8 +69,7 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
   dnnl::memory dst_mem;
 
   // If zero point exists and we are NOT dequantizing int32, then substract zp from x and scale
-  if (node.Input(IN_X_ZERO_POINT).Exists() &&
-      (x_mem.get_desc().data_type() != dnnl::memory::data_type::s32)) {
+  if (isZeroPointUseful && (x_mem.get_desc().data_type() != dnnl::memory::data_type::s32)) {
     // Get Zero point
     auto x_zp_mem = sp.GetMemory(node.Input(IN_X_ZERO_POINT));
     // Get mds for operands
@@ -66,16 +81,18 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
       Padd(&x_zp_md, x_dims, false);
     } else {
       // For N-D zp
-      Padd(&x_zp_md, static_cast<uint64_t>(axis) + 1, x_md.dims().size());
+      Padd(&x_zp_md, static_cast<uint64_t>(axis) + 1, x_dims);
     }
 
     // Create binary desc
     auto binary_d = dnnl::binary::desc(dnnl::algorithm::binary_sub, x_md, x_zp_md, dst_md);
     // Add post op scale
-    dnnl::post_ops binary_ops;
     dnnl::primitive_attr binary_attr;
-    binary_ops.append_binary(dnnl::algorithm::binary_mul, x_scale_md);
-    binary_attr.set_post_ops(binary_ops);
+    {
+      dnnl::post_ops binary_ops;
+      binary_ops.append_binary(dnnl::algorithm::binary_mul, x_scale_md);
+      binary_attr.set_post_ops(binary_ops);
+    }
     // Add post op to scale result
     auto binary_pd = dnnl::binary::primitive_desc(binary_d, binary_attr, dnnl_engine);
     // Move to GPU if available
@@ -99,7 +116,6 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
     dst_mem = dnnl::memory(binary_pd.dst_desc(), dnnl_engine);
     auto binary_prim = dnnl::binary(binary_pd);
 
-    // We recycle the x_mem
     sp.AddPrimitive(binary_prim, {{DNNL_ARG_SRC_0, x_mem},
                                   {DNNL_ARG_SRC_1, x_scale_mem},
                                   {DNNL_ARG_DST, dst_mem}});
@@ -113,6 +129,25 @@ void DnnlDequantizeLinear::CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode&
   }
 }
 
+bool DnnlDequantizeLinear::isZeroPointNonZero(dnnl::memory* zp_mem) {
+  // Because zp will always be int8, uint8 or int32, this cast is always valid
+  auto zp_data = static_cast<uint8_t*>(zp_mem->get_data_handle());
+  //  Adjust the iteration num
+  auto topline = zp_mem->get_desc().dims().size();
+  if (zp_mem->get_desc().data_type() == dnnl::memory::data_type::s32) {
+    topline *= 4;
+  }
+  // ZP is either a scalar or a 1-D vector so iterate over all the dimensions
+  // and search for a zp != 0
+  for (size_t i = 0; i < topline; ++i) {
+    if (zp_data[i] != 0) {
+      return true;
+    }
+  }
+  // If ZP is full of zeros then it is not needed
+  return false;
+}
+
 void DnnlDequantizeLinear::Padd(dnnl::memory::desc* target_md, size_t front_pad, size_t back_pad) {
   // Pads an input to broadcast the op correctly
   auto target_dims = target_md->dims();

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

@@ -24,6 +24,7 @@ class DnnlDequantizeLinear {
   void CreatePrimitive(DnnlSubgraphPrimitive& sp, DnnlNode& node);
 
  private:
+  bool isZeroPointNonZero(dnnl::memory* zp_mem);
   int64_t GetAxis(DnnlNode& node, size_t x_dims);
   void Padd(dnnl::memory::desc* target, size_t front_pad, size_t back_pad);
   void ValidateDims(DnnlSubgraphPrimitive& sp, DnnlNode& node);

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

@@ -9,7 +9,7 @@ namespace ort_dnnl {
 
 DnnlTensor DnnlNode::empty_tensor_ = DnnlTensor("");
 
-DnnlTensor::DnnlTensor(const NodeArg* arg) {
+DnnlTensor::DnnlTensor(const NodeArg* arg, bool isConstantInitializer) {
   if (!arg || !arg->Exists()) {
     tensor_name_ = "";
   } else {
@@ -20,6 +20,7 @@ DnnlTensor::DnnlTensor(const NodeArg* arg) {
   arg_type_ = arg->Type();
   arg_type_proto_ = ONNX_NAMESPACE::TypeProto::Create();
   arg_type_proto_->copy_from(arg->TypeAsProto());
+  isConstant_ = isConstantInitializer;
 }
 
 DnnlTensor::DnnlTensor(std::string name) {
@@ -124,6 +125,10 @@ bool DnnlTensor::IsDynamic() {
   return false;
 }
 
+bool DnnlTensor::IsConstant() {
+  return isConstant_;
+}
+
 bool DnnlTensor::Exists() {
   return !(tensor_name_ == "");
 }
@@ -355,7 +360,9 @@ void DnnlSubgraph::Build(const GraphViewer& graph_viewer) {
     for (auto input : node->InputDefs()) {
       if (input && input->Exists() && input->Name() != "") {
         if (!dnnl_tensors_.count(input->Name())) {
-          dnnl_tensors_[input->Name()] = std::make_unique<DnnlTensor>(input);
+          dnnl_tensors_[input->Name()] = 
+            std::make_unique<DnnlTensor>(input,
+                                         graph_viewer.IsConstantInitializer(input->Name(), true));
         }
         dnnl_tensors_[input->Name()]->AddConsumer(DnnlNodeArg(dnnl_node, index, false));
         inputs.push_back(dnnl_tensors_[input->Name()].get());

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

@@ -36,16 +36,18 @@ class DnnlNodeArg {
 
 class DnnlTensor {
  public:
-  DnnlTensor(const NodeArg* arg);
+  DnnlTensor(const NodeArg* arg, bool isConstantInitializer = false);
   DnnlTensor(std::string name);
   DnnlTensor() = default;
   std::string Name() const;
   dnnl::memory::dims Dim() const;
   dnnl::memory::data_type Type() const;
   dnnl::memory::format_tag Format();
-  //check whether the tensor is dynamic, e.g. contains unspecified dimension
+  // Check whether the tensor is dynamic, e.g. contains unspecified dimension
   bool IsDynamic();
-  //check whether the tensor exsits for optional input output
+  // Check whether the tensor is constant initializer
+  bool IsConstant();
+  // Check whether the tensor exsits for optional input output
   bool Exists();
   std::vector<DnnlNodeArg>& GetConsumers() { return consumers_; };
   DnnlNodeArg& GetProducer() { return producer_; };
@@ -64,6 +66,7 @@ class DnnlTensor {
   //a tensor can have no producer (input.initializer) or no consumer (output for subgraph)
   DnnlNodeArg producer_;
   std::vector<DnnlNodeArg> consumers_;
+  bool isConstant_;
 };
 
 class DnnlNode {

onnxruntime/test/providers/dnnl/math/quantize_linear_ops_test.cc

@@ -0,0 +1,63 @@
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+
+
+namespace onnxruntime {
+namespace test {
+
+#ifdef USE_DNNL
+// The same as the default provider, but in this case with constant initializers to test optimization
+
+TEST(DequantizeLinearOpTest, DNNL_Uint8_ConstantInitializer) {
+  OpTester test("DequantizeLinear", 10);
+  std::vector<int64_t> dims{4};
+  test.AddInput<uint8_t>("x", dims, {0, 3, 128, 255});
+  test.AddInput<float>("x_scale", {}, {2.0f});
+  test.AddInput<uint8_t>("x_zero_point", {}, {128}, true);
+  test.AddOutput<float>("y", dims, {-256.0f, -250.0f, 0.0f, 254.0f});
+  test.Run();
+}
+
+// scalar zero & scale with int8
+TEST(DequantizeLinearOpTest, DNNL_Int8_ConstantInitializer) {
+  OpTester test("DequantizeLinear", 10);
+  std::vector<int64_t> dims{4};
+  test.AddInput<int8_t>("x", dims, {-30, -3, 100, 127});
+  test.AddInput<float>("x_scale", {}, {2.0f});
+  test.AddInput<int8_t>("x_zero_point", {}, {-10}, true);
+  test.AddOutput<float>("y", dims, {-40.0f, 14.0f, 220.0f, 274.0f});
+  // Disable Tensorrt EP due to error:node1_quantize_scale_node: out of bounds channel axis 1. Number of input dimensions is 1.
+  test.Run();
+}
+
+// scalar zero & scale with int8
+TEST(DequantizeLinearOpTest, DNNL_Int32_ConstantInitializer) {
+  OpTester test("DequantizeLinear", 10);
+  std::vector<int64_t> dims{4};
+  test.AddInput<int32_t>("x", dims, {-30, -3, 100, 127});
+  test.AddInput<float>("x_scale", {}, {2.0f}, true);
+  test.AddOutput<float>("y", dims, {-60.f, -6.f, 200.f, 254.f});
+  test.Run();
+}
+
+// 2d inputs
+TEST(DequantizeLinearOpTest, DNNL_2D_ConstantInitializer) {
+  OpTester test("DequantizeLinear", 10);
+  std::vector<int64_t> dims{3, 4};
+  test.AddInput<uint8_t>("X", dims,
+                         {0, 1, 2, 3,
+                          0, 1, 2, 3,
+                          0, 10, 20, 30});
+  test.AddInput<float>("scale", {}, {1.0f});
+  test.AddInput<uint8_t>("zero_point", {}, {0}, true);
+  test.AddOutput<float>("Y", dims,
+                        {0, 1, 2, 3,
+                         0, 1, 2, 3,
+                         0, 10, 20, 30});
+  test.Run();
+}
+
+#endif  // USE_DNNL
+
+}  // namespace test
+}  // namespace onnxruntime