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   std::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;  // TF_GUARDED_BY(mu_) but annotalysis doesn't
81                                 // like 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 TF_SCOPED_LOCKABLE ScopedUnlockUnrefVar {
101  public:
ScopedUnlockUnrefVar(Var * var)102   explicit ScopedUnlockUnrefVar(Var* var) TF_EXCLUSIVE_LOCK_FUNCTION(var_->mu())
103       : var_(var) {
104     if (var_) {
105       var_->mu()->lock();
106     }
107   }
Release()108   void Release() TF_UNLOCK_FUNCTION() {
109     if (var_) {
110       var_->mu()->unlock();
111       var_->Unref();
112       var_ = nullptr;
113     }
114   }
TF_UNLOCK_FUNCTION()115   ~ScopedUnlockUnrefVar() TF_UNLOCK_FUNCTION() { Release(); }
116 
117  private:
118   Var* var_;
119 
120   ScopedUnlockUnrefVar(const ScopedUnlockUnrefVar&) = delete;
121   ScopedUnlockUnrefVar(ScopedUnlockUnrefVar&&) = delete;
122   ScopedUnlockUnrefVar& operator=(const ScopedUnlockUnrefVar&) = delete;
123   ScopedUnlockUnrefVar& operator=(ScopedUnlockUnrefVar&&) = delete;
124 };
125 
126 }  //  end namespace tensorflow
127 
128 #endif  // TENSORFLOW_CORE_FRAMEWORK_RESOURCE_VAR_H_
129