chore: import upstream snapshot with attribution

tensorflow/lite/kernels/internal/utils/BUILD

@@ -0,0 +1,42 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "sparsity_format_converter",
+    srcs = ["sparsity_format_converter.cc"],
+    hdrs = ["sparsity_format_converter.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:common",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+cc_test(
+    name = "sparsity_format_converter_test",
+    srcs = ["sparsity_format_converter_test.cc"],
+    data = ["//tensorflow/lite:testdata/sparse_tensor.bin"],
+    tags = [
+        "tflite_not_portable",
+    ],
+    deps = [
+        ":sparsity_format_converter",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:framework",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/kernels/internal/utils/sparsity_format_converter.cc

@@ -0,0 +1,402 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+namespace internal {
+namespace sparsity {
+
+namespace {
+uint64_t GetFlattenedIndex(const std::vector<int>& indices,
+                           const std::vector<int>& shape) {
+  uint64_t index = 0;
+  int sub_elements = 1;
+  for (int i = shape.size() - 1; i >= 0; i--) {
+    index += indices[i] * sub_elements;
+    sub_elements *= shape[i];
+  }
+  return index;
+}
+
+std::vector<int> TfLiteIntArrayToVector(const TfLiteIntArray* int_array) {
+  std::vector<int> values;
+  if (!int_array) {
+    return values;
+  }
+
+  values.resize(int_array->size);
+  for (size_t i = 0; i < int_array->size; i++) {
+    values[i] = int_array->data[i];
+  }
+
+  return values;
+}
+
+}  // namespace
+
+template <typename T>
+FormatConverter<T>::FormatConverter(
+    const std::vector<int>& shape, const std::vector<int>& traversal_order,
+    const std::vector<TfLiteDimensionType>& format,
+    const std::vector<int>& block_size, const std::vector<int>& block_map)
+    : dense_shape_(shape),
+      traversal_order_(traversal_order),
+      block_size_(block_size),
+      block_map_(block_map) {
+  dense_size_ = 1;
+  int block_dim = 0;
+  blocked_shape_.resize(shape.size());
+  format_.resize(shape.size() + block_map.size());
+  for (int i = 0; i < shape.size(); i++) {
+    format_[i] = format[traversal_order[i]];
+    dense_size_ *= shape[i];
+    if (block_dim < block_map.size() && block_map[block_dim] == i) {
+      blocked_shape_[i] = shape[i] / block_size[block_dim];
+      block_dim++;
+    } else {
+      blocked_shape_[i] = shape[i];
+    }
+  }
+
+  // Only dense blocks are supported.
+  for (int i = 0; i < block_map.size(); i++) {
+    format_[i + shape.size()] = kTfLiteDimDense;
+  }
+}
+
+template <typename T>
+TfLiteStatus FormatConverter<T>::DenseToSparse(const T* src_data) {
+  int num_original_dims = dense_shape_.size();
+  int num_block_dims = block_map_.size();
+  int num_expanded_dims = num_original_dims + num_block_dims;
+  std::vector<int> expanded_shape(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; i++) {
+    if (i < num_original_dims) {
+      expanded_shape[i] = blocked_shape_[i];
+    } else {
+      expanded_shape[i] = block_size_[i - num_original_dims];
+    }
+  }
+
+  std::vector<int> shape_offset(num_original_dims);
+  shape_offset[shape_offset.size() - 1] = 1;
+  for (int i = num_original_dims - 1; i > 0; --i) {
+    shape_offset[i - 1] = shape_offset[i] * dense_shape_[i];
+  }
+
+  std::vector<int> expanded_shape_offset(num_expanded_dims);
+  for (int i = 0; i < num_original_dims; ++i) {
+    expanded_shape_offset[i] = shape_offset[i];
+  }
+  for (int i = 0; i < num_block_dims; ++i) {
+    int mapped_dim = block_map_[i];
+    expanded_shape_offset[num_original_dims + i] = shape_offset[mapped_dim];
+    expanded_shape_offset[mapped_dim] *= block_size_[i];
+  }
+
+  std::vector<int> dst_ordered_offset(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; ++i) {
+    dst_ordered_offset[i] = expanded_shape_offset[traversal_order_[i]];
+  }
+
+  std::vector<bool> dst_dim_has_nonzeroes(num_expanded_dims);
+  std::fill(dst_dim_has_nonzeroes.begin(), dst_dim_has_nonzeroes.end(), false);
+  std::vector<int> inner_compressed_dim(num_expanded_dims);
+  int most_recent_compressed_dim = -1;
+  std::vector<int> num_segments_of_next_compressed_dim(num_expanded_dims);
+  int segment_count = 1;
+  for (int i = num_expanded_dims - 1; i >= 0; --i) {
+    inner_compressed_dim[i] = most_recent_compressed_dim;
+    if (format_[i] == kTfLiteDimSparseCSR) {
+      most_recent_compressed_dim = i;
+      num_segments_of_next_compressed_dim[i] = segment_count;
+      segment_count = 1;
+    } else {
+      num_segments_of_next_compressed_dim[i] = -1;
+      segment_count *= expanded_shape[traversal_order_[i]];
+    }
+  }
+
+  dim_metadata_.resize(num_expanded_dims * 2);
+  std::vector<int> dst_sparse_dims;
+  dst_sparse_dims.reserve(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; ++i) {
+    dim_metadata_[i * 2].clear();
+    dim_metadata_[i * 2 + 1].clear();
+    if (format_[i] == kTfLiteDimDense) {
+      // If dimension is dense, just store the shape.
+      dim_metadata_[i * 2].push_back(expanded_shape[traversal_order_[i]]);
+    } else {
+      dim_metadata_[i * 2].push_back(0);  // Segment array always begins with 0.
+      dst_sparse_dims.push_back(i);       // Add dimension to the sparse list.
+    }
+  }
+
+  // This algorithm assumes that the block size is small enough for all the
+  // elements to fit in cache, so the strided accesses from different traversal
+  // order and the write-first-erase-later strategy shouldn't be too slow
+  int dst_dim_idx = num_expanded_dims;
+  std::vector<int> coordinate(num_expanded_dims, 0);
+  int dense_tensor_idx = 0;
+  while (dst_dim_idx >= 0) {
+    if (dst_dim_idx == num_expanded_dims) {
+      // We have a complete coordinate. Add the element to the value array if it
+      // is not zero, or if the last dimension is dense.
+      if (!IsZero(src_data[dense_tensor_idx])) {
+        data_.push_back(src_data[dense_tensor_idx]);
+        // Mark all sparse dimensions that their current indices have nonzeroes.
+        for (auto dst_dim : dst_sparse_dims) {
+          if (!dst_dim_has_nonzeroes[dst_dim]) {
+            // Only add the index to the indices array if the current nonzero
+            // is the first nonzero of the block.
+            dim_metadata_[2 * dst_dim + 1].push_back(coordinate[dst_dim]);
+            dst_dim_has_nonzeroes[dst_dim] = true;
+          }
+        }
+      } else if (format_[num_expanded_dims - 1] == kTfLiteDimDense) {
+        data_.push_back(src_data[dense_tensor_idx]);
+      }
+      --dst_dim_idx;
+    } else {
+      int original_dim_idx = traversal_order_[dst_dim_idx];
+      int dim_size = expanded_shape[original_dim_idx];
+      if (dst_dim_has_nonzeroes[dst_dim_idx]) {
+        // If the previous block has nonzeroes, reset the flag to false since
+        // we have just moved to a new block.
+        dst_dim_has_nonzeroes[dst_dim_idx] = false;
+      } else if (format_[dst_dim_idx] == kTfLiteDimSparseCSR) {
+        // This block is empty. Delete unnecessary values if compressed.
+        int next_compressed_dim = inner_compressed_dim[dst_dim_idx];
+        int erase_offset = dim_metadata_[2 * dst_dim_idx + 1].size() *
+                           num_segments_of_next_compressed_dim[dst_dim_idx];
+        if (next_compressed_dim >= 0) {
+          auto& segments = dim_metadata_[2 * inner_compressed_dim[dst_dim_idx]];
+          segments.erase(segments.begin() + 1 + erase_offset, segments.end());
+        } else {
+          data_.erase(data_.begin() + erase_offset, data_.end());
+        }
+      }
+      if (++coordinate[dst_dim_idx] < dim_size) {
+        // The current dst_dim_idx is valid (not out of bound).
+        dense_tensor_idx += dst_ordered_offset[dst_dim_idx];
+        ++dst_dim_idx;
+      } else {
+        // dst_dim_idx has reached its dim size. Update segment array and go
+        // back to incrementing the previous dimension (dst_dim_idx - 1).
+        if (format_[dst_dim_idx] == kTfLiteDimSparseCSR) {
+          dim_metadata_[2 * dst_dim_idx].push_back(
+              dim_metadata_[2 * dst_dim_idx + 1].size());
+        }
+        coordinate[dst_dim_idx] = -1;
+        dense_tensor_idx -= dst_ordered_offset[dst_dim_idx] * dim_size;
+        --dst_dim_idx;
+      }
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+template <typename T>
+FormatConverter<T>::FormatConverter(
+    const std::vector<int>& shape, const std::vector<int>& traversal_order,
+    const std::vector<TfLiteDimensionType>& format,
+    const std::vector<int>& dense_size,
+    const std::vector<std::vector<int>>& segments,
+    const std::vector<std::vector<int>>& indices,
+    const std::vector<int>& block_map) {
+  InitSparseToDenseConverter(shape, traversal_order, format, dense_size,
+                             segments, indices, block_map);
+}
+
+template <typename T>
+FormatConverter<T>::FormatConverter(const std::vector<int>& shape,
+                                    const TfLiteSparsity& sparsity) {
+  auto traversal_order = TfLiteIntArrayToVector(sparsity.traversal_order);
+  auto block_map = TfLiteIntArrayToVector(sparsity.block_map);
+
+  std::vector<TfLiteDimensionType> format(sparsity.dim_metadata_size);
+  std::vector<int> dense_size(sparsity.dim_metadata_size);
+  std::vector<std::vector<int>> segments(sparsity.dim_metadata_size);
+  std::vector<std::vector<int>> indices(sparsity.dim_metadata_size);
+  for (int i = 0; i < sparsity.dim_metadata_size; i++) {
+    format[i] = sparsity.dim_metadata[i].format;
+    dense_size[i] = sparsity.dim_metadata[i].dense_size;
+    segments[i] =
+        TfLiteIntArrayToVector(sparsity.dim_metadata[i].array_segments);
+    indices[i] = TfLiteIntArrayToVector(sparsity.dim_metadata[i].array_indices);
+  }
+
+  InitSparseToDenseConverter(shape, std::move(traversal_order),
+                             std::move(format), std::move(dense_size),
+                             std::move(segments), std::move(indices),
+                             std::move(block_map));
+}
+
+template <typename T>
+void FormatConverter<T>::InitSparseToDenseConverter(
+    std::vector<int> shape, std::vector<int> traversal_order,
+    std::vector<TfLiteDimensionType> format, std::vector<int> dense_size,
+    std::vector<std::vector<int>> segments,
+    std::vector<std::vector<int>> indices, std::vector<int> block_map) {
+  dense_shape_ = std::move(shape);
+  traversal_order_ = std::move(traversal_order);
+  block_map_ = std::move(block_map);
+  format_ = std::move(format);
+
+  dense_size_ = 1;
+  for (int i = 0; i < dense_shape_.size(); i++) {
+    dense_size_ *= dense_shape_[i];
+  }
+
+  dim_metadata_.resize(2 * format_.size());
+  for (int i = 0; i < format_.size(); i++) {
+    if (format_[i] == kTfLiteDimDense) {
+      dim_metadata_[2 * i] = {dense_size[i]};
+    } else {
+      dim_metadata_[2 * i] = std::move(segments[i]);
+      dim_metadata_[2 * i + 1] = std::move(indices[i]);
+    }
+  }
+
+  int original_rank = dense_shape_.size();
+  int block_dim = 0;
+
+  blocked_shape_.resize(original_rank);
+  block_size_.resize(block_map_.size());
+  for (int i = 0; i < original_rank; i++) {
+    if (block_dim < block_map_.size() && block_map_[block_dim] == i) {
+      if (original_rank + block_dim < traversal_order_.size()) {
+        int orig_dim = traversal_order_[original_rank + block_dim];
+        block_size_[block_dim] = dense_size[orig_dim];
+        blocked_shape_[i] = dense_shape_[i] / dense_size[orig_dim];
+        block_dim++;
+      }
+    } else {
+      blocked_shape_[i] = dense_shape_[i];
+    }
+  }
+}
+
+template <typename T>
+void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
+                                  int level, int prev_idx, int* src_data_ptr,
+                                  T* dest_data) {
+  if (level == indices.size()) {
+    int orig_rank = dense_shape_.size();
+    std::vector<int> orig_idx;
+    orig_idx.resize(orig_rank);
+    int i = 0;
+    for (; i < orig_idx.size(); i++) {
+      int orig_dim = traversal_order_[i];
+      orig_idx[orig_dim] = indices[i];
+    }
+
+    for (; i < indices.size(); i++) {
+      const int block_idx = traversal_order_[i] - orig_rank;
+      const int orig_dim = block_map_[block_idx];
+      orig_idx[orig_dim] =
+          orig_idx[orig_dim] * block_size_[block_idx] + indices[i];
+    }
+
+    dest_data[GetFlattenedIndex(orig_idx, dense_shape_)] =
+        src_data[*src_data_ptr];
+
+    *src_data_ptr = *src_data_ptr + 1;
+    return;
+  }
+
+  const int metadata_idx = 2 * level;
+  const int shape_of_level = dim_metadata_[metadata_idx][0];
+  if (format_[level] == kTfLiteDimDense) {
+    for (int i = 0; i < shape_of_level; i++) {
+      indices[level] = i;
+      Populate(src_data, indices, level + 1, prev_idx * shape_of_level + i,
+               src_data_ptr, dest_data);
+    }
+  } else if (prev_idx + 1 < dim_metadata_[metadata_idx].size()) {
+    const auto& array_segments = dim_metadata_[metadata_idx];
+    const auto& array_indices = dim_metadata_[metadata_idx + 1];
+    for (int i = array_segments[prev_idx]; i < array_segments[prev_idx + 1];
+         i++) {
+      if (i < array_indices.size() && level < indices.size()) {
+        indices[level] = array_indices[i];
+        Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data);
+      }
+    }
+  }
+}
+
+template <typename T>
+TfLiteStatus FormatConverter<T>::SparseToDense(const T* src_data) {
+  data_.resize(dense_size_);
+  std::fill(data_.begin(), data_.end(), T(0));
+
+  int total_rank = traversal_order_.size();
+  int src_data_ptr = 0;
+  std::vector<int> indices(total_rank);
+  Populate(src_data, indices, 0, 0, &src_data_ptr, data_.data());
+
+  return kTfLiteOk;
+}
+
+template <typename T>
+TfLiteStatus FormatConverter<T>::SparseToDense(const T* src_data,
+                                               const size_t dest_size,
+                                               T* dest_data,
+                                               TfLiteContext* context) {
+  if (dest_size != dense_size_) {
+    TF_LITE_MAYBE_KERNEL_LOG(
+        context, "unexpected buffer size for densified data, expected %zu.\n",
+        dense_size_);
+    return kTfLiteError;
+  }
+
+  // For types like Eigen::half, we cannot do a simple memset() with 0 values.
+  for (auto i = 0; i < dest_size; i++) {
+    dest_data[i] = T(0);
+  }
+
+  const int total_rank = traversal_order_.size();
+  int src_data_ptr = 0;
+  std::vector<int> indices(total_rank);
+  Populate(src_data, indices, 0, 0, &src_data_ptr, dest_data);
+
+  return kTfLiteOk;
+}
+
+template <typename T>
+bool FormatConverter<T>::IsZero(const T val) {
+  return (val == static_cast<T>(0));
+}
+
+template class FormatConverter<int32_t>;
+template class FormatConverter<int8_t>;
+template class FormatConverter<float>;
+template class FormatConverter<Eigen::half>;
+
+}  // namespace sparsity
+}  // namespace internal
+}  // namespace tflite

tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h

@@ -0,0 +1,150 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_
+
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace internal {
+namespace sparsity {
+
+// A converter that keeps an internal representation of sparse tensor parameters
+// and converts tensors between dense and sparse formats.
+template <typename T>
+class FormatConverter {
+ public:
+  /*
+   * Creates a dense to sparse converter.
+   * @param shape             Shape of the dense tensor.
+   * @param traversal_order   In what order to traverse all dimensions,
+   *                          including block dimensions.
+   * @param format            Whether each dimension in the dense tensor is
+   *                          dense or sparse (not in the traversal order).
+   * @param block_size        Size of each block dimension.
+   * @param block_map         Map from block dimension to original tensor
+   *                          dimension.
+   */
+  FormatConverter(const std::vector<int>& shape,
+                  const std::vector<int>& traversal_order,
+                  const std::vector<TfLiteDimensionType>& format,
+                  const std::vector<int>& block_size = {},
+                  const std::vector<int>& block_map = {});
+
+  /*
+   * Creates a sparse to dense converter.
+   * @param shape             Shape of the target dense tensor.
+   * @param traversal_order   In what order to traverse all dimensions,
+   *                          including block dimensions.
+   * @param format            Whether each dimension in the dense tensor is
+   *                          dense or sparse (not in the traversal order).
+   * @param dense_size        Size of each dense dimension in the sparse tensor.
+   *                          Should be 0 for sparse dimensions.
+   * @param segments          Segments of each dimension in the sparse tensor.
+   *                          Should be empty for dense dimensions.
+   * @param indices           Indices in the dense tensor for each dimension.
+   *                          Should be empty for dense dimensions.
+   * @param block_map         Map from block dimension to original tensor
+   *                          dimension.
+   */
+  FormatConverter(const std::vector<int>& shape,
+                  const std::vector<int>& traversal_order,
+                  const std::vector<TfLiteDimensionType>& format,
+                  const std::vector<int>& dense_size,
+                  const std::vector<std::vector<int>>& segments,
+                  const std::vector<std::vector<int>>& indices,
+                  const std::vector<int>& block_map = {});
+
+  /* Creates a sparse to dense converter.
+   * @param shape      Shape of the target dense tensor.
+   * @param sparsity   Sparsity parameter of the sparse TfLiteTensor.
+   */
+  FormatConverter(const std::vector<int>& shape,
+                  const TfLiteSparsity& sparsity);
+
+  const std::vector<T>& GetData() { return data_; }
+  const std::vector<std::vector<int>>& GetDimMetadata() {
+    return dim_metadata_;
+  }
+
+  // Method for dense to sparse conversion. Need to call GetData() method to get
+  // the compressed data.
+  TfLiteStatus DenseToSparse(const T* src_data);
+
+  // Method for sparse to dense conversion. Need to call GetData() method to get
+  // the decompressed data.
+  TfLiteStatus SparseToDense(const T* src_data);
+  // Method for sparse to dense conversion with caller provided buffer. No need
+  // to call GetData() with this method.
+  TfLiteStatus SparseToDense(const T* src_data, const size_t dest_size,
+                             T* dest_data, TfLiteContext* context = nullptr);
+
+ private:
+  // Helper function for initializing this converter for sparse to dense
+  // conversion.
+  void InitSparseToDenseConverter(std::vector<int> shape,
+                                  std::vector<int> traversal_order,
+                                  std::vector<TfLiteDimensionType> format,
+                                  std::vector<int> dense_size,
+                                  std::vector<std::vector<int>> segments,
+                                  std::vector<std::vector<int>> indices,
+                                  std::vector<int> block_map);
+
+  // A recursive function to fetch data from the compressed src_data buffer and
+  // populate the dense buffer.
+  void Populate(const T* src_data, std::vector<int> indices, int level,
+                int prev_idx, int* src_data_ptr, T* dest_data);
+
+  // Check if val is equal to zero.
+  bool IsZero(const T val);
+
+  // Shape of the conceptual dense tensor.
+  std::vector<int> dense_shape_;
+  // Shape of the dense tensor with inner blocks reduced. For example, a (4, 4)
+  // tensor with (2, 2) block has blocked_shape (2, 2).
+  std::vector<int> blocked_shape_;
+  // Total number of elements in the dense tensor.
+  size_t dense_size_;
+  // Has n(original dimension)+k(block_dimension) elements.
+  std::vector<int> traversal_order_;
+  // Format of each dimension in the traversal order.
+  std::vector<TfLiteDimensionType> format_;
+  // Size of each block dimension, in the same order as block map.
+  std::vector<int> block_size_;
+  // Map from block dimension to the original tensor dimension.
+  std::vector<int> block_map_;
+  // Metadata of each dimension in the traversal order.
+  // Each dimension needs two vectors. For dense dimensions, the first vector
+  // stores the size of that dimension, and the second vector is empty. For
+  // sparse dimensions, the first vector stores the segments and the second one
+  // stores the indices.
+  std::vector<std::vector<int>> dim_metadata_;
+  // Actual buffer holding data after conversion. Could be sparse buffer or
+  // dense buffer.
+  std::vector<T> data_;
+};
+
+extern template class FormatConverter<int32_t>;
+extern template class FormatConverter<int8_t>;
+extern template class FormatConverter<float>;
+extern template class FormatConverter<Eigen::half>;
+}  // namespace sparsity
+}  // namespace internal
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_

tensorflow/lite/kernels/internal/utils/sparsity_format_converter_test.cc

@@ -0,0 +1,648 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h"
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+namespace internal {
+namespace sparsity {
+namespace {
+TEST(FormatConverterTest, SimpleTestD0D1) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {0, 1};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimDense};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {3};
+  const std::vector<int> dm1 = {4};
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestS0D1) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {0, 1};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimSparseCSR,
+                                                   kTfLiteDimDense};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 2};
+  const std::vector<int> dm0_1 = {0, 2};
+  const std::vector<int> dm1 = {4};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 0, 9, 8, 5, 0, 0, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestD0S1) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {0, 1};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {3};
+  const std::vector<int> dm1_0 = {0, 3, 3, 5};
+  const std::vector<int> dm1_1 = {0, 2, 3, 0, 3};
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 9, 8, 5, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestS0S1) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {0, 1};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimSparseCSR,
+                                                   kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 2};
+  const std::vector<int> dm0_1 = {0, 2};
+  const std::vector<int> dm1_0 = {0, 3, 5};
+  const std::vector<int> dm1_1 = {0, 2, 3, 0, 3};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 9, 8, 5, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestD1D0) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {1, 0};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimDense};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {4};
+  const std::vector<int> dm1 = {3};
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 0, 5, 0, 0, 0, 9, 0, 0, 8, 0, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestS1D0) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {1, 0};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 3};
+  const std::vector<int> dm0_1 = {0, 2, 3};
+  const std::vector<int> dm1 = {3};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 0, 5, 9, 0, 0, 8, 0, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestD1S0) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {1, 0};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimSparseCSR,
+                                                   kTfLiteDimDense};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {4};
+  const std::vector<int> dm1_0 = {0, 2, 2, 3, 5};
+  const std::vector<int> dm1_1 = {0, 2, 0, 0, 2};
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 5, 9, 8, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, SimpleTestS1S0) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 4};
+  const std::vector<int> traversal_order = {1, 0};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimSparseCSR,
+                                                   kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 3};
+  const std::vector<int> dm0_1 = {0, 2, 3};
+  const std::vector<int> dm1_0 = {0, 2, 3, 5};
+  const std::vector<int> dm1_1 = {0, 2, 0, 0, 2};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 5, 9, 8, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, 3DTestS0D1S2) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 2, 2};
+  const std::vector<int> traversal_order = {0, 1, 2};
+  const std::vector<TfLiteDimensionType> format = {
+      kTfLiteDimSparseCSR, kTfLiteDimDense, kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 2};
+  const std::vector<int> dm0_1 = {0, 2};
+  const std::vector<int> dm1 = {2};
+  const std::vector<int> dm2_0 = {0, 1, 3, 4, 5};
+  const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
+
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+  EXPECT_EQ(dm2_0, dim_metadata[4]);
+  EXPECT_EQ(dm2_1, dim_metadata[5]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 9, 8, 5, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, 3DTestD0D1S2) {
+  const std::vector<int> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  const std::vector<int> dense_shape = {3, 2, 2};
+  const std::vector<int> traversal_order = {0, 1, 2};
+  const std::vector<TfLiteDimensionType> format = {
+      kTfLiteDimDense, kTfLiteDimDense, kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {3};
+  const std::vector<int> dm1 = {2};
+  const std::vector<int> dm2_0 = {0, 1, 3, 3, 3, 4, 5};
+  const std::vector<int> dm2_1 = {0, 0, 1, 0, 1};
+
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1, dim_metadata[2]);
+  EXPECT_EQ(dm2_0, dim_metadata[4]);
+  EXPECT_EQ(dm2_1, dim_metadata[5]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {6, 9, 8, 5, 7};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, 3DTestS0S1S2) {
+  const std::vector<int> dense_values = {1, 7, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0,
+                                         0, 0, 0, 0, 0, 2, 0, 0, 4, 8, 3, 9};
+  const std::vector<int> dense_shape = {3, 4, 2};
+  const std::vector<int> traversal_order = {0, 1, 2};
+  const std::vector<TfLiteDimensionType> format = {
+      kTfLiteDimSparseCSR, kTfLiteDimSparseCSR, kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 2};
+  const std::vector<int> dm0_1 = {0, 2};
+  const std::vector<int> dm1_0 = {0, 2, 5};
+  const std::vector<int> dm1_1 = {0, 2, 0, 2, 3};
+  const std::vector<int> dm2_0 = {0, 2, 3, 4, 6, 8};
+  const std::vector<int> dm2_1 = {0, 1, 1, 1, 0, 1, 0, 1};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm2_0, dim_metadata[4]);
+  EXPECT_EQ(dm2_1, dim_metadata[5]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 7, 5, 2, 4, 8, 3, 9};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, 3DTestS0S2S1) {
+  const std::vector<int> dense_values = {1, 0, 0, 0, 7, 0, 5, 0, 0, 0, 0, 0,
+                                         0, 0, 0, 0, 0, 0, 4, 3, 2, 0, 8, 9};
+  const std::vector<int> dense_shape = {3, 2, 4};
+  const std::vector<int> traversal_order = {0, 2, 1};
+  const std::vector<TfLiteDimensionType> format = {
+      kTfLiteDimSparseCSR, kTfLiteDimSparseCSR, kTfLiteDimSparseCSR};
+  FormatConverter<int> converter(dense_shape, traversal_order, format);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0_0 = {0, 2};
+  const std::vector<int> dm0_1 = {0, 2};
+  const std::vector<int> dm1_0 = {0, 2, 5};
+  const std::vector<int> dm1_1 = {0, 2, 0, 2, 3};
+  const std::vector<int> dm2_0 = {0, 2, 3, 4, 6, 8};
+  const std::vector<int> dm2_1 = {0, 1, 1, 1, 0, 1, 0, 1};
+  EXPECT_EQ(dm0_0, dim_metadata[0]);
+  EXPECT_EQ(dm0_1, dim_metadata[1]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm2_0, dim_metadata[4]);
+  EXPECT_EQ(dm2_1, dim_metadata[5]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 7, 5, 2, 4, 8, 3, 9};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, BlockTestD0D1) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0,
+                                         0, 0, 5, 0, 0, 0, 0, 6};
+  const std::vector<int> dense_shape = {4, 4};
+  const std::vector<int> traversal_order = {0, 1, 2, 3};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimDense};
+  const std::vector<int> block_size = {2, 2};
+  const std::vector<int> block_map = {0, 1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm = {2};
+  EXPECT_EQ(dm, dim_metadata[0]);
+  EXPECT_EQ(dm, dim_metadata[2]);
+  EXPECT_EQ(dm, dim_metadata[4]);
+  EXPECT_EQ(dm, dim_metadata[6]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 0, 0, 4, 2, 3, 0, 0,
+                                          0, 0, 0, 0, 5, 0, 0, 6};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+// BCSR
+TEST(FormatConverterTest, BlockTestD0S11DBlock) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0,
+                                         0, 0, 5, 0, 0, 0, 0, 6};
+  const std::vector<int> dense_shape = {4, 4};
+  const std::vector<int> traversal_order = {0, 1, 2};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  const std::vector<int> block_size = {2};
+  const std::vector<int> block_map = {1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm0 = {4};
+  const std::vector<int> dm2 = {2};
+  const std::vector<int> dm1_0 = {0, 2, 3, 4, 5};
+  const std::vector<int> dm1_1 = {0, 1, 0, 1, 1};
+  EXPECT_EQ(dm0, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm2, dim_metadata[4]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 0, 2, 3, 0, 4, 5, 0, 0, 6};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+// BCSR
+TEST(FormatConverterTest, BlockTestD0S12DBlock) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0,
+                                         0, 0, 5, 0, 0, 0, 0, 6};
+  const std::vector<int> dense_shape = {4, 4};
+  const std::vector<int> traversal_order = {0, 1, 2, 3};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  const std::vector<int> block_size = {2, 2};
+  const std::vector<int> block_map = {0, 1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm = {2};
+  const std::vector<int> dm1_0 = {0, 2, 3};
+  const std::vector<int> dm1_1 = {0, 1, 1};
+  EXPECT_EQ(dm, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm, dim_metadata[4]);
+  EXPECT_EQ(dm, dim_metadata[6]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 0, 0, 4, 2, 3, 0, 0, 5, 0, 0, 6};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+// BCSC
+TEST(FormatConverterTest, BlockTestD1S0) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0,
+                                         0, 0, 5, 0, 0, 0, 0, 6};
+  const std::vector<int> dense_shape = {4, 4};
+  const std::vector<int> traversal_order = {1, 0, 3, 2};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimSparseCSR,
+                                                   kTfLiteDimDense};
+  const std::vector<int> block_size = {2, 2};
+  const std::vector<int> block_map = {0, 1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm = {2};
+  const std::vector<int> dm1_0 = {0, 1, 3};
+  const std::vector<int> dm1_1 = {0, 0, 1};
+  EXPECT_EQ(dm, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm, dim_metadata[4]);
+  EXPECT_EQ(dm, dim_metadata[6]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 0, 0, 4, 2, 0, 3, 0, 5, 0, 0, 6};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+// BCSR with last block being empty
+TEST(FormatConverterTest, BlockTestD0S1LastBlockEmpty) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0,
+                                         0, 0, 0, 0, 0, 0, 0, 0};
+  const std::vector<int> dense_shape = {4, 4};
+  const std::vector<int> traversal_order = {0, 1, 2, 3};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  const std::vector<int> block_size = {2, 2};
+  const std::vector<int> block_map = {0, 1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm = {2};
+  const std::vector<int> dm1_0 = {0, 2, 2};
+  const std::vector<int> dm1_1 = {0, 1};
+  EXPECT_EQ(dm, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm, dim_metadata[4]);
+  EXPECT_EQ(dm, dim_metadata[6]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 0, 0, 4, 2, 3, 0, 0};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+
+TEST(FormatConverterTest, BlockTestD0S1ColMajorBlock) {
+  const std::vector<int> dense_values = {1, 0, 2, 3, 0, 4, 0, 0, 1, 0, 2,
+                                         3, 0, 4, 0, 0, 0, 0, 5, 0, 0, 0,
+                                         0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
+  const std::vector<int> dense_shape = {4, 8};
+  const std::vector<int> traversal_order = {0, 1, 3, 2};
+  const std::vector<TfLiteDimensionType> format = {kTfLiteDimDense,
+                                                   kTfLiteDimSparseCSR};
+  const std::vector<int> block_size = {2, 2};
+  const std::vector<int> block_map = {0, 1};
+  FormatConverter<int> converter(dense_shape, traversal_order, format,
+                                 block_size, block_map);
+  converter.DenseToSparse(dense_values.data());
+
+  const auto& dim_metadata = converter.GetDimMetadata();
+  const std::vector<int> dm = {2};
+  const std::vector<int> dm1_0 = {0, 3, 4};
+  const std::vector<int> dm1_1 = {0, 1, 2, 1};
+  EXPECT_EQ(dm, dim_metadata[0]);
+  EXPECT_EQ(dm1_0, dim_metadata[2]);
+  EXPECT_EQ(dm1_1, dim_metadata[3]);
+  EXPECT_EQ(dm, dim_metadata[4]);
+  EXPECT_EQ(dm, dim_metadata[6]);
+
+  const auto& data = converter.GetData();
+  const std::vector<int> expected_data = {1, 1, 0, 0, 2, 2, 3, 3,
+                                          0, 0, 4, 4, 5, 0, 0, 0};
+  EXPECT_EQ(expected_data, data);
+
+  converter.SparseToDense(expected_data.data());
+  const auto& data_back = converter.GetData();
+  EXPECT_EQ(data_back, dense_values);
+
+  std::vector<int> dense_data(dense_values.size());
+  converter.SparseToDense(expected_data.data(), dense_data.size(),
+                          dense_data.data(), nullptr);
+  EXPECT_EQ(dense_data, dense_values);
+}
+}  // namespace
+}  // namespace sparsity
+}  // namespace internal
+}  // namespace tflite