Copilot commented on code in PR #3144:
URL: https://github.com/apache/brpc/pull/3144#discussion_r4217694142
##########
src/brpc/rdma_transport.cpp:
##########
@@ -30,10 +30,19 @@ DECLARE_bool(usercode_in_pthread);
extern SocketVarsCollector *g_vars;
-void RdmaTransport::Init(Socket *socket, const SocketOptions &options) {
+void RdmaTransport::DoInit(Socket *socket, const SocketOptions &options, bool
use_gdr) {
CHECK(_rdma_ep == nullptr);
- if (options.socket_mode == SOCKET_MODE_RDMA) {
- _rdma_ep = new rdma::RdmaEndpoint(socket);
+ // gdr mode is a special mode of rdma mode.
+ // both rdma mode and gdr mode need init rdma::RdmaEndpoint.
+ if (options.socket_mode == SOCKET_MODE_RDMA ||
+ options.socket_mode == SOCKET_MODE_GDR) {
+ _rdma_ep = new(std::nothrow)rdma::RdmaEndpoint(socket, use_gdr);
+ if (!_rdma_ep) {
+ const int saved_errno = errno;
+ PLOG(ERROR) << "Fail to create RdmaEndpoint";
+ socket->SetFailed(
+ saved_errno, "Fail to create RdmaEndpoint: %s",
berror(saved_errno));
+ }
_rdma_state = RDMA_UNKNOWN;
} else {
_rdma_state = RDMA_OFF;
Review Comment:
If `new(std::nothrow)` fails, `_rdma_ep` stays null but `_rdma_state` is
still set to `RDMA_UNKNOWN`, which can make later code paths assume RDMA is
active and dereference `_rdma_ep`. Suggestion: on allocation failure, set
`_rdma_state = RDMA_OFF` (or `FALLBACK_TCP` if that’s the intended behavior)
and return early after `socket->SetFailed(...)`.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch 512B from gpu to cpu.
+// if header_size + meta_size + body_size(without attachment) is less than
512, then one
+// d2h is enough for one rpc.
+
+const int prefetch_d2h_size = 512;
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ LOG(FATAL) << "RpcMessage is not in gpu!!!";
+ return MakeParseError(PARSE_ERROR_TRY_OTHERS);
+ }
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
+ if (prefetch_d2h_data == NULL) {
+ // Pool exhausted or single-block limit hit: ask the caller to retry
+ // instead of crashing the whole server.
+ LOG(ERROR) << "alloc host data for prefetch d2h failed, size="
+ << prefetch_d2h_size;
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+
+ // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size;
+ n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size);
+ size_t copy_size = n > 12 ? 12 : n;
+ memcpy(header_buf, prefetch_d2h_data, copy_size);
+
+ do {
+ if (n >= 4) {
+ void* dummy = header_buf;
+ if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+ pe = PARSE_ERROR_TRY_OTHERS;
+ break;
+ }
+ } else {
+ if (memcmp(header_buf, "PRPC", n) != 0) {
+ pe = PARSE_ERROR_TRY_OTHERS;
+ break;
+ }
+ }
+ if (n < sizeof(header_buf)) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ butil::RawUnpacker(header_buf +
4).unpack32(body_size).unpack32(meta_size);
+ if (body_size > FLAGS_max_body_size) {
+ // We need this log to report the body_size to give users some
clues
+ // which is not printed in InputMessenger.
+ LOG(ERROR) << "body_size=" << body_size << " from "
+ << socket->remote_side() << " is too large";
+ pe = PARSE_ERROR_TOO_BIG_DATA;
+ break;
+ } else if (source->length() < sizeof(header_buf) + body_size) {
+ pe = PARSE_ERROR_NOT_ENOUGH_DATA;
+ break;
+ }
+ if (meta_size > body_size) {
+ LOG(ERROR) << "meta_size=" << meta_size << " is bigger than
body_size="
+ << body_size;
+ // Pop the message
+ source->pop_front(sizeof(header_buf) + body_size);
+ pe = PARSE_ERROR_TRY_OTHERS;
+ break;
+ }
+ } while (0);
+
+ if (pe != PARSE_OK) {
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ return MakeParseError(pe);
+ }
+
+ source->pop_front(sizeof(header_buf));
+ MostCommonMessage* msg = MostCommonMessage::Get();
+
+ if (header_size + meta_size <= n) {
+ auto deleter = [host_allocator, prefetch_d2h_data](void* data) {
host_allocator->DeallocateRaw(prefetch_d2h_data); };
+ // `n` is the number of bytes actually prefetched from GPU. It is
stored
+ // in the meta's `data_meta' so that FillReqBufGpu/FillResBufGpu can
+ // decide whether the prefetched host buffer also covers the body (and
+ // thus avoid a second D2H copy). This block is host memory
(prefetched),
+ // NOT GPU memory, so we deliberately use append_user_data_with_meta
+ // (not append_user_data_gpu) to keep is_gpu_memory() returning false.
+ msg->meta.append_user_data_with_meta((char*)prefetch_d2h_data +
header_size, meta_size, deleter, n);
+ source->pop_front(meta_size);
+ } else {
+ host_allocator->DeallocateRaw(prefetch_d2h_data);
+ source->cutn_from_gpu(&msg->meta, meta_size);
+ }
+ source->cutn(&msg->payload, body_size - meta_size);
+ return MakeMessage(msg);
+}
+
+
+void FillReqBufGpu(butil::IOBuf* req_buf, MostCommonMessage* msg, int
body_without_attachment_size) {
+ int meta_size = msg->meta.size();
+ bool is_gpu_memory = msg->payload.is_gpu_memory();
+ if (!is_gpu_memory) {
+ LOG(FATAL) << "message is not on gpu!!!";
+ }
Review Comment:
Same issue as the parser: `LOG(FATAL)` here can crash the server if the
message is unexpectedly host-backed (or partially host-backed). Consider
turning this into a controller failure / error log + fallback behavior (e.g.,
call the normal CPU path) to keep the server resilient.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch 512B from gpu to cpu.
+// if header_size + meta_size + body_size(without attachment) is less than
512, then one
+// d2h is enough for one rpc.
+
+const int prefetch_d2h_size = 512;
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ LOG(FATAL) << "RpcMessage is not in gpu!!!";
Review Comment:
Using `LOG(FATAL)` in the protocol parser will terminate the whole process
on unexpected input/state. Since this function already returns a parse error,
consider downgrading to `LOG(ERROR)` (or `DLOG`), and rely on the returned
`ParseResult` to recover/try other protocols.
##########
src/brpc/rdma/rdma_endpoint.cpp:
##########
@@ -63,7 +66,7 @@ DEFINE_bool(rdma_recv_zerocopy, true, "Enable zerocopy for
receive side");
DEFINE_int32(rdma_zerocopy_min_size, 512, "The minimal size for receive
zerocopy");
DEFINE_int32(rdma_cqe_poll_once, 32, "The maximum of cqe number polled once.");
DEFINE_int32(rdma_prepared_qp_size, 128, "SQ and RQ size for prepared QP.");
-DEFINE_int32(rdma_prepared_qp_cnt, 1024, "Initial count of prepared QP.");
+DEFINE_int32(rdma_prepared_qp_cnt, 256, "Initial count of prepared QP.");
Review Comment:
This PR changes the default `rdma_prepared_qp_cnt` from 1024 to 256, which
is a behavior/performance tuning change but isn’t mentioned in the PR
description’s change list/side-effects. Either document this as an intentional
side effect (and why), or keep the default unchanged and make the new value
part of GDR-only configuration.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch 512B from gpu to cpu.
+// if header_size + meta_size + body_size(without attachment) is less than
512, then one
+// d2h is enough for one rpc.
+
+const int prefetch_d2h_size = 512;
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ LOG(FATAL) << "RpcMessage is not in gpu!!!";
+ return MakeParseError(PARSE_ERROR_TRY_OTHERS);
+ }
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
+ if (prefetch_d2h_data == NULL) {
+ // Pool exhausted or single-block limit hit: ask the caller to retry
+ // instead of crashing the whole server.
+ LOG(ERROR) << "alloc host data for prefetch d2h failed, size="
+ << prefetch_d2h_size;
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+
+ // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size;
+ n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size);
+ size_t copy_size = n > 12 ? 12 : n;
+ memcpy(header_buf, prefetch_d2h_data, copy_size);
+
+ do {
+ if (n >= 4) {
+ void* dummy = header_buf;
+ if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+ pe = PARSE_ERROR_TRY_OTHERS;
+ break;
+ }
+ } else {
Review Comment:
This uses type-punning and potentially unaligned `uint32_t` loads from
`header_buf`/string literal, which is undefined behavior on some architectures
and can violate strict-aliasing rules. Suggestion: use `memcmp(header_buf,
\"PRPC\", 4)` (or copy into a `uint32_t` via `memcpy`) to avoid
alignment/aliasing issues.
##########
src/butil/gpu/gpu_block_pool.cpp:
##########
@@ -0,0 +1,498 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <iostream>
+#include <chrono>
+#include <gflags/gflags.h>
+#include <sys/stat.h>
+#include "butil/fast_rand.h"
+#include "gpu_block_pool.h"
+namespace butil {
+namespace gdr {
+DEFINE_int32(gdr_block_size_kb, 512, "gdr block size in KB");
+DEFINE_int32(max_gdr_regions, 32, "max num of gdr regions");
+
+#define CHECK_CUDA(call) \
+do { \
+ auto _sts = (call); \
+ if (_sts != cudaSuccess) { \
+ LOG(FATAL) << " cuda error:" \
+ << (cudaGetErrorString(_sts)) << std::string(" at ") \
+ << __FILE__ << ": " << __LINE__; \
+ } \
+} while (0);
+
+size_t GetGdrBlockSize() {
+ return FLAGS_gdr_block_size_kb * 1024;
+}
+
+bool verify_same_context() {
+ static std::once_flag init_flag;
+ static int original_device = -1;
+
+ int current_device;
+ CHECK_CUDA(cudaGetDevice(¤t_device));
+
+ std::call_once(init_flag, [&]() {
+ original_device = current_device;
+ });
+
+ return (current_device == original_device);
+}
+
+void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ void *d_data;
+
+ LOG(INFO) << "try to alloc " << gpu_mem_size << " bytes from gpu " <<
gpu_id;
+
+ CHECK_CUDA(cudaMalloc(&d_data, gpu_mem_size));
+ cudaDeviceSynchronize();
+ return (void *)d_data;
+}
+
+void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+
+ LOG(INFO) << "try to alloc " << cpu_mem_size << " bytes from gpu " <<
gpu_id << "on host";
+
+ void* mem = NULL;
+
+ CHECK_CUDA(cudaMallocHost(&mem, cpu_mem_size));
+
+ cudaDeviceSynchronize();
+
+ return mem;
+}
+
+
+BlockPoolAllocators* BlockPoolAllocators::GetInstance() {
+ return ::Singleton<BlockPoolAllocators,
+ LeakySingletonTraits<BlockPoolAllocators>>::get();
+}
+
+void BlockPoolAllocators::init(int gpu_id, ibv_pd* pd) {
+ LOG(INFO) << "set GPU BlockPoolAllocator for " << gpu_id;
+ size_t region_size = 1024LL * 1024 * 1024;
+ size_t block_size = FLAGS_gdr_block_size_kb * 1024;
+ gpu_mem_alloc = new BlockPoolAllocator(gpu_id, true, pd, block_size,
region_size);
+
+ region_size = 32LL * 1024 * 1024;
+ block_size = 512;
+ cpu_mem_alloc = new BlockPoolAllocator(gpu_id, false, pd, block_size,
region_size);
+
+ gpu_stream_pool = new GPUStreamPool(gpu_id);
+}
+
+bool InitGPUBlockPool(int gpu_id, ibv_pd* pd) {
+ BlockPoolAllocators::GetInstance()->init(gpu_id, pd);
+ return true;
+}
+
+class BlockHeaderList {
+ public:
+ BlockHeaderList() {
+ objects_.reserve(kMaxObjects);
+ }
+ virtual ~BlockHeaderList() {
+ for (size_t i = 0; i < objects_.size(); i++) {
+ delete objects_[i];
+ }
+ }
+
+ BlockHeader* New() {
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (!objects_.empty()) {
+ BlockHeader* result = objects_.back();
+ objects_.pop_back();
+ return result;
+ }
+ }
+ return new BlockHeader;
+ }
+ void Release(BlockHeader* obj) {
+ obj->Reset();
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (objects_.size() < kMaxObjects) {
+ objects_.push_back(obj);
+ return;
+ }
+ }
+ delete obj;
+ }
+
+ private:
+ static const int kMaxObjects = 100000;
+
+ std::mutex mu_;
+ std::vector<BlockHeader*> objects_;
+};
+
+static BlockHeaderList* get_bh_list() {
+ static BlockHeaderList* bh_list = new BlockHeaderList();
+ return bh_list;
+}
+
+BlockPoolAllocator::BlockPoolAllocator(int gpuId, bool onGpu, ibv_pd* ibvPd,
+ size_t blockSize, size_t regionSize) :
+ gpu_id(gpuId)
+ , on_gpu(onGpu)
+ , pd(ibvPd)
+ , BLOCK_SIZE(std::max(blockSize, sizeof(BlockHeader)))
+ , REGION_SIZE((regionSize / blockSize) * blockSize) // 对齐到块大小的倍数
+ , freeList(nullptr)
+ , g_region_num(0)
+ , totalAllocated(0)
+ , totalDeallocated(0)
+ , peakUsage(0) {
+ g_regions.resize(FLAGS_max_gdr_regions);
+ LOG(INFO) << "Memory Pool initialized: block_size=" << BLOCK_SIZE
+ << ", region_size=" << REGION_SIZE << ", max_gdr_regions=" <<
FLAGS_max_gdr_regions
+ << ", gpu_id=" << gpu_id << ", on_gpu=" << on_gpu << ", pd=" << pd;
+
+ // The very first region is required at construction time; failure here is
+ // fatal because the pool would be unusable.
+ if (extendRegion() != 0) {
+ LOG(FATAL) << "Failed to initialize GDR memory pool (first region) on
gpu "
+ << gpu_id << ", on_gpu " << on_gpu;
+ }
+}
+
+BlockPoolAllocator::~BlockPoolAllocator() {
+#ifdef DEBUG
+ printStatistics();
+#endif
+
+ for (int i = 0; i < FLAGS_max_gdr_regions; i++) {
+ Region* r = &g_regions[i];
+ if (!r->mr) {
+ continue;
+ }
+
+ LOG(INFO) << "try to free " << r->size << " bytes from gpu " << gpu_id
<< ", on_gpu " << on_gpu;
+ ibv_dereg_mr(r->mr);
+ if (on_gpu) {
+ CHECK_CUDA(cudaFree(reinterpret_cast<void*>(r->start)));
+ } else {
+ CHECK_CUDA(cudaFreeHost(reinterpret_cast<void*>(r->start)));
+ }
+ }
+}
+
+Region* BlockPoolAllocator::GetRegion(const void* buf) {
+ if (!buf) {
+ errno = EINVAL;
+ return NULL;
+ }
+ Region* r = NULL;
+ uintptr_t addr = (uintptr_t)buf;
+ for (int i = 0; i < FLAGS_max_gdr_regions; ++i) {
+ if (g_regions[i].aligned_start == 0) {
+ break;
+ }
+ if (addr >= g_regions[i].aligned_start &&
+ addr < g_regions[i].aligned_start + g_regions[i].aligned_size)
{
+ r = &g_regions[i];
+ break;
+ }
+ }
+ return r;
+}
+
+uint32_t BlockPoolAllocator::get_lkey(const void* buf) {
+ Region* r = GetRegion(buf);
+ if (!r) {
+ LOG(ERROR) << "can not get a region for buf " << buf;
+ return 0;
+ }
+
+ if (!r->mr) {
+ LOG(FATAL) << "region has not been registered into rdma yet, addr:" <<
r->start;
+ return 0;
+ }
+
+ return r->mr->lkey;
+}
+
+void* BlockPoolAllocator::AllocateRaw(size_t num_bytes) {
+ if (num_bytes == 0) {
+ return nullptr;
+ }
+ if (num_bytes > BLOCK_SIZE) {
+ // A single block cannot satisfy the request; fail gracefully instead
of
+ // aborting the whole server.
+ LOG(ERROR) << "try to alloc " << num_bytes
+ << " bytes, bigger than block_size " << BLOCK_SIZE;
+ errno = ENOMEM;
+ return nullptr;
+ }
+
+ auto startTime = std::chrono::high_resolution_clock::now();
+
+ std::lock_guard<std::mutex> lock(poolMutex);
+
+ if (!freeList) {
+ if (extendRegion() != 0) {
+ // Out of regions or registration failed; let the caller handle it.
+ errno = ENOMEM;
+ return nullptr;
+ }
+ }
+
+ BlockHeader* block = freeList;
+ freeList = freeList->next;
+
+ void* addr = block->addr;
+ get_bh_list()->Release(block);
+
+ totalAllocated++;
+ peakUsage = std::max(peakUsage, totalAllocated - totalDeallocated);
+
+ auto endTime = std::chrono::high_resolution_clock::now();
+ auto duration =
std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+#ifdef DEBUG
+ if (duration.count() > 1000) {
+ LOG(INFO) << "Slow allocation: " << duration.count() << " ns";
+ }
+#endif
+
+ return addr;
+}
+
+void BlockPoolAllocator::DeallocateRaw(void* ptr) {
+ if (!ptr) return;
+
+ std::lock_guard<std::mutex> lock(poolMutex);
+
+ BlockHeader* block = get_bh_list()->New();
+ block->addr = ptr;
+ block->next = freeList;
+ freeList = block;
+
+ totalDeallocated++;
+}
+
+void BlockPoolAllocator::printStatistics() const {
+ LOG(INFO) << "=== Memory Pool Statistics ===";
+ LOG(INFO) << "Total regions: " << g_region_num
+ << ", Total blocks allocated: " << totalAllocated
+ << ", Total blocks deallocated: " << totalDeallocated
+ << ", Current usage: " << (totalAllocated - totalDeallocated) << "
blocks"
+ << ", Peak usage: " << peakUsage << " blocks"
+ << ", Memory efficiency: "
+ << (static_cast<double>(totalAllocated - totalDeallocated) /
+ (g_region_num * (REGION_SIZE / BLOCK_SIZE)) * 100)
+ << "%";
+}
+
+int BlockPoolAllocator::extendRegion() {
+ if (g_region_num == FLAGS_max_gdr_regions) {
+ LOG(ERROR) << "Gdr Memory pool reaches max regions (" <<
FLAGS_max_gdr_regions << ")";
+ errno = ENOMEM;
+ return -1;
+ }
+
+ auto startTime = std::chrono::high_resolution_clock::now();
+ void* ptr = nullptr;
+ void* aligned_ptr = nullptr;
+ int alignment = 4096;
+
+ if (on_gpu) {
+ ptr = get_gpu_mem(gpu_id, REGION_SIZE);
+ } else {
+ ptr = get_cpu_mem(gpu_id, REGION_SIZE);
+ }
+ if (ptr == nullptr) {
+ LOG(ERROR) << "Failed to allocate region memory on gpu " << gpu_id
+ << ", on_gpu " << on_gpu;
+ return -1;
+ }
+
+ aligned_ptr = (void*)(((uintptr_t)ptr + alignment - 1) & ~(alignment - 1));
+
+ int64_t aligned_bytes = REGION_SIZE;
+ if (ptr != aligned_ptr) {
+ uintptr_t region_end = uintptr_t(ptr) + REGION_SIZE;
+ uintptr_t aligned_end_ptr = region_end & ~(alignment - 1);
+ aligned_bytes = uintptr_t(aligned_end_ptr) - uintptr_t(aligned_ptr);
+ LOG(WARNING) << "addr is not aligned with 4096: " << ptr << ",
aligned_bytes: " << aligned_bytes
+ << ", region_size: " << REGION_SIZE;
+ }
+
+ LOG(INFO) << "reg_mr for ptr: " << aligned_ptr << ", size:" <<
aligned_bytes;
+ // IBV_ACCESS_RELAXED_ORDERING is only available on newer libibverbs; guard
+ // it so the build does not break on older toolchains.
+ int access_flags =
+ IBV_ACCESS_LOCAL_WRITE |
+ IBV_ACCESS_REMOTE_READ |
+ IBV_ACCESS_REMOTE_WRITE;
+#ifdef IBV_ACCESS_RELAXED_ORDERING
+ access_flags |= IBV_ACCESS_RELAXED_ORDERING;
+#endif
+ auto mr = ibv_reg_mr(pd, aligned_ptr, aligned_bytes, access_flags);
+
+ if (!mr) {
+ LOG(ERROR) << "Failed to register MR: " << strerror(errno)
+ << ", pd " << pd << ", aligned_ptr:" << aligned_ptr;
+ // Free the region memory we just allocated so it is not leaked.
+ if (on_gpu) {
+ cudaFree(reinterpret_cast<void*>(ptr));
+ } else {
+ cudaFreeHost(reinterpret_cast<void*>(ptr));
+ }
+ return -1;
+ } else {
+ LOG(INFO) << "Success to register MR: "
+ << ", pd " << pd << ", aligned_ptr:" << aligned_ptr;
+ }
+
+ LOG(INFO) << "try to init region, g_region_num:" << g_region_num;
+ size_t blockCount = aligned_bytes / BLOCK_SIZE;
+ Region* region = &g_regions[g_region_num++];
+ region->start = (uintptr_t)ptr;
+ region->aligned_start = (uintptr_t)aligned_ptr;
+ region->mr = mr;
+ region->size = REGION_SIZE;
+ region->aligned_size = aligned_bytes;
+ region->blockCount = blockCount;
+
+
+ LOG(INFO) << "try to insert list, freeList:" << freeList << ",
blockCount:" << blockCount;
+ BlockHeader* lastBlock = nullptr;
+ for (size_t i = 0; i < blockCount; ++i) {
+ BlockHeader* block = get_bh_list()->New();
+ block->addr = reinterpret_cast<void*>(static_cast<char*>(aligned_ptr)
+ i * BLOCK_SIZE);
+ if (lastBlock != nullptr) {
+ lastBlock->next = block;
+ } else {
+ freeList = block;
+ }
+ lastBlock = block;
+ }
+
+ if (lastBlock) {
+ lastBlock->next = nullptr;
+ }
+
+ auto endTime = std::chrono::high_resolution_clock::now();
+ auto duration =
std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+ LOG(INFO) << "Extended region #" << g_region_num << ": " << blockCount
+ << " blocks (" << (REGION_SIZE / (1024 * 1024)) << " MB)" << ", on_gpu
" << on_gpu
+ << ", cost " << duration.count() << " ns";
+ return 0;
+}
+
+GPUStreamPool::GPUStreamPool(int gpu_id) :
+ gpu_id_(gpu_id) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ d2d_streams_.resize(kMaxConcurrent);
+ d2h_streams_.resize(kMaxConcurrent);
+ for (int i = 0; i < kMaxConcurrent; i++) {
+ CHECK_CUDA(cudaStreamCreate(&d2d_streams_[i]));
+ CHECK_CUDA(cudaStreamCreate(&d2h_streams_[i]));
+ }
+ CHECK_CUDA(cudaDeviceSynchronize());
+}
+
+GPUStreamPool::~GPUStreamPool() {
+ CHECK_CUDA(cudaDeviceSynchronize());
+ for (int i = 0; i < kMaxConcurrent; i++) {
+ CHECK_CUDA(cudaStreamDestroy(d2d_streams_[i]));
+ CHECK_CUDA(cudaStreamDestroy(d2h_streams_[i]));
+ }
+}
+
+void GPUStreamPool::fast_d2d(std::vector<void*>& src_list,
+ std::vector<int64_t>& length_list,
+ void* dst) {
+#ifdef DEBUG
+ if (!verify_same_context()) {
+ LOG(FATAL) << "Context mismatch!";
+ return;
+ }
+#endif
+ int64_t offset = 0;
+ int segs = src_list.size();
+ if (segs == 0) return;
+ if (segs != length_list.size()) {
+ LOG(FATAL) << "src list size is not equal with length list size!!!";
+ }
+
+ int stream_idx = 0;
+ {
+ std::lock_guard<std::mutex> stream_lb_lock(d2d_lb_lock_);
+ d2d_cnt_.fetch_add(1);
+ stream_idx = d2d_cnt_ % kMaxConcurrent;
+ }
+ std::lock_guard<std::mutex> stream_lock(d2d_locks_[stream_idx]);
+ CHECK_CUDA(cudaStreamSynchronize(d2d_streams_[stream_idx]));
+ for (int i = 0; i < segs; i++) {
+ if (length_list[i] == 0) {
+ continue;
+ }
+ CHECK_CUDA(cudaMemcpyAsync(static_cast<char*>(dst) + offset,
src_list[i], length_list[i],
+ cudaMemcpyDeviceToDevice, d2d_streams_[stream_idx]));
+ offset += length_list[i];
+ }
+ CHECK_CUDA(cudaStreamSynchronize(d2d_streams_[stream_idx]));
Review Comment:
The initial `cudaStreamSynchronize()` before enqueueing copies is redundant
if you already synchronize at the end: the final synchronize will wait for both
any previous work in the stream and the newly enqueued memcpy operations.
Removing the first synchronize reduces latency on hot paths.
##########
src/butil/gpu/gpu_block_pool.cpp:
##########
@@ -0,0 +1,498 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <iostream>
+#include <chrono>
+#include <gflags/gflags.h>
+#include <sys/stat.h>
+#include "butil/fast_rand.h"
+#include "gpu_block_pool.h"
+namespace butil {
+namespace gdr {
+DEFINE_int32(gdr_block_size_kb, 512, "gdr block size in KB");
+DEFINE_int32(max_gdr_regions, 32, "max num of gdr regions");
+
+#define CHECK_CUDA(call) \
+do { \
+ auto _sts = (call); \
+ if (_sts != cudaSuccess) { \
+ LOG(FATAL) << " cuda error:" \
+ << (cudaGetErrorString(_sts)) << std::string(" at ") \
+ << __FILE__ << ": " << __LINE__; \
+ } \
+} while (0);
+
+size_t GetGdrBlockSize() {
+ return FLAGS_gdr_block_size_kb * 1024;
+}
+
+bool verify_same_context() {
+ static std::once_flag init_flag;
+ static int original_device = -1;
+
+ int current_device;
+ CHECK_CUDA(cudaGetDevice(¤t_device));
+
+ std::call_once(init_flag, [&]() {
+ original_device = current_device;
+ });
+
+ return (current_device == original_device);
+}
+
+void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ void *d_data;
+
+ LOG(INFO) << "try to alloc " << gpu_mem_size << " bytes from gpu " <<
gpu_id;
+
+ CHECK_CUDA(cudaMalloc(&d_data, gpu_mem_size));
+ cudaDeviceSynchronize();
+ return (void *)d_data;
+}
+
+void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+
+ LOG(INFO) << "try to alloc " << cpu_mem_size << " bytes from gpu " <<
gpu_id << "on host";
+
+ void* mem = NULL;
+
+ CHECK_CUDA(cudaMallocHost(&mem, cpu_mem_size));
+
+ cudaDeviceSynchronize();
+
+ return mem;
+}
+
+
+BlockPoolAllocators* BlockPoolAllocators::GetInstance() {
+ return ::Singleton<BlockPoolAllocators,
+ LeakySingletonTraits<BlockPoolAllocators>>::get();
+}
+
+void BlockPoolAllocators::init(int gpu_id, ibv_pd* pd) {
+ LOG(INFO) << "set GPU BlockPoolAllocator for " << gpu_id;
+ size_t region_size = 1024LL * 1024 * 1024;
+ size_t block_size = FLAGS_gdr_block_size_kb * 1024;
+ gpu_mem_alloc = new BlockPoolAllocator(gpu_id, true, pd, block_size,
region_size);
+
+ region_size = 32LL * 1024 * 1024;
+ block_size = 512;
+ cpu_mem_alloc = new BlockPoolAllocator(gpu_id, false, pd, block_size,
region_size);
+
+ gpu_stream_pool = new GPUStreamPool(gpu_id);
+}
+
+bool InitGPUBlockPool(int gpu_id, ibv_pd* pd) {
+ BlockPoolAllocators::GetInstance()->init(gpu_id, pd);
+ return true;
+}
+
+class BlockHeaderList {
+ public:
+ BlockHeaderList() {
+ objects_.reserve(kMaxObjects);
+ }
+ virtual ~BlockHeaderList() {
+ for (size_t i = 0; i < objects_.size(); i++) {
+ delete objects_[i];
+ }
+ }
+
+ BlockHeader* New() {
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (!objects_.empty()) {
+ BlockHeader* result = objects_.back();
+ objects_.pop_back();
+ return result;
+ }
+ }
+ return new BlockHeader;
+ }
+ void Release(BlockHeader* obj) {
+ obj->Reset();
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (objects_.size() < kMaxObjects) {
+ objects_.push_back(obj);
+ return;
+ }
+ }
+ delete obj;
+ }
+
+ private:
+ static const int kMaxObjects = 100000;
+
+ std::mutex mu_;
+ std::vector<BlockHeader*> objects_;
+};
+
+static BlockHeaderList* get_bh_list() {
+ static BlockHeaderList* bh_list = new BlockHeaderList();
+ return bh_list;
+}
+
+BlockPoolAllocator::BlockPoolAllocator(int gpuId, bool onGpu, ibv_pd* ibvPd,
+ size_t blockSize, size_t regionSize) :
+ gpu_id(gpuId)
+ , on_gpu(onGpu)
+ , pd(ibvPd)
+ , BLOCK_SIZE(std::max(blockSize, sizeof(BlockHeader)))
+ , REGION_SIZE((regionSize / blockSize) * blockSize) // 对齐到块大小的倍数
+ , freeList(nullptr)
+ , g_region_num(0)
+ , totalAllocated(0)
+ , totalDeallocated(0)
+ , peakUsage(0) {
+ g_regions.resize(FLAGS_max_gdr_regions);
+ LOG(INFO) << "Memory Pool initialized: block_size=" << BLOCK_SIZE
+ << ", region_size=" << REGION_SIZE << ", max_gdr_regions=" <<
FLAGS_max_gdr_regions
+ << ", gpu_id=" << gpu_id << ", on_gpu=" << on_gpu << ", pd=" << pd;
+
+ // The very first region is required at construction time; failure here is
+ // fatal because the pool would be unusable.
+ if (extendRegion() != 0) {
+ LOG(FATAL) << "Failed to initialize GDR memory pool (first region) on
gpu "
+ << gpu_id << ", on_gpu " << on_gpu;
+ }
+}
+
+BlockPoolAllocator::~BlockPoolAllocator() {
+#ifdef DEBUG
+ printStatistics();
+#endif
+
+ for (int i = 0; i < FLAGS_max_gdr_regions; i++) {
+ Region* r = &g_regions[i];
+ if (!r->mr) {
+ continue;
+ }
+
+ LOG(INFO) << "try to free " << r->size << " bytes from gpu " << gpu_id
<< ", on_gpu " << on_gpu;
+ ibv_dereg_mr(r->mr);
+ if (on_gpu) {
+ CHECK_CUDA(cudaFree(reinterpret_cast<void*>(r->start)));
+ } else {
+ CHECK_CUDA(cudaFreeHost(reinterpret_cast<void*>(r->start)));
+ }
+ }
+}
+
+Region* BlockPoolAllocator::GetRegion(const void* buf) {
+ if (!buf) {
+ errno = EINVAL;
+ return NULL;
+ }
+ Region* r = NULL;
+ uintptr_t addr = (uintptr_t)buf;
+ for (int i = 0; i < FLAGS_max_gdr_regions; ++i) {
+ if (g_regions[i].aligned_start == 0) {
+ break;
+ }
+ if (addr >= g_regions[i].aligned_start &&
+ addr < g_regions[i].aligned_start + g_regions[i].aligned_size)
{
+ r = &g_regions[i];
+ break;
+ }
+ }
+ return r;
+}
+
+uint32_t BlockPoolAllocator::get_lkey(const void* buf) {
+ Region* r = GetRegion(buf);
+ if (!r) {
+ LOG(ERROR) << "can not get a region for buf " << buf;
+ return 0;
+ }
+
+ if (!r->mr) {
+ LOG(FATAL) << "region has not been registered into rdma yet, addr:" <<
r->start;
+ return 0;
+ }
+
+ return r->mr->lkey;
+}
+
+void* BlockPoolAllocator::AllocateRaw(size_t num_bytes) {
+ if (num_bytes == 0) {
+ return nullptr;
+ }
+ if (num_bytes > BLOCK_SIZE) {
+ // A single block cannot satisfy the request; fail gracefully instead
of
+ // aborting the whole server.
+ LOG(ERROR) << "try to alloc " << num_bytes
+ << " bytes, bigger than block_size " << BLOCK_SIZE;
+ errno = ENOMEM;
+ return nullptr;
+ }
+
+ auto startTime = std::chrono::high_resolution_clock::now();
+
+ std::lock_guard<std::mutex> lock(poolMutex);
+
+ if (!freeList) {
+ if (extendRegion() != 0) {
+ // Out of regions or registration failed; let the caller handle it.
+ errno = ENOMEM;
+ return nullptr;
+ }
+ }
+
+ BlockHeader* block = freeList;
+ freeList = freeList->next;
+
+ void* addr = block->addr;
+ get_bh_list()->Release(block);
+
+ totalAllocated++;
+ peakUsage = std::max(peakUsage, totalAllocated - totalDeallocated);
+
+ auto endTime = std::chrono::high_resolution_clock::now();
+ auto duration =
std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+#ifdef DEBUG
+ if (duration.count() > 1000) {
+ LOG(INFO) << "Slow allocation: " << duration.count() << " ns";
+ }
+#endif
Review Comment:
The allocation fast-path always takes two timestamps and computes
`duration`, even in non-DEBUG builds. This adds overhead on every allocation.
Suggestion: guard the timing collection itself under `#ifdef DEBUG` (or another
profiling macro), not just the logging.
##########
example/rdma_performance/client.cpp:
##########
@@ -278,6 +322,12 @@ int main(int argc, char* argv[]) {
if (FLAGS_use_rdma) {
brpc::rdma::GlobalRdmaInitializeOrDie();
}
+#ifdef BRPC_WITH_GDR
+ else if (FLAGS_use_gdr) {
+ brpc::rdma::GlobalRdmaInitializeOrDie();
Review Comment:
With the current `if/else if`, `GlobalGdrInitializeOrDie()` won’t run when
both `--use_rdma=true` (default) and `--use_gdr=true`, because the first branch
wins. Suggestion: initialize RDMA when either flag is set, and initialize GDR
additionally when `--use_gdr` is set.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
