Add verifier check for sparse tensors.

PiperOrigin-RevId: 286286868
Change-Id: I92bb40fb8eff9a16e5b382fec19befd39c93aeaa

tensorflow/lite/testdata/sparse_tensor.json

 {
   "version": 3,
-  "operator_codes": [
-  ],
+  "operator_codes": [{"builtin_code": "CUSTOM", "custom_code": "FakeOp"}],
   "subgraphs": [
     {
       "tensors": [
@@ -12,7 +11,7 @@
           ],
           "name": "sparse_tensor",
           "buffer": 1,
-          "type": "FLOAT32",
+          "type": "INT8",
           "quantization": {
           },
           "is_variable": "false",
@@ -41,12 +40,9 @@
           }
         }
       ],
-      "inputs": [
-      ],
-      "outputs": [
-      ],
-      "operators": [
-      ]
+      "inputs": [0],
+      "outputs": [0],
+      "operators": [{"inputs":[-1], "outputs":[-1]}]
     }
   ],
   "buffers": [

tensorflow/lite/tools/BUILD

@@ -84,6 +84,9 @@ cc_test(
     name = "verifier_test",
     size = "small",
     srcs = ["verifier_test.cc"],
+    data = [
+        "//tensorflow/lite:testdata/sparse_tensor.bin",
+    ],
     tags = [
         "tflite_not_portable",
     ],

tensorflow/lite/tools/verifier.cc

@@ -14,7 +14,10 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/lite/tools/verifier.h"
+
 #include <climits>
+#include <cstdint>
+
 #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<uint64_t> VerifyAndCountElements(
+    const SparsityParameters& sparsity, const std::vector<int>& 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<uint64_t> 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<int> 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);

tensorflow/lite/tools/verifier_test.cc

@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/tools/verifier.h"
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -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<ModelT> 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<ModelT> 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<ModelT> 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<ModelT> 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<ModelT> 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