Commit message: Fix one division by zero

PiperOrigin-RevId: 369474832
Change-Id: I1082858ed78d9b2e4738ce30b231955973d49e1e
CVE: CVE-2021-29528
CWE: CWE-369
Vulnerability description: TensorFlow is an end-to-end open source platform for machine learning. An attacker can trigger a division by 0 in `tf.raw_ops.QuantizedMul`. This is because the implementation(https://github.com/tensorflow/tensorflow/blob/55900e961ed4a23b438392024912154a2c2f5e85/tensorflow/core/kernels/quantized_mul_op.cc#L188-L198) does a division by a quantity that is controlled by the caller. 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_mul_op.cc

void VectorMultiply<quint8, qint32>(OpKernelContext* context,
                                    const quint8* x_data, int32 offset_x,
                                    const quint8* y_data, int32 offset_y,
                                    int64 num_elements, qint32* output) {
  const uint8x8_t offset_x_8x8 = vmov_n_u8(offset_x);
  const uint8x8_t offset_y_8x8 = vmov_n_u8(offset_y);
  int i;
  // Go through the results in 16-element chunks for NEON acceleration.
  for (i = 0; i < (num_elements - 15); i += 16) {
    // Load the vector inputs.
    const uint8* x_data_ptr = &(x_data->value) + i;
    const uint8x16_t x_8x16 = vld1q_u8(x_data_ptr);
    const uint8* y_data_ptr = &(y_data->value) + i;
    const uint8x16_t y_8x16 = vld1q_u8(y_data_ptr);

    // Break into two sets of vectors so we can do further calculations easily.
    const uint8x8_t x_high_8x8 = vget_high_u8(x_8x16);
    const uint8x8_t x_low_8x8 = vget_low_u8(x_8x16);
    const uint8x8_t y_high_8x8 = vget_high_u8(y_8x16);
    const uint8x8_t y_low_8x8 = vget_low_u8(y_8x16);

    // Subtract off the offset values to get 16-bit results.
    const int16x8_t x_minus_offset_high_16x8 =
        vreinterpretq_s16_u16(vsubl_u8(x_high_8x8, offset_x_8x8));
    const int16x8_t x_minus_offset_low_16x8 =
        vreinterpretq_s16_u16(vsubl_u8(x_low_8x8, offset_x_8x8));
    const int16x8_t y_minus_offset_high_16x8 =
        vreinterpretq_s16_u16(vsubl_u8(y_high_8x8, offset_y_8x8));
    const int16x8_t y_minus_offset_low_16x8 =
        vreinterpretq_s16_u16(vsubl_u8(y_low_8x8, offset_y_8x8));

    // We have to work with 4-wide vectors, so extract them.
    const int16x4_t x_high_high_16x4 = vget_high_s16(x_minus_offset_high_16x8);
    const int16x4_t x_high_low_16x4 = vget_low_s16(x_minus_offset_high_16x8);
    const int16x4_t x_low_high_16x4 = vget_high_s16(x_minus_offset_low_16x8);
    const int16x4_t x_low_low_16x4 = vget_low_s16(x_minus_offset_low_16x8);
    const int16x4_t y_high_high_16x4 = vget_high_s16(y_minus_offset_high_16x8);
    const int16x4_t y_high_low_16x4 = vget_low_s16(y_minus_offset_high_16x8);
    const int16x4_t y_low_high_16x4 = vget_high_s16(y_minus_offset_low_16x8);
    const int16x4_t y_low_low_16x4 = vget_low_s16(y_minus_offset_low_16x8);

    // Perform the multiplication.
    const int32x4_t z_high_high_32x4 =
        vmull_s16(x_high_high_16x4, y_high_high_16x4);
    const int32x4_t z_high_low_32x4 =
        vmull_s16(x_high_low_16x4, y_high_low_16x4);
    const int32x4_t z_low_high_32x4 =
        vmull_s16(x_low_high_16x4, y_low_high_16x4);
    const int32x4_t z_low_low_32x4 = vmull_s16(x_low_low_16x4, y_low_low_16x4);

    // Write out the results.
    int32* output_ptr = &(output->value) + i;
    vst1q_s32(output_ptr + 0, z_low_low_32x4);
    vst1q_s32(output_ptr + 4, z_low_high_32x4);
    vst1q_s32(output_ptr + 8, z_high_low_32x4);
    vst1q_s32(output_ptr + 12, z_high_high_32x4);
  }
  for (; i < num_elements; ++i) {
    output[i] = (static_cast<int32>(x_data[i]) - offset_x) *
                (static_cast<int32>(y_data[i]) - offset_y);
  }
}

Project: tensorflow
File: quantized_mul_op.cc

void ScalarMultiply<quint8, qint32>(OpKernelContext* context,
                                    const quint8* full_input,
                                    int32 full_input_offset, int64 num_elements,
                                    quint8 scalar_input,
                                    int32 scalar_input_offset, qint32* output) {
  const int16 scalar_minus_offset =
      static_cast<int16>(scalar_input) - scalar_input_offset;
  const int16x4_t scalar_minus_offset_16x4 = vmov_n_s16(scalar_minus_offset);
  const uint8x8_t full_input_offset_8x8 = vmov_n_u8(full_input_offset);
  // Go through the results in 16-element chunks for NEON acceleration.
  int i;
  for (i = 0; i < (num_elements - 15); i += 16) {
    // Load the tensor inputs.
    const uint8* full_input_ptr = &(full_input->value) + i;
    const uint8x16_t full_input_8x16 = vld1q_u8(full_input_ptr);

    // Break into two sets of vectors so we can do further calculations
    // easily.
    const uint8x8_t full_input_high_8x8 = vget_high_u8(full_input_8x16);
    const uint8x8_t full_input_low_8x8 = vget_low_u8(full_input_8x16);

    // Subtract off the offset value to get 16-bit results.
    const int16x8_t full_input_minus_offset_high_16x8 = vreinterpretq_s16_u16(
        vsubl_u8(full_input_high_8x8, full_input_offset_8x8));
    const int16x8_t full_input_minus_offset_low_16x8 = vreinterpretq_s16_u16(
        vsubl_u8(full_input_low_8x8, full_input_offset_8x8));

    // We have to work with 4-wide vectors, so extract them.
    const int16x4_t x_high_high_16x4 =
        vget_high_s16(full_input_minus_offset_high_16x8);
    const int16x4_t x_high_low_16x4 =
        vget_low_s16(full_input_minus_offset_high_16x8);
    const int16x4_t x_low_high_16x4 =
        vget_high_s16(full_input_minus_offset_low_16x8);
    const int16x4_t x_low_low_16x4 =
        vget_low_s16(full_input_minus_offset_low_16x8);

    // Perform the multiplication.
    const int32x4_t z_high_high_32x4 =
        vmull_s16(x_high_high_16x4, scalar_minus_offset_16x4);
    const int32x4_t z_high_low_32x4 =
        vmull_s16(x_high_low_16x4, scalar_minus_offset_16x4);
    const int32x4_t z_low_high_32x4 =
        vmull_s16(x_low_high_16x4, scalar_minus_offset_16x4);
    const int32x4_t z_low_low_32x4 =
        vmull_s16(x_low_low_16x4, scalar_minus_offset_16x4);

    // Write out the results.
    int32* output_ptr = &(output->value) + i;
    vst1q_s32(output_ptr + 0, z_low_low_32x4);
    vst1q_s32(output_ptr + 4, z_low_high_32x4);
    vst1q_s32(output_ptr + 8, z_high_low_32x4);
    vst1q_s32(output_ptr + 12, z_high_high_32x4);
  }
  // Finish up any remaining elements that weren't a multiple of 16.
  for (; i < num_elements; ++i) {
    output[i] = (static_cast<int32>(full_input[i]) - full_input_offset) *
                scalar_minus_offset;
  }
}

Project: tensorflow
File: quantized_mul_op.cc

void VectorMultiply(OpKernelContext* context, const T* x_data, int32 offset_x,
                    const T* y_data, int32 offset_y, int64 num_elements,
                    Toutput* output) {
  for (int i = 0; i < num_elements; ++i) {
    output[i] = (static_cast<int32>(x_data[i]) - offset_x) *
                (static_cast<int32>(y_data[i]) - offset_y);
  }
}

Project: tensorflow
File: quantized_mul_op.cc

void VectorTensorMultiply(const T* vector_data, int32 vector_offset,
                          int64 vector_num_elements, const T* tensor_data,
                          int32 tensor_offset, int64 tensor_num_elements,
                          Toutput* output) {
  for (int i = 0; i < tensor_num_elements; ++i) {
    const int64 vector_i = i % vector_num_elements;
    output[i] = (static_cast<int32>(vector_data[vector_i]) - vector_offset) *
                (static_cast<int32>(tensor_data[i]) - tensor_offset);
  }
}

Project: tensorflow
File: quantized_mul_op.cc

void ScalarMultiply(OpKernelContext* context, const T* full_input,
                    int32 full_input_offset, int64 num_elements, T scalar_input,
                    int32 scalar_input_offset, Toutput* output) {
  const int32 scalar_minus_offset =
      static_cast<int32>(scalar_input) - scalar_input_offset;
  for (int i = 0; i < num_elements; ++i) {
    output[i] = (static_cast<int32>(full_input[i]) - full_input_offset) *
                scalar_minus_offset;
  }
}

Project: tensorflow
File: quantized_mul_op.cc

void VectorTensorMultiply<quint8, qint32>(
    const quint8* vector_data, int32 vector_offset, int64 vector_num_elements,
    const quint8* tensor_data, int32 tensor_offset, int64 tensor_num_elements,
    qint32* output) {
  const uint8x8_t offset_x_8x8 = vmov_n_u8(vector_offset);
  const uint8x8_t offset_y_8x8 = vmov_n_u8(tensor_offset);
  CHECK_EQ(0, tensor_num_elements % vector_num_elements);
  for (int base_i = 0; base_i < tensor_num_elements;
       base_i += vector_num_elements) {
    int i = base_i;
    const int end_i = base_i + vector_num_elements;
    // Go through the results in 16-element chunks for NEON acceleration.
    int vector_i;
    for (vector_i = 0; vector_i < (vector_num_elements - 15);
         vector_i += 16, i += 16) {
      // Load the vector inputs.
      const uint8* x_data_ptr = &(vector_data->value) + vector_i;
      const uint8x16_t x_8x16 = vld1q_u8(x_data_ptr);
      const uint8* y_data_ptr = &(tensor_data->value) + i;
      const uint8x16_t y_8x16 = vld1q_u8(y_data_ptr);

      // Break into two sets of vectors so we can do further calculations
      // easily.
      const uint8x8_t x_high_8x8 = vget_high_u8(x_8x16);
      const uint8x8_t x_low_8x8 = vget_low_u8(x_8x16);
      const uint8x8_t y_high_8x8 = vget_high_u8(y_8x16);
      const uint8x8_t y_low_8x8 = vget_low_u8(y_8x16);

      // Subtract off the offset values to get 16-bit results.
      const int16x8_t x_minus_offset_high_16x8 =
          vreinterpretq_s16_u16(vsubl_u8(x_high_8x8, offset_x_8x8));
      const int16x8_t x_minus_offset_low_16x8 =
          vreinterpretq_s16_u16(vsubl_u8(x_low_8x8, offset_x_8x8));
      const int16x8_t y_minus_offset_high_16x8 =
          vreinterpretq_s16_u16(vsubl_u8(y_high_8x8, offset_y_8x8));
      const int16x8_t y_minus_offset_low_16x8 =
          vreinterpretq_s16_u16(vsubl_u8(y_low_8x8, offset_y_8x8));

      // We have to work with 4-wide vectors, so extract them.
      const int16x4_t x_high_high_16x4 =
          vget_high_s16(x_minus_offset_high_16x8);
      const int16x4_t x_high_low_16x4 = vget_low_s16(x_minus_offset_high_16x8);
      const int16x4_t x_low_high_16x4 = vget_high_s16(x_minus_offset_low_16x8);
      const int16x4_t x_low_low_16x4 = vget_low_s16(x_minus_offset_low_16x8);
      const int16x4_t y_high_high_16x4 =
          vget_high_s16(y_minus_offset_high_16x8);
      const int16x4_t y_high_low_16x4 = vget_low_s16(y_minus_offset_high_16x8);
      const int16x4_t y_low_high_16x4 = vget_high_s16(y_minus_offset_low_16x8);
      const int16x4_t y_low_low_16x4 = vget_low_s16(y_minus_offset_low_16x8);

      // Perform the multiplication.
      const int32x4_t z_high_high_32x4 =
          vmull_s16(x_high_high_16x4, y_high_high_16x4);
      const int32x4_t z_high_low_32x4 =
          vmull_s16(x_high_low_16x4, y_high_low_16x4);
      const int32x4_t z_low_high_32x4 =
          vmull_s16(x_low_high_16x4, y_low_high_16x4);
      const int32x4_t z_low_low_32x4 =
          vmull_s16(x_low_low_16x4, y_low_low_16x4);

      // Write out the results.
      int32* output_ptr = &(output->value) + i;
      vst1q_s32(output_ptr + 0, z_low_low_32x4);
      vst1q_s32(output_ptr + 4, z_low_high_32x4);
      vst1q_s32(output_ptr + 8, z_high_low_32x4);
      vst1q_s32(output_ptr + 12, z_high_high_32x4);
    }
    for (; i < end_i; ++i, ++vector_i) {
      output[i] = (static_cast<int32>(vector_data[vector_i]) - vector_offset) *
                  (static_cast<int32>(tensor_data[i]) - tensor_offset);
    }
  }
}

Project: tensorflow
File: quantized_mul_op.cc

  explicit QuantizedMulOp(OpKernelConstruction* context) : OpKernel(context) {}