2 #include <c10/core/Device.h>
3 #include <c10/core/Stream.h>
4 #include <c10/core/impl/DeviceGuardImplInterface.h>
5 #include <c10/util/Logging.h>
10 constexpr c10::DeviceType kMTIADeviceType = c10::DeviceType::MTIA;
11 constexpr c10::DeviceIndex kMTIADeviceCount = 2;
12 static thread_local c10::DeviceIndex current_device = 0;
13 static thread_local std::array<c10::Stream, kMTIADeviceCount> current_streams =
14     {c10::Stream::unpack3(0, 0, c10::DeviceType::MTIA),
15      c10::Stream::unpack3(0, 1, c10::DeviceType::MTIA)};
18 static std::array<c10::Stream, kMTIADeviceCount> default_streams = {
19     c10::Stream::unpack3(0, 0, c10::DeviceType::MTIA),
20     c10::Stream::unpack3(0, 1, c10::DeviceType::MTIA)};
21 struct MTIAGuardImpl final : public c10::impl::DeviceGuardImplInterface {
23   explicit MTIAGuardImpl(c10::DeviceType t) {  in MTIAGuardImpl()
26   c10::DeviceType type() const override {  in type()
29   c10::Device exchangeDevice(c10::Device d) const override {  in exchangeDevice()
30     c10::Device old_device = getDevice();  in exchangeDevice()
36   c10::Device getDevice() const override {  in getDevice()
37     return c10::Device(kMTIADeviceType, current_device);  in getDevice()
40   void setDevice(c10::Device d) const override {  in setDevice()
41     c10::Device current_device = getDevice();  in setDevice()
46   void uncheckedSetDevice(c10::Device d) const noexcept override {  in uncheckedSetDevice()
49   c10::Stream getStream(c10::Device d) const noexcept override {  in getStream()
52   c10::Stream getNewStream(c10::Device d, int priority = 0) const override {  in getNewStream()
54     return c10::Stream::unpack3(stream_id_gen++, d.index(), d.type());  in getNewStream()
56   c10::Stream getDefaultStream(c10::Device d) const override {  in getDefaultStream()
59   c10::Stream getStreamFromGlobalPool(  in getStreamFromGlobalPool()
60       c10::Device d,  in getStreamFromGlobalPool()
62     return c10::Stream::unpack3(stream_id_gen++, d.index(), d.type());  in getStreamFromGlobalPool()
65   c10::Stream exchangeStream(c10::Stream s) const noexcept override {  in exchangeStream()
66     c10::Stream old_stream = getStream(s.device());  in exchangeStream()
69   c10::DeviceIndex deviceCount() const noexcept override {  in deviceCount()
73   void destroyEvent(void* event, const c10::DeviceIndex device_index)  in destroyEvent()
80       const c10::Stream& stream,  in record()
81       const c10::DeviceIndex device_index,  in record()
82       const c10::EventFlag flag) const override {  in record()
101   void block(void* event, const c10::Stream& stream) const override {  in block()
113   bool queryStream(const c10::Stream& stream) const override {  in queryStream()
118   void synchronizeStream(const c10::Stream& stream) const override {  in synchronizeStream()
123       const c10::DataPtr& data_ptr,  in recordDataPtrOnStream()
124       const c10::Stream& stream) const override {  in recordDataPtrOnStream()
129 …double elapsedTime(void* event1, void* event2, const c10::DeviceIndex device_index) const override…  in elapsedTime()
148   c10::DeviceIndex deviceCount() const override {  in deviceCount()
150     return c10::DeviceIndex(2);  in deviceCount()
153   void deviceSynchronize(c10::DeviceIndex device_index) const override {  in deviceSynchronize()
162   c10::DeviceIndex exchangeDevice(c10::DeviceIndex device) const override {  in exchangeDevice()
171   c10::DeviceIndex maybeExchangeDevice(c10::DeviceIndex device) const override {  in maybeExchangeDevice()
181   c10::Stream getDefaultStream(c10::DeviceIndex device) const override {  in getDefaultStream()
187   c10::Stream getCurrentStream(c10::DeviceIndex device) const override {  in getCurrentStream()
193   void setCurrentStream(const c10::Stream& stream) const override {  in setCurrentStream()
199   c10::DeviceIndex getCurrentDevice() const override {  in getCurrentDevice()
205   void setCurrentDevice(c10::DeviceIndex device) const override {  in setCurrentDevice()