blob: 23feed9aab92bc22f2d6c6d5b642f8abd3939c07 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
|
#ifndef TENSORFLOW_COMMON_RUNTIME_GPU_VISITABLE_ALLOCATOR_H_
#define TENSORFLOW_COMMON_RUNTIME_GPU_VISITABLE_ALLOCATOR_H_
#include <functional>
#include "tensorflow/core/framework/allocator.h"
namespace tensorflow {
// Subclass VisitableAllocator instead of Allocator when a memory
// allocator needs to enable some kind of registration/deregistration
// of memory areas.
class VisitableAllocator : public Allocator {
public:
// Visitor gets called with a pointer to a memory area and its
// size in bytes.
typedef std::function<void(void*, size_t)> Visitor;
// Register a visitor guaranteed to be called exactly once on each
// chunk of memory newly allocated from the underlying device.
// Typically, chunks will be reused and possibly sub-divided by a
// pool manager, so the calls will happen only once per process
// execution, not once per tensor (re)allocation.
virtual void AddAllocVisitor(Visitor visitor) = 0;
// Register a visitor guaranteed to be called on each chunk of
// memory returned to the underlying device.
virtual void AddFreeVisitor(Visitor visitor) = 0;
};
} // namespace tensorflow
#endif // TENSORFLOW_COMMON_RUNTIME_GPU_VISITABLE_ALLOCATOR_H_
|
