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