1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file pipe.h |
22 | * \brief Platform independent pipe, used for IPC. |
23 | */ |
24 | #ifndef TVM_SUPPORT_PIPE_H_ |
25 | #define TVM_SUPPORT_PIPE_H_ |
26 | |
27 | #include <dmlc/io.h> |
28 | #include <tvm/runtime/logging.h> |
29 | |
30 | #ifdef _WIN32 |
31 | #include <windows.h> |
32 | #else |
33 | #include <errno.h> |
34 | #include <unistd.h> |
35 | |
36 | #include <cstdlib> |
37 | #include <cstring> |
38 | #endif |
39 | |
40 | namespace tvm { |
41 | namespace support { |
42 | |
43 | /*! \brief Platform independent pipe */ |
44 | class Pipe : public dmlc::Stream { |
45 | public: |
46 | #ifdef _WIN32 |
47 | using PipeHandle = HANDLE; |
48 | #else |
49 | using PipeHandle = int; |
50 | #endif |
51 | /*! \brief Construct a pipe from system handle. */ |
52 | explicit Pipe(int64_t handle) : handle_(static_cast<PipeHandle>(handle)) {} |
53 | /*! \brief destructor */ |
54 | ~Pipe() { Flush(); } |
55 | using Stream::Read; |
56 | using Stream::Write; |
57 | /*! |
58 | * \brief reads data from a file descriptor |
59 | * \param ptr pointer to a memory buffer |
60 | * \param size block size |
61 | * \return the size of data read |
62 | */ |
63 | size_t Read(void* ptr, size_t size) final { |
64 | if (size == 0) return 0; |
65 | #ifdef _WIN32 |
66 | DWORD nread; |
67 | ICHECK(ReadFile(handle_, static_cast<TCHAR*>(ptr), &nread, nullptr)) |
68 | << "Read Error: " << GetLastError(); |
69 | #else |
70 | ssize_t nread; |
71 | nread = read(handle_, ptr, size); |
72 | ICHECK_GE(nread, 0) << "Write Error: " << strerror(errno); |
73 | #endif |
74 | return static_cast<size_t>(nread); |
75 | } |
76 | /*! |
77 | * \brief write data to a file descriptor |
78 | * \param ptr pointer to a memory buffer |
79 | * \param size block size |
80 | * \return the size of data read |
81 | */ |
82 | void Write(const void* ptr, size_t size) final { |
83 | if (size == 0) return; |
84 | #ifdef _WIN32 |
85 | DWORD nwrite; |
86 | ICHECK(WriteFile(handle_, static_cast<const TCHAR*>(ptr), &nwrite, nullptr) && |
87 | static_cast<size_t>(nwrite) == size) |
88 | << "Write Error: " << GetLastError(); |
89 | #else |
90 | ssize_t nwrite; |
91 | nwrite = write(handle_, ptr, size); |
92 | ICHECK_EQ(static_cast<size_t>(nwrite), size) << "Write Error: " << strerror(errno); |
93 | #endif |
94 | } |
95 | /*! |
96 | * \brief Flush the pipe; |
97 | */ |
98 | void Flush() { |
99 | #ifdef _WIN32 |
100 | FlushFileBuffers(handle_); |
101 | #endif |
102 | } |
103 | /*! \brief close the pipe */ |
104 | void Close() { |
105 | #ifdef _WIN32 |
106 | CloseHandle(handle_); |
107 | #else |
108 | close(handle_); |
109 | #endif |
110 | } |
111 | |
112 | private: |
113 | PipeHandle handle_; |
114 | }; |
115 | } // namespace support |
116 | } // namespace tvm |
117 | |
118 | #endif // TVM_SUPPORT_PIPE_H_ |
119 | |