reid from https://github.com/michuanhaohao/reid-strong-baseline
zhangmeng
2020-01-17 f7c4a3cfd07adede3308f8d9d3d7315427d90a7c
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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#pragma once
 
#include <TH/THGeneral.h>
 
#include <c10/core/Allocator.h>
 
#define TH_ALLOCATOR_MAPPED_SHARED 1
#define TH_ALLOCATOR_MAPPED_SHAREDMEM 2
#define TH_ALLOCATOR_MAPPED_EXCLUSIVE 4
#define TH_ALLOCATOR_MAPPED_NOCREATE 8
#define TH_ALLOCATOR_MAPPED_KEEPFD 16
#define TH_ALLOCATOR_MAPPED_FROMFD 32
#define TH_ALLOCATOR_MAPPED_UNLINK 64
 
/* default malloc/free allocator. malloc and realloc raise an error (using
 * THError) on allocation failure.
 */
TH_API c10::Allocator* getTHDefaultAllocator(void);
 
// Sentinel value/type to help distinguish the file descriptor constructor from
// the non-file descriptor constructor
enum WithFd { WITH_FD };
 
class CAFFE2_API THMapAllocator {
 public:
  THMapAllocator(const char *filename, int flags, size_t size);
  THMapAllocator(WithFd, const char *filename, int fd, int flags, size_t size);
  THMapAllocator(const THMapAllocator&) = delete;
  THMapAllocator& operator=(const THMapAllocator&) = delete;
  THMapAllocator(THMapAllocator&&) = delete;
  THMapAllocator& operator=(THMapAllocator&&) = delete;
 
  const char* filename() const { return filename_.c_str(); }
  int fd() const {
#ifdef _WIN32
    AT_ERROR("THMapAllocator::fd() is unsupported on Windows");
#else
    return fd_;
#endif
  }
  ptrdiff_t size() const { return size_; }
  // Return a pointer to the actual data for this allocator
  // (in the case of the refcounted allocator, this is offset
  // from the base pointer.)
  virtual void* data() const { return base_ptr_; }
 
  static THMapAllocator* fromDataPtr(const at::DataPtr&);
  static at::DataPtr makeDataPtr(const char *filename, int flags, size_t size, size_t* actual_size_out);
  static at::DataPtr makeDataPtr(WithFd, const char *filename, int fd, int flags, size_t size, size_t* actual_size_out);
 
  // Closes the data.  Helps us avoid destructor shenanigans
  virtual void close();
 
  // This is very dangerous.  You have to redefine this destructor for each
  // subclass
  virtual ~THMapAllocator() { close(); }
 
protected:
  bool closed_ = false;
  std::string filename_;
  int flags_ = 0;
  ptrdiff_t size_; /* mapped size */
#ifdef _WIN32
  void* handle_;
  void* event_;
  std::string eventname_;
#else
  int fd_ = -1;
#endif
  void *base_ptr_ = nullptr;
};
 
// Base-from-member idiom
struct CAFFE2_API THRefcountedMapAllocatorArgCheck {
  THRefcountedMapAllocatorArgCheck(int flags);
};
 
class CAFFE2_API THRefcountedMapAllocator
    : private THRefcountedMapAllocatorArgCheck,
      public THMapAllocator {
 public:
  THRefcountedMapAllocator(const char *filename, int flags, size_t size);
  THRefcountedMapAllocator(WithFd, const char *filename, int fd, int flags, size_t size);
 
  static THRefcountedMapAllocator* fromDataPtr(const at::DataPtr&);
  static at::DataPtr makeDataPtr(const char *filename, int flags, size_t size, size_t* actual_size_out);
  static at::DataPtr makeDataPtr(WithFd, const char *filename, int fd, int flags, size_t size, size_t* actual_size_out);
 
  void* data() const override;
 
  void incref();
  int decref();
  void close() override;
 
  virtual ~THRefcountedMapAllocator() { close(); }
 
protected:
  void checkFlags();
  void initializeAlloc();
};