diff --git a/tensorflow/core/kernels/ops_util.h b/tensorflow/core/kernels/ops_util.h index 2d81e682ca2..68a9c37406a 100644 --- a/tensorflow/core/kernels/ops_util.h +++ b/tensorflow/core/kernels/ops_util.h @@ -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 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(s); } diff --git a/tensorflow/core/kernels/ops_util_test.cc b/tensorflow/core/kernels/ops_util_test.cc index 04a42a9921e..42ffef6735b 100644 --- a/tensorflow/core/kernels/ops_util_test.cc +++ b/tensorflow/core/kernels/ops_util_test.cc @@ -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(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(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(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(shape, start, end); + EXPECT_EQ(output, false); +} } // namespace } // namespace tensorflow