UnifyDTypeScale

参考：tvm/src/relay/quantize/realize.cc

float ChooseDomScale(const std::vector<const QRealizeIntExprNode*>& nptrs) {
  if (nptrs.size() == 2) {
    // x = a * s1, y = b * s2
    // x + y = (a * s1 / s2 + b) * s2, if s1 > s2
    //       = (a + b * s2 / s1) * s1, if s2 > s1
    float s1 = GetScalarFromConstant<float>(nptrs[0]->dom_scale);
    float s2 = GetScalarFromConstant<float>(nptrs[1]->dom_scale);
    return s1 > s2 ? s2 : s1;
  } else {
    const QConfig& cfg = QConfig::Current();
    float scale = cfg->global_scale;
    return scale / std::pow(2.0, cfg->nbit_activation - 1);
  }
}

这段代码定义了一个名为 ChooseDomScale 的函数，用于选择两个节点中较小的一个作为它们的共同量化比例。

函数接收 QRealizeIntExprNode 指针的向量 nptrs 作为参数。如果向量的大小为 2，则根据两个节点的量化比例计算它们的和，并返回较小的那个比例。具体来说，如果 s1 > s2，则返回 s2；否则返回 s1。

如果向量的大小不为 2，则获取当前的量化配置（QConfig::Current()），并返回全局比例除以 2 的 cfg->nbit_activation - 1 次方的结果。

/* \brief Unify the dom scale of arguments */
Array<Expr> UnifyDTypeScale(const Array<Expr>& ref_args, const Array<Expr>& args,
                            DataType* dtype_ptr, Expr* scale_ptr,
                            DataType dtype = DataType::Void()) {
  static const Op& simulated_quantize = Op::Get("relay.op.annotation.simulated_quantize");
  const QConfig& cfg = QConfig::Current();

  std::vector<const QRealizeIntExprNode*> nptrs;
  Array<Expr> ret;
  for (auto arg : args) {
    const auto* nptr = arg.as<QRealizeIntExprNode>();
    ICHECK(nptr);
    nptrs.push_back(nptr);
    ret.push_back(nptr->data);
  }

  // unify the data type
  ICHECK_EQ(ref_args.size(), args.size());

  if (dtype.is_void()) {
    if (ret.size() == 2 && nptrs[1]->dtype == cfg->dtype_input) {
      dtype = cfg->dtype_input;
    } else {
      dtype = cfg->dtype_activation;
    }
  }

  for (size_t i = 0; i < ret.size(); ++i) {
    auto ref_arg = ref_args[i].as<CallNode>();
    if (nptrs[i]->dtype != dtype) {
      ret.Set(i, Cast(ret[i], dtype));
    } else if (ref_arg && ref_arg->op.same_as(simulated_quantize) &&
               ref_arg->attrs.as<SimulatedQuantizeAttrs>()->kind == kQInput) {
      auto new_arg = Cast(ret[i], cfg->dtype_input);
      new_arg = StopFusion(new_arg);
      ret.Set(i, Cast(new_arg, dtype));
    }
  }

  // unify the dom_scale
  float s = ChooseDomScale(nptrs);
  Expr dom_scale = MakeConstantScalar(DataType::Float(32), s);
  for (size_t i = 0; i < ret.size(); ++i) {
    float cur_s = GetScalarFromConstant<float>(nptrs[i]->dom_scale);
    ret.Set(i, MulAndDiv(ret[i], cur_s, s, dtype, ref_args[i]->type_as<TensorTypeNode>()->shape));
  }

  *dtype_ptr = dtype;
  *scale_ptr = dom_scale;
  return ret;
}

这段代码定义了一个名为UnifyDTypeScale的函数，用于统一参数的数据类型和量化比例。

函数接收5个参数：

• ref_args：参考参数；

• args：需要处理的参数；

• dtype_ptr：指向数据类型的指针；

• scale_ptr：指向量化比例的指针；

• dtype：默认为空的数据类型。

函数首先获取当前的量化配置（QConfig::Current()），然后遍历args中的每个参数，将其转换为QRealizeIntExprNode类型，并将其添加到nptrs向量中。同时，将每个参数的数据部分添加到ret数组中。

接下来，函数检查是否需要统一数据类型。如果dtype为空，则根据ret的大小和cfg->dtype_input的值来设置dtype。否则，使用给定的dtype值。

然后，函数遍历ret中的每个元素，并根据以下条件进行转换：

• 如果当前参数的数据类型与dtype不同，则将其转换为dtype类型；

• 如果当前参数是模拟量化节点且其属性为输入类型，则将其转换为cfg->dtype_input类型，并停止融合操作。

最后，函数调用ChooseDomScale函数来选择两个节点中较小的一个作为它们的共同量化比例，并将结果存储在dom_scale变量中。接着，遍历ret中的每个元素，将其乘以当前比例和共同比例之间的比值，并将结果存储回ret数组中。

最后，函数将dtype和dom_scale分别存储到dtype_ptr和scale_ptr指向的位置，并返回ret数组。