1c8b59c046895fa5b6d79f73e0b5817330fcfbfc1A. Unique TensorFlower/* Copyright 2015 The TensorFlow Authors. All Rights Reserved. 29c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 39c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurLicensed under the Apache License, Version 2.0 (the "License"); 49c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudluryou may not use this file except in compliance with the License. 59c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurYou may obtain a copy of the License at 69c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 79c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur http://www.apache.org/licenses/LICENSE-2.0 89c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 99c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurUnless required by applicable law or agreed to in writing, software 109c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurdistributed under the License is distributed on an "AS IS" BASIS, 119c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 129c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath KudlurSee the License for the specific language governing permissions and 139c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlurlimitations under the License. 149c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur==============================================================================*/ 159c3043ff3bf31a6a81810b4ce9e87ef936f1f529Manjunath Kudlur 16f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/op_segment.h" 17f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 18b481783fe0e00a86f6feb20a8dcad5fc4fc936a4Josh Levenberg#include <vector> 19f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/allocator.h" 20f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/node_def_builder.h" 21f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/framework/op_kernel.h" 22f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/kernels/ops_util.h" 23f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/lib/core/errors.h" 24f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur#include "tensorflow/core/lib/core/status_test_util.h" 2556313def004795f75ef8281a0294c958d28f1e06Vijay Vasudevan#include "tensorflow/core/lib/strings/strcat.h" 2656313def004795f75ef8281a0294c958d28f1e06Vijay Vasudevan#include "tensorflow/core/platform/logging.h" 27c8eaac926c929e07ac8db69f67803a2223ff2d93Josh Levenberg#include "tensorflow/core/platform/test.h" 28d9cfc64a2ddf05c0b093c8fb6704c67452ee3ea0Vijay Vasudevan#include "tensorflow/core/public/version.h" 29f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 30f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurnamespace tensorflow { 31f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 32f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlurclass OpSegmentTest : public ::testing::Test { 33f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur protected: 34f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur DeviceBase device_; 35f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur std::vector<NodeDef> int32_nodedefs_; 36f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur std::vector<NodeDef> float_nodedefs_; 37f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 38f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegmentTest() : device_(Env::Default()) { 39f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur for (int i = 0; i < 10; ++i) { 40f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur NodeDef def; 41f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur TF_CHECK_OK(NodeDefBuilder(strings::StrCat("op", i), "Mul") 42f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Input("x", 0, DT_INT32) 43f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Input("y", 0, DT_INT32) 44f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Finalize(&def)); 45f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur int32_nodedefs_.push_back(def); 46f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur TF_CHECK_OK(NodeDefBuilder(strings::StrCat("op", i), "Mul") 47f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Input("x", 0, DT_FLOAT) 48f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Input("y", 0, DT_FLOAT) 49f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur .Finalize(&def)); 50f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur float_nodedefs_.push_back(def); 51f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 52f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 53f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 54f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur void ValidateOpAndTypes(OpKernel* op, const NodeDef& expected, DataType dt) { 55f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ASSERT_NE(op, nullptr); 56f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(expected.DebugString(), op->def().DebugString()); 57f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(2, op->num_inputs()); 58f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(dt, op->input_type(0)); 59f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(dt, op->input_type(1)); 60f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(1, op->num_outputs()); 61f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_EQ(dt, op->output_type(0)); 62f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 63f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 64f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegment::CreateKernelFn GetFn(const NodeDef* ndef) { 65f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return [this, ndef](OpKernel** kernel) { 66f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Status s; 67d9cfc64a2ddf05c0b093c8fb6704c67452ee3ea0Vijay Vasudevan auto created = CreateOpKernel(DEVICE_CPU, &device_, cpu_allocator(), 68d9cfc64a2ddf05c0b093c8fb6704c67452ee3ea0Vijay Vasudevan *ndef, TF_GRAPH_DEF_VERSION, &s); 69f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur if (s.ok()) { 70f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur *kernel = created.release(); 71f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 72f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return s; 73f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur }; 74f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 75f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur}; 76f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 77f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST_F(OpSegmentTest, Basic) { 78f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegment opseg; 79f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpKernel* op; 80f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 81f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.AddHold("A"); 82f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.AddHold("B"); 83f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur for (int i = 0; i < 10; ++i) { 84f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Register in session A. 85f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur auto* ndef = &float_nodedefs_[i]; 86c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan TF_EXPECT_OK(opseg.FindOrCreate("A", ndef->name(), &op, GetFn(ndef))); 87f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ValidateOpAndTypes(op, *ndef, DT_FLOAT); 88f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 89f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Register in session B. 90f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ndef = &int32_nodedefs_[i]; 91c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan TF_EXPECT_OK(opseg.FindOrCreate("B", ndef->name(), &op, GetFn(ndef))); 92f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ValidateOpAndTypes(op, *ndef, DT_INT32); 93f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 94f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 95f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur auto reterr = [](OpKernel** kernel) { 96f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur return errors::Internal("Should not be called"); 97f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur }; 98f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur for (int i = 0; i < 10; ++i) { 99f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Lookup op in session A. 100c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan TF_EXPECT_OK( 101c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan opseg.FindOrCreate("A", strings::StrCat("op", i), &op, reterr)); 102f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ValidateOpAndTypes(op, float_nodedefs_[i], DT_FLOAT); 103f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 104f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Lookup op in session B. 105c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan TF_EXPECT_OK( 106c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan opseg.FindOrCreate("B", strings::StrCat("op", i), &op, reterr)); 107f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ValidateOpAndTypes(op, int32_nodedefs_[i], DT_INT32); 108f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur } 109f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 110f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("A"); 111f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("B"); 112f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 113f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 114f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST_F(OpSegmentTest, SessionNotFound) { 115f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegment opseg; 116f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpKernel* op; 117f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur NodeDef def = float_nodedefs_[0]; 118f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Status s = opseg.FindOrCreate("A", def.name(), &op, GetFn(&def)); 119f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_TRUE(errors::IsNotFound(s)) << s; 120f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 121f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 122f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST_F(OpSegmentTest, CreateFailure) { 123f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegment opseg; 124f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpKernel* op; 125f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur NodeDef def = float_nodedefs_[0]; 126f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur def.set_op("nonexistop"); 127f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.AddHold("A"); 128f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur Status s = opseg.FindOrCreate("A", def.name(), &op, GetFn(&def)); 129f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur EXPECT_TRUE(errors::IsNotFound(s)) << s; 130f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("A"); 131f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 132f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 133f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath KudlurTEST_F(OpSegmentTest, AddRemoveHolds) { 134f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpSegment opseg; 135f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur OpKernel* op; 136f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur const auto& ndef = int32_nodedefs_[0]; 137f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 138f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // No op. 139f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("null"); 140f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 141f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Thread1 register the op and wants to ensure it alive. 142f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.AddHold("foo"); 143c3c27f275f94af6e32ecaff66a6bf439abdaff5bVijay Vasudevan TF_EXPECT_OK(opseg.FindOrCreate("foo", ndef.name(), &op, GetFn(&ndef))); 144f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 145f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Thread2 starts some execution needs "op" to be alive. 146f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.AddHold("foo"); 147f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 148f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Thread1 clears session "foo". E.g., a master sends CleanupGraph 149f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // before an execution finishes. 150f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("foo"); 151f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 152f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Thread2 should still be able to access "op". 153f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur ValidateOpAndTypes(op, ndef, DT_INT32); 154f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 155f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur // Thread2 then remove its hold on "foo". 156f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur opseg.RemoveHold("foo"); 157f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} 158f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur 159f41959ccb2d9d4c722fe8fc3351401d53bcf490Manjunath Kudlur} // namespace tensorflow 160