Commit message: Prevent division by 0 in `QuantizedBiasAdd`.

PiperOrigin-RevId: 370117454
Change-Id: I3804e2ac8dcc6d3afcc92e27853e2325a017ca4d
CVE: CVE-2021-29546
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can trigger an integer division by zero undefined behavior in `tf.raw_ops.QuantizedBiasAdd`. This is because the implementation of the Eigen kernel(https://github.com/tensorflow/tensorflow/blob/61bca8bd5ba8a68b2d97435ddfafcdf2b85672cd/tensorflow/core/kernels/quantization_utils.h#L812-L849) does a division by the number of elements of the smaller input (based on shape) without checking that this is not zero. The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.4, as these are also affected and still in supported range.

Project: tensorflow
File: quantized_bias_add_op.cc

  void Compute(OpKernelContext* context) override {
    const Tensor& input = context->input(0);
    const Tensor& bias = context->input(1);
    const float input_min = context->input(2).flat<float>()(0);
    const float input_max = context->input(3).flat<float>()(0);
    const float bias_min = context->input(4).flat<float>()(0);
    const float bias_max = context->input(5).flat<float>()(0);

    OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input.shape()),
                errors::InvalidArgument("Input tensor must be at least 2D: ",
                                        input.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
                errors::InvalidArgument("Biases must be 1D: ",
                                        bias.shape().DebugString()));
    const auto last_dim = input.shape().dims() - 1;
    OP_REQUIRES(
        context, bias.shape().dim_size(0) == input.shape().dim_size(last_dim),
        errors::InvalidArgument(
            "Must provide as many biases as the last dimension "
            "of the input tensor: ",
            bias.shape().DebugString(), " vs. ", input.shape().DebugString()));

    Tensor* output = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input.shape(), &output));

    float total_min;
    float total_max;

    if (meta::IsSupportedAndEnabled() && std::is_same<T1, quint8>() &&
        std::is_same<T2, quint8>() && std::is_same<T3, qint32>()) {
      auto input_ui8_array = input.flat<quint8>();
      auto bias_ui8_array = bias.flat<quint8>();
      GetOutputMinAndMaxForQuantizedAdd(input_min, input_max, bias_min,
                                        bias_max, &total_min, &total_max);
      meta::QuantizedBiasAdd(context, input_ui8_array.data(),
                             input_ui8_array.size(), bias_ui8_array.data(),
                             bias_ui8_array.size(), input_min, input_max,
                             bias_min, bias_max, total_min, total_max,
                             output->flat<qint32>().data());
    } else {
      QuantizedAddUsingEigen<T1, T2, T3>(
          context->template eigen_device<CPUDevice>(), input, input_min,
          input_max, bias, bias_min, bias_max, output, &total_min, &total_max);
    }

    Tensor* output_min = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(1, {}, &output_min));
    output_min->flat<float>()(0) = total_min;

    Tensor* output_max = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(2, {}, &output_max));
    output_max->flat<float>()(0) = total_max;
  }