Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix tf.raw_ops. QuantizedMatMul vulnerability from non scalar min/max…
… a/b arguments.

PiperOrigin-RevId: 465364257
  • Loading branch information
pak-laura authored and tensorflower-gardener committed Aug 4, 2022
1 parent a0f6cdc commit aca766a
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 20 deletions.
15 changes: 15 additions & 0 deletions tensorflow/core/kernels/quantized_matmul_op.cc
Expand Up @@ -20,11 +20,14 @@ limitations under the License.
#define GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK
#include "public/gemmlowp.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/kernels/meta_support.h"
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/kernels/reference_gemm.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/errors.h"

namespace tensorflow {

Expand Down Expand Up @@ -75,9 +78,21 @@ class QuantizedMatMulOp : public OpKernel {
void Compute(OpKernelContext* context) override {
const Tensor& a = context->input(0);
const Tensor& b = context->input(1);
OP_REQUIRES(context, TensorShapeUtils::IsScalar(context->input(2).shape()),
errors::InvalidArgument("min_a must be a scalar, but got shape",
context->input(2).shape()));
const float min_a = context->input(2).flat<float>()(0);
OP_REQUIRES(context, context->input(3).NumElements() == 1,
errors::InvalidArgument("max_a must be a scalar, but got shape",
context->input(3).shape()));
const float max_a = context->input(3).flat<float>()(0);
OP_REQUIRES(context, context->input(4).NumElements() == 1,
errors::InvalidArgument("min_b must be a scalar, but got shape",
context->input(4).shape()));
const float min_b = context->input(4).flat<float>()(0);
OP_REQUIRES(context, context->input(5).NumElements() == 1,
errors::InvalidArgument("max_b must be a scalar, but got shape",
context->input(5).shape()));
const float max_b = context->input(5).flat<float>()(0);

// Make sure that we have valid quantization ranges for the input buffers.
Expand Down
78 changes: 58 additions & 20 deletions tensorflow/core/kernels/quantized_matmul_op_test.cc
Expand Up @@ -62,10 +62,10 @@ TEST_F(QuantizedMatMulTest, Small_NoParams) {
// | 15 | 16 | 17 | 18 |
AddInputFromArray<quint8>(TensorShape({3, 4}),
{7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18});
AddInputFromArray<float>(TensorShape({1}), {0});
AddInputFromArray<float>(TensorShape({1}), {255.0f});
AddInputFromArray<float>(TensorShape({1}), {0});
AddInputFromArray<float>(TensorShape({1}), {255.0f});
AddInputFromArray<float>(TensorShape({}), {0});
AddInputFromArray<float>(TensorShape({}), {255.0f});
AddInputFromArray<float>(TensorShape({}), {0});
AddInputFromArray<float>(TensorShape({}), {255.0f});

TF_ASSERT_OK(RunOpKernel());
// Here are the results we expect, from hand calculations:
Expand Down Expand Up @@ -118,10 +118,10 @@ TEST_F(QuantizedMatMulTest, VerySmall_WithParams) {
// The B matrix is:
// | 1 |
AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
AddInputFromArray<float>(TensorShape({1}), {-12.0f});
AddInputFromArray<float>(TensorShape({1}), {243.0f});
AddInputFromArray<float>(TensorShape({1}), {1.0f});
AddInputFromArray<float>(TensorShape({1}), {256.0f});
AddInputFromArray<float>(TensorShape({}), {-12.0f});
AddInputFromArray<float>(TensorShape({}), {243.0f});
AddInputFromArray<float>(TensorShape({}), {1.0f});
AddInputFromArray<float>(TensorShape({}), {256.0f});
TF_ASSERT_OK(RunOpKernel());
// We're requesting C = A.transposed() * B,
// so we expect to get these results:
Expand Down Expand Up @@ -162,12 +162,50 @@ TEST_F(QuantizedMatMulTest, VerySmall_BadRange) {
// The B matrix is:
// | 1 |
AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
AddInputFromArray<float>(TensorShape({1}), {-12.0f});
AddInputFromArray<float>(TensorShape({1}), {243.0f});
AddInputFromArray<float>(TensorShape({}), {-12.0f});
AddInputFromArray<float>(TensorShape({}), {243.0f});
// Here we set the range so that the min and max are equal, so we expect to
// see an error when we run.
AddInputFromArray<float>(TensorShape({1}), {1.0f});
AddInputFromArray<float>(TensorShape({1}), {1.0f});
AddInputFromArray<float>(TensorShape({}), {1.0f});
AddInputFromArray<float>(TensorShape({}), {1.0f});
EXPECT_EQ(::tensorflow::error::INVALID_ARGUMENT, RunOpKernel().code());
}

// This test multiplies two 1x1 8bit matrices, but sets invalid quantized min
// and max values, so we expect to get an error
TEST_F(QuantizedMatMulTest, VerySmall_BadMinMax) {
// These parameters reflect a typical production usage of eight-bit matmuls
// in an Inception-style network.
const bool transpose_a = true;
const int a_rows = 1;
const int a_cols = 1;
const int b_rows = 1;
const int b_cols = 1;
const bool transpose_b = false;
TF_ASSERT_OK(NodeDefBuilder("quantized_mat_mul_op", "QuantizedMatMul")
.Input(FakeInput(DT_QUINT8))
.Input(FakeInput(DT_QUINT8))
.Input(FakeInput(DT_FLOAT))
.Input(FakeInput(DT_FLOAT))
.Input(FakeInput(DT_FLOAT))
.Input(FakeInput(DT_FLOAT))
.Attr("Toutput", DataTypeToEnum<qint32>::v())
.Attr("transpose_a", transpose_a)
.Attr("transpose_b", transpose_b)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
// The A matrix is:
// | -1 |
AddInputFromArray<quint8>(TensorShape({a_rows, a_cols}), {11});
// The B matrix is:
// | 1 |
AddInputFromArray<quint8>(TensorShape({b_rows, b_cols}), {0});
// Here we set the error of a non scalar min_a value, so we expect to see an
// error when we run.
AddInputFromArray<float>(TensorShape({1}), {2});
AddInputFromArray<float>(TensorShape({}), {243.0f});
AddInputFromArray<float>(TensorShape({}), {1.0f});
AddInputFromArray<float>(TensorShape({}), {256.0f});
EXPECT_EQ(::tensorflow::error::INVALID_ARGUMENT, RunOpKernel().code());
}

Expand Down Expand Up @@ -233,10 +271,10 @@ TEST_F(QuantizedMatMulTest, Small_WithParams) {
3,
6,
});
AddInputFromArray<float>(TensorShape({1}), {-12.0f});
AddInputFromArray<float>(TensorShape({1}), {243.0f});
AddInputFromArray<float>(TensorShape({1}), {0});
AddInputFromArray<float>(TensorShape({1}), {255.0f});
AddInputFromArray<float>(TensorShape({}), {-12.0f});
AddInputFromArray<float>(TensorShape({}), {243.0f});
AddInputFromArray<float>(TensorShape({}), {0});
AddInputFromArray<float>(TensorShape({}), {255.0f});
TF_ASSERT_OK(RunOpKernel());
// We're requesting C = A.transposed() * B,
// so we expect to get these results:
Expand Down Expand Up @@ -326,10 +364,10 @@ TEST_F(QuantizedMatMulTest, Medium_WithParams) {

AddInputFromArray<quint8>(a_quantized.shape(), a_quantized.flat<quint8>());
AddInputFromArray<quint8>(b_quantized.shape(), b_quantized.flat<quint8>());
AddInputFromArray<float>(TensorShape({1}), {a_min});
AddInputFromArray<float>(TensorShape({1}), {a_max});
AddInputFromArray<float>(TensorShape({1}), {b_min});
AddInputFromArray<float>(TensorShape({1}), {b_max});
AddInputFromArray<float>(TensorShape({}), {a_min});
AddInputFromArray<float>(TensorShape({}), {a_max});
AddInputFromArray<float>(TensorShape({}), {b_min});
AddInputFromArray<float>(TensorShape({}), {b_max});
TF_ASSERT_OK(RunOpKernel());

Tensor expected_float(DT_FLOAT, {a_cols, b_cols});
Expand Down

0 comments on commit aca766a

Please sign in to comment.