Fix a bug in the format converter.

PiperOrigin-RevId: 309160555
Change-Id: I3fe5c899ebba18f7b19c49c6551a9641a1cc9502
This commit is contained in:
Yunlu Li 2020-04-29 22:00:59 -07:00 committed by TensorFlower Gardener
parent 3c520614a3
commit 532920d246
2 changed files with 89 additions and 55 deletions

View File

@ -250,7 +250,7 @@ FormatConverter<T>::FormatConverter(const std::vector<int>& shape,
for (int i = 0; i < original_rank; i++) {
if (block_dim < block_map_.size() && block_map_[block_dim] == i) {
int orig_dim = traversal_order_[original_rank + block_dim];
block_size_[i] = sparsity.dim_metadata[orig_dim].dense_size;
block_size_[block_dim] = sparsity.dim_metadata[orig_dim].dense_size;
blocked_shape_[i] = shape[i] / sparsity.dim_metadata[orig_dim].dense_size;
block_dim++;
} else {
@ -273,9 +273,10 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
}
for (; i < indices.size(); i++) {
int orig_dim = block_map_[traversal_order_[i] - orig_rank];
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_[orig_dim] + indices[i];
orig_idx[orig_dim] * block_size_[block_idx] + indices[i];
}
data_[GetFlattenedIndex(orig_idx, dense_shape_)] = src_data[*src_data_ptr];

View File

@ -31,18 +31,18 @@ TEST(FormatConverterTest, SimpleTestD0D1) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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 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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -55,7 +55,7 @@ TEST(FormatConverterTest, SimpleTestS0D1) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -63,12 +63,12 @@ TEST(FormatConverterTest, SimpleTestS0D1) {
EXPECT_EQ(dm0_1, dim_metadata[1]);
EXPECT_EQ(dm1, dim_metadata[2]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -81,7 +81,7 @@ TEST(FormatConverterTest, SimpleTestD0S1) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -89,12 +89,12 @@ TEST(FormatConverterTest, SimpleTestD0S1) {
EXPECT_EQ(dm1_0, dim_metadata[2]);
EXPECT_EQ(dm1_1, dim_metadata[3]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -107,7 +107,7 @@ TEST(FormatConverterTest, SimpleTestS0S1) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -117,12 +117,12 @@ TEST(FormatConverterTest, SimpleTestS0S1) {
EXPECT_EQ(dm1_0, dim_metadata[2]);
EXPECT_EQ(dm1_1, dim_metadata[3]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -135,18 +135,18 @@ TEST(FormatConverterTest, SimpleTestD1D0) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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 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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -159,7 +159,7 @@ TEST(FormatConverterTest, SimpleTestS1D0) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -167,12 +167,12 @@ TEST(FormatConverterTest, SimpleTestS1D0) {
EXPECT_EQ(dm0_1, dim_metadata[1]);
EXPECT_EQ(dm1, dim_metadata[2]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -185,7 +185,7 @@ TEST(FormatConverterTest, SimpleTestD1S0) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -193,12 +193,12 @@ TEST(FormatConverterTest, SimpleTestD1S0) {
EXPECT_EQ(dm1_0, dim_metadata[2]);
EXPECT_EQ(dm1_1, dim_metadata[3]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -211,7 +211,7 @@ TEST(FormatConverterTest, SimpleTestS1S0) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -221,12 +221,12 @@ TEST(FormatConverterTest, SimpleTestS1S0) {
EXPECT_EQ(dm1_0, dim_metadata[2]);
EXPECT_EQ(dm1_1, dim_metadata[3]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -239,7 +239,7 @@ TEST(FormatConverterTest, 3DTestS0D1S2) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -252,12 +252,12 @@ TEST(FormatConverterTest, 3DTestS0D1S2) {
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -270,7 +270,7 @@ TEST(FormatConverterTest, 3DTestD0D1S2) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -281,12 +281,12 @@ TEST(FormatConverterTest, 3DTestD0D1S2) {
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -300,7 +300,7 @@ TEST(FormatConverterTest, 3DTestS0S1S2) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -314,12 +314,12 @@ TEST(FormatConverterTest, 3DTestS0S1S2) {
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -333,7 +333,7 @@ TEST(FormatConverterTest, 3DTestS0S2S1) {
FormatConverter<int> converter(dense_shape, traversal_order, format);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -347,12 +347,12 @@ TEST(FormatConverterTest, 3DTestS0S2S1) {
EXPECT_EQ(dm2_0, dim_metadata[4]);
EXPECT_EQ(dm2_1, dim_metadata[5]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -369,25 +369,58 @@ TEST(FormatConverterTest, BlockTestD0D1) {
block_size, block_map);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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 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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
// BCSR
TEST(FormatConverterTest, BlockTestD0S1) {
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);
}
// 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};
@ -400,7 +433,7 @@ TEST(FormatConverterTest, BlockTestD0S1) {
block_size, block_map);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -410,12 +443,12 @@ TEST(FormatConverterTest, BlockTestD0S1) {
EXPECT_EQ(dm, dim_metadata[4]);
EXPECT_EQ(dm, dim_metadata[6]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -433,7 +466,7 @@ TEST(FormatConverterTest, BlockTestD1S0) {
block_size, block_map);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -443,12 +476,12 @@ TEST(FormatConverterTest, BlockTestD1S0) {
EXPECT_EQ(dm, dim_metadata[4]);
EXPECT_EQ(dm, dim_metadata[6]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -466,7 +499,7 @@ TEST(FormatConverterTest, BlockTestD0S1LastBlockEmpty) {
block_size, block_map);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -476,12 +509,12 @@ TEST(FormatConverterTest, BlockTestD0S1LastBlockEmpty) {
EXPECT_EQ(dm, dim_metadata[4]);
EXPECT_EQ(dm, dim_metadata[6]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
@ -499,7 +532,7 @@ TEST(FormatConverterTest, BlockTestD0S1ColMajorBlock) {
block_size, block_map);
converter.DenseToSparse(dense_values.data());
const auto& dim_metadata = converter.GetDimMetadata();
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};
@ -509,13 +542,13 @@ TEST(FormatConverterTest, BlockTestD0S1ColMajorBlock) {
EXPECT_EQ(dm, dim_metadata[4]);
EXPECT_EQ(dm, dim_metadata[6]);
const auto& data = converter.GetData();
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();
const auto data_back = converter.GetData();
EXPECT_EQ(data_back, dense_values);
}
} // namespace