Repository: microsoft/onnxruntime
File: onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc (C++)

Code before:
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/optimizer/qdq_transformer/weight_bias_quantization.h"

#include "core/common/common.h"
#include "core/util/qmath.h"
#include "core/graph/graph_utils.h"
#include "core/graph/graph_viewer.h"
#include "core/optimizer/initializer.h"
#include "core/optimizer/utils.h"
#include "core/optimizer/qdq_transformer/qdq_util.h"

namespace onnxruntime {

/**
 * Checks whether or not the output path from a given node leads to a `QuantizeLinear` node,
 * with type and shape are preserved and no branching in between.
 *
 * @param node The starting node to check the output path from.
 * @param graph The graph containing the nodes.
 *
 * @return true if the path exist, false otherwise.
 */
static bool IsTypePreserveNoBranchPathToQuantizeLinear(const Node& node, const Graph& graph) {
  const Node* current = &node;
  while (true) {
    // Conv / ConvTranspose / Gemm should produce single output
    if (current->OutputDefs().empty()) {
      return false;
    }
    if (current->OutputDefs().size() != 1) {
      return false;
    }
    const std::vector<const Node*>& consumers = graph.GetConsumerNodes(current->OutputDefs()[0]->Name());
    // Branching or no consumer: not eligible
    if (consumers.size() != 1) {
      return false;
    }
    const Node* consumer = consumers[0];
    if (consumer->OpType() == QDQ::QOpName) {
      return true;
    }
    // Must be single-input, single-output op
    const auto& i_defs = consumer->InputDefs();
    const auto& o_defs = consumer->OutputDefs();
    if (i_defs.size() != 1 || o_defs.size() != 1) {
      return false;
    }
    // Element type must match

Review comment (performance):
Instead of making this generically handle any operator type between the Conv and the Q (e.g., `Conv -> Any* -> Q`), could we instead limit this explicitly to only `Clip` and `Relu` operator types?

The reason I originally added the check for the final Q is to ensure that this optimizer does not generate a non-valid QDQ node unit as defined by the rules here: https://github.com/microsoft/onnxruntime/blob/5d77b73c4ea102f5d4ea653e5c6461244f553def/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc#L361. These rules only allow a Relu/Clip to appear between a Conv (or Gemm) and the output Q. Any other operator type would result in an invalid QDQ node unit.

Also, this was originally added to fix certain customer models. We should ensure that still works.


Code after:
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/optimizer/qdq_transformer/weight_bias_quantization.h"

#include "core/common/common.h"
#include "core/util/qmath.h"
#include "core/graph/graph_utils.h"
#include "core/graph/graph_viewer.h"
#include "core/optimizer/initializer.h"
#include "core/optimizer/utils.h"
#include "core/optimizer/qdq_transformer/qdq_util.h"

namespace onnxruntime {

/**
 * Checks whether or not the output path from a given node leads to a QuantizeLinear op, optionally, with no
 * branching ReLU or Clip op in between. See also: NodeGroupSelector::GetQDQSelection() in qdq_selectors.cc.
 *
 * @param node The starting node to check the output path from.
 * @param graph The graph containing the nodes.
 *
 * @return true if the path exist, false otherwise.
 */
static bool IsNoBranchPathToQuantizeLinear(const Node& node, const Graph& graph) {
  const Node* current = &node;
  while (true) {
    // Conv / ConvTranspose / Gemm produces single output
    if (current->OutputDefs().size() != 1) {
      return false;
    }
    const std::vector<const Node*>& consumers = graph.GetConsumerNodes(current->OutputDefs()[0]->Name());
    // Branching or no consumer: not eligible
    if (consumers.size() != 1) {
      return false;
    }
    const Node* consumer = consumers[0];
    if (consumer->OpType() == QDQ::QOpName) {
      return true;
    }
    // Allow ReLU or Clip, see also: NodeGroupSelector::GetQDQSelection() in qdq_selectors.cc.
    if (consumer->OpType() != "Relu" && consumer->OpType() != "Clip") {
      return false;
    }
    current = consumer;
  }
}

Status WeightBiasQuantization::ApplyImpl(Graph& graph, bool& modified, int graph_level,
                                         const logging::Logger& logger) const {