Skip to content

Commit

Permalink
Add benchmark for util/SparseTensor's SparseReorder, and optimize Spa…
Browse files Browse the repository at this point in the history 
…rseTensor::Reorder.

Change: 134425452
  • Loading branch information
@ebrevdo @tensorflower-gardener
ebrevdo authored and tensorflower-gardener committed Sep 27, 2016
1 parent 86226d5 commit 1a295f8
Show file tree
Hide file tree
Showing 3 changed files with 163 additions and 17 deletions.
34 changes: 29 additions & 5 deletions tensorflow/core/util/sparse/dim_comparator.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,14 +46,14 @@ class DimComparator {
public:
typedef typename gtl::ArraySlice<int64> VarDimArray;

inline DimComparator(const TTypes<int64>::Matrix& ix,
const VarDimArray& order, int dims)
: ix_(ix), order_(order), dims_(dims) {
DimComparator(const TTypes<int64>::Matrix& ix, const VarDimArray& order,
const TensorShape& shape)
: ix_(ix), order_(order), dims_(shape.dims()) {
CHECK_GT(order.size(), size_t{0}) << "Must order using at least one index";
CHECK_LE(order.size(), dims_) << "Can only sort up to dims";
CHECK_LE(order.size(), shape.dims()) << "Can only sort up to dims";
for (size_t d = 0; d < order.size(); ++d) {
CHECK_GE(order[d], 0);
CHECK_LT(order[d], dims);
CHECK_LT(order[d], shape.dims());
}
}

Expand Down Expand Up @@ -84,9 +84,33 @@ class DimComparator {
return 0;
}

protected:
const TTypes<int64>::Matrix ix_;
const VarDimArray order_;
const int dims_;
const std::vector<int64>* ix_order_;
};

template <int ORDER_DIM>
class FixedDimComparator : DimComparator {
public:
FixedDimComparator(const TTypes<int64>::Matrix& ix, const VarDimArray& order,
const TensorShape& shape)
: DimComparator(ix, order, shape) {
CHECK_EQ(order.size(), ORDER_DIM);
}
inline bool operator()(const int64 i, const int64 j) const {
bool value = false;
for (int di = 0; di < ORDER_DIM; ++di) {
const int64 d = order_[di];
if (ix_(i, d) < ix_(j, d)) {
value = true;
break;
}
if (ix_(i, d) > ix_(j, d)) break;
}
return value;
}
};

} // namespace sparse
Expand Down
24 changes: 20 additions & 4 deletions tensorflow/core/util/sparse/sparse_tensor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ class SparseTensor {

// Returns the tensor shape (the dimensions of the "densified"
// tensor this tensor represents).
const TensorShape shape() const { return shape_; }
const TensorShape& shape() const { return shape_; }

const VarDimArray order() const { return order_; }

Expand Down Expand Up @@ -271,13 +271,29 @@ void SparseTensor::Reorder(const VarDimArray& order) {
auto ix_t = ix_.matrix<int64>();
auto vals_t = vals_.vec<T>();

DimComparator sorter(ix_t, order, dims_);

std::vector<int64> reorder(num_entries());
std::iota(reorder.begin(), reorder.end(), 0);

// Sort to get order of indices
std::sort(reorder.begin(), reorder.end(), sorter);
switch (order.size()) {
#define CASE_SORT(ORDER_SIZE) \
case ORDER_SIZE: { \
FixedDimComparator<ORDER_SIZE> sorter(ix_t, order, shape()); \
std::sort(reorder.begin(), reorder.end(), sorter); \
break; \
}
CASE_SORT(0);
CASE_SORT(1);
CASE_SORT(2);
CASE_SORT(3);
CASE_SORT(4);
CASE_SORT(5);
#undef CASE_SORT
default: {
DimComparator sorter(ix_t, order, shape());
std::sort(reorder.begin(), reorder.end(), sorter);
}
}

// We have a forward reordering, but what we'll need is a
// permutation (the inverse). This can be calculated with O(1)
Expand Down
122 changes: 114 additions & 8 deletions tensorflow/core/util/sparse/sparse_tensor_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@ limitations under the License.
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/random/simple_philox.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"

namespace tensorflow {
namespace sparse {
Expand Down Expand Up @@ -65,17 +67,26 @@ TEST(SparseTensorTest, DimComparatorSorts) {

// new order should be: {0, 4, 3, 2, 1}
std::vector<int64> order{0, 1, 2};
DimComparator sorter(map, order, NDIM);
TensorShape shape;
for (int i = 0; i < NDIM; ++i) shape.AddDim(N);
DimComparator sorter(map, order, shape);
std::sort(sorting.begin(), sorting.end(), sorter);
EXPECT_EQ(sorting, std::vector<int64>({0, 4, 3, 2, 1}));

FixedDimComparator<3> sorter_fixed(map, order, shape);
std::sort(sorting.begin(), sorting.end(), sorter_fixed);
EXPECT_EQ(sorting, std::vector<int64>({0, 4, 3, 2, 1}));

// new order should be: {0, 3, 2, 1, 4}
std::vector<int64> order1{2, 0, 1};
DimComparator sorter1(map, order1, NDIM);
DimComparator sorter1(map, order1, shape);
for (std::size_t n = 0; n < N; ++n) sorting[n] = n;
std::sort(sorting.begin(), sorting.end(), sorter1);
EXPECT_EQ(sorting, std::vector<int64>({0, 3, 2, 1, 4}));

FixedDimComparator<3> sorter1_fixed(map, order1, shape);
for (std::size_t n = 0; n < N; ++n) sorting[n] = n;
std::sort(sorting.begin(), sorting.end(), sorter1_fixed);
EXPECT_EQ(sorting, std::vector<int64>({0, 3, 2, 1, 4}));
}

Expand Down Expand Up @@ -304,18 +315,16 @@ TEST(SparseTensorTest, SparseTensorCheckBoundaries) {
Status st_indices_valid = st.IndicesValid();
EXPECT_FALSE(st_indices_valid.ok());
// Error message references index 4 because of the call to Reorder.
EXPECT_EQ(
"indices[4] = [11,0,0] is out of bounds: need 0 <= index < [10,10,10]",
st_indices_valid.error_message());
EXPECT_EQ("[11,0,0] is out of bounds: need 0 <= index < [10,10,10]",
st_indices_valid.error_message().substr(13));

ix_t(0, 0) = -1;
ix.matrix<int64>() = ix_t;
st.Reorder<string>(order);
st_indices_valid = st.IndicesValid();
EXPECT_FALSE(st_indices_valid.ok());
EXPECT_EQ(
"indices[0] = [-1,0,0] is out of bounds: need 0 <= index < [10,10,10]",
st_indices_valid.error_message());
EXPECT_EQ("[-1,0,0] is out of bounds: need 0 <= index < [10,10,10]",
st_indices_valid.error_message().substr(13));

ix_t(0, 0) = 0;
ix.matrix<int64>() = ix_t;
Expand Down Expand Up @@ -616,6 +625,103 @@ TEST(SparseTensorTest, Dim0SparseTensorToDenseTensor) {
EXPECT_EQ(dense.scalar<int32>()(), 5);
}

static void BM_SparseReorderFloat(int iters, int N32, int NDIM32) {
random::PhiloxRandom philox(301, 17);
random::SimplePhilox rnd(&philox);
const int64 NDIM = static_cast<int64>(NDIM32);
const int64 N = static_cast<int64>(N32);
Tensor ix(DT_INT64, TensorShape({N, NDIM}));
Tensor vals(DT_FLOAT, TensorShape({N}));
TensorShape shape;
std::vector<int64> order;
for (int d = 0; d < NDIM32; ++d) {
shape.AddDim(1000);
order.push_back(d);
}
std::vector<int64> reorder;
reorder.push_back(1);
reorder.push_back(0);
for (int d = 2; d < NDIM32; ++d) {
reorder.push_back(d);
}
auto ix_t = ix.matrix<int64>();
testing::UseRealTime();

while (--iters) {
testing::StopTiming();
for (int64 i = 0; i < N; ++i) {
for (int d = 0; d < NDIM32; ++d) {
ix_t(i, d) = rnd.Rand64() % 1000;
}
}
SparseTensor st(ix, vals, shape, order);

testing::StartTiming();
st.Reorder<float>(reorder);
}
}

static void BM_SparseReorderString(int iters, int N32, int NDIM32) {
random::PhiloxRandom philox(301, 17);
random::SimplePhilox rnd(&philox);
const int64 NDIM = static_cast<int64>(NDIM32);
const int64 N = static_cast<int64>(N32);
Tensor ix(DT_INT64, TensorShape({N, NDIM}));
Tensor vals(DT_STRING, TensorShape({N}));
TensorShape shape;
std::vector<int64> order;
auto ix_t = ix.matrix<int64>();
auto vals_t = vals.vec<string>();
for (int i = 0; i < N32; ++i) {
int len = rnd.Rand32() % 1000;
vals_t(i).resize(len);
}
for (int d = 0; d < NDIM32; ++d) {
shape.AddDim(1000);
order.push_back(d);
}
std::vector<int64> reorder;
reorder.push_back(1);
reorder.push_back(0);
for (int d = 2; d < NDIM32; ++d) {
reorder.push_back(d);
}
testing::UseRealTime();

while (--iters) {
testing::StopTiming();
for (int64 i = 0; i < N; ++i) {
for (int d = 0; d < NDIM32; ++d) {
ix_t(i, d) = rnd.Rand64() % 1000;
}
}
SparseTensor st(ix, vals, shape, order);

testing::StartTiming();
st.Reorder<string>(reorder);
}
}

BENCHMARK(BM_SparseReorderFloat)->ArgPair(10, 2);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(100, 2);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(1000, 2);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(10000, 2);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(100000, 2);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(10, 3);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(100, 3);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(1000, 3);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(10000, 3);
BENCHMARK(BM_SparseReorderFloat)->ArgPair(100000, 3);

BENCHMARK(BM_SparseReorderString)->ArgPair(10, 2);
BENCHMARK(BM_SparseReorderString)->ArgPair(100, 2);
BENCHMARK(BM_SparseReorderString)->ArgPair(1000, 2);
BENCHMARK(BM_SparseReorderString)->ArgPair(10000, 2);
BENCHMARK(BM_SparseReorderString)->ArgPair(10, 3);
BENCHMARK(BM_SparseReorderString)->ArgPair(100, 3);
BENCHMARK(BM_SparseReorderString)->ArgPair(1000, 3);
BENCHMARK(BM_SparseReorderString)->ArgPair(10000, 3);

} // namespace
} // namespace sparse
} // namespace tensorflow

0 comments on commit 1a295f8

Please sign in to comment.