Coverage Report

Created: 2025-06-13 07:02

/src/tesseract/src/lstm/networkscratch.h
Line
Count
Source (jump to first uncovered line)
1
///////////////////////////////////////////////////////////////////////
2
// File:        networkscratch.h
3
// Description: Scratch space for Network layers that hides distinction
4
//              between float/int implementations.
5
// Author:      Ray Smith
6
//
7
// (C) Copyright 2014, Google Inc.
8
// Licensed under the Apache License, Version 2.0 (the "License");
9
// you may not use this file except in compliance with the License.
10
// You may obtain a copy of the License at
11
// http://www.apache.org/licenses/LICENSE-2.0
12
// Unless required by applicable law or agreed to in writing, software
13
// distributed under the License is distributed on an "AS IS" BASIS,
14
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
// See the License for the specific language governing permissions and
16
// limitations under the License.
17
///////////////////////////////////////////////////////////////////////
18
19
#ifndef TESSERACT_LSTM_NETWORKSCRATCH_H_
20
#define TESSERACT_LSTM_NETWORKSCRATCH_H_
21
22
#include <mutex>
23
#include "matrix.h"
24
#include "networkio.h"
25
26
namespace tesseract {
27
28
// Generic scratch space for network layers. Provides NetworkIO that can store
29
// a complete set (over time) of intermediates, and vector<float>
30
// scratch space that auto-frees after use. The aim here is to provide a set
31
// of temporary buffers to network layers that can be reused between layers
32
// and don't have to be reallocated on each call.
33
class NetworkScratch {
34
public:
35
2
  NetworkScratch() : int_mode_(false) {}
36
0
  ~NetworkScratch() = default;
37
38
  // Sets the network representation. If the representation is integer, then
39
  // default (integer) NetworkIOs are separated from the always-float variety.
40
  // This saves memory by having separate int-specific and float-specific
41
  // stacks. If the network representation is float, then all NetworkIOs go
42
  // to the float stack.
43
0
  void set_int_mode(bool int_mode) {
44
0
    int_mode_ = int_mode;
45
0
  }
46
47
  // Class that acts like a NetworkIO (by having an implicit cast operator),
48
  // yet actually holds a pointer to NetworkIOs in the source NetworkScratch,
49
  // and knows how to unstack the borrowed pointers on destruction.
50
  class IO {
51
  public:
52
    // The NetworkIO should be sized after construction.
53
    IO(const NetworkIO &src, NetworkScratch *scratch)
54
690k
        : int_mode_(scratch->int_mode_ && src.int_mode()), scratch_space_(scratch) {
55
690k
      network_io_ =
56
690k
          int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow();
57
690k
    }
58
    // Default constructor for arrays. Use one of the Resize functions
59
    // below to initialize and size.
60
345k
    IO() : int_mode_(false), network_io_(nullptr), scratch_space_(nullptr) {}
61
62
1.03M
    ~IO() {
63
1.03M
      if (scratch_space_ == nullptr) {
64
345k
        ASSERT_HOST(network_io_ == nullptr);
65
690k
      } else if (int_mode_) {
66
0
        scratch_space_->int_stack_.Return(network_io_);
67
690k
      } else {
68
690k
        scratch_space_->float_stack_.Return(network_io_);
69
690k
      }
70
1.03M
    }
71
    // Resizes the array (and stride), avoiding realloc if possible, to the
72
    // size from various size specs:
73
    // Same time size, given number of features.
74
0
    void Resize(const NetworkIO &src, int num_features, NetworkScratch *scratch) {
75
0
      if (scratch_space_ == nullptr) {
76
0
        int_mode_ = scratch->int_mode_ && src.int_mode();
77
0
        scratch_space_ = scratch;
78
0
        network_io_ =
79
0
            int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow();
80
0
      }
81
0
      network_io_->Resize(src, num_features);
82
0
    }
83
    // Resizes to a specific size as a temp buffer. No batches, no y-dim.
84
0
    void Resize2d(bool int_mode, int width, int num_features, NetworkScratch *scratch) {
85
0
      if (scratch_space_ == nullptr) {
86
0
        int_mode_ = scratch->int_mode_ && int_mode;
87
0
        scratch_space_ = scratch;
88
0
        network_io_ =
89
0
            int_mode_ ? scratch_space_->int_stack_.Borrow() : scratch_space_->float_stack_.Borrow();
90
0
      }
91
0
      network_io_->Resize2d(int_mode, width, num_features);
92
0
    }
93
    // Resize forcing a float representation with the width of src and the given
94
    // number of features.
95
0
    void ResizeFloat(const NetworkIO &src, int num_features, NetworkScratch *scratch) {
96
0
      if (scratch_space_ == nullptr) {
97
0
        int_mode_ = false;
98
0
        scratch_space_ = scratch;
99
0
        network_io_ = scratch_space_->float_stack_.Borrow();
100
0
      }
101
0
      network_io_->ResizeFloat(src, num_features);
102
0
    }
103
104
    // Returns a ref to a NetworkIO that enables *this to be treated as if
105
    // it were just a NetworkIO*.
106
1.03M
    NetworkIO &operator*() {
107
1.03M
      return *network_io_;
108
1.03M
    }
109
0
    NetworkIO *operator->() {
110
0
      return network_io_;
111
0
    }
112
1.03M
    operator NetworkIO *() {
113
1.03M
      return network_io_;
114
1.03M
    }
115
116
  private:
117
    // True if this is from the always-float stack, otherwise the default stack.
118
    bool int_mode_;
119
    // The NetworkIO that we have borrowed from the scratch_space_.
120
    NetworkIO *network_io_;
121
    // The source scratch_space_. Borrowed pointer, used to free the
122
    // NetworkIO. Don't delete!
123
    NetworkScratch *scratch_space_;
124
  }; // class IO.
125
126
  // Class that acts like a fixed array of float, yet actually uses space
127
  // from a vector<float> in the source NetworkScratch, and knows how
128
  // to unstack the borrowed vector on destruction.
129
  class FloatVec {
130
  public:
131
    // The array will have size elements in it, uninitialized.
132
0
    FloatVec(int size, NetworkScratch *scratch) : vec_(nullptr), scratch_space_(scratch) {
133
0
      Init(size, scratch);
134
0
    }
135
    // Default constructor is for arrays. Use Init to setup.
136
3.45M
    FloatVec() : vec_(nullptr), data_(nullptr), scratch_space_(nullptr) {}
137
3.45M
    ~FloatVec() {
138
3.45M
      if (scratch_space_ != nullptr) {
139
3.10M
        scratch_space_->vec_stack_.Return(vec_);
140
3.10M
      }
141
3.45M
    }
142
143
3.10M
    void Init(int /*size*/, int reserve, NetworkScratch *scratch) {
144
3.10M
      if (scratch_space_ != nullptr && vec_ != nullptr) {
145
0
        scratch_space_->vec_stack_.Return(vec_);
146
0
      }
147
3.10M
      scratch_space_ = scratch;
148
3.10M
      vec_ = scratch_space_->vec_stack_.Borrow();
149
      // TODO: optimize.
150
3.10M
      vec_->resize(reserve);
151
3.10M
      data_ = &(*vec_)[0];
152
3.10M
    }
153
154
1.38M
    void Init(int size, NetworkScratch *scratch) {
155
1.38M
      Init(size, size, scratch);
156
1.38M
    }
157
158
    // Use the cast operator instead of operator[] so the FloatVec can be used
159
    // as a TFloat* argument to a function call.
160
295M
    operator TFloat *() const {
161
295M
      return data_;
162
295M
    }
163
0
    TFloat *get() {
164
0
      return data_;
165
0
    }
166
167
  private:
168
    // Vector borrowed from the scratch space. Use Return to free it.
169
    std::vector<TFloat> *vec_;
170
    // Short-cut pointer to the underlying array.
171
    TFloat *data_;
172
    // The source scratch_space_. Borrowed pointer, used to free the
173
    // vector. Don't delete!
174
    NetworkScratch *scratch_space_;
175
  }; // class FloatVec
176
177
  // Class that acts like a 2-D array of TFloat, yet actually uses space
178
  // from the source NetworkScratch, and knows how to unstack the borrowed
179
  // array on destruction.
180
  class GradientStore {
181
  public:
182
    // Default constructor is for arrays. Use Init to setup.
183
0
    GradientStore() : array_(nullptr), scratch_space_(nullptr) {}
184
0
    ~GradientStore() {
185
0
      if (scratch_space_ != nullptr) {
186
0
        scratch_space_->array_stack_.Return(array_);
187
0
      }
188
0
    }
189
190
0
    void Init(int size1, int size2, NetworkScratch *scratch) {
191
0
      if (scratch_space_ != nullptr && array_ != nullptr) {
192
0
        scratch_space_->array_stack_.Return(array_);
193
0
      }
194
0
      scratch_space_ = scratch;
195
0
      array_ = scratch_space_->array_stack_.Borrow();
196
0
      array_->Resize(size1, size2, 0.0);
197
0
    }
198
199
    // Accessors to get to the underlying TransposedArray.
200
0
    TransposedArray *get() const {
201
0
      return array_;
202
0
    }
203
0
    const TransposedArray &operator*() const {
204
0
      return *array_;
205
0
    }
206
207
  private:
208
    // Array borrowed from the scratch space. Use Return to free it.
209
    TransposedArray *array_;
210
    // The source scratch_space_. Borrowed pointer, used to free the
211
    // vector. Don't delete!
212
    NetworkScratch *scratch_space_;
213
  }; // class GradientStore
214
215
  // Class that does the work of holding a stack of objects, a stack pointer
216
  // and a vector of in-use flags, so objects can be returned out of order.
217
  // It is safe to attempt to Borrow/Return in multiple threads.
218
  template <typename T>
219
  class Stack {
220
  public:
221
8
    Stack() = default;
tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Stack()
Line
Count
Source
221
4
    Stack() = default;
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Stack()
Line
Count
Source
221
2
    Stack() = default;
tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Stack()
Line
Count
Source
221
2
    Stack() = default;
222
223
0
    ~Stack() {
224
0
      for (auto data : stack_) {
225
0
        delete data;
226
0
      }
227
0
    }
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::~Stack()
Unexecuted instantiation: tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::~Stack()
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::~Stack()
228
229
    // Lends out the next free item, creating one if none available, sets
230
    // the used flags and increments the stack top.
231
3.79M
    T *Borrow() {
232
3.79M
      std::lock_guard<std::mutex> lock(mutex_);
233
3.79M
      if (stack_top_ == stack_.size()) {
234
12
        stack_.push_back(new T);
235
12
        flags_.push_back(false);
236
12
      }
237
3.79M
      flags_[stack_top_] = true;
238
3.79M
      return stack_[stack_top_++];
239
3.79M
    }
tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Borrow()
Line
Count
Source
231
690k
    T *Borrow() {
232
690k
      std::lock_guard<std::mutex> lock(mutex_);
233
690k
      if (stack_top_ == stack_.size()) {
234
4
        stack_.push_back(new T);
235
4
        flags_.push_back(false);
236
4
      }
237
690k
      flags_[stack_top_] = true;
238
690k
      return stack_[stack_top_++];
239
690k
    }
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Borrow()
Line
Count
Source
231
3.10M
    T *Borrow() {
232
3.10M
      std::lock_guard<std::mutex> lock(mutex_);
233
3.10M
      if (stack_top_ == stack_.size()) {
234
8
        stack_.push_back(new T);
235
8
        flags_.push_back(false);
236
8
      }
237
3.10M
      flags_[stack_top_] = true;
238
3.10M
      return stack_[stack_top_++];
239
3.10M
    }
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Borrow()
240
    // Takes back the given item, and marks it free. Item does not have to be
241
    // the most recently lent out, but free slots don't get re-used until the
242
    // blocking item is returned. The assumption is that there will only be
243
    // small, temporary variations from true stack use. (Determined by the order
244
    // of destructors within a local scope.)
245
3.79M
    void Return(T *item) {
246
3.79M
      std::lock_guard<std::mutex> lock(mutex_);
247
      // Linear search will do.
248
3.79M
      int index = stack_top_;
249
3.79M
      while (--index >= 0 && stack_[index] != item) {
250
0
      }
251
3.79M
      if (index >= 0) {
252
3.79M
        flags_[index] = false;
253
3.79M
      }
254
7.59M
      while (stack_top_ > 0 && !flags_[stack_top_ - 1]) {
255
3.79M
        --stack_top_;
256
3.79M
      }
257
3.79M
    }
tesseract::NetworkScratch::Stack<tesseract::NetworkIO>::Return(tesseract::NetworkIO*)
Line
Count
Source
245
690k
    void Return(T *item) {
246
690k
      std::lock_guard<std::mutex> lock(mutex_);
247
      // Linear search will do.
248
690k
      int index = stack_top_;
249
690k
      while (--index >= 0 && stack_[index] != item) {
250
0
      }
251
690k
      if (index >= 0) {
252
690k
        flags_[index] = false;
253
690k
      }
254
1.38M
      while (stack_top_ > 0 && !flags_[stack_top_ - 1]) {
255
690k
        --stack_top_;
256
690k
      }
257
690k
    }
tesseract::NetworkScratch::Stack<std::__1::vector<float, std::__1::allocator<float> > >::Return(std::__1::vector<float, std::__1::allocator<float> >*)
Line
Count
Source
245
3.10M
    void Return(T *item) {
246
3.10M
      std::lock_guard<std::mutex> lock(mutex_);
247
      // Linear search will do.
248
3.10M
      int index = stack_top_;
249
3.10M
      while (--index >= 0 && stack_[index] != item) {
250
0
      }
251
3.10M
      if (index >= 0) {
252
3.10M
        flags_[index] = false;
253
3.10M
      }
254
6.21M
      while (stack_top_ > 0 && !flags_[stack_top_ - 1]) {
255
3.10M
        --stack_top_;
256
3.10M
      }
257
3.10M
    }
Unexecuted instantiation: tesseract::NetworkScratch::Stack<tesseract::TransposedArray>::Return(tesseract::TransposedArray*)
258
259
  private:
260
    std::vector<T *> stack_;
261
    std::vector<bool> flags_;
262
    unsigned stack_top_ = 0;
263
    std::mutex mutex_;
264
  }; // class Stack.
265
266
private:
267
  // If true, the network weights are int8_t, if false, float.
268
  bool int_mode_;
269
  // Stacks of NetworkIO and vector<float>. Once allocated, they are not
270
  // deleted until the NetworkScratch is deleted.
271
  Stack<NetworkIO> int_stack_;
272
  Stack<NetworkIO> float_stack_;
273
  Stack<std::vector<TFloat>> vec_stack_;
274
  Stack<TransposedArray> array_stack_;
275
};
276
277
} // namespace tesseract.
278
279
#endif // TESSERACT_LSTM_NETWORKSCRATCH_H_