[IO] Support cross-endian

dmlc-core @ 9b3f9753

-Subproject commit d3f7fbb53e5b037c0f5bf6bd21871ccc720690cc
+Subproject commit 9b3f9753ae81d657743c555e0cacc4e43f0bed2d

include/tvm/runtime/serializer.h

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file tvm/runtime/serializer.h
+ * \brief Serializer extension to support TVM data types
+ *  Include this file to enable serialization of DLDataType, DLContext
+ */
+#ifndef TVM_RUNTIME_SERIALIZER_H_
+#define TVM_RUNTIME_SERIALIZER_H_
+
+#include <dmlc/io.h>
+#include <dmlc/serializer.h>
+#include "./c_runtime_api.h"
+
+namespace dmlc {
+namespace serializer {
+
+template<>
+struct Handler<DLDataType> {
+  inline static void Write(Stream *strm, const DLDataType& dtype) {
+    Handler<uint8_t>::Write(strm, dtype.code);
+    Handler<uint8_t>::Write(strm, dtype.bits);
+    Handler<uint16_t>::Write(strm, dtype.lanes);
+  }
+  inline static bool Read(Stream *strm, DLDataType* dtype) {
+    if (!Handler<uint8_t>::Read(strm, &(dtype->code))) return false;
+    if (!Handler<uint8_t>::Read(strm, &(dtype->bits))) return false;
+    if (!Handler<uint16_t>::Read(strm, &(dtype->lanes))) return false;
+    return true;
+  }
+};
+
+template<>
+struct Handler<DLContext> {
+  inline static void Write(Stream *strm, const DLContext& ctx) {
+    int32_t device_type = static_cast<int32_t>(ctx.device_type);
+    Handler<int32_t>::Write(strm, device_type);
+    Handler<int32_t>::Write(strm, ctx.device_id);
+  }
+  inline static bool Read(Stream *strm, DLContext* ctx) {
+    int32_t device_type = 0;
+    if (!Handler<int32_t>::Read(strm, &(device_type))) return false;
+    ctx->device_type = static_cast<DLDeviceType>(device_type);
+    if (!Handler<int32_t>::Read(strm, &(ctx->device_id))) return false;
+    return true;
+  }
+};
+
+}  // namespace serializer
+}  // namespace dmlc
+#endif  // TVM_RUNTIME_SERIALIZER_H_

nnvm/src/compiler/graph_runtime.cc

@@ -7,6 +7,7 @@
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/runtime/serializer.h>
 #include "./graph_runtime.h"
 
 namespace nnvm {
@@ -38,46 +39,53 @@ NNVM_REGISTER_OP(tvm_op)
 
 bool SaveDLTensor(dmlc::Stream* strm, DLTensor* tensor) {
   uint64_t header = kTVMNDArrayMagic, reserved = 0;
-  strm->Write(&header, sizeof(header));
-  strm->Write(&reserved, sizeof(reserved));
-
-  strm->Write(&tensor->ctx, sizeof(tensor->ctx));
-  strm->Write(&tensor->ndim, sizeof(tensor->ndim));
-  strm->Write(&tensor->dtype, sizeof(tensor->dtype));
-
+  strm->Write(header);
+  strm->Write(reserved);
+  strm->Write(tensor->ctx);
+  strm->Write(tensor->ndim);
+  strm->Write(tensor->dtype);
   int ndim = tensor->ndim;
-  strm->Write(tensor->shape, sizeof(int64_t) * ndim);
+  strm->WriteArray(tensor->shape, ndim);
 
-  int type_size = tensor->dtype.bits / 8;
-  int64_t size = 1;
+  int type_bytes = tensor->dtype.bits / 8;
+  int64_t num_elems = 1;
   for (int i = 0; i < ndim; ++i) {
-    size *= tensor->shape[i];
+    num_elems *= tensor->shape[i];
   }
-  int64_t data_byte_size = type_size * size;
-  strm->Write(&data_byte_size, sizeof(data_byte_size));
+  int64_t data_byte_size = type_bytes * num_elems;
+  strm->Write(data_byte_size);
+  // handle endianness of data correctly.
+  if (DMLC_IO_NO_ENDIAN_SWAP) {
     strm->Write(tensor->data, data_byte_size);
+  } else {
+    uint8_t* dptr = reinterpret_cast<uint8_t*>(tensor->data);
+    std::vector<uint8_t> bytes(dptr, dptr + data_byte_size);
+    dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
+    strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
+  }
   return true;
 }
 
 DLTensor* LoadDLTensor(dmlc::Stream* strm) {
   uint64_t header, reserved;
-  CHECK(strm->Read(&header, sizeof(header)))
+  CHECK(strm->Read(&header))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&reserved, sizeof(reserved)))
+  CHECK(strm->Read(&reserved))
       << "Invalid DLTensor file format";
   CHECK(header == kTVMNDArrayMagic)
       << "Invalid DLTensor file format";
-
   DLTensor tensor;
-  CHECK(strm->Read(&tensor.ctx, sizeof(tensor.ctx)))
+  CHECK(strm->Read(&(tensor.ctx)))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&tensor.ndim, sizeof(tensor.ndim)))
+  CHECK(strm->Read(&(tensor.ndim)))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&tensor.dtype, sizeof(tensor.dtype)))
+  CHECK(strm->Read(&(tensor.dtype)))
       << "Invalid DLTensor file format";
   std::vector<int64_t> shape(tensor.ndim);
-  CHECK(strm->Read(&shape[0], sizeof(int64_t) * tensor.ndim))
+  if (tensor.ndim != 0) {
+    CHECK(strm->ReadArray(&shape[0], tensor.ndim))
         << "Invalid DLTensor file format";
+  }
   DLTensor* ret;
   CHECK_EQ(TVMArrayAlloc(shape.data(),
                          tensor.ndim,
@@ -87,18 +95,21 @@ DLTensor* LoadDLTensor(dmlc::Stream* strm) {
                          static_cast<int>(tensor.ctx.device_type),
                          tensor.ctx.device_id,
                          &ret), 0) << TVMGetLastError();
-  int64_t size = 1;
-  int type_size = ret->dtype.bits / 8;
+  int64_t num_elems = 1;
+  int elem_bytes = (ret->dtype.bits + 7) / 8;
   for (int i = 0; i < ret->ndim; ++i) {
-    size *= ret->shape[i];
+    num_elems *= ret->shape[i];
   }
   int64_t data_byte_size;
-  CHECK(strm->Read(&data_byte_size, sizeof(data_byte_size)))
+  CHECK(strm->Read(&data_byte_size))
       << "Invalid DLTensor file format";
-  CHECK(data_byte_size == type_size * size)
+  CHECK(data_byte_size == num_elems * elem_bytes)
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(ret->data, type_size * size))
+  CHECK(strm->Read(ret->data, data_byte_size))
       << "Invalid DLTensor file format";
+  if (!DMLC_IO_NO_ENDIAN_SWAP) {
+    dmlc::ByteSwap(ret->data, elem_bytes, num_elems);
+  }
   return ret;
 }
 
@@ -118,12 +129,12 @@ TVM_REGISTER_GLOBAL("nnvm.compiler._save_param_dict")
     dmlc::MemoryStringStream strm(&bytes);
     dmlc::Stream* fo = &strm;
     uint64_t header = kTVMNDArrayListMagic, reserved = 0;
-    fo->Write(&header, sizeof(header));
-    fo->Write(&reserved, sizeof(reserved));
+    fo->Write(header);
+    fo->Write(reserved);
     fo->Write(names);
     {
       uint64_t sz = static_cast<uint64_t>(arrays.size());
-      fo->Write(&sz, sizeof(sz));
+      fo->Write(sz);
       for (size_t i = 0; i < sz; ++i) {
         SaveDLTensor(fo, arrays[i]);
       }
@@ -150,7 +161,6 @@ TVM_REGISTER_GLOBAL("nnvm.compiler._load_param_dict")
         << "Invalid parameters file format";
     CHECK(strm->Read(&reserved))
         << "Invalid parameters file format";
-
     CHECK(strm->Read(&names))
         << "Invalid parameters file format";
     uint64_t sz;

python/tvm/contrib/rpc/base.py

@@ -73,7 +73,7 @@ def sendjson(sock, data):
         Python value to be sent.
     """
     data = json.dumps(data)
-    sock.sendall(struct.pack("@i", len(data)))
+    sock.sendall(struct.pack("<i", len(data)))
     sock.sendall(data.encode("utf-8"))
 
 
@@ -90,7 +90,7 @@ def recvjson(sock):
     value : object
         The value received.
     """
-    size = struct.unpack("@i", recvall(sock, 4))[0]
+    size = struct.unpack("<i", recvall(sock, 4))[0]
     data = json.loads(py_str(recvall(sock, size)))
     return data
 

python/tvm/contrib/rpc/client.py

@@ -192,8 +192,8 @@ class TrackerSession(object):
 
     def _connect(self):
         self._sock = base.connect_with_retry(self._addr)
-        self._sock.sendall(struct.pack("@i", base.RPC_TRACKER_MAGIC))
-        magic = struct.unpack("@i", base.recvall(self._sock, 4))[0]
+        self._sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+        magic = struct.unpack("<i", base.recvall(self._sock, 4))[0]
         if magic != base.RPC_TRACKER_MAGIC:
             raise RuntimeError("%s is not RPC Tracker" % str(self._addr))
 

python/tvm/contrib/rpc/proxy.py

@@ -58,14 +58,14 @@ class ForwardHandler(object):
     def _init_step(self, message):
         if self._magic is None:
             assert len(message) == 4
-            self._magic = struct.unpack('@i', message)[0]
+            self._magic = struct.unpack('<i', message)[0]
             if self._magic != base.RPC_MAGIC:
                 logging.info("Invalid RPC magic from %s", self.name())
                 self.close()
             self._init_req_nbytes = 4
         elif self._rpc_key_length is None:
             assert len(message) == 4
-            self._rpc_key_length = struct.unpack('@i', message)[0]
+            self._rpc_key_length = struct.unpack('<i', message)[0]
             self._init_req_nbytes = self._rpc_key_length
         elif self.rpc_key is None:
             assert len(message) == self._rpc_key_length
@@ -269,12 +269,12 @@ class ProxyServerHandler(object):
         lhs.forward_proxy = rhs
         rhs.forward_proxy = lhs
 
-        lhs.send_data(struct.pack('@i', base.RPC_CODE_SUCCESS))
-        lhs.send_data(struct.pack('@i', len(rhs.rpc_key)))
+        lhs.send_data(struct.pack('<i', base.RPC_CODE_SUCCESS))
+        lhs.send_data(struct.pack('<i', len(rhs.rpc_key)))
         lhs.send_data(rhs.rpc_key.encode("utf-8"))
 
-        rhs.send_data(struct.pack('@i', base.RPC_CODE_SUCCESS))
-        rhs.send_data(struct.pack('@i', len(lhs.rpc_key)))
+        rhs.send_data(struct.pack('<i', base.RPC_CODE_SUCCESS))
+        rhs.send_data(struct.pack('<i', len(lhs.rpc_key)))
         rhs.send_data(lhs.rpc_key.encode("utf-8"))
         logging.info("Pairup connect %s  and %s", lhs.name(), rhs.name())
 
@@ -299,8 +299,8 @@ class ProxyServerHandler(object):
             if self._tracker_conn is None:
                 self._tracker_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 self._tracker_conn.connect(self._tracker_addr)
-                self._tracker_conn.sendall(struct.pack("@i", base.RPC_TRACKER_MAGIC))
-                magic = struct.unpack("@i", base.recvall(self._tracker_conn, 4))[0]
+                self._tracker_conn.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+                magic = struct.unpack("<i", base.recvall(self._tracker_conn, 4))[0]
                 if magic != base.RPC_TRACKER_MAGIC:
                     self.loop.stop()
                     raise RuntimeError("%s is not RPC Tracker" % str(self._tracker_addr))
@@ -371,7 +371,7 @@ class ProxyServerHandler(object):
             if handler.match_key in self._server_pool:
                 self._pair_up(self._server_pool.pop(handler.match_key), handler)
             else:
-                handler.send_data(struct.pack('@i', base.RPC_CODE_MISMATCH))
+                handler.send_data(struct.pack('<i', base.RPC_CODE_MISMATCH))
                 handler.signal_close()
 
     def _handler_ready_proxy_mode(self, handler):
@@ -395,12 +395,12 @@ class ProxyServerHandler(object):
                     logging.info("Timeout client connection %s, cannot find match key=%s",
                                  handler.name(), key)
                     pool_dst.pop(key)
-                    handler.send_data(struct.pack('@i', base.RPC_CODE_MISMATCH))
+                    handler.send_data(struct.pack('<i', base.RPC_CODE_MISMATCH))
                     handler.signal_close()
             self.loop.call_later(timeout, cleanup)
         else:
             logging.info("Duplicate connection with same key=%s", key)
-            handler.send_data(struct.pack('@i', base.RPC_CODE_DUPLICATE))
+            handler.send_data(struct.pack('<i', base.RPC_CODE_DUPLICATE))
             handler.signal_close()
 
     def handler_ready(self, handler):
@@ -538,13 +538,13 @@ def websocket_proxy_server(url, key=""):
         on_message = create_on_message(conn)
         temp = _server_env(None)
         # Start connecton
-        conn.write_message(struct.pack('@i', base.RPC_MAGIC), binary=True)
+        conn.write_message(struct.pack('<i', base.RPC_MAGIC), binary=True)
         key = "server:" + key
-        conn.write_message(struct.pack('@i', len(key)), binary=True)
+        conn.write_message(struct.pack('<i', len(key)), binary=True)
         conn.write_message(key.encode("utf-8"), binary=True)
         msg = yield conn.read_message()
         assert len(msg) >= 4
-        magic = struct.unpack('@i', msg[:4])[0]
+        magic = struct.unpack('<i', msg[:4])[0]
         if magic == base.RPC_CODE_DUPLICATE:
             raise RuntimeError("key: %s has already been used in proxy" % key)
         elif magic == base.RPC_CODE_MISMATCH:

python/tvm/contrib/rpc/server.py

@@ -124,23 +124,23 @@ def _listen_loop(sock, port, rpc_key, tracker_addr, load_library):
                         unmatch_period_count = 0
                     continue
             conn, addr = listen_sock.accept()
-            magic = struct.unpack("@i", base.recvall(conn, 4))[0]
+            magic = struct.unpack("<i", base.recvall(conn, 4))[0]
             if magic != base.RPC_MAGIC:
                 conn.close()
                 continue
-            keylen = struct.unpack("@i", base.recvall(conn, 4))[0]
+            keylen = struct.unpack("<i", base.recvall(conn, 4))[0]
             key = py_str(base.recvall(conn, keylen))
             arr = key.split()
             expect_header = "client:" + matchkey
             server_key = "server:" + rpc_key
             if arr[0] != expect_header:
-                conn.sendall(struct.pack("@i", base.RPC_CODE_MISMATCH))
+                conn.sendall(struct.pack("<i", base.RPC_CODE_MISMATCH))
                 conn.close()
                 logging.info("RPCServer: mismatch key from %s", addr)
                 continue
             else:
-                conn.sendall(struct.pack("@i", base.RPC_CODE_SUCCESS))
-                conn.sendall(struct.pack("@i", len(server_key)))
+                conn.sendall(struct.pack("<i", base.RPC_CODE_SUCCESS))
+                conn.sendall(struct.pack("<i", len(server_key)))
                 conn.sendall(server_key.encode("utf-8"))
                 return conn, addr, _parse_server_opt(arr[1:])
 
@@ -151,8 +151,8 @@ def _listen_loop(sock, port, rpc_key, tracker_addr, load_library):
             # step 1: setup tracker and report to tracker
             if tracker_addr and tracker_conn is None:
                 tracker_conn = base.connect_with_retry(tracker_addr)
-                tracker_conn.sendall(struct.pack("@i", base.RPC_TRACKER_MAGIC))
-                magic = struct.unpack("@i", base.recvall(tracker_conn, 4))[0]
+                tracker_conn.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+                magic = struct.unpack("<i", base.recvall(tracker_conn, 4))[0]
                 if magic != base.RPC_TRACKER_MAGIC:
                     raise RuntimeError("%s is not RPC Tracker" % str(tracker_addr))
                 # report status of current queue
@@ -193,17 +193,17 @@ def _connect_proxy_loop(addr, key, load_library):
         try:
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             sock.connect(addr)
-            sock.sendall(struct.pack("@i", base.RPC_MAGIC))
-            sock.sendall(struct.pack("@i", len(key)))
+            sock.sendall(struct.pack("<i", base.RPC_MAGIC))
+            sock.sendall(struct.pack("<i", len(key)))
             sock.sendall(key.encode("utf-8"))
-            magic = struct.unpack("@i", base.recvall(sock, 4))[0]
+            magic = struct.unpack("<i", base.recvall(sock, 4))[0]
             if magic == base.RPC_CODE_DUPLICATE:
                 raise RuntimeError("key: %s has already been used in proxy" % key)
             elif magic == base.RPC_CODE_MISMATCH:
                 logging.info("RPCProxy do not have matching client key %s", key)
             elif magic != base.RPC_CODE_SUCCESS:
                 raise RuntimeError("%s is not RPC Proxy" % str(addr))
-            keylen = struct.unpack("@i", base.recvall(sock, 4))[0]
+            keylen = struct.unpack("<i", base.recvall(sock, 4))[0]
             remote_key = py_str(base.recvall(sock, keylen))
             opts = _parse_server_opt(remote_key.split()[1:])
             logging.info("RPCProxy connected to %s", str(addr))

python/tvm/contrib/rpc/tracker.py

@@ -143,11 +143,11 @@ class TCPEventHandler(tornado_util.TCPHandler):
         if len(message) != 4:
             logging.info("Invalid connection from %s", self.name())
             self.close()
-        magic = struct.unpack('@i', message)[0]
+        magic = struct.unpack('<i', message)[0]
         if magic != RPC_TRACKER_MAGIC:
             logging.info("Invalid magic from %s", self.name())
             self.close()
-        self.write_message(struct.pack('@i', RPC_TRACKER_MAGIC), binary=True)
+        self.write_message(struct.pack('<i', RPC_TRACKER_MAGIC), binary=True)
         self._init_req_nbytes = 0
 
     def on_message(self, message):
@@ -168,7 +168,7 @@ class TCPEventHandler(tornado_util.TCPHandler):
         while True:
             if self._msg_size == 0:
                 if len(self._data) >= 4:
-                    self._msg_size = struct.unpack('@i', self._data[:4])[0]
+                    self._msg_size = struct.unpack('<i', self._data[:4])[0]
                 else:
                     return
             if self._msg_size != 0 and len(self._data) >= self._msg_size + 4:
@@ -184,7 +184,7 @@ class TCPEventHandler(tornado_util.TCPHandler):
         """return value to the output"""
         data = json.dumps(data)
         self.write_message(
-            struct.pack('@i', len(data)), binary=True)
+            struct.pack('<i', len(data)), binary=True)
         self.write_message(data.encode("utf-8"), binary=True)
 
     def call_handler(self, args):
@@ -355,8 +355,8 @@ class Tracker(object):
     def _stop_tracker(self):
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.connect((self.host, self.port))
-        sock.sendall(struct.pack("@i", base.RPC_TRACKER_MAGIC))
-        magic = struct.unpack("@i", base.recvall(sock, 4))[0]
+        sock.sendall(struct.pack("<i", base.RPC_TRACKER_MAGIC))
+        magic = struct.unpack("<i", base.recvall(sock, 4))[0]
         assert magic == base.RPC_TRACKER_MAGIC
         base.sendjson(sock, [TrackerCode.STOP, self.stop_key])
         assert base.recvjson(sock) == TrackerCode.SUCCESS

src/runtime/file_util.cc

@@ -4,6 +4,7 @@
  */
 #include <dmlc/json.h>
 #include <dmlc/logging.h>
+#include <tvm/runtime/serializer.h>
 #include <fstream>
 
 #include "./file_util.h"

src/runtime/graph/graph_runtime.cc

@@ -4,6 +4,7 @@
  */
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/runtime/serializer.h>
 #include <dmlc/memory_io.h>
 #include <dmlc/json.h>
 #include <numeric>
@@ -397,24 +398,25 @@ class GraphRuntime : public ModuleNode {
 
 
 void GraphRuntime::LoadDLTensor(dmlc::Stream* strm, DLTensor* dst) {
+  // always use strm->Read to maintain endianness conversion
   uint64_t header, reserved;
-  CHECK(strm->Read(&header, sizeof(header)))
+  CHECK(strm->Read(&header))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&reserved, sizeof(reserved)))
+  CHECK(strm->Read(&reserved))
       << "Invalid DLTensor file format";
   CHECK(header == kTVMNDArrayMagic)
       << "Invalid DLTensor file format";
 
   DLTensor tensor;
-  CHECK(strm->Read(&tensor.ctx, sizeof(tensor.ctx)))
+  CHECK(strm->Read(&(tensor.ctx)))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&tensor.ndim, sizeof(tensor.ndim)))
+  CHECK(strm->Read(&(tensor.ndim)))
       << "Invalid DLTensor file format";
-  CHECK(strm->Read(&tensor.dtype, sizeof(tensor.dtype)))
+  CHECK(strm->Read(&(tensor.dtype)))
       << "Invalid DLTensor file format";
   std::vector<int64_t> shape(tensor.ndim);
   if (tensor.ndim != 0) {
-    CHECK(strm->Read(&shape[0], sizeof(int64_t) * tensor.ndim))
+    CHECK(strm->ReadArray(&shape[0], tensor.ndim))
         << "Invalid DLTensor file format";
   }
   CHECK_EQ(tensor.ndim, dst->ndim) << "param dimension mismatch";
@@ -425,18 +427,23 @@ void GraphRuntime::LoadDLTensor(dmlc::Stream* strm, DLTensor* dst) {
     CHECK_EQ(shape[i], dst->shape[i]) << "param shape mismatch";
   }
   size_t bits = dst->dtype.bits * dst->dtype.lanes;
-  size_t size = (bits + 7) / 8;
+  size_t elem_bytes = (bits + 7) / 8;
+  size_t num_elems = 1;
   for (int i = 0; i < dst->ndim; ++i) {
-    size *= dst->shape[i];
+    num_elems *= dst->shape[i];
   }
   uint64_t data_byte_size;
-  CHECK(strm->Read(&data_byte_size, sizeof(data_byte_size)))
+  CHECK(strm->Read(&data_byte_size))
       << "Invalid DLTensor file format";
-  CHECK(data_byte_size == size)
+  CHECK_EQ(data_byte_size, elem_bytes * num_elems)
       << "Invalid DLTensor file format";
   std::vector<uint8_t> bytes(data_byte_size + 1);
   CHECK(strm->Read(&bytes[0], data_byte_size))
       << "Invalid DLTensor file format";
+  // explicitly swap endian when necessary.
+  if (!DMLC_IO_NO_ENDIAN_SWAP) {
+    dmlc::ByteSwap(&bytes[0], elem_bytes, num_elems);
+  }
   TVM_CCALL(TVMArrayCopyFromBytes(dst, &bytes[0], data_byte_size));
 }
 
@@ -453,9 +460,8 @@ void GraphRuntime::LoadParams(dmlc::Stream* strm) {
   CHECK(strm->Read(&names))
       << "Invalid parameters file format";
   uint64_t sz;
-  strm->Read(&sz, sizeof(sz));
+  strm->Read(&sz);
   size_t size = static_cast<size_t>(sz);
-
   CHECK(size == names.size())
       << "Invalid parameters file format";
   for (size_t i = 0; i < size; ++i) {

src/runtime/rpc/rpc_session.cc

@@ -6,6 +6,7 @@
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/runtime/serializer.h>
 #include <memory>
 #include <array>
 #include <string>
@@ -44,7 +45,7 @@ struct RPCArgBuffer {
 };
 
 // Event handler for RPC events.
-class RPCSession::EventHandler {
+class RPCSession::EventHandler : public dmlc::Stream {
  public:
   EventHandler(common::RingBuffer* reader,
                common::RingBuffer* writer,
@@ -71,6 +72,15 @@ class RPCSession::EventHandler {
       return 0;
     }
   }
+  // Request number of bytes from reader.
+  void RequestBytes(size_t nbytes) {
+    pending_request_bytes_ += nbytes;
+    reader_->Reserve(pending_request_bytes_);
+  }
+  // Whether we are ready to handle next request.
+  bool Ready() {
+    return reader_->bytes_available() >= pending_request_bytes_;
+  }
   bool CanCleanShutdown() const {
     return state_ == kRecvCode;
   }
@@ -86,12 +96,12 @@ class RPCSession::EventHandler {
         case kInitHeader: HandleInitHeader(); break;
         case kRecvCode: HandleRecvCode(); break;
         case kRecvCallHandle: {
-          this->Read(&call_handle_, sizeof(call_handle_));
+          CHECK(this->Read(&call_handle_));
           this->SwitchToState(kRecvPackedSeqNumArgs);
           break;
         }
         case kRecvPackedSeqNumArgs: {
-          this->Read(&num_packed_args_, sizeof(num_packed_args_));
+          CHECK(this->Read(&num_packed_args_));
           arg_buf_.reset(new RPCArgBuffer());
           arg_buf_->value.resize(num_packed_args_);
           arg_buf_->tcode.resize(num_packed_args_);
@@ -100,7 +110,7 @@ class RPCSession::EventHandler {
         }
         case kRecvPackedSeqTypeCode: {
           if (num_packed_args_ != 0) {
-            this->Read(arg_buf_->tcode.data(), sizeof(int) * num_packed_args_);
+            this->ReadArray(arg_buf_->tcode.data(), num_packed_args_);
           }
           arg_index_ = 0;
           arg_recv_stage_ = 0;
@@ -164,8 +174,8 @@ class RPCSession::EventHandler {
   }
   // send Packed sequence to writer.
   void SendPackedSeq(const TVMValue* arg_values, const int* type_codes, int n) {
-    writer_->Write(&n, sizeof(n));
-    writer_->Write(type_codes, sizeof(int) * n);
+    this->Write(n);
+    this->WriteArray(type_codes, n);
     // Argument packing.
     for (int i = 0; i < n; ++i) {
       int tcode = type_codes[i];
@@ -173,14 +183,20 @@ class RPCSession::EventHandler {
       switch (tcode) {
         case kDLInt:
         case kDLUInt:
-        case kDLFloat:
+        case kDLFloat: {
+          this->Write<int64_t>(value.v_int64);
+          break;
+        }
         case kTVMType: {
-          writer_->Write(&value, sizeof(TVMValue));
+          this->Write(value.v_type);
+          // padding
+          int32_t padding = 0;
+          this->Write<int32_t>(padding);
           break;
         }
         case kTVMContext: {
           value.v_ctx = StripSessMask(value.v_ctx);
-          writer_->Write(&value, sizeof(TVMValue));
+          this->Write(value.v_ctx);
           break;
         }
         case kFuncHandle:
@@ -188,7 +204,7 @@ class RPCSession::EventHandler {
         case kHandle: {
           // always send handle in 64 bit.
           uint64_t handle = reinterpret_cast<uint64_t>(value.v_handle);
-          writer_->Write(&handle, sizeof(uint64_t));
+          this->Write(handle);
           break;
         }
         case kArrayHandle: {
@@ -196,11 +212,11 @@ class RPCSession::EventHandler {
           TVMContext ctx = StripSessMask(arr->ctx);
           uint64_t data = reinterpret_cast<uint64_t>(
               static_cast<RemoteSpace*>(arr->data)->data);
-          writer_->Write(&data, sizeof(uint64_t));
-          writer_->Write(&ctx, sizeof(ctx));
-          writer_->Write(&(arr->ndim), sizeof(int));
-          writer_->Write(&(arr->dtype), sizeof(DLDataType));
-          writer_->Write(arr->shape, sizeof(int64_t) * arr->ndim);
+          this->Write(data);
+          this->Write(ctx);
+          this->Write(arr->ndim);
+          this->Write(arr->dtype);
+          this->WriteArray(arr->shape, arr->ndim);
           CHECK(arr->strides == nullptr)
               << "Donot support strided remote array";
           CHECK_EQ(arr->byte_offset, 0)
@@ -211,15 +227,15 @@ class RPCSession::EventHandler {
         case kStr: {
           const char* s = value.v_str;
           uint64_t len = strlen(s);
-          writer_->Write(&len, sizeof(len));
-          writer_->Write(s, sizeof(char) * len);
+          this->Write(len);
+          this->WriteArray(s, len);
           break;
         }
         case kBytes: {
           TVMByteArray* bytes = static_cast<TVMByteArray*>(arg_values[i].v_handle);
           uint64_t len = bytes->size;
-          writer_->Write(&len, sizeof(len));
-          writer_->Write(bytes->data, sizeof(char) * len);
+          this->Write(len);
+          this->WriteArray(bytes->data, len);
           break;
         }
         default: {
@@ -230,6 +246,23 @@ class RPCSession::EventHandler {
     }
   }
 
+  // Endian aware IO handling
+  using Stream::Read;
+  using Stream::Write;
+  using Stream::ReadArray;
+  using Stream::WriteArray;
+
+  inline bool Read(RPCCode* code) {
+    int cdata;
+    if (!this->Read(&cdata)) return false;
+    *code = static_cast<RPCCode>(cdata);
+    return true;
+  }
+  inline void Write(RPCCode code) {
+    int cdata = static_cast<int>(code);
+    this->Write(cdata);
+  }
+
  protected:
   enum State {
     kInitHeader,
@@ -370,10 +403,22 @@ class RPCSession::EventHandler {
       switch (tcode) {
         case kDLInt:
         case kDLUInt:
-        case kDLFloat:
-        case kTVMType:
+        case kDLFloat: {
+          this->Read<int64_t>(&(value.v_int64));
+          ++arg_index_;
+          this->SwitchToState(kRecvPackedSeqArg);
+          break;
+        }
+        case kTVMType: {
+          this->Read(&(value.v_type));
+          int32_t padding = 0;
+          this->Read<int32_t>(&padding);
+          ++arg_index_;
+          this->SwitchToState(kRecvPackedSeqArg);
+          break;
+        }
         case kTVMContext: {
-          this->Read(&value, sizeof(TVMValue));
+          this->Read(&(value.v_ctx));
           ++arg_index_;
           this->SwitchToState(kRecvPackedSeqArg);
           break;
@@ -383,7 +428,7 @@ class RPCSession::EventHandler {
         case kHandle: {
           // always send handle in 64 bit.
           uint64_t handle;
-          this->Read(&handle, sizeof(handle));
+          this->Read(&handle);
           value.v_handle = reinterpret_cast<void*>(handle);
           ++arg_index_;
           this->SwitchToState(kRecvPackedSeqArg);
@@ -398,7 +443,7 @@ class RPCSession::EventHandler {
         case kStr:
         case kBytes: {
           uint64_t len;
-          this->Read(&len, sizeof(len));
+          this->Read(&len);
           temp_bytes_.reset( new RPCByteArrayBuffer());
           temp_bytes_->data.resize(len);
           arg_recv_stage_ = 1;
@@ -409,12 +454,12 @@ class RPCSession::EventHandler {
       case kArrayHandle: {
         temp_array_.reset(new RPCDataArrayBuffer());
         uint64_t handle;
-        this->Read(&handle, sizeof(handle));
+        this->Read(&handle);
         DLTensor& tensor = temp_array_->tensor;
         tensor.data = reinterpret_cast<void*>(handle);
-        this->Read(&(tensor.ctx), sizeof(TVMContext));
-        this->Read(&(tensor.ndim), sizeof(int));
-        this->Read(&(tensor.dtype), sizeof(DLDataType));
+        this->Read(&(tensor.ctx));
+        this->Read(&(tensor.ndim));
+        this->Read(&(tensor.dtype));
         temp_array_->shape.resize(tensor.ndim);
         tensor.shape = temp_array_->shape.data();
         arg_recv_stage_ = 1;
@@ -432,7 +477,7 @@ class RPCSession::EventHandler {
       CHECK_EQ(arg_recv_stage_, 1);
       if (tcode == kStr || tcode == kBytes) {
         if (temp_bytes_->data.size() != 0) {
-          this->Read(&(temp_bytes_->data[0]), temp_bytes_->data.size());
+          this->ReadArray(&(temp_bytes_->data[0]), temp_bytes_->data.size());
         }
         if (tcode == kStr) {
           value.v_str = temp_bytes_->data.c_str();
@@ -445,7 +490,7 @@ class RPCSession::EventHandler {
       } else {
         CHECK_EQ(tcode, kArrayHandle);
         DLTensor& tensor = temp_array_->tensor;
-        this->Read(tensor.shape, tensor.ndim * sizeof(int64_t));
+        this->ReadArray(tensor.shape, tensor.ndim);
         value.v_handle = &tensor;
         arg_buf_->temp_array.emplace_back(std::move(temp_array_));
       }
@@ -458,20 +503,20 @@ class RPCSession::EventHandler {
   void HandleInitHeader() {
     if (init_header_step_ == 0) {
       int32_t len;
-      this->Read(&len, sizeof(len));
+      this->Read(&len);
       remote_key_->resize(len);
       init_header_step_ = 1;
       this->RequestBytes(len);
       return;
     } else {
       CHECK_EQ(init_header_step_, 1);
-      this->Read(dmlc::BeginPtr(*remote_key_), remote_key_->length());
+      this->ReadArray(dmlc::BeginPtr(*remote_key_), remote_key_->length());
       this->SwitchToState(kRecvCode);
     }
   }
   // Handler for read code.
   void HandleRecvCode() {
-    this->Read(&code_, sizeof(code_));
+    this->Read(&code_);
     if (code_ > RPCCode::kSystemFuncStart) {
       SwitchToState(kRecvPackedSeqNumArgs);
       return;
@@ -511,14 +556,14 @@ class RPCSession::EventHandler {
   void HandleCopyFromRemote() {
     uint64_t handle, offset, size;
     TVMContext ctx;
-    this->Read(&handle, sizeof(handle));
-    this->Read(&offset, sizeof(offset));
-    this->Read(&size, sizeof(size));
-    this->Read(&ctx, sizeof(ctx));
+    this->Read(&handle);
+    this->Read(&offset);
+    this->Read(&size);
+    this->Read(&ctx);
     if (ctx.device_type == kDLCPU) {
       RPCCode code = RPCCode::kCopyAck;
-      writer_->Write(&code, sizeof(code));
-      writer_->Write(reinterpret_cast<char*>(handle) + offset, size);
+      this->Write(code);
+      this->WriteArray(reinterpret_cast<char*>(handle) + offset, size);
     } else {
       temp_data_.resize(size + 1);
       try {
@@ -530,11 +575,11 @@ class RPCSession::EventHandler {
             dmlc::BeginPtr(temp_data_), 0,
             size, ctx, cpu_ctx, nullptr);
         RPCCode code = RPCCode::kCopyAck;
-        writer_->Write(&code, sizeof(code));
-        writer_->Write(&temp_data_[0], size);
+        this->Write(code);
+        this->WriteArray(&temp_data_[0], size);
       } catch (const std::runtime_error &e) {
         RPCCode code = RPCCode::kException;
-        writer_->Write(&code, sizeof(code));
+        this->Write(code);
         TVMValue ret_value;
         ret_value.v_str = e.what();
         int ret_tcode = kStr;
@@ -548,10 +593,10 @@ class RPCSession::EventHandler {
     // use static variable to persist state.
     // This only works if next stage is immediately after this.
     if (arg_recv_stage_ == 0) {
-      this->Read(&copy_handle_, sizeof(uint64_t));
-      this->Read(&copy_offset_, sizeof(uint64_t));
-      this->Read(&copy_size_, sizeof(uint64_t));
-      this->Read(&copy_ctx_, sizeof(TVMContext));
+      CHECK(this->Read(&copy_handle_));
+      CHECK(this->Read(&copy_offset_));
+      CHECK(this->Read(&copy_size_));
+      CHECK(this->Read(&copy_ctx_));
       arg_recv_stage_ = 1;
       CHECK_EQ(pending_request_bytes_, 0U);
       this->RequestBytes(copy_size_);
@@ -563,11 +608,11 @@ class RPCSession::EventHandler {
       RPCCode code = RPCCode::kReturn;
       std::string errmsg;
       if (copy_ctx_.device_type == kDLCPU) {
-        this->Read(
+        this->ReadArray(
             reinterpret_cast<char*>(copy_handle_) + copy_offset_, copy_size_);
       } else {
         temp_data_.resize(copy_size_ + 1);
-        this->Read(&temp_data_[0], copy_size_);
+        this->ReadArray(&temp_data_[0], copy_size_);
         try {
           TVMContext cpu_ctx;
           cpu_ctx.device_type = kDLCPU;
@@ -583,7 +628,7 @@ class RPCSession::EventHandler {
           ret_tcode = kStr;
         }
       }
-      writer_->Write(&code, sizeof(code));
+      this->Write(code);
       SendPackedSeq(&ret_value, &ret_tcode, 1);
       arg_recv_stage_ = 0;
       this->SwitchToState(kRecvCode);
@@ -603,7 +648,7 @@ class RPCSession::EventHandler {
       std::unique_ptr<RPCArgBuffer> args = std::move(arg_buf_);
       f(args->AsTVMArgs(), &rv);
       RPCCode code = RPCCode::kReturn;
-      writer_->Write(&code, sizeof(code));
+      this->Write(code);
       if (rv.type_code() == kStr) {
         ret_value.v_str = rv.ptr<std::string>()->c_str();
         ret_tcode = kStr;
@@ -630,7 +675,7 @@ class RPCSession::EventHandler {
       }
     } catch (const std::runtime_error& e) {
       RPCCode code = RPCCode::kException;
-      writer_->Write(&code, sizeof(code));
+      this->Write(code);
       ret_value.v_str = e.what();
       ret_tcode = kStr;
       SendPackedSeq(&ret_value, &ret_tcode, 1);
@@ -640,19 +685,14 @@ class RPCSession::EventHandler {
  private:
   // Utility functions
   // Internal read function, update pending_request_bytes_
-  void Read(void* data, size_t size) {
+  size_t Read(void* data, size_t size) final {
     CHECK_LE(size, pending_request_bytes_);
     reader_->Read(data, size);
     pending_request_bytes_ -= size;
+    return size;
   }
-  // Request number of bytes from reader.
-  void RequestBytes(size_t nbytes) {
-    pending_request_bytes_ += nbytes;
-    reader_->Reserve(pending_request_bytes_);
-  }
-  // Whether we are ready to handle next request.
-  bool Ready() {
-    return reader_->bytes_available() >= pending_request_bytes_;
+  void Write(const void* data, size_t size) final {
+    writer_->Write(data, size);
   }
   // Number of pending bytes requests
   size_t pending_request_bytes_;
@@ -766,7 +806,7 @@ RPCSession::~RPCSession() {
 void RPCSession::Shutdown() {
   if (channel_ != nullptr) {
     RPCCode code = RPCCode::kShutdown;
-    writer_.Write(&code, sizeof(code));
+    handler_->Write(code);
     // flush all writing buffer to output channel.
     try {
       while (writer_.bytes_available() != 0) {
@@ -788,7 +828,6 @@ void RPCSession::ServerLoop() {
   }
   TVMRetValue rv;
   CHECK(HandleUntilReturnEvent(&rv, false, nullptr) == RPCCode::kShutdown);
-  LOG(INFO) << "Shutdown...";
   if (const auto* f = Registry::Get("tvm.contrib.rpc.server.shutdown")) {
     (*f)();
   }
@@ -821,9 +860,9 @@ void RPCSession::CallFunc(void* h,
                           const PackedFunc* fwrap) {
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   RPCCode code = RPCCode::kCallFunc;
-  writer_.Write(&code, sizeof(code));
+  handler_->Write(code);
   uint64_t handle = reinterpret_cast<uint64_t>(h);
-  writer_.Write(&handle, sizeof(handle));
+  handler_->Write(handle);
   handler_->SendPackedSeq(args.values, args.type_codes, args.num_args);
   code = HandleUntilReturnEvent(rv, true, fwrap);
   CHECK(code == RPCCode::kReturn) << "code=" << static_cast<int>(code);
@@ -838,15 +877,15 @@ void RPCSession::CopyToRemote(void* from,
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   ctx_to = handler_->StripSessMask(ctx_to);
   RPCCode code = RPCCode::kCopyToRemote;
-  writer_.Write(&code, sizeof(code));
+  handler_->Write(code);
   uint64_t handle = reinterpret_cast<uint64_t>(to);
-  writer_.Write(&handle, sizeof(handle));
+  handler_->Write(handle);
   uint64_t offset = static_cast<uint64_t>(to_offset);
-  writer_.Write(&offset, sizeof(offset));
+  handler_->Write(offset);
   uint64_t size = static_cast<uint64_t>(data_size);
-  writer_.Write(&size, sizeof(size));
-  writer_.Write(&ctx_to, sizeof(ctx_to));
-  writer_.Write(reinterpret_cast<char*>(from) + from_offset, data_size);
+  handler_->Write(size);
+  handler_->Write(ctx_to);
+  handler_->WriteArray(reinterpret_cast<char*>(from) + from_offset, data_size);
   TVMRetValue rv;
   CHECK(HandleUntilReturnEvent(&rv, true, nullptr) == RPCCode::kReturn);
 }
@@ -860,26 +899,27 @@ void RPCSession::CopyFromRemote(void* from,
   std::lock_guard<std::recursive_mutex> lock(mutex_);
   ctx_from = handler_->StripSessMask(ctx_from);
   RPCCode code = RPCCode::kCopyFromRemote;
-  writer_.Write(&code, sizeof(code));
+  handler_->Write(code);
   uint64_t handle = reinterpret_cast<uint64_t>(from);
-  writer_.Write(&handle, sizeof(handle));
+  handler_->Write(handle);
   uint64_t offset = static_cast<uint64_t>(from_offset);
-  writer_.Write(&offset, sizeof(offset));
+  handler_->Write(offset);
   uint64_t size = static_cast<uint64_t>(data_size);
-  writer_.Write(&size, sizeof(size));
-  writer_.Write(&ctx_from, sizeof(ctx_from));
+  handler_->Write(size);
+  handler_->Write(ctx_from);
   TVMRetValue rv;
   CHECK(HandleUntilReturnEvent(&rv, true, nullptr) == RPCCode::kCopyAck);
   reader_.Reserve(data_size);
-  while (reader_.bytes_available() < data_size) {
-    size_t bytes_needed = data_size - reader_.bytes_available();
+  handler_->RequestBytes(data_size);
+  while (!handler_->Ready()) {
+    size_t bytes_needed = handler_->BytesNeeded();
     reader_.WriteWithCallback([this](void* data, size_t size) {
         size_t n = channel_->Recv(data, size);
         CHECK_NE(n, 0U) << "Channel closes before we get neded bytes";
         return n;
       }, bytes_needed);
   }
-  reader_.Read(reinterpret_cast<char*>(to) + to_offset, data_size);
+  handler_->ReadArray(reinterpret_cast<char*>(to) + to_offset, data_size);
   handler_->FinishCopyAck();
 }