1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #include "tensorflow/core/kernels/immutable_constant_op.h" 16 17 #include <algorithm> 18 #include <tuple> 19 20 #include "tensorflow/cc/ops/standard_ops.h" 21 #include "tensorflow/core/framework/allocator.h" 22 #include "tensorflow/core/framework/tensor.h" 23 #include "tensorflow/core/graph/graph_def_builder.h" 24 #include "tensorflow/core/lib/core/status_test_util.h" 25 #include "tensorflow/core/lib/io/path.h" 26 #include "tensorflow/core/platform/test.h" 27 #include "tensorflow/core/platform/test_benchmark.h" 28 #include "tensorflow/core/public/session.h" 29 30 namespace tensorflow { 31 namespace { 32 // A safe alignment that equal to memmapped page alignment on many modern 33 // architectures. 34 constexpr size_t kTestAlignment = 4096; 35 constexpr size_t kTestTensorSize = 4; 36 constexpr size_t kTestTensorSizeBytes = kTestTensorSize * sizeof(float); 37 38 // A test ReadOnlyMemoryRegion implementation. 39 class TestReadOnlyMemoryRegion : public ReadOnlyMemoryRegion { 40 public: 41 TestReadOnlyMemoryRegion() = delete; 42 explicit TestReadOnlyMemoryRegion(uint64 length) 43 : memptr_(cpu_allocator()->AllocateRaw(kTestAlignment, length)), 44 length_(length) {} 45 ~TestReadOnlyMemoryRegion() override { 46 cpu_allocator()->DeallocateRaw(memptr_); 47 } 48 const void* data() override { return memptr_; } 49 float* GetWritableDataStart() { return reinterpret_cast<float*>(memptr_); } 50 uint64 length() override { return length_; } 51 52 protected: 53 void* memptr_; 54 uint64 length_; 55 }; 56 57 // A mock file system and environment class that creates ReadOnlyMemoryRegion 58 // from allocated memory. 59 class TestFileSystem : public NullFileSystem { 60 public: 61 ~TestFileSystem() override = default; 62 Status NewReadOnlyMemoryRegionFromFile( 63 const string& fname, 64 std::unique_ptr<ReadOnlyMemoryRegion>* result) override { 65 float val = 0; 66 StringPiece scheme, host, path; 67 io::ParseURI(fname, &scheme, &host, &path); 68 // For the tests create in-memory regions with float values equal to the 69 // region name. 70 if (path == "/2") { 71 val = 2.0f; 72 } else if (path == "/3") { 73 val = 3.0f; 74 } else { 75 val = 0.0f; 76 } 77 78 auto region = new TestReadOnlyMemoryRegion(kTestTensorSizeBytes); 79 std::fill_n(region->GetWritableDataStart(), kTestTensorSize, val); 80 result->reset(region); 81 return Status::OK(); 82 } 83 }; 84 85 REGISTER_FILE_SYSTEM("test", TestFileSystem); 86 87 struct ImmutableConstantOpTest {}; 88 89 TEST(ImmutableConstantOpTest, Simple) { 90 const TensorShape kTestTensorShape({4, 1}); 91 const TensorShape kTestTensorShapeT({1, 4}); 92 auto root = Scope::NewRootScope().ExitOnError(); 93 auto node1 = 94 ops::ImmutableConst(root, DT_FLOAT, kTestTensorShape, "test:///2"); 95 auto node2 = 96 ops::ImmutableConst(root, DT_FLOAT, kTestTensorShapeT, "test:///3"); 97 auto result = ops::MatMul(root, node1, node2); 98 GraphDef graph_def; 99 TF_ASSERT_OK(root.ToGraphDef(&graph_def)); 100 SessionOptions session_options; 101 session_options.env = Env::Default(); 102 session_options.config.mutable_graph_options() 103 ->mutable_optimizer_options() 104 ->set_opt_level(OptimizerOptions_Level_L0); 105 std::unique_ptr<Session> session(NewSession(session_options)); 106 ASSERT_TRUE(session != nullptr) << "Failed to create session"; 107 TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph"; 108 std::vector<Tensor> outputs; 109 TF_ASSERT_OK(session->Run({}, {result.node()->name() + ":0"}, {}, &outputs)); 110 ASSERT_EQ(outputs.size(), 1); 111 EXPECT_EQ(outputs.front().flat<float>()(0), 2.0f * 3.0f); 112 EXPECT_EQ(outputs.front().flat<float>()(1), 2.0f * 3.0f); 113 EXPECT_EQ(outputs.front().flat<float>()(2), 2.0f * 3.0f); 114 EXPECT_EQ(outputs.front().flat<float>()(kTestTensorSize - 1), 2.0f * 3.0f); 115 } 116 117 // Creates a test graph with two immutable_const tensors and a simple math 118 // operation, one of nodes has wrong size, check that error properly reported. 119 120 TEST(ImmutableConstantOpTest, ExecutionError) { 121 const TensorShape kBadTensorShape({40, 100}); 122 const TensorShape kTestTensorShapeT({1, 4}); 123 124 auto root = Scope::DisabledShapeInferenceScope().ExitOnError(); 125 auto node1 = 126 ops::ImmutableConst(root, DT_FLOAT, kBadTensorShape, "test:///2"); 127 auto node2 = 128 ops::ImmutableConst(root, DT_FLOAT, kTestTensorShapeT, "test:///3"); 129 auto result = ops::MatMul(root, node1, node2); 130 GraphDef graph_def; 131 TF_ASSERT_OK(root.ToGraphDef(&graph_def)); 132 SessionOptions session_options; 133 session_options.env = Env::Default(); 134 std::unique_ptr<Session> session(NewSession(session_options)); 135 ASSERT_TRUE(session != nullptr) << "Failed to create session"; 136 TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph"; 137 std::vector<Tensor> outputs; 138 // Check that the run returned error. 139 EXPECT_EQ( 140 session->Run({}, {result.node()->name() + ":0"}, {}, &outputs).code(), 141 error::INTERNAL); 142 } 143 144 Status CreateTempFile(Env* env, float value, uint64 size, string* filename) { 145 const string dir = testing::TmpDir(); 146 *filename = io::JoinPath(dir, strings::StrCat("file_", value)); 147 std::unique_ptr<WritableFile> file; 148 TF_RETURN_IF_ERROR(env->NewWritableFile(*filename, &file)); 149 for (uint64 i = 0; i < size; ++i) { 150 StringPiece sp(static_cast<char*>(static_cast<void*>(&value)), 151 sizeof(value)); 152 TF_RETURN_IF_ERROR(file->Append(sp)); 153 } 154 TF_RETURN_IF_ERROR(file->Close()); 155 return Status::OK(); 156 } 157 158 TEST(ImmutableConstantOpTest, FromFile) { 159 const TensorShape kFileTensorShape({1000, 1}); 160 Env* env = Env::Default(); 161 auto root = Scope::NewRootScope().ExitOnError(); 162 163 string two_file, three_file; 164 TF_ASSERT_OK(CreateTempFile(env, 2.0f, 1000, &two_file)); 165 TF_ASSERT_OK(CreateTempFile(env, 3.0f, 1000, &three_file)); 166 auto node1 = ops::ImmutableConst(root, DT_FLOAT, kFileTensorShape, two_file); 167 auto node2 = 168 ops::ImmutableConst(root, DT_FLOAT, kFileTensorShape, three_file); 169 auto result = ops::MatMul(root, node1, node2, ops::MatMul::TransposeB(true)); 170 171 GraphDef graph_def; 172 TF_ASSERT_OK(root.ToGraphDef(&graph_def)); 173 SessionOptions session_options; 174 session_options.config.mutable_graph_options() 175 ->mutable_optimizer_options() 176 ->set_opt_level(OptimizerOptions_Level_L0); 177 std::unique_ptr<Session> session(NewSession(session_options)); 178 ASSERT_TRUE(session != nullptr) << "Failed to create session"; 179 TF_ASSERT_OK(session->Create(graph_def)) << "Can't create test graph"; 180 std::vector<Tensor> outputs; 181 TF_ASSERT_OK(session->Run({}, {result.node()->name() + ":0"}, {}, &outputs)); 182 ASSERT_EQ(outputs.size(), 1); 183 EXPECT_EQ(outputs.front().flat<float>()(0), 2.0f * 3.0f); 184 EXPECT_EQ(outputs.front().flat<float>()(1), 2.0f * 3.0f); 185 EXPECT_EQ(outputs.front().flat<float>()(2), 2.0f * 3.0f); 186 } 187 188 } // namespace 189 } // namespace tensorflow 190
