Add a no-copy SparseToDense method.
PiperOrigin-RevId: 346602648 Change-Id: Id3c953222a2c1b7d8810d9b08d6427720b0e5cc8
This commit is contained in:
parent
698571054b
commit
b43f88a832
tensorflow/lite
delegates/xnnpack
tools/optimize/sparsity
@ -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: {
|
||||
|
@ -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;
|
||||
}
|
||||
|
@ -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.
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user