18f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower/* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 28f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 38f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerLicensed under the Apache License, Version 2.0 (the "License"); 48f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFloweryou may not use this file except in compliance with the License. 58f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerYou may obtain a copy of the License at 68f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 78f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower http://www.apache.org/licenses/LICENSE-2.0 88f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 98f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerUnless required by applicable law or agreed to in writing, software 108f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerdistributed under the License is distributed on an "AS IS" BASIS, 118f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 128f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerSee the License for the specific language governing permissions and 138f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerlimitations under the License. 148f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower==============================================================================*/ 158f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 168f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/resource_op_kernel.h" 178f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 188f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include <memory> 198f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 208f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/allocator.h" 218f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/node_def_builder.h" 228f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/node_def_util.h" 238f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/op_kernel.h" 248f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/framework/types.h" 258f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/lib/core/errors.h" 268f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/lib/core/status_test_util.h" 278f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/lib/strings/strcat.h" 288f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/platform/mutex.h" 298f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/platform/test.h" 308f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/platform/thread_annotations.h" 318f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower#include "tensorflow/core/public/version.h" 328f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 338f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowernamespace tensorflow { 348f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowernamespace { 358f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 368f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower// Stub DeviceBase subclass which only returns allocators. 378f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerclass StubDevice : public DeviceBase { 388f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower public: 398f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubDevice() : DeviceBase(nullptr) {} 408f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 418f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Allocator* GetAllocator(AllocatorAttributes) override { 428f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return cpu_allocator(); 438f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 448f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower}; 458f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 468f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower// Stub resource for testing resource op kernel. 478f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerclass StubResource : public ResourceBase { 488f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower public: 498f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower string DebugString() override { return ""; } 508f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower int code; 518f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower}; 528f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 538f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerclass StubResourceOpKernel : public ResourceOpKernel<StubResource> { 548f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower public: 558f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower using ResourceOpKernel::ResourceOpKernel; 568f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 578f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubResource* resource() LOCKS_EXCLUDED(mu_) { 588f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower mutex_lock lock(mu_); 598f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return resource_; 608f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 618f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 628f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower private: 638f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Status CreateResource(StubResource** resource) override { 648f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower *resource = CHECK_NOTNULL(new StubResource); 658f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return GetNodeAttr(def(), "code", &(*resource)->code); 668f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 678f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 688f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Status VerifyResource(StubResource* resource) override { 698f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower int code; 708f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_RETURN_IF_ERROR(GetNodeAttr(def(), "code", &code)); 718f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower if (code != resource->code) { 728f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return errors::InvalidArgument("stub has code ", resource->code, 738f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower " but requested code ", code); 748f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 758f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return Status::OK(); 768f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 778f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower}; 788f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 798f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerREGISTER_OP("StubResourceOp") 808f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Attr("code: int") 818f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Attr("container: string = ''") 828f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Attr("shared_name: string = ''") 838f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Output("output: Ref(string)"); 848f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 858f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerREGISTER_KERNEL_BUILDER(Name("StubResourceOp").Device(DEVICE_CPU), 868f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubResourceOpKernel); 878f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 888f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerclass ResourceOpKernelTest : public ::testing::Test { 898f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower protected: 908f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower std::unique_ptr<StubResourceOpKernel> CreateOp(int code, 918f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower const string& shared_name) { 928f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower NodeDef node_def; 938f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_CHECK_OK( 948f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower NodeDefBuilder(strings::StrCat("test-node", count_++), "StubResourceOp") 958f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Attr("code", code) 968f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Attr("shared_name", shared_name) 978f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower .Finalize(&node_def)); 988f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Status status; 998f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower std::unique_ptr<OpKernel> op(CreateOpKernel( 1008f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower DEVICE_CPU, &device_, device_.GetAllocator(AllocatorAttributes()), 1018f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower node_def, TF_GRAPH_DEF_VERSION, &status)); 1028f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(status) << status; 1038f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_TRUE(op != nullptr); 1048f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1058f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower // Downcast to StubResourceOpKernel to call resource() later. 1068f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower std::unique_ptr<StubResourceOpKernel> resource_op( 1078f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower dynamic_cast<StubResourceOpKernel*>(op.get())); 1088f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_TRUE(resource_op != nullptr); 1098f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower if (resource_op != nullptr) { 1108f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower op.release(); 1118f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 1128f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return resource_op; 1138f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 1148f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1158f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Status RunOpKernel(OpKernel* op) { 1168f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower OpKernelContext::Params params; 1178f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1188f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower params.device = &device_; 1198f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower params.resource_manager = &mgr_; 1208f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower params.op_kernel = op; 1218f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1228f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower OpKernelContext context(¶ms); 1238f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower op->Compute(&context); 1248f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower return context.status(); 1258f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower } 1268f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1278f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubDevice device_; 1288f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ResourceMgr mgr_; 1298f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower int count_ = 0; 1308f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower}; 1318f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1328f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerTEST_F(ResourceOpKernelTest, PrivateResource) { 1338f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower // Empty shared_name means private resource. 1348f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower const int code = -100; 1358f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op = CreateOp(code, ""); 1368f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op != nullptr); 1378f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(RunOpKernel(op.get())); 1388f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1398f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower // Default non-shared name provided from ContainerInfo. 1408f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower const string key = "_0_" + op->name(); 1418f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1428f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubResource* resource; 1438f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_ASSERT_OK( 1448f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower mgr_.Lookup<StubResource>(mgr_.default_container(), key, &resource)); 1458f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_EQ(op->resource(), resource); // Check resource identity. 1468f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_EQ(code, resource->code); // Check resource stored information. 1478f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower resource->Unref(); 1488f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1498f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower // Destroy the op kernel. Expect the resource to be released. 1508f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower op = nullptr; 1518f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower Status s = 1528f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower mgr_.Lookup<StubResource>(mgr_.default_container(), key, &resource); 1538f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1548f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_FALSE(s.ok()); 1558f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} 1568f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1578f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerTEST_F(ResourceOpKernelTest, SharedResource) { 1588f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower const string shared_name = "shared_stub"; 1598f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower const int code = -201; 1608f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op = CreateOp(code, shared_name); 1618f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op != nullptr); 1628f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(RunOpKernel(op.get())); 1638f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1648f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower StubResource* resource; 1658f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_ASSERT_OK(mgr_.Lookup<StubResource>(mgr_.default_container(), shared_name, 1668f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower &resource)); 1678f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_EQ(op->resource(), resource); // Check resource identity. 1688f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_EQ(code, resource->code); // Check resource stored information. 1698f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower resource->Unref(); 1708f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1718f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower // Destroy the op kernel. Expect the resource not to be released. 1728f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower op = nullptr; 1738f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_ASSERT_OK(mgr_.Lookup<StubResource>(mgr_.default_container(), shared_name, 1748f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower &resource)); 1758f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower resource->Unref(); 1768f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} 1778f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1788f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerTEST_F(ResourceOpKernelTest, LookupShared) { 1798f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op1 = CreateOp(-333, "shared_stub"); 1808f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op2 = CreateOp(-333, "shared_stub"); 1818f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op1 != nullptr); 1828f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op2 != nullptr); 1838f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1848f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(RunOpKernel(op1.get())); 1858f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(RunOpKernel(op2.get())); 1868f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_EQ(op1->resource(), op2->resource()); 1878f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} 1888f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1898f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlowerTEST_F(ResourceOpKernelTest, VerifyResource) { 1908f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op1 = CreateOp(-444, "shared_stub"); 1918f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower auto op2 = CreateOp(0, "shared_stub"); // Different resource code. 1928f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op1 != nullptr); 1938f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower ASSERT_TRUE(op2 != nullptr); 1948f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 1958f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower TF_EXPECT_OK(RunOpKernel(op1.get())); 1968f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_FALSE(RunOpKernel(op2.get()).ok()); 1978f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_TRUE(op1->resource() != nullptr); 1988f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower EXPECT_TRUE(op2->resource() == nullptr); 1998f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} 2008f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower 2018f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} // namespace 2028f706abe75391bc5bd877761e244b2cf57816326A. Unique TensorFlower} // namespace tensorflow 203