From 1640ac6ad682027212431e08c247dcf99a005eeb Mon Sep 17 00:00:00 2001 From: Yunlu Li Date: Wed, 18 Dec 2019 16:15:47 -0800 Subject: [PATCH] Add verifier check for sparse tensors. PiperOrigin-RevId: 286286868 Change-Id: I92bb40fb8eff9a16e5b382fec19befd39c93aeaa --- tensorflow/lite/testdata/sparse_tensor.bin | Bin 412 -> 504 bytes tensorflow/lite/testdata/sparse_tensor.json | 14 +- tensorflow/lite/tools/BUILD | 3 + tensorflow/lite/tools/verifier.cc | 132 ++++++++++++++++++- tensorflow/lite/tools/verifier_test.cc | 137 ++++++++++++++++++++ 5 files changed, 275 insertions(+), 11 deletions(-) diff --git a/tensorflow/lite/testdata/sparse_tensor.bin b/tensorflow/lite/testdata/sparse_tensor.bin index d1445ac648065da9918a1ba72ab8b53374273b5e..497ce68a3ace0dd23c01d95beef93be0b23347cb 100644 GIT binary patch delta 188 zcmbQk{DXOd0^^2>isnEfQo25Xfq}sVh%JDa1&A4eSOAC_7`PaCfFuh82ap93gUSB~ z0+2cw4P=7#!DtB(gF%Qv4yczMXo>)X2#_xU6afRyJ3tDg{{RES#NBr_KunNT96*vC XNOAz#K%fA!jtz+260=kN3m6yxiu4wE delta 85 zcmeytJcoIL0%OHQMRUd%6C +#include + #include "absl/container/flat_hash_set.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/string_util.h" @@ -110,6 +113,115 @@ bool VerifyStringTensorBuffer(const Tensor& tensor, const Buffer& buffer, return true; } +// The sparsity parameter defines a tree structure to map each non-zero element +// stored in the flattened buffer back to its index in the conceptual dense +// tensor. +// Traverse the tree level by level, count total number of elements, and +// validate the sparsity parameters along the way. +absl::optional VerifyAndCountElements( + const SparsityParameters& sparsity, const std::vector& dim_sizes) { + const int total_level = sparsity.traversal_order()->size(); + uint64_t num_elements = 1; + for (int i = 0; i < total_level; i++) { + const int original_dim = sparsity.traversal_order()->Get(i); + const auto* dim_metadata = sparsity.dim_metadata()->Get(i); + if (dim_metadata->format() == DimensionType_DENSE) { + if (dim_metadata->dense_size() != dim_sizes[original_dim]) { + return absl::nullopt; + } + + // Each index in a dense dimension is stored implicitly. + num_elements *= dim_metadata->dense_size(); + } else { + const auto* array_segments = dim_metadata->array_segments(); + const auto* array_indices = dim_metadata->array_indices(); + if (array_segments == nullptr || array_indices == nullptr) { + return absl::nullopt; + } + + for (int j = 0; j < array_segments->size() - 1; j++) { + if (array_segments->Get(j) < 0 || array_segments->Get(j + 1) < 0 || + array_segments->Get(j) > array_segments->Get(j + 1)) { + return absl::nullopt; + } + } + + if (num_elements != array_segments->size() - 1) { + return absl::nullopt; + } + + if (array_indices->size() != + array_segments->Get(array_segments->size() - 1)) { + return absl::nullopt; + } + + for (int j = 0; j < array_indices->size(); j++) { + if (array_indices->Get(j) < 0 || + array_indices->Get(j) >= dim_sizes[original_dim]) { + return absl::nullopt; + } + } + + // Need to reset num_elements when seeing a sparse dimension. + num_elements = array_indices->size(); + } + } + + return num_elements; +} + +absl::optional VerifyAndCountSparseElements(const Tensor& tensor) { + const auto* sparsity = tensor.sparsity(); + if (sparsity->traversal_order() == nullptr || + sparsity->dim_metadata() == nullptr) { + return absl::nullopt; + } + + const int total_dims = sparsity->traversal_order()->size(); + + if (sparsity->dim_metadata()->size() != total_dims) { + return absl::nullopt; + } + + const int block_rank = total_dims - tensor.shape()->size(); + if (block_rank > 0 && (sparsity->block_map() == nullptr || + sparsity->block_map()->size() != block_rank)) { + return absl::nullopt; + } + + // For a n-dimensional tensor (d0, ..., dn-1) with k-dimensional block (dn, + // ..., dn+k-1), the expanded_dim_sizes holds the size of each dimension in + // the order of (d0, ..., dn-1, dn, ..., dn+k-1), not the traversal order. + // For example, a 4x4 tensor with 2x2 block has expanded_dim_sizes = {2, 2, 2, + // 2}. + std::vector expanded_dim_sizes; + expanded_dim_sizes.resize(total_dims); + const int original_rank = tensor.shape()->size(); + // First go through the original tensor dimensions, populate their sizes. + for (int i = 0; i < original_rank; i++) { + expanded_dim_sizes[i] = tensor.shape()->Get(i); + } + // Then go through the block dimensions, and + // 1. populate block dimension size. + // 2. block_map[i] has the original dimension that block dimension i maps + // to. Divide the size of the original dimension by the size of the ith + // block dimension. + for (int i = 0; i < block_rank; i++) { + int original_block_dim = + sparsity->traversal_order()->Get(i + original_rank); + int block_dim_size = + sparsity->dim_metadata()->Get(i + original_rank)->dense_size(); + if (block_dim_size == 0) { + return absl::nullopt; + } + + expanded_dim_sizes[original_block_dim] = block_dim_size; + expanded_dim_sizes[sparsity->block_map()->Get(i)] /= block_dim_size; + } + + return VerifyAndCountElements(*sparsity, expanded_dim_sizes); +} + // Verifies numeric tensor has legit buffer. bool VerifyNumericTensorBuffer(const Tensor& tensor, const Buffer& buffer, ErrorReporter* error_reporter) { @@ -118,14 +230,30 @@ bool VerifyNumericTensorBuffer(const Tensor& tensor, const Buffer& buffer, // Empty tensor. Avoid further checks. return true; } - for (int dim : *tensor.shape()) { - bytes_required *= dim; + if (tensor.sparsity() != nullptr) { + const auto num_elements = VerifyAndCountSparseElements(tensor); + if (!num_elements.has_value()) { + ReportError(error_reporter, "Tensor %s has invalid sparsity parameters", + tensor.name()->c_str()); + return false; + } + bytes_required = num_elements.value(); if (bytes_required > UINT_MAX) { ReportError(error_reporter, "Tensor %s dimension overflow", tensor.name()->c_str()); return false; } + } else { + for (int dim : *tensor.shape()) { + bytes_required *= dim; + if (bytes_required > UINT_MAX) { + ReportError(error_reporter, "Tensor %s dimension overflow", + tensor.name()->c_str()); + return false; + } + } } + switch (tensor.type()) { case TensorType_FLOAT32: bytes_required *= sizeof(float); diff --git a/tensorflow/lite/tools/verifier_test.cc b/tensorflow/lite/tools/verifier_test.cc index be3a06f2eb2..a945e980030 100644 --- a/tensorflow/lite/tools/verifier_test.cc +++ b/tensorflow/lite/tools/verifier_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/lite/tools/verifier.h" +#include #include #include @@ -25,6 +26,8 @@ limitations under the License. #include "tensorflow/lite/allocation.h" #include "tensorflow/lite/core/api/flatbuffer_conversions.h" #include "tensorflow/lite/error_reporter.h" +#include "tensorflow/lite/model.h" +#include "tensorflow/lite/mutable_op_resolver.h" #include "tensorflow/lite/op_resolver.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/testing/util.h" @@ -33,6 +36,11 @@ limitations under the License. namespace tflite { +namespace { +static const char* kSparseTensorTestModel = + "tensorflow/lite/testdata/sparse_tensor.bin"; +} // namespace + class MockErrorReporter : public ErrorReporter { public: MockErrorReporter() : buffer_size_(0) {} @@ -552,6 +560,135 @@ TEST(VerifyModel, TypedTensorShapeMatchesTensorBufferSize) { } } +TEST(VerifyModel, SimpleValidSparseTensor) { + const auto model = FlatBufferModel::BuildFromFile(kSparseTensorTestModel); + ASSERT_TRUE(model); + + std::unique_ptr scoped_model; + scoped_model.reset(model->GetModel()->UnPack()); + + flatbuffers::FlatBufferBuilder builder; + auto model_ = Model::Pack(builder, scoped_model.get()); + + ::tflite::FinishModelBuffer(builder, model_); + MockErrorReporter mock_reporter; + MutableOpResolver resolver; + TfLiteRegistration fake_op; + resolver.AddCustom("FakeOp", &fake_op); + Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter); + ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter)); +} + +TEST(VerifyModel, InvalidSparseTensorMissingBlockMap) { + const auto model = FlatBufferModel::BuildFromFile(kSparseTensorTestModel); + ASSERT_TRUE(model); + + std::unique_ptr scoped_model; + scoped_model.reset(model->GetModel()->UnPack()); + + auto* tensor = scoped_model->subgraphs[0]->tensors[0].get(); + tensor->sparsity->block_map = {}; + + flatbuffers::FlatBufferBuilder builder; + auto model_ = Model::Pack(builder, scoped_model.get()); + + ::tflite::FinishModelBuffer(builder, model_); + MockErrorReporter mock_reporter; + MutableOpResolver resolver; + TfLiteRegistration fake_op; + resolver.AddCustom("FakeOp", &fake_op); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter)); + EXPECT_THAT(mock_reporter.GetAsString(), + ::testing::ContainsRegex("invalid sparsity parameters")); +} + +TEST(VerifyModel, InvalidSparseTensorIndexOutOfBound) { + const auto model = FlatBufferModel::BuildFromFile(kSparseTensorTestModel); + ASSERT_TRUE(model); + + std::unique_ptr scoped_model; + scoped_model.reset(model->GetModel()->UnPack()); + + auto* tensor = scoped_model->subgraphs[0]->tensors[0].get(); + tensor->sparsity->dim_metadata[1]->array_indices[1] = 5; + + flatbuffers::FlatBufferBuilder builder; + auto model_ = Model::Pack(builder, scoped_model.get()); + + ::tflite::FinishModelBuffer(builder, model_); + MockErrorReporter mock_reporter; + MutableOpResolver resolver; + TfLiteRegistration fake_op; + resolver.AddCustom("FakeOp", &fake_op); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter)); + EXPECT_THAT(mock_reporter.GetAsString(), + ::testing::ContainsRegex("invalid sparsity parameters")); +} + +TEST(VerifyModel, InvalidSparseTensorInvalidBuffer) { + const auto model = FlatBufferModel::BuildFromFile(kSparseTensorTestModel); + ASSERT_TRUE(model); + + std::unique_ptr scoped_model; + scoped_model.reset(model->GetModel()->UnPack()); + + // Expected to have 12 numbers in buffer. + scoped_model->buffers[1]->data = {0, 1, 2, 3, 4, 5, 6, 7}; + + flatbuffers::FlatBufferBuilder builder; + auto model_ = Model::Pack(builder, scoped_model.get()); + + ::tflite::FinishModelBuffer(builder, model_); + MockErrorReporter mock_reporter; + MutableOpResolver resolver; + TfLiteRegistration fake_op; + resolver.AddCustom("FakeOp", &fake_op); + ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter)); + EXPECT_THAT(mock_reporter.GetAsString(), + ::testing::ContainsRegex( + "requires 12 bytes, but is allocated with 8 bytes buffer")); +} + +TEST(VerifyModel, ValidSparseTensorBCSC) { + const auto model = FlatBufferModel::BuildFromFile(kSparseTensorTestModel); + ASSERT_TRUE(model); + + std::unique_ptr scoped_model; + scoped_model.reset(model->GetModel()->UnPack()); + + auto* tensor = scoped_model->subgraphs[0]->tensors[0].get(); + tensor->sparsity->traversal_order = {1, 0, 3, 2}; + tensor->sparsity->block_map = {0, 1}; + tensor->sparsity->dim_metadata[0]->format = DimensionType_DENSE; + tensor->sparsity->dim_metadata[0]->dense_size = 2; + + tensor->sparsity->dim_metadata[1]->format = DimensionType_SPARSE_CSR; + tensor->sparsity->dim_metadata[1]->array_segments = {0, 1, 3}; + tensor->sparsity->dim_metadata[1]->array_indices = {0, 0, 1}; + + tensor->sparsity->dim_metadata[2]->format = DimensionType_DENSE; + tensor->sparsity->dim_metadata[2]->dense_size = 2; + tensor->sparsity->dim_metadata[3]->format = DimensionType_DENSE; + tensor->sparsity->dim_metadata[3]->dense_size = 2; + + flatbuffers::FlatBufferBuilder builder; + auto model_ = Model::Pack(builder, scoped_model.get()); + + ::tflite::FinishModelBuffer(builder, model_); + MockErrorReporter mock_reporter; + MutableOpResolver resolver; + TfLiteRegistration fake_op; + resolver.AddCustom("FakeOp", &fake_op); + ASSERT_TRUE(Verify(builder.GetBufferPointer(), builder.GetSize(), resolver, + &mock_reporter)); +} + +// TODO(b/145614687): Add more tricky test cases for sparse tensor verification. // TODO(yichengfan): make up malicious files to test with. } // namespace tflite -- GitLab