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
101
102
103
104
105
#pragma once
 
#include <c10/core/impl/DeviceGuardImplInterface.h>
 
#include <array>
 
namespace c10 {
namespace impl {
 
// FakeGuardImpl is hardcoded to have eight devices.  Not for
// any good reason, just to simplify code.
constexpr size_t kFakeGuardImplMaxDevices = 8;
 
/**
 * A fake implementation of DeviceGuardImplInterface suitable for testing.
 * The current device is modeled as a mutable field in the guard implementation
 * class.  See DeviceGuard_test.cpp for an example use.
 */
template <DeviceType T>
struct FakeGuardImpl final : public DeviceGuardImplInterface {
  static constexpr DeviceType static_type = T;
  // Runtime device type is not used
  FakeGuardImpl(DeviceType) {}
  FakeGuardImpl() {}
  DeviceType type() const override {
    return T;
  }
  Device exchangeDevice(Device d) const override {
    AT_ASSERT(d.type() == type());
    AT_ASSERT(d.index() < kFakeGuardImplMaxDevices);
    Device old_device = getDevice();
    if (old_device.index() != d.index()) {
      current_device_ = d.index();
    }
    return old_device;
  }
  Device getDevice() const override {
    return Device(type(), current_device_);
  }
  void setDevice(Device d) const override {
    AT_ASSERT(d.type() == type());
    AT_ASSERT(d.index() >= 0);
    AT_ASSERT(d.index() < kFakeGuardImplMaxDevices);
    current_device_ = d.index();
  }
  void uncheckedSetDevice(Device d) const noexcept override {
    current_device_ = d.index();
  }
  Stream getStream(Device d) const noexcept override {
    return Stream(Stream::UNSAFE, d, current_streams_[d.index()]);
  }
  Stream exchangeStream(Stream s) const noexcept override {
    auto old_id = current_streams_[s.device_index()];
    current_streams_[s.device_index()] = s.id();
    return Stream(Stream::UNSAFE, s.device(), old_id);
  }
  DeviceIndex deviceCount() const noexcept override {
    return kFakeGuardImplMaxDevices;
  }
 
  // Event-related functions
  void record(
    void** event,
    const Stream& stream,
    const DeviceIndex device_index,
    const EventFlag flag) const override { }
  void block(
    void* event,
    const Stream& stream) const override { }
  bool queryEvent(void* event) const override { return true; }
  void destroyEvent(
    void* event,
    const DeviceIndex device_index) const noexcept override { }
 
  // Convenience methods for testing
  static DeviceIndex getDeviceIndex() {
    return current_device_;
  }
  static void setDeviceIndex(DeviceIndex i) {
    AT_ASSERT(i >= 0);
    AT_ASSERT(i < kFakeGuardImplMaxDevices);
    current_device_ = i;
  }
  static StreamId getCurrentStreamIdFor(DeviceIndex i) {
    return current_streams_.at(i);
  }
  static void resetStreams() {
    current_streams_.fill(0);
  }
private:
  thread_local static DeviceIndex current_device_;
  thread_local static std::array<StreamId, kFakeGuardImplMaxDevices> current_streams_;
};
 
template <DeviceType T>
thread_local DeviceIndex FakeGuardImpl<T>::current_device_ = 0;
 
template <DeviceType T>
constexpr DeviceType FakeGuardImpl<T>::static_type;
 
template <DeviceType T>
thread_local std::array<StreamId, kFakeGuardImplMaxDevices> FakeGuardImpl<T>::current_streams_ = {0,0,0,0,0,0,0,0};
 
 
}} // namespace c10::impl