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

Code before:
      if (!weight_scale_shape || weight_scale_shape->dim_size() != 1 || !weight_scale_shape->dim(0).has_dim_value() ||
          weight_scale_shape->dim(0).dim_value() != bias_size) {
        continue;
      }

      const auto& dq_attrs = dq_1->GetAttributes();
      auto attr_it = dq_attrs.find("block_size");
      // Default value of block_size=0 has no significance. Don't skip weight_bias_quantization.
      if (attr_it != dq_attrs.end() && attr_it->second.i() != 0) {
        continue;
      }

      int64_t axis = 1;
      if (auto axis_iter = dq_attrs.find("axis"); axis_iter != dq_attrs.end()) {
        axis = axis_iter->second.i();
      }

      int64_t expected_axis = 0;
      if (node.OpType() == "Gemm") {
        int64_t transB = 0;
        const auto& gemm_attrs = node.GetAttributes();
        if (auto trans_b_iter = gemm_attrs.find("transB"); trans_b_iter != gemm_attrs.end()) {
          transB = trans_b_iter->second.i();
        }
        expected_axis = transB == 0 ? 1 : 0;
        axis = (axis == -1) ? 1 : axis;  // normalize for comparison
      }

      if (axis != expected_axis) {
        continue;
      }
    }

    NodeArg* weight_scale_arg = nullptr;
    if (!dq_1) {
      Initializer initializer(graph, *weight_proto, graph.ModelPath());
      const float* weight_data = initializer.data<float>();

      // Quantize float32 weight to int8_t (per-tensor, symmetric).
      // int8_t quantization of input[1] works with input[0] of all types.
      float scale;
      int8_t zp;
      GetQuantizationParameter(weight_data, static_cast<int64_t>(initializer.size()), scale, zp, nullptr);

      // Weight scale initializer.
      ONNX_NAMESPACE::TensorProto weight_scale_proto;
      weight_scale_proto.set_name(graph.GenerateNodeArgName(node.Name() + "_weight_scale"));
      weight_scale_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
      weight_scale_proto.mutable_float_data()->Add(scale);
      weight_scale_arg = &graph_utils::AddInitializer(graph, weight_scale_proto);


Review comment (suggestion):
can the normalization be done more generally like this?

https://github.com/microsoft/onnxruntime/blob/88f2652dee465faa485cf68cb39c4489f28945bb/onnxruntime/core/providers/common.h#L24-L35

maybe it can also move out of this Gemm-only block.

Code after:
      auto weight_scale_shape = dq_1->InputDefs()[1]->Shape();
      if (!weight_scale_shape || weight_scale_shape->dim_size() != 1 || !weight_scale_shape->dim(0).has_dim_value() ||
          weight_scale_shape->dim(0).dim_value() != bias_size) {
        continue;
      }

      const auto& dq_attrs = dq_1->GetAttributes();
      auto attr_it = dq_attrs.find("block_size");
      // Default value of block_size=0 has no significance. Don't skip weight_bias_quantization.
      if (attr_it != dq_attrs.end() && attr_it->second.i() != 0) {
        continue;
      }

      int64_t axis = 1;
      if (auto axis_iter = dq_attrs.find("axis"); axis_iter != dq_attrs.end()) {
        axis = axis_iter->second.i();
        const ONNX_NAMESPACE::TensorShapeProto* weight_shape = weight_arg->Shape();
        if (!weight_shape && dq_1->InputDefs()[0]) {
          weight_shape = dq_1->InputDefs()[0]->Shape();
        }
        if (axis < 0 && !weight_shape) {
          continue;
        }
        axis = HandleNegativeAxis(axis, weight_shape->dim_size());
      }

      int64_t expected_axis = 0;
      if (node.OpType() == "Gemm") {
        int64_t transB = 0;
        const auto& gemm_attrs = node.GetAttributes();
        if (auto trans_b_iter = gemm_attrs.find("transB"); trans_b_iter != gemm_attrs.end()) {
          transB = trans_b_iter->second.i();
        }
        expected_axis = transB == 0 ? 1 : 0;
      }

      if (axis != expected_axis) {
        continue;
      }
    }

    NodeArg* weight_scale_arg = nullptr;
    if (!dq_1) {
      Initializer initializer(graph, *weight_proto, graph.ModelPath());
      const float* weight_data = initializer.data<float>();

      // Quantize float32 weight to int8_t (per-tensor, symmetric).
      // int8_t quantization of input[1] works with input[0] of all types.
      float scale;
      int8_t zp;
      GetQuantizationParameter(weight_data, static_cast<int64_t>(initializer.size()), scale, zp, nullptr);