Add a no-copy SparseToDense method.

PiperOrigin-RevId: 346602648
Change-Id: Id3c953222a2c1b7d8810d9b08d6427720b0e5cc8
This commit is contained in:
Yunlu Li 2020-12-09 11:52:59 -08:00 committed by TensorFlower Gardener
parent 698571054b
commit b43f88a832
4 changed files with 143 additions and 14 deletions
tensorflow/lite

View File

@ -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<float*>(unpacked_data);
tflite::optimize::sparsity::FormatConverter<float> converter(
vector_shape, *input_tensor.sparsity);
converter.SparseToDense(
static_cast<const float*>(input_tensor.data.data));
const std::vector<float> out = converter.GetData();
std::memcpy(unpacked_data, out.data(), out.size() * sizeof(float));
static_cast<const float*>(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<Eigen::half*>(unpacked_data);
tflite::optimize::sparsity::FormatConverter<Eigen::half> converter(
vector_shape, *input_tensor.sparsity);
converter.SparseToDense(
static_cast<const Eigen::half*>(input_tensor.data.data));
const std::vector<Eigen::half> out = converter.GetData();
std::memcpy(unpacked_data, out.data(),
out.size() * sizeof(Eigen::half));
static_cast<const Eigen::half*>(input_tensor.data.data),
dense_size, unpacked_fp16_data, context);
break;
}
default: {

View File

@ -261,7 +261,8 @@ FormatConverter<T>::FormatConverter(const std::vector<int>& shape,
template <typename T>
void FormatConverter<T>::Populate(const T* src_data, std::vector<int> 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<int> orig_idx;
@ -279,7 +280,8 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> 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<T>::Populate(const T* src_data, std::vector<int> 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<T>::Populate(const T* src_data, std::vector<int> 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<T>::SparseToDense(const T* src_data) {
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);
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 %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<int> indices(total_rank);
Populate(src_data, indices, 0, 0, &src_data_ptr, dest_data);
return kTfLiteOk;
}

View File

@ -54,18 +54,27 @@ class FormatConverter {
FormatConverter(const std::vector<int>& shape,
const TfLiteSparsity& sparsity);
// TODO(b/175040247): Return const reference to avoid copy.
std::vector<T> GetData() { return data_; }
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);
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<int> 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<int> 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<int> traversal_order_;
// Format of each dimension in the traversal order.

View File

@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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) {
@ -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<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
@ -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<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
@ -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<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
@ -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<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
@ -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<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) {
@ -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<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