From b43f88a832f06a85d49b1fd962d98504e917ba94 Mon Sep 17 00:00:00 2001 From: Yunlu Li Date: Wed, 9 Dec 2020 11:52:59 -0800 Subject: [PATCH] Add a no-copy SparseToDense method. PiperOrigin-RevId: 346602648 Change-Id: Id3c953222a2c1b7d8810d9b08d6427720b0e5cc8 --- .../delegates/xnnpack/xnnpack_delegate.cc | 17 ++-- .../optimize/sparsity/format_converter.cc | 37 ++++++-- .../optimize/sparsity/format_converter.h | 13 ++- .../sparsity/format_converter_test.cc | 90 +++++++++++++++++++ 4 files changed, 143 insertions(+), 14 deletions(-) diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc index 4b177da8113..d3b073cd6e3 100644 --- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc +++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc @@ -3178,22 +3178,25 @@ TfLiteIntArray* Delegate::PrepareOpsToDelegate(TfLiteContext* context) { switch (input_tensor.type) { case kTfLiteFloat32: { + const size_t dense_size = context->tensors[t].bytes / sizeof(float); + float* unpacked_fp32_data = reinterpret_cast(unpacked_data); tflite::optimize::sparsity::FormatConverter converter( vector_shape, *input_tensor.sparsity); converter.SparseToDense( - static_cast(input_tensor.data.data)); - const std::vector out = converter.GetData(); - std::memcpy(unpacked_data, out.data(), out.size() * sizeof(float)); + static_cast(input_tensor.data.data), dense_size, + unpacked_fp32_data, context); break; } case kTfLiteFloat16: { + const size_t dense_size = + context->tensors[t].bytes / sizeof(Eigen::half); + Eigen::half* unpacked_fp16_data = + reinterpret_cast(unpacked_data); tflite::optimize::sparsity::FormatConverter converter( vector_shape, *input_tensor.sparsity); converter.SparseToDense( - static_cast(input_tensor.data.data)); - const std::vector out = converter.GetData(); - std::memcpy(unpacked_data, out.data(), - out.size() * sizeof(Eigen::half)); + static_cast(input_tensor.data.data), + dense_size, unpacked_fp16_data, context); break; } default: { diff --git a/tensorflow/lite/tools/optimize/sparsity/format_converter.cc b/tensorflow/lite/tools/optimize/sparsity/format_converter.cc index d6a80f585d5..47f63e77cbb 100644 --- a/tensorflow/lite/tools/optimize/sparsity/format_converter.cc +++ b/tensorflow/lite/tools/optimize/sparsity/format_converter.cc @@ -261,7 +261,8 @@ FormatConverter::FormatConverter(const std::vector& shape, template void FormatConverter::Populate(const T* src_data, std::vector indices, - int level, int prev_idx, int* src_data_ptr) { + int level, int prev_idx, int* src_data_ptr, + T* dest_data) { if (level == indices.size()) { int orig_rank = dense_shape_.size(); std::vector orig_idx; @@ -279,7 +280,8 @@ void FormatConverter::Populate(const T* src_data, std::vector indices, orig_idx[orig_dim] * block_size_[block_idx] + indices[i]; } - data_[GetFlattenedIndex(orig_idx, dense_shape_)] = src_data[*src_data_ptr]; + dest_data[GetFlattenedIndex(orig_idx, dense_shape_)] = + src_data[*src_data_ptr]; *src_data_ptr = *src_data_ptr + 1; return; @@ -291,7 +293,7 @@ void FormatConverter::Populate(const T* src_data, std::vector indices, 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); + src_data_ptr, dest_data); } } else { const auto& array_segments = dim_metadata_[metadata_idx]; @@ -299,7 +301,7 @@ void FormatConverter::Populate(const T* src_data, std::vector indices, for (int i = array_segments[prev_idx]; i < array_segments[prev_idx + 1]; i++) { indices[level] = array_indices[i]; - Populate(src_data, indices, level + 1, i, src_data_ptr); + Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data); } } } @@ -312,7 +314,32 @@ TfLiteStatus FormatConverter::SparseToDense(const T* src_data) { int total_rank = traversal_order_.size(); int src_data_ptr = 0; std::vector indices(total_rank); - Populate(src_data, indices, 0, 0, &src_data_ptr); + Populate(src_data, indices, 0, 0, &src_data_ptr, data_.data()); + + return kTfLiteOk; +} + +template +TfLiteStatus FormatConverter::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 %lld.\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 indices(total_rank); + Populate(src_data, indices, 0, 0, &src_data_ptr, dest_data); return kTfLiteOk; } diff --git a/tensorflow/lite/tools/optimize/sparsity/format_converter.h b/tensorflow/lite/tools/optimize/sparsity/format_converter.h index 46e7d93b0b7..5edeb223c86 100644 --- a/tensorflow/lite/tools/optimize/sparsity/format_converter.h +++ b/tensorflow/lite/tools/optimize/sparsity/format_converter.h @@ -54,18 +54,27 @@ class FormatConverter { FormatConverter(const std::vector& shape, const TfLiteSparsity& sparsity); + // TODO(b/175040247): Return const reference to avoid copy. std::vector GetData() { return data_; } std::vector> 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); private: // 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 indices, int level, - int prev_idx, int* src_data_ptr); + int prev_idx, int* src_data_ptr, T* dest_data); // Check if val is equal to zero. bool IsZero(const T val); @@ -76,7 +85,7 @@ class FormatConverter { // tensor with (2, 2) block has blocked_shape (2, 2). std::vector blocked_shape_; // Total number of elements in the dense tensor. - uint64_t dense_size_; + size_t dense_size_; // Has n(original dimension)+k(block_dimension) elements. std::vector traversal_order_; // Format of each dimension in the traversal order. diff --git a/tensorflow/lite/tools/optimize/sparsity/format_converter_test.cc b/tensorflow/lite/tools/optimize/sparsity/format_converter_test.cc index 96919d22d4a..f3889fb1063 100644 --- a/tensorflow/lite/tools/optimize/sparsity/format_converter_test.cc +++ b/tensorflow/lite/tools/optimize/sparsity/format_converter_test.cc @@ -44,6 +44,11 @@ TEST(FormatConverterTest, SimpleTestD0D1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -70,6 +75,11 @@ TEST(FormatConverterTest, SimpleTestS0D1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -96,6 +106,11 @@ TEST(FormatConverterTest, SimpleTestD0S1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -124,6 +139,11 @@ TEST(FormatConverterTest, SimpleTestS0S1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -148,6 +168,11 @@ TEST(FormatConverterTest, SimpleTestD1D0) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -174,6 +199,11 @@ TEST(FormatConverterTest, SimpleTestS1D0) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -200,6 +230,11 @@ TEST(FormatConverterTest, SimpleTestD1S0) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -228,6 +263,11 @@ TEST(FormatConverterTest, SimpleTestS1S0) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -259,6 +299,11 @@ TEST(FormatConverterTest, 3DTestS0D1S2) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -288,6 +333,11 @@ TEST(FormatConverterTest, 3DTestD0D1S2) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -321,6 +371,11 @@ TEST(FormatConverterTest, 3DTestS0S1S2) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -354,6 +409,11 @@ TEST(FormatConverterTest, 3DTestS0S2S1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -384,6 +444,11 @@ TEST(FormatConverterTest, BlockTestD0D1) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector dense_data(dense_values.size()); + converter.SparseToDense(expected_data.data(), dense_data.size(), + dense_data.data(), nullptr); + EXPECT_EQ(dense_data, dense_values); } // BCSR @@ -417,6 +482,11 @@ TEST(FormatConverterTest, BlockTestD0S11DBlock) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector dense_data(dense_values.size()); + converter.SparseToDense(expected_data.data(), dense_data.size(), + dense_data.data(), nullptr); + EXPECT_EQ(dense_data, dense_values); } // BCSR @@ -450,6 +520,11 @@ TEST(FormatConverterTest, BlockTestD0S12DBlock) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector dense_data(dense_values.size()); + converter.SparseToDense(expected_data.data(), dense_data.size(), + dense_data.data(), nullptr); + EXPECT_EQ(dense_data, dense_values); } // BCSC @@ -483,6 +558,11 @@ TEST(FormatConverterTest, BlockTestD1S0) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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 @@ -516,6 +596,11 @@ TEST(FormatConverterTest, BlockTestD0S1LastBlockEmpty) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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) { @@ -550,6 +635,11 @@ TEST(FormatConverterTest, BlockTestD0S1ColMajorBlock) { converter.SparseToDense(expected_data.data()); const auto data_back = converter.GetData(); EXPECT_EQ(data_back, dense_values); + + std::vector 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