experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Check EIGEN_MAX_ALIGN_BYTES to prevent mod-by-0 (#10380)

Browse Source 
* Check EIGEN_MAX_ALIGN_BYTES to prevent mod-by-0

If EIGEN_MAX_ALIGN_BYTES is set to 0, alignment checks that mod by EIGEN_MAX_ALIGN_BYTES fail at runtime.

* Returns true, as in tensorflow/core/framework/tensor.h
* Update unit tests

* Enable tests only if EIGEN_MAX_ALIGN_BYTES > 0
This commit is contained in:
Yan (Asta) Li 2017-06-07 21:59:27 -07:00 committed by Benoit Steiner
parent 91cb809bd6
commit 390e57a75c
2 changed files with 47 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/core/kernels/ops_util.h
View File

@ -50,8 +50,12 @@ bool IsInnerDimsSizeAligned(const TensorShape& s) {
if (s.dims() == 0) return false;
const int64 dim0_size = s.dim_size(0);
if (dim0_size == 0) return false;
#if EIGEN_MAX_ALIGN_BYTES == 0
return true;
#else
const int64 bytes_per_dim0 = (s.num_elements() / dim0_size) * sizeof(T);
return bytes_per_dim0 % EIGEN_MAX_ALIGN_BYTES == 0;
#endif
}
// Given a shape 's' of a tensor of type T and the `start` and `end` index of a
@ -61,6 +65,9 @@ bool IsInnerDimsSizeAligned(const TensorShape& s) {
template <typename T>
bool IsDim0SliceAligned(const TensorShape& s, int64 start, int64 end_or_size) {
if (s.dims() == 1) {
#if EIGEN_MAX_ALIGN_BYTES == 0
return true;
#else
bool start_aligned = (start * sizeof(T)) % EIGEN_MAX_ALIGN_BYTES == 0;
// End is aligned if either the explicit end index is passed and is a
// a multiple of EIGEN_MAX_ALIGN_BYTES, or the start index is aligned and
@ -68,6 +75,7 @@ bool IsDim0SliceAligned(const TensorShape& s, int64 start, int64 end_or_size) {
// index, or start and size.
bool end_aligned = (end_or_size * sizeof(T)) % EIGEN_MAX_ALIGN_BYTES == 0;
return start_aligned && end_aligned;
#endif
} else {
return IsInnerDimsSizeAligned<T>(s);
}

39
tensorflow/core/kernels/ops_util_test.cc
View File

@ -286,6 +286,14 @@ TEST_F(OpsUtilTest, SanitizeThreadSuffix) {
}
TEST_F(OpsUtilTest, Aligned1DSlice) {
#if EIGEN_MAX_ALIGN_BYTES == 0
// When EIGEN_MAX_ALIGN_BYTES is 0, a 1D tensor is always aligned.
Tensor t(DT_FLOAT, TensorShape({3}));
int64 start = 0;
int64 end = 1;
bool output = IsDim0SliceAligned<float>(t.shape(), start, end);
EXPECT_EQ(output, true);
#else
Tensor t(DT_FLOAT, TensorShape({EIGEN_MAX_ALIGN_BYTES * 2}));
int64 start = 0;
int64 end = EIGEN_MAX_ALIGN_BYTES;
@ -295,8 +303,10 @@ TEST_F(OpsUtilTest, Aligned1DSlice) {
Tensor sliced;
CHECK(sliced.CopyFrom(t.Slice(start, end), TensorShape({end - start})));
EXPECT_EQ(sliced.IsAligned(), true);
#endif
}
#if EIGEN_MAX_ALIGN_BYTES > 0
TEST_F(OpsUtilTest, Misaligned1DSlice) {
Tensor t(DT_FLOAT, TensorShape({EIGEN_MAX_ALIGN_BYTES * 2}));
int64 start = 1;
@ -308,8 +318,18 @@ TEST_F(OpsUtilTest, Misaligned1DSlice) {
CHECK(sliced.CopyFrom(t.Slice(start, end), TensorShape({end - start})));
EXPECT_EQ(sliced.IsAligned(), false);
}
#endif
TEST_F(OpsUtilTest, Aligned2DSliceOfDim0) {
#if EIGEN_MAX_ALIGN_BYTES == 0
// When EIGEN_MAX_ALIGN_BYTES is 0 and the size of the first dimension is nonzero,
// a multidimensional tensor is always aligned.
Tensor t(DT_FLOAT, TensorShape({3, 4}));
int64 start = 1;
int64 end = 2;
bool output = IsDim0SliceAligned<float>(t.shape(), start, end);
EXPECT_EQ(output, true);
#else
// For multidimensional tensors, alignment is dictated by inner_dim_size.
int64 inner_dim_size = EIGEN_MAX_ALIGN_BYTES;
Tensor t(DT_FLOAT, TensorShape({3, inner_dim_size}));
@ -321,8 +341,10 @@ TEST_F(OpsUtilTest, Aligned2DSliceOfDim0) {
Tensor sliced;
CHECK(sliced.CopyFrom(t.Slice(start, end), TensorShape({1, inner_dim_size})));
EXPECT_EQ(sliced.IsAligned(), true);
#endif
}
#if EIGEN_MAX_ALIGN_BYTES > 0
TEST_F(OpsUtilTest, Misaligned2DSliceOfDim0) {
// For multidimensional tensors, alignment is dictated by inner_dim_size.
int64 inner_dim_size = EIGEN_MAX_ALIGN_BYTES + 1;
@ -336,6 +358,23 @@ TEST_F(OpsUtilTest, Misaligned2DSliceOfDim0) {
CHECK(sliced.CopyFrom(t.Slice(start, end), TensorShape({1, inner_dim_size})));
EXPECT_EQ(sliced.IsAligned(), false);
}
#endif
TEST_F(OpsUtilTest, MisalignedEmptyShape) {
TensorShape shape({});
int64 start = 1;
int64 end = 2;
bool output = IsDim0SliceAligned<float>(shape, start, end);
EXPECT_EQ(output, false);
}
TEST_F(OpsUtilTest, MisalignedEmptyDim0) {
TensorShape shape({0, 1, 2});
int64 start = 0;
int64 end = 1;
bool output = IsDim0SliceAligned<float>(shape, start, end);
EXPECT_EQ(output, false);
}
} // namespace
} // namespace tensorflow