1 /* Copyright 2018 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 16 #ifndef TENSORFLOW_CORE_FRAMEWORK_RESOURCE_VAR_H_ 17 #define TENSORFLOW_CORE_FRAMEWORK_RESOURCE_VAR_H_ 18 19 #include "tensorflow/core/framework/resource_mgr.h" 20 #include "tensorflow/core/lib/core/status.h" 21 22 // Forward declarations to avoid introducing a dependency on headers in 23 // "tensorflow/core/graph/...". 24 class GraphDefBuilder; 25 26 namespace tensorflow { 27 28 // Resource stored by variables in the resource manager (new, resource-style 29 // version). 30 // 31 // These variables have a mixed access mode: they can operate on copy-on-write 32 // mode (the default) or copy-on-read mode (used only for sparse access). 33 // 34 // When copy-on-write mode is enabled reading the value of the variable involves 35 // grabbing its mutex in shared mode and aliasing the internal tensor as the 36 // output of the read operation, increasing its reference count. Writing, 37 // conversely, works by, under an exclusive lock, detecting whether there are 38 // outstanding aliases of the tensor, using the reference count, copying the 39 // tensor if they exist, and writing to either the original or a copy with no 40 // outstanding aliases. Sparse operations are not supported in copy-on-write 41 // mode. 42 // 43 // When a variable is accessed sparsely it switches to copy-on-read mode. To 44 // switch we need to grab an exclusive lock and might (if there are aliases) 45 // need to copy the entire tensor. Once copy-on-read mode is enabled, no tensor 46 // is allowed to alias the variable's internal tensor. This means dense reads 47 // must return a copy of the variable, done while holding a shared lock. Dense 48 // writes do not need to check whether aliases exist, and can always write 49 // directly to the buffer without making a copy, while holding an exclusive 50 // lock. Sparse reads and sparse writes, on the other hand, can be done under a 51 // shared or exclusive mutex (the damage from writes under a shared mutex is 52 // limited since no other buffer is allowed to alias the variable's 53 // buffer). Using an exclusive mutex disallows concurrent writes and concurrent 54 // sparse reads, providing some extra safety at the expense of performance, 55 // while shared mutex allow for "hogwild" behavior. Doing sparse writes under a 56 // shared mutex prevents them from overlapping with dense writes, which is 57 // necessary as dense writes can change the shape the of the tensor. 58 // 59 // Transitioning a variable from copy-on-read mode to copy-on-write mode is 60 // currently not supported. To upgrade a variable from copy-on-write to 61 // copy-on-read use `EnsureSparseVariableAccess()`, and then grab the variable's 62 // mutex as desired. To access the variable in dense mode grab the mutex either 63 // directly or via `MaybeLockVariableInputMutexesInOrder` on all variables being 64 // modified and then call `PrepareToUpdateVariable` on them in any order. 65 class Var : public ResourceBase { 66 public: Var(DataType dtype)67 explicit Var(DataType dtype) : tensor_(dtype) {} 68 69 // When locking multiple variables, the locks must be acquired in order of 70 // increasing mu() address. 71 // TODO(ebrevdo): Use LockSet instead of exposing mu. mu()72 mutex* mu() { return &mu_; } tensor()73 Tensor* tensor() { return &tensor_; } 74 75 Status AsGraphDef(GraphDefBuilder* builder, Node** out) const override; 76 DebugString()77 std::string DebugString() const override { 78 return strings::StrCat(DataTypeString(tensor_.dtype()), "/", 79 tensor_.shape().DebugString()); 80 } 81 82 // Only used in the resource variable path. In resource variables, 83 // tensor.IsInitialized() can be true (i.e. have memory allocated to it) while 84 // there is not a good value there due to a race condition, and it's possible 85 // to stumble upon this during variable.initialized_value(). So it's best to 86 // just store directly whether the variable is initialized. 87 bool is_initialized = false; // TF_GUARDED_BY(mu_) but annotalysis doesn't 88 // like it. 89 90 // Also fake-guarded by mu_. Should be set to True whenever any sparse 91 // operation uses the variable. Once this is true no tensor is allowed to 92 // alias the memory of the variable, and we always copy the variable on 93 // reads. This allows sparse operations to happen with only a shared lock if 94 // so desired. 95 std::atomic<bool> copy_on_read_mode{false}; 96 97 private: 98 mutex mu_; 99 Tensor tensor_; 100 ~Var()101 ~Var() override {} 102 TF_DISALLOW_COPY_AND_ASSIGN(Var); 103 }; 104 105 // Does unlock and unref automatically when going out of scope, and also 106 // supports early manual release. 107 class TF_SCOPED_LOCKABLE ScopedUnlockUnrefVar { 108 public: ScopedUnlockUnrefVar(Var * var)109 explicit ScopedUnlockUnrefVar(Var* var) TF_EXCLUSIVE_LOCK_FUNCTION(var_->mu()) 110 : var_(var) { 111 if (var_) { 112 var_->mu()->lock(); 113 } 114 } Release()115 void Release() TF_UNLOCK_FUNCTION() { 116 if (var_) { 117 var_->mu()->unlock(); 118 var_->Unref(); 119 var_ = nullptr; 120 } 121 } TF_UNLOCK_FUNCTION()122 ~ScopedUnlockUnrefVar() TF_UNLOCK_FUNCTION() { Release(); } 123 124 private: 125 Var* var_; 126 127 ScopedUnlockUnrefVar(const ScopedUnlockUnrefVar&) = delete; 128 ScopedUnlockUnrefVar(ScopedUnlockUnrefVar&&) = delete; 129 ScopedUnlockUnrefVar& operator=(const ScopedUnlockUnrefVar&) = delete; 130 ScopedUnlockUnrefVar& operator=(ScopedUnlockUnrefVar&&) = delete; 131 }; 132 133 } // end namespace tensorflow 134 135 #endif // TENSORFLOW_CORE_FRAMEWORK_RESOURCE_VAR_H_ 136