| /* Copyright 2018 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 <vector> |
| |
| #include "tensorflow/core/common_runtime/dma_helper.h" |
| #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h" |
| #include "tensorflow/core/common_runtime/scoped_allocator.h" |
| #include "tensorflow/core/common_runtime/scoped_allocator_mgr.h" |
| #include "tensorflow/core/framework/fake_input.h" |
| #include "tensorflow/core/framework/node_def_builder.h" |
| #include "tensorflow/core/framework/op.h" |
| #include "tensorflow/core/framework/op_kernel.h" |
| #include "tensorflow/core/framework/tensor.h" |
| #include "tensorflow/core/framework/tensor_shape.h" |
| #include "tensorflow/core/framework/tensor_types.h" |
| #include "tensorflow/core/graph/graph.h" |
| #include "tensorflow/core/graph/node_builder.h" |
| #include "tensorflow/core/graph/testlib.h" |
| #include "tensorflow/core/kernels/ops_testutil.h" |
| #include "tensorflow/core/platform/test_benchmark.h" |
| #include "tensorflow/core/platform/types.h" |
| |
| namespace tensorflow { |
| |
| class ScopedAllocatorOpTest : public OpsTestBase { |
| protected: |
| void MakeOp(const TensorShape& shape, |
| const gtl::ArraySlice<TensorShape>& shapes, DataType dtype, |
| const string& name, int32 id, int32 expected_call_count) { |
| TF_EXPECT_OK(NodeDefBuilder("scoped_allocator_op", "_ScopedAllocator") |
| .Attr("T", dtype) |
| .Attr("shape", shape) |
| .Attr("shapes", shapes) |
| .Attr("sa_name", name) |
| .Attr("id", id) |
| .Attr("expected_call_count", expected_call_count) |
| .Finalize(node_def())); |
| TF_EXPECT_OK(InitOp()); |
| TF_ASSERT_OK(RunOpKernel()); |
| |
| // Allocate and Deallocate the tensors so that memory is not leaked |
| AllocatorAttributes attr; |
| Allocator* allocator; |
| for (size_t i = 0; i < shapes.size(); i++) { |
| attr.scope_id = id + i + 1; |
| allocator = device_->GetScopedAllocator(attr, context_->step_id()); |
| Tensor temp(allocator, dtype, shapes[i]); |
| } |
| } |
| }; |
| |
| TEST_F(ScopedAllocatorOpTest, Simple) { |
| MakeOp(TensorShape({8}), {TensorShape({8})}, DT_FLOAT, "test", 120, 1); |
| MakeOp(TensorShape({1024}), {TensorShape({32, 32})}, DT_DOUBLE, "test1", 130, |
| 1); |
| MakeOp(TensorShape({204}), |
| {TensorShape({64}), TensorShape({3, 3}), TensorShape({5, 5, 5})}, |
| DT_HALF, "test2", 140, 3); |
| MakeOp(TensorShape({1024}), {TensorShape({512}), TensorShape({64, 8})}, |
| DT_UINT32, "test3", 150, 2); |
| } |
| |
| // PrepOp is common to ConcatOp tests and SplitOpTests. |
| // It allocates a backing tensor that is large enough to hold all slices defined |
| // by fields, creates ScopedAllocatorInstances for each field, allocates the |
| // tensors, and assigns them as inputs to the op. |
| // We won't use the AddInput* suite of functions from ops_testutil.h because |
| // they allocate new tensors for each input. We need to mimic what a |
| // ScopedAllocator would do. |
| void PrepOp(DataType dtype, int32 id, |
| const std::vector<TensorShape>& fields_shapes, |
| std::vector<ScopedAllocator::Field>* fields, |
| Tensor** backing_tensor, Allocator* allocator, |
| ScopedAllocatorMgr* sam, const string& op_name, |
| std::vector<Tensor>* tensors, |
| gtl::InlinedVector<TensorValue, 4>* inputs, |
| const DataTypeVector& input_types) { |
| ScopedAllocatorMgr::PopulateFields(id, fields_shapes, dtype, fields); |
| // We don't simply allocate a tensor with shape as backing_tensor_shape, |
| // because we need to account for padding in the fields. We actually need a |
| // tensor of size at least (fields[-1].offset + fields[-1].bytes). |
| size_t num_bytes = fields->back().offset + fields->back().bytes; |
| int32_t num_elements = num_bytes / DataTypeSize(dtype); |
| CHECK_EQ(num_bytes % DataTypeSize(dtype), 0); |
| |
| *backing_tensor = new Tensor(allocator, dtype, {num_elements}); |
| int64 step_id = 10; |
| Status s = sam->AddScopedAllocator(**backing_tensor, step_id, id, |
| "sa_" + op_name + "_test", *fields, |
| fields_shapes.size()); |
| TF_ASSERT_OK(s); |
| |
| ScopedAllocatorContainer* sac = sam->GetContainer(step_id); |
| std::vector<ScopedAllocatorInstance*> sa_instances(fields_shapes.size(), |
| nullptr); |
| for (size_t i = 0; i < fields_shapes.size(); i++) { |
| sa_instances[i] = sac->GetInstance(id + i + 1); |
| tensors->push_back(Tensor(sa_instances[i], dtype, fields_shapes[i])); |
| } |
| // Now add the tensor as an input to ScopedAllocator<op_name>Op. |
| // Order matters here, so first add the backing tensor, then the slices. |
| inputs->reserve(1 + tensors->size()); |
| CHECK_GT(input_types.size(), inputs->size()); |
| CHECK_EQ(input_types[inputs->size()], dtype); |
| inputs->push_back({nullptr, *backing_tensor}); |
| for (size_t i = 0; i < tensors->size(); i++) { |
| CHECK_EQ(input_types[inputs->size()], dtype); |
| inputs->push_back({nullptr, &((*tensors)[i])}); |
| } |
| } |
| |
| class ScopedAllocatorConcatOpTest : public OpsTestBase { |
| protected: |
| void BuildNodeDef(const TensorShape& shape, DataType dtype, |
| const string& name, int32 id, int32 num_tensors) { |
| TF_EXPECT_OK( |
| NodeDefBuilder("scoped_allocator_concat_op", "_ScopedAllocatorConcat") |
| .Attr("shape", shape) |
| .Attr("T", dtype) |
| .Attr("N", num_tensors) |
| .Attr("sa_name", name) |
| .Attr("id", id) |
| .Input(FakeInput(dtype)) // backing tensor |
| .Input(FakeInput(num_tensors, dtype)) // list of tensors |
| .Finalize(node_def())); |
| shape_ = shape; |
| reshape_ = false; |
| } |
| |
| void BuildNodeDefWithReshape(const TensorShape& shape, DataType dtype, |
| bool reshape, const string& name, int32 id, |
| int32 num_tensors) { |
| TF_EXPECT_OK( |
| NodeDefBuilder("scoped_allocator_concat_op", "_ScopedAllocatorConcat") |
| .Attr("shape", shape) |
| .Attr("T", dtype) |
| .Attr("reshape", reshape) |
| .Attr("N", num_tensors) |
| .Attr("sa_name", name) |
| .Attr("id", id) |
| .Input(FakeInput(dtype)) // backing tensor |
| .Input(FakeInput(num_tensors, dtype)) // list of tensors |
| .Finalize(node_def())); |
| shape_ = shape; |
| reshape_ = reshape; |
| } |
| |
| void MakeOp(const TensorShape& shape, DataType dtype, bool reshape, |
| const string& name, int32 id, int32 num_tensors) { |
| BuildNodeDefWithReshape(shape, dtype, reshape, name, id, num_tensors); |
| TF_EXPECT_OK(InitOp()); |
| } |
| |
| void ExecOp(DataType dtype, int32 id, |
| const std::vector<TensorShape>& fields_shapes) { |
| Tensor* backing_tensor = nullptr; |
| std::vector<Tensor> tensors; |
| std::vector<ScopedAllocator::Field> fields; |
| PrepOp(dtype, id, fields_shapes, &fields, &backing_tensor, allocator(), |
| device_->GetScopedAllocatorMgr(), "concat", &tensors, &inputs_, |
| input_types_); |
| |
| TF_ASSERT_OK(RunOpKernel()); |
| |
| // Check input and output are same tensor. |
| const Tensor& input = context_->input(0); |
| OpOutputList output_list; |
| Status s = context_->output_list("output", &output_list); |
| TF_ASSERT_OK(s); |
| const Tensor& output = *(output_list[0]); |
| CHECK_EQ(DMAHelper::base(&input), DMAHelper::base(&output)); |
| CHECK_EQ(input.dtype(), output.dtype()); |
| CHECK_EQ(input.NumElements(), output.NumElements()); |
| if (reshape_) { |
| CHECK_EQ(shape_, output.shape()); |
| } else { |
| TensorShape expected_shape({input.NumElements()}); |
| CHECK_EQ(expected_shape, output.shape()); |
| } |
| |
| // Free the backing tensor which was allocated in PrepOp. |
| delete backing_tensor; |
| } |
| |
| private: |
| TensorShape shape_; |
| bool reshape_; |
| }; |
| |
| TEST_F(ScopedAllocatorConcatOpTest, Success1) { |
| MakeOp({32}, DT_FLOAT, false, "test", 120, 2); |
| ExecOp(DT_FLOAT, 120, {{16}, {16}}); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, Success2) { |
| MakeOp({2, 2, 2}, DT_DOUBLE, false, "test", 120, 2); |
| ExecOp(DT_DOUBLE, 120, {{2, 2}, {2, 2}}); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, Success3) { |
| MakeOp({3, 3, 3}, DT_HALF, false, "test", 120, 3); |
| ExecOp(DT_HALF, 120, {{3, 3}, {3, 3}, {3, 3}}); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, Reshape) { |
| MakeOp({2, 2, 4}, DT_DOUBLE, true, "test", 120, 2); |
| |
| // The elements of the third parameter to ExecOp must be multiples of |
| // Allocator::kAllocatorAlignment in size. If they are not, the backing |
| // tensor allocated by PrepOp will have too many elements and reshaping |
| // will fail. |
| ExecOp(DT_DOUBLE, 120, {{2, 4}, {2, 4}}); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, NoReshapeAttr) { |
| BuildNodeDef({3, 4, 4}, DT_HALF, "test", 120, 3); |
| TF_EXPECT_OK(InitOp()); |
| ExecOp(DT_HALF, 120, {{4, 4}, {4, 4}, {4, 4}}); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, FailDtypeCheck) { |
| MakeOp({8}, DT_FLOAT, false, "test", 120, 2); |
| EXPECT_DEATH(ExecOp(DT_DOUBLE, 120, {{4}, {4}}), ""); |
| } |
| |
| TEST_F(ScopedAllocatorConcatOpTest, FailNumElementsCheck) { |
| MakeOp({32}, DT_FLOAT, false, "test", 120, 2); |
| AddInputFromArray<float>({8}, {0, 1, 2, 3, 4, 5, 6, 7}); |
| AddInputFromArray<float>({4}, {0, 1, 2, 3}); |
| AddInputFromArray<float>({4}, {4, 5, 6, 7}); |
| Status s = RunOpKernel(); |
| EXPECT_EQ(s.code(), error::INVALID_ARGUMENT); |
| } |
| |
| // This test should fail because the backing tensor and the input tensors are |
| // unrelated, i.e. the inputs are not slices of the backing tensor. |
| TEST_F(ScopedAllocatorConcatOpTest, FailBounds) { |
| MakeOp({8}, DT_DOUBLE, false, "test", 120, 2); |
| AddInputFromArray<double>({8}, {0, 1, 2, 3, 4, 5, 6, 7}); |
| AddInputFromArray<double>({4}, {0, 1, 2, 3}); |
| AddInputFromArray<double>({4}, {4, 5, 6, 7}); |
| Status s = RunOpKernel(); |
| EXPECT_EQ(s.code(), error::INVALID_ARGUMENT); |
| } |
| |
| class ScopedAllocatorSplitOpTest : public OpsTestBase { |
| protected: |
| void BuildNodeDef(const TensorShape& in_shape, DataType dtype, |
| const string& name, int32 id, int32 num_tensors, |
| const std::vector<TensorShape>& out_shapes) { |
| TF_EXPECT_OK( |
| NodeDefBuilder("scoped_allocator_split_op", "_ScopedAllocatorSplit") |
| .Attr("T", dtype) |
| .Attr("N", num_tensors) |
| .Attr("sa_name", name) |
| .Attr("id", id) |
| .Attr("shapes", out_shapes) |
| .Input(FakeInput(dtype)) // backing tensor and input |
| .Input( |
| FakeInput(num_tensors, dtype)) // list of subtensors to forward |
| .Finalize(node_def())); |
| } |
| |
| void MakeOp(const TensorShape& in_shape, DataType dtype, const string& name, |
| int32 id, int32 num_tensors, |
| const std::vector<TensorShape>& out_shapes) { |
| BuildNodeDef(in_shape, dtype, name, id, num_tensors, out_shapes); |
| TF_EXPECT_OK(InitOp()); |
| } |
| |
| // Similar to ConcatOpTest, we add inputs that are allocated from |
| // ScopedAllocator so that the memory lines up nicely. |
| void ExecOp(DataType dtype, int32 id, |
| const std::vector<TensorShape>& fields_shapes) { |
| Tensor* backing_tensor = nullptr; |
| std::vector<Tensor> tensors; |
| std::vector<ScopedAllocator::Field> fields; |
| PrepOp(dtype, id, fields_shapes, &fields, &backing_tensor, allocator(), |
| device_->GetScopedAllocatorMgr(), "split", &tensors, &inputs_, |
| input_types_); |
| |
| TF_ASSERT_OK(RunOpKernel()); |
| |
| // Check that outputs are slices of backing tensor. |
| const Tensor& input = context_->input(0); |
| const void* lower_limit = DMAHelper::base(&input); |
| const char* lower_limit_c = |
| static_cast<const char*>(lower_limit); // for pointer arithmetic |
| OpOutputList output_list; |
| Status s = context_->output_list("output", &output_list); |
| TF_ASSERT_OK(s); |
| for (int i = 0; i < output_list.size(); i++) { |
| const Tensor& output = *(output_list[i]); |
| const void* expected_base = |
| static_cast<const void*>(lower_limit_c + fields[i].offset); |
| CHECK_EQ(output.dtype(), input.dtype()); |
| CHECK_EQ(expected_base, DMAHelper::base(&output)); |
| CHECK_EQ(output.NumElements(), fields_shapes[i].num_elements()); |
| } |
| |
| // Free the backing tensor which was allocated in PrepOp. |
| delete backing_tensor; |
| } |
| }; |
| |
| TEST_F(ScopedAllocatorSplitOpTest, Success1) { |
| MakeOp({32}, DT_FLOAT, "test", 120, 2, {{16}, {16}}); |
| ExecOp(DT_FLOAT, 120, {{16}, {16}}); |
| } |
| |
| TEST_F(ScopedAllocatorSplitOpTest, Success2) { |
| MakeOp({2, 2, 2}, DT_DOUBLE, "test", 120, 2, {{2, 2}, {2, 2}}); |
| ExecOp(DT_DOUBLE, 120, {{2, 2}, {2, 2}}); |
| } |
| |
| TEST_F(ScopedAllocatorSplitOpTest, Success3) { |
| MakeOp({3, 3, 3}, DT_HALF, "test", 120, 3, {{3, 3}, {3, 3}, {3, 3}}); |
| ExecOp(DT_HALF, 120, {{3, 3}, {3, 3}, {3, 3}}); |
| } |
| |
| TEST_F(ScopedAllocatorSplitOpTest, FailNLessThan2) { |
| BuildNodeDef({4, 4}, DT_FLOAT, "test", 120, 1, {{4, 4}}); |
| Status s = InitOp(); |
| EXPECT_EQ(s.code(), error::INVALID_ARGUMENT); |
| } |
| |
| TEST_F(ScopedAllocatorSplitOpTest, FailDtypeCheck) { |
| MakeOp({8}, DT_FLOAT, "test", 120, 2, {{4}, {4}}); |
| EXPECT_DEATH(ExecOp(DT_HALF, 120, {{4}, {4}}), ""); |
| } |
| |
| TEST_F(ScopedAllocatorSplitOpTest, FailBounds) { |
| MakeOp({8}, DT_DOUBLE, "test", 120, 2, {{4}, {4}}); |
| AddInputFromArray<double>({8}, {0, 1, 2, 3, 4, 5, 6, 7}); |
| AddInputFromArray<double>({4}, {0, 1, 2, 3}); |
| AddInputFromArray<double>({4}, {4, 5, 6, 7}); |
| Status s = RunOpKernel(); |
| EXPECT_EQ(s.code(), error::INVALID_ARGUMENT); |
| } |
| |
| } // end namespace tensorflow |
