Change GPU memory detection logics in baidu_rpc_protocol

Author: sunce4t

src/brpc/policy/baidu_rpc_protocol.cpp

@@ -115,8 +115,8 @@ ParseResult ParseRpcMessage(butil::IOBuf* source, Socket* socket,
 
 #if BRPC_WITH_GDR
     void* prefetch_d2h_data = NULL;
-    uint64_t data_meta = source->get_first_data_meta();
-    bool is_gpu_memory = (data_meta > 0 && data_meta <= UINT_MAX);
+    uint32_t data_meta = source->get_first_data_meta_high32();
+    bool is_gpu_memory = (data_meta == static_cast<uint32_t>(butil::IOBuf::GPU_MEMORY));
     butil::gdr::BlockPoolAllocator* host_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_cpu_allocator();
     if (is_gpu_memory) {
         prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
@@ -859,8 +859,8 @@ void ProcessRpcRequest(InputMessageBase* msg_base) {
             int body_without_attachment_size = req_size - meta.attachment_size();
 #if BRPC_WITH_GDR
             int meta_size = msg->meta.size();
-            uint64_t data_meta = msg->payload.get_first_data_meta();
-            bool is_gpu_memory = (data_meta > 0 && data_meta <= UINT_MAX);
+            uint32_t data_meta = msg->payload.get_first_data_meta_high32();
+            bool is_gpu_memory = (data_meta == static_cast<uint32_t>(butil::IOBuf::GPU_MEMORY));
             if(is_gpu_memory) {
                 int64_t real_prefetch_d2h_size = msg->meta.get_first_data_meta();
                 if (header_size + meta_size + body_without_attachment_size <= real_prefetch_d2h_size) {
@@ -1054,8 +1054,8 @@ void ProcessRpcResponse(InputMessageBase* msg_base) {
         butil::IOBuf* res_buf_ptr = &msg->payload;
 
 #if BRPC_WITH_GDR
-        uint64_t data_meta = msg->payload.get_first_data_meta();
-        bool is_gpu_memory = (data_meta > 0 && data_meta <= UINT_MAX);
+        uint32_t data_meta = msg->payload.get_first_data_meta_high32();
+        bool is_gpu_memory = (data_meta == static_cast<uint32_t>(butil::IOBuf::GPU_MEMORY));
 #endif  // BRPC_WITH_GDR
         if (meta.has_attachment_size()) {
             if (meta.attachment_size() > res_size) {

src/brpc/rdma/rdma_endpoint.cpp

@@ -1056,7 +1056,8 @@ int RdmaEndpoint::PostRecv(uint32_t num, bool zerocopy) {
             void* device_ptr = device_allocator->AllocateRaw(g_rdma_recv_block_size);
             auto deleter = [device_allocator](void* data) { device_allocator->DeallocateRaw(data); };
             lkey = device_allocator->get_lkey(device_ptr);
-            _rbuf[_rq_received].append_user_data_with_meta(device_ptr, g_rdma_recv_block_size, deleter , lkey);
+            uint64_t data_meta = (static_cast<uint64_t>(butil::IOBuf::GPU_MEMORY) << 32) | lkey;
+            _rbuf[_rq_received].append_user_data_with_meta(device_ptr, g_rdma_recv_block_size, deleter , data_meta);
             _rbuf_data[_rq_received] = device_ptr;
 #else
             butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received],

src/butil/iobuf.cpp

@@ -1193,7 +1193,21 @@ uint64_t IOBuf::get_first_data_meta() {
     if (!(r.block->flags & IOBUF_BLOCK_FLAGS_USER_DATA)) {
         return 0;
     }
-    return r.block->u.data_meta;
+    return (r.block->u.data_meta & 0x00000000FFFFFFFF);
+}
+
+// only when user use append_user_data_with_meta(), lkey is stored in data_meta
+// We add this function for GDR, we want to know whether the data is in Host memory or GPU memory
+// since lkey is uint32_t type, thus we use the high 32 bit to store
+uint32_t IOBuf::get_first_data_meta_high32() {
+    if (_ref_num() == 0) {
+        return 0;
+    }
+    IOBuf::BlockRef const& r = _ref_at(0);
+    if (!(r.block->flags & IOBUF_BLOCK_FLAGS_USER_DATA)) {
+        return 0;
+    }
+    return (uint32_t)(r.block->u.data_meta >> 32);
 }
 
 void* IOBuf::get_first_data_ptr() {

src/butil/iobuf.h

@@ -70,6 +70,11 @@ friend class SingleIOBuf;
     static const size_t DEFAULT_BLOCK_SIZE = 8192;
     static const size_t INITIAL_CAP = 32; // must be power of 2
 
+    enum MemoryMeta {
+        HOST_MEMORY = 0,
+        GPU_MEMORY = 1
+    };
+
     struct Block;
 
     // can't directly use `struct iovec' here because we also need to access the
@@ -265,6 +270,11 @@ friend class SingleIOBuf;
     // The meta is specified with append_user_data_with_meta before.
     // 0 means the meta is invalid.
     uint64_t get_first_data_meta();
+
+    // Get the high 32 bits of the data meta of the first byte in this IOBuf.
+    // The meta is specified with append_user_data_with_meta before.
+    // we use 0 to specify host memory, 1 to specify GPU memory
+    uint32_t get_first_data_meta_high32();
     void* get_first_data_ptr();
 
     // Resizes the buf to a length of n characters.