From a67ee929f5aa2e16478d10e3287a248f34078cb5 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 27 May 2020 10:21:19 -0700 Subject: [PATCH] add a tensorflow::batch_util::CopyContiguousSlices utility function for slicing out a contiguous pieces of tensors along the batch dimension and copying them to another tensor. PiperOrigin-RevId: 313414257 Change-Id: I2530c58ed53ad8e92e5f976f2dd1728296d12185 --- tensorflow/core/framework/BUILD | 1 + tensorflow/core/framework/batch_util_test.cc | 61 ++++++++++++++++ tensorflow/core/framework/tensor.h | 5 ++ tensorflow/core/util/batch_util.cc | 73 ++++++++++++++++++++ tensorflow/core/util/batch_util.h | 11 +++ 5 files changed, 151 insertions(+) create mode 100644 tensorflow/core/framework/batch_util_test.cc diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD index 0e923bd1236..52f15dcb5c2 100644 --- a/tensorflow/core/framework/BUILD +++ b/tensorflow/core/framework/BUILD @@ -1008,6 +1008,7 @@ tf_cc_tests( srcs = [ "allocator_test.cc", "attr_value_util_test.cc", + "batch_util_test.cc", "bfloat16_test.cc", "cancellation_test.cc", "common_shape_fns_test.cc", diff --git a/tensorflow/core/framework/batch_util_test.cc b/tensorflow/core/framework/batch_util_test.cc new file mode 100644 index 00000000000..4e98371bda7 --- /dev/null +++ b/tensorflow/core/framework/batch_util_test.cc @@ -0,0 +1,61 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { + +TEST(CopyContiguousSlicesTest, CompatibleShape) { + Tensor src(DT_FLOAT, {7, 1, 2}); + Tensor dst(DT_FLOAT, {9, 2, 1}); + auto s = batch_util::CopyContiguousSlices( + src, /*src_offset=*/2, /*dst_offset=*/0, /*num_slices=*/5, &dst); + ASSERT_EQ(error::OK, s.code()); +} + +TEST(CopyContiguousSlicesTest, SourceOffsetOutOfRange) { + Tensor src(DT_FLOAT, {7, 1, 2}); + Tensor dst(DT_FLOAT, {9, 2, 1}); + auto s = batch_util::CopyContiguousSlices( + src, /*src_offset=*/7, /*dst_offset=*/0, /*num_slices=*/5, &dst); + ASSERT_EQ(error::FAILED_PRECONDITION, s.code()); +} + +TEST(CopyContiguousSlicesTest, DstOffsetOutOfRange) { + Tensor src(DT_FLOAT, {7, 1, 2}); + Tensor dst(DT_FLOAT, {9, 2, 1}); + auto s = batch_util::CopyContiguousSlices( + src, /*src_offset=*/0, /*dst_offset=*/0, /*num_slices=*/8, &dst); + ASSERT_EQ(error::FAILED_PRECONDITION, s.code()); +} + +TEST(CopyContiguousSlicesTest, CheckDstWithExpectedValues) { + auto src = test::AsTensor({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, + TensorShape({5, 2})); + Tensor dst(DT_FLOAT, {9, 2, 1}); + auto s = batch_util::CopyContiguousSlices( + src, /*src_offset=*/1, /*dst_offset=*/5, /*num_slices=*/3, &dst); + ASSERT_EQ(error::OK, s.code()); + test::ExpectTensorEqual( + test::AsTensor({2, 3, 4, 5, 6, 7}, TensorShape({3, 2, 1})), + dst.Slice(5, 8)); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h index 744a14e007e..28eab3ab1e0 100644 --- a/tensorflow/core/framework/tensor.h +++ b/tensorflow/core/framework/tensor.h @@ -53,6 +53,8 @@ namespace batch_util { Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index); Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index); Status MaybeMoveSliceToElement(Tensor* parent, Tensor* element, int64 index); +Status CopyContiguousSlices(const Tensor& src, int64 src_offset, + int64 dst_offset, int64 num_slices, Tensor* dst); } // namespace batch_util /// @ingroup core @@ -679,6 +681,9 @@ class Tensor { friend Status batch_util::MaybeMoveSliceToElement( Tensor* parent, Tensor* element, int64 index); // For access to base(). + friend Status batch_util::CopyContiguousSlices( + const Tensor& src, int64 src_offset, int64 dst_offset, int64 num_slices, + Tensor* dst); // For access to base(). bool CanUseDMA() const; diff --git a/tensorflow/core/util/batch_util.cc b/tensorflow/core/util/batch_util.cc index 0aff6b00f1c..b88c365ced0 100644 --- a/tensorflow/core/util/batch_util.cc +++ b/tensorflow/core/util/batch_util.cc @@ -216,6 +216,79 @@ Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index) { } } +Status CopyContiguousSlices(const Tensor& src, int64 src_offset, + int64 dst_offset, int64 num_slices, Tensor* dst) { + if (src.dtype() != dst->dtype()) { + return errors::FailedPrecondition( + "CopyContiguousSlices cannot perform copy: src and dst have different " + "dtypes. Source dtype: ", + src.dtype(), " dstination dtype: ", dst->dtype(), "."); + } + if (src.dims() < 1) { + return errors::FailedPrecondition( + "CopyContiguousSlices cannot perform copy: src has to be a tensor with " + "rank >= 1. Source shape: ", + src.shape().DebugString()); + } + + if (dst->dims() < 1) { + return errors::FailedPrecondition( + "CopyContiguousSlices cannot perform copy: dst has to be a tensor " + "with rank >= 1. Dest shape: ", + dst->shape().DebugString()); + } + + const int64 src_dim0 = src.dim_size(0); + const int64 dst_dim0 = dst->dim_size(0); + int64 src_chip_size = 1; + int64 dst_chip_size = 1; + for (int i = 1; i < src.dims(); ++i) { + src_chip_size *= src.dim_size(i); + } + for (int i = 1; i < dst->dims(); ++i) { + dst_chip_size *= dst->dim_size(i); + } + + if (src_chip_size != dst_chip_size) { + return errors::FailedPrecondition( + "CopyContiguousSlices cannot perform copy: source and dst shapes are" + "not compatible. Source shape: ", + src.shape().DebugString(), ", dst shape: ", dst->shape().DebugString()); + } + + if (src_chip_size == 0 && dst_chip_size == 0) { + return Status::OK(); + } + + if (src_offset < 0 || src_offset + num_slices > src_dim0 || dst_offset < 0 || + dst_offset + num_slices > dst_dim0) { + return errors::FailedPrecondition( + "CopyContiguousSlices cannot perform copy: index out of range. " + "src_offset: ", + src_offset, ", num_slices: ", num_slices, ", src_dim0: ", src_dim0, + ", dst_offset: ", dst_offset, ", dst_dim0: ", dst_dim0, "."); + } + +#define HANDLE_TYPE(T) \ + case DataTypeToEnum::value: { \ + const T* src_p = src.base() + (src_chip_size * src_offset); \ + T* dst_p = dst->base() + (dst_chip_size * dst_offset); \ + HandleSliceToElement(src_p, dst_p, src_chip_size * num_slices); \ + return Status::OK(); \ + } + + switch (src.dtype()) { + TF_CALL_ALL_TYPES(HANDLE_TYPE); + TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE); + TF_CALL_uint32(HANDLE_TYPE); + TF_CALL_uint64(HANDLE_TYPE); +#undef HANDLE_TYPE + default: + return errors::Unimplemented("CopyContiguousSlices unhandled data type: ", + src.dtype()); + } +} + // Copies the index^th slice of parent (in the 0th dimension) into element. // // NOTE(mrry): The implementation may be able to optimize the copy to a move. diff --git a/tensorflow/core/util/batch_util.h b/tensorflow/core/util/batch_util.h index eee0309fbc4..d44d82ec0a2 100644 --- a/tensorflow/core/util/batch_util.h +++ b/tensorflow/core/util/batch_util.h @@ -32,6 +32,17 @@ Status CopyElementToSlice(Tensor element, Tensor* parent, int64 index); // Copies the index^th slice of parent (in the 0th dimension) into element. Status CopySliceToElement(const Tensor& parent, Tensor* element, int64 index); +// Copies 'num_slices' contiguous slices from 'src' tensor starting from index +// 'src_offset' into target tensor 'dst', and places them into slices +// starting from 'dst_offset'. +// +// This function requires 'src' and 'dst' to have compatible shapes. That is it +// requires cum_prod(src.shape[1:] == cum_prod(dst->shape[1:]). For example if +// source is of shape [x, 2, 1] and dst is a tensor of shape [y, 1, 2], this +// function can still proceed successfully. +Status CopyContiguousSlices(const Tensor& src, int64 src_offset, + int64 dst_offset, int64 num_slices, Tensor* dst); + // Copies the index^th slice of parent (in the 0th dimension) into element. // // NOTE(mrry): The implementation may be able to optimize the copy to a move.