From 2b497cb3ba7afed7cf7e3e893b7567996d2a3ca5 Mon Sep 17 00:00:00 2001 From: randomkang <550941794@qq.com> Date: Wed, 3 Jun 2026 00:36:19 +0800 Subject: [PATCH 1/4] support gpu direct rdma 1) recv all data on gpu first 2) the gpu block is alloced from a gpu block pool 3) brpc header, meta and body will be copied from gpu to cpu to process. 4) To decrease the d2h counts, we will prefetch 512B to memory Co-authored-by: sunce4t <847480001@qq.com> --- BUILD.bazel | 9 + Makefile | 1 + bazel/config/BUILD.bazel | 8 +- config_brpc.sh | 17 +- docs/cn/gdr.md | 43 ++ docs/en/gdr.md | 44 ++ example/BUILD.bazel | 5 +- example/rdma_performance/client.cpp | 58 ++- example/rdma_performance/server.cpp | 7 + src/brpc/acceptor.h | 2 +- src/brpc/channel.cpp | 2 + src/brpc/gdr_transport.cpp | 35 ++ src/brpc/gdr_transport.h | 32 ++ src/brpc/policy/baidu_rpc_protocol.cpp | 44 +- src/brpc/policy/baidu_rpc_protocol.h | 13 + src/brpc/policy/baidu_rpc_protocol_gpu.cpp | 228 ++++++++++ src/brpc/rdma/rdma_endpoint.cpp | 104 ++++- src/brpc/rdma/rdma_endpoint.h | 24 +- src/brpc/rdma/rdma_helper.cpp | 38 +- src/brpc/rdma/rdma_helper.h | 7 + src/brpc/rdma_transport.cpp | 13 +- src/brpc/rdma_transport.h | 3 +- src/brpc/socket_mode.h | 5 +- src/brpc/transport_factory.cpp | 17 +- src/butil/gpu/gpu_block_pool.cpp | 475 +++++++++++++++++++++ src/butil/gpu/gpu_block_pool.h | 175 ++++++++ src/butil/iobuf.cpp | 102 +++++ src/butil/iobuf.h | 10 + 28 files changed, 1475 insertions(+), 46 deletions(-) create mode 100644 docs/cn/gdr.md create mode 100644 docs/en/gdr.md create mode 100644 src/brpc/gdr_transport.cpp create mode 100644 src/brpc/gdr_transport.h create mode 100644 src/brpc/policy/baidu_rpc_protocol_gpu.cpp create mode 100644 src/butil/gpu/gpu_block_pool.cpp create mode 100644 src/butil/gpu/gpu_block_pool.h diff --git a/BUILD.bazel b/BUILD.bazel index 5dc5fcf726..ef4bf402de 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -51,6 +51,9 @@ DEFINES = [ }) + select({ "//bazel/config:brpc_with_rdma": ["BRPC_WITH_RDMA=1"], "//conditions:default": [], +}) + select({ + "//bazel/config:brpc_with_gdr": ["-DBRPC_WITH_GDR=1"], + "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_debug_bthread_sche_safety": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1"], "//conditions:default": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=0"], @@ -94,6 +97,11 @@ LINKOPTS = [ "-libverbs", ], "//conditions:default": [], +}) + select({ + "//bazel/config:brpc_with_gdr": [ + "-lcuda -lcudart", + ], + "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_asan": ["-fsanitize=address"], "//conditions:default": [], @@ -237,6 +245,7 @@ BUTIL_SRCS = [ "src/butil/iobuf.cpp", "src/butil/single_iobuf.cpp", "src/butil/iobuf_profiler.cpp", + "src/butil/gpu/gpu_block_pool.cpp", "src/butil/binary_printer.cpp", "src/butil/recordio.cc", "src/butil/popen.cpp", diff --git a/Makefile b/Makefile index 86de388448..c45b2441cf 100644 --- a/Makefile +++ b/Makefile @@ -97,6 +97,7 @@ BUTIL_SOURCES = \ src/butil/files/scoped_temp_dir.cc \ src/butil/file_util.cc \ src/butil/file_util_posix.cc \ + src/butil/gpu/gpu_block_pool.cpp \ src/butil/guid.cc \ src/butil/guid_posix.cc \ src/butil/hash.cc \ diff --git a/bazel/config/BUILD.bazel b/bazel/config/BUILD.bazel index eec551da8b..12319f0dbe 100644 --- a/bazel/config/BUILD.bazel +++ b/bazel/config/BUILD.bazel @@ -104,6 +104,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "brpc_with_gdr", + define_values = {"BRPC_WITH_GDR": "true"}, + visibility = ["//visibility:public"], +) + config_setting( name = "brpc_with_boringssl", define_values = {"BRPC_WITH_BORINGSSL": "true"}, @@ -149,4 +155,4 @@ config_setting( name = "with_babylon_counter", define_values = {"with_babylon_counter": "true"}, visibility = ["//visibility:public"], -) \ No newline at end of file +) diff --git a/config_brpc.sh b/config_brpc.sh index 1c05942068..ed7d8f1872 100755 --- a/config_brpc.sh +++ b/config_brpc.sh @@ -54,10 +54,11 @@ else LDD=ldd fi -TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,nodebugsymbols,werror -n 'config_brpc' -- "$@"` +TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-gdr,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,nodebugsymbols,werror -n 'config_brpc' -- "$@"` WITH_GLOG=0 WITH_THRIFT=0 WITH_RDMA=0 +WITH_GDR=0 WITH_MESALINK=0 WITH_BTHREAD_TRACER=0 WITH_ASAN=0 @@ -89,6 +90,7 @@ while true; do --with-glog ) WITH_GLOG=1; shift 1 ;; --with-thrift) WITH_THRIFT=1; shift 1 ;; --with-rdma) WITH_RDMA=1; shift 1 ;; + --with-gdr) WITH_GDR=1; shift 1 ;; --with-mesalink) WITH_MESALINK=1; shift 1 ;; --with-bthread-tracer) WITH_BTHREAD_TRACER=1; shift 1 ;; --with-debug-bthread-sche-safety ) BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1; shift 1 ;; @@ -536,6 +538,18 @@ if [ $WITH_RDMA != 0 ]; then append_to_output "WITH_RDMA=1" fi +if [ $WITH_GDR != 0 ]; then + CUDA_LIB="/usr/local/cuda/lib64" + CUDA_HDR="/usr/local/cuda/include" + append_to_output_libs "$CUDA_LIB" + append_to_output_headers "$CUDA_HDR" + + CPPFLAGS="${CPPFLAGS} -DBRPC_WITH_GDR" + + append_to_output "DYNAMIC_LINKINGS+=-lcuda -lcudart" + append_to_output "WITH_GDR=1" +fi + if [ $WITH_MESALINK != 0 ]; then CPPFLAGS="${CPPFLAGS} -DUSE_MESALINK" fi @@ -667,6 +681,7 @@ print_info "System: $SYSTEM" if [ $WITH_GLOG -ne 0 ]; then print_info "With glog: yes"; fi if [ $WITH_THRIFT -ne 0 ]; then print_info "With thrift: yes"; fi if [ $WITH_RDMA -ne 0 ]; then print_info "With RDMA: yes"; fi +if [ $WITH_GDR -ne 0 ]; then print_info "With GDR: yes"; fi if [ $WITH_MESALINK -ne 0 ]; then print_info "With MesaLink: yes"; fi if [ $WITH_BTHREAD_TRACER -ne 0 ]; then print_info "With bthread tracer: yes"; fi if [ $WITH_ASAN -ne 0 ]; then print_info "With ASAN: yes"; fi diff --git a/docs/cn/gdr.md b/docs/cn/gdr.md new file mode 100644 index 0000000000..7e46f19d2f --- /dev/null +++ b/docs/cn/gdr.md @@ -0,0 +1,43 @@ +# 编译 + +GDR: GPU Direct Rdma, gdr 是rdma的一种特殊模式，其通过rdma将数据直接收到了gpu的显存上。 + +由于GDR对驱动与硬件有要求，目前仅支持在Linux系统编译并运行GDR功能。 + +目前GDR只支持baidu std protocol。 + +使用config_brpc： +```bash +sh config_brpc.sh --with-rdma --with-gdr --headers="/usr/include" --libs="/usr/lib64 /usr/bin" +make + +cd example/rdma_performance # 示例程序 +make +``` + +使用bazel: +```bash +# Server +bazel build --define=BRPC_WITH_RDMA=true --define=BRPC_WITH_GDR=true example:rdma_performance_server +# Client +bazel build --define=BRPC_WITH_RDMA=true --define=BRPC_WITH_GDR=true example:rdma_performance_client +``` + +# 基本实现 + +GDR是RDMA的一种特殊形式，在使用GDR之前，必须对RDMA和GDR都进行Global Init。 +GDR新增了一个显存池，类似于RDMA内存池，显存池的数据也是按照block进行组织的。 +当打开GDR功能后，框架通过DoPostRecvGDR来发起显存上的WQE。 +在接收到数据后，我们将header、meta、body（不包括attachment）copy回内存进行处理。 +AttachMent位于显存上，用户可以调用IOBuf::copy_from_gpu接口将attachment从brpc框架层copy到应用层进行处理。 + + +注意： +1. 在使用gdr功能时，需要将环境变量MLX5_SCATTER_TO_CQE设置为0. + + +# 参数 + +可配置参数说明： +* gdr_block_size_kb: 使用gdr传送数据时，block的大小（单位为KB），默认为512； +* max_gdr_regions: gdr显存池所使用Region的最大个数，每个Region大小为1GB； diff --git a/docs/en/gdr.md b/docs/en/gdr.md new file mode 100644 index 0000000000..2e968f3e69 --- /dev/null +++ b/docs/en/gdr.md @@ -0,0 +1,44 @@ +Compile GDR: + +GPU Direct RDMA. GDR is a special mode of RDMA that allows data to be received directly into the GPU’s memory through RDMA. +Because GDR requires specific drivers and hardware support, it is currently only available for compilation and execution on Linux systems. +At present, GDR only supports the Baidu STD protocol. + +To use config_brpc: + +sh config_brpc.sh --with-rdma --with-gdr --headers="/usr/include" --libs="/usr/lib64 /usr/bin" +make +cd example/rdma_performance # Example program +make + +To use Bazel: + +# Server +bazel build --define=BRPC_WITH_RDMA=true --define=BRPC_WITH_GDR=true example:rdma_performance_server + +# Client +bazel build --define=BRPC_WITH_RDMA=true --define=BRPC_WITH_GDR=true example:rdma_performance_client + + +Basic Implementation: + +GDR is a special form of RDMA. Before using GDR, both RDMA and GDR must be globally initialized. + +GDR introduces a GPU memory pool, similar to the RDMA memory pool. Data in the GPU memory pool is also organized in blocks. + +When GDR is enabled, the framework initiates WQEs on GPU memory through DoPostRecvGDR. + +After receiving data, the header, meta, and body (excluding attachments) are copied back to host memory for processing. +Attachments remain in GPU memory, and users can call IOBuf::copy_from_gpu to copy attachments from the brpc framework layer to the application layer. + +Note: + +When using GDR, the environment variable MLX5_SCATTER_TO_CQE must be set to 0. + +Parameters + +Configurable parameters: + +gdr_block_size_kb: The block size (in KB) used when transferring data via GDR. Default is 512. + +max_gdr_regions: The maximum number of regions used by the GDR GPU memory pool. Each region is 1 GB. diff --git a/example/BUILD.bazel b/example/BUILD.bazel index 4ee7cb140f..098d283973 100644 --- a/example/BUILD.bazel +++ b/example/BUILD.bazel @@ -34,6 +34,9 @@ COPTS = [ }) + select({ "//bazel/config:brpc_with_rdma": ["-DBRPC_WITH_RDMA=1"], "//conditions:default": [""], +}) + select({ + "//bazel/config:brpc_with_gdr": ["-DBRPC_WITH_GDR=1"], + "//conditions:default": [""], }) brpc_proto_library( @@ -119,4 +122,4 @@ cc_binary( deps = [ "//:brpc", ], -) \ No newline at end of file +) diff --git a/example/rdma_performance/client.cpp b/example/rdma_performance/client.cpp index 2e8acc4051..40944d65a1 100644 --- a/example/rdma_performance/client.cpp +++ b/example/rdma_performance/client.cpp @@ -15,6 +15,10 @@ // specific language governing permissions and limitations // under the License. +#ifdef BRPC_WITH_GDR +#include +#include +#endif #include #include #include @@ -42,6 +46,7 @@ DEFINE_string(connection_type, "single", "Connection type of the channel"); DEFINE_string(protocol, "baidu_std", "Protocol type."); DEFINE_string(servers, "0.0.0.0:8002+0.0.0.0:8002", "IP Address of servers"); DEFINE_bool(use_rdma, true, "Use RDMA or not"); +DEFINE_bool(use_gdr, false, "Use GDR or not"); DEFINE_int32(rpc_timeout_ms, 2000, "RPC call timeout"); DEFINE_int32(test_seconds, 20, "Test running time"); DEFINE_int32(test_iterations, 0, "Test iterations"); @@ -84,16 +89,46 @@ class PerformanceTest { , _stop(false) { if (attachment_size > 0) { - _addr = malloc(attachment_size); - butil::fast_rand_bytes(_addr, attachment_size); - _attachment.append(_addr, attachment_size); +#ifdef BRPC_WITH_GDR + if (FLAGS_use_gdr) { + int gpu_id = 0; + cudaSetDevice(gpu_id); + cudaMalloc(&_addr, attachment_size); + auto pd = brpc::rdma::GetRdmaPd(); + mr = ibv_reg_mr(pd, _addr, attachment_size, + IBV_ACCESS_LOCAL_WRITE | + IBV_ACCESS_REMOTE_READ | + IBV_ACCESS_REMOTE_WRITE); + if (!mr) { + LOG(FATAL) << "Failed to register MR:" << strerror(errno) + << ", addr:" << _addr; + } + auto deleter = [](void* data) {}; + _attachment.append_user_data_with_meta(_addr, attachment_size, deleter, mr->lkey); + } + else +#endif + { + _addr = malloc(attachment_size); + butil::fast_rand_bytes(_addr, attachment_size); + _attachment.append(_addr, attachment_size); + } } _echo_attachment = echo_attachment; } ~PerformanceTest() { if (_addr) { - free(_addr); +#ifdef BRPC_WITH_GDR + if (FLAGS_use_gdr) { + ibv_dereg_mr(mr); + cudaFree(_addr); + } + else +#endif + { + free(_addr); + } } delete _channel; } @@ -103,6 +138,11 @@ class PerformanceTest { int Init() { brpc::ChannelOptions options; options.socket_mode = FLAGS_use_rdma? brpc::SOCKET_MODE_RDMA : brpc::SOCKET_MODE_TCP; +#ifdef BRPC_WITH_GDR + if (FLAGS_use_gdr) { + options.socket_mode = brpc::SOCKET_MODE_GDR; + } +#endif options.protocol = FLAGS_protocol; options.connection_type = FLAGS_connection_type; options.timeout_ms = FLAGS_rpc_timeout_ms; @@ -203,6 +243,9 @@ class PerformanceTest { } private: +#ifdef BRPC_WITH_GDR + ibv_mr* mr; +#endif void* _addr; brpc::Channel* _channel; uint64_t _start_time; @@ -223,6 +266,7 @@ void Test(int thread_num, int attachment_size) { << ", Depth: " << FLAGS_queue_depth << ", Attachment: " << attachment_size << "B" << ", RDMA: " << (FLAGS_use_rdma ? "yes" : "no") + << ", GDR: " << (FLAGS_use_gdr ? "yes" : "no") << ", Echo: " << (FLAGS_echo_attachment ? "yes]" : "no]") << std::endl; g_total_bytes.store(0, butil::memory_order_relaxed); @@ -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(); + brpc::rdma::GlobalGdrInitializeOrDie(); + } +#endif brpc::StartDummyServerAt(FLAGS_dummy_port); diff --git a/example/rdma_performance/server.cpp b/example/rdma_performance/server.cpp index 2e93e1eec7..eca2641513 100644 --- a/example/rdma_performance/server.cpp +++ b/example/rdma_performance/server.cpp @@ -28,6 +28,7 @@ DEFINE_int32(port, 8002, "TCP Port of this server"); DEFINE_bool(use_rdma, true, "Use RDMA or not"); +DEFINE_bool(use_gdr, false, "Use GDR or not"); butil::atomic g_last_time(0); @@ -77,6 +78,12 @@ int main(int argc, char* argv[]) { brpc::ServerOptions options; options.socket_mode = FLAGS_use_rdma? brpc::SOCKET_MODE_RDMA : brpc::SOCKET_MODE_TCP; +#ifdef BRPC_WITH_GDR + if (FLAGS_use_gdr) { + options.socket_mode = brpc::SOCKET_MODE_GDR; + } +#endif + if (server.Start(FLAGS_port, &options) != 0) { LOG(ERROR) << "Fail to start EchoServer"; return -1; diff --git a/src/brpc/acceptor.h b/src/brpc/acceptor.h index 77942beca2..f28d3f5bce 100644 --- a/src/brpc/acceptor.h +++ b/src/brpc/acceptor.h @@ -111,7 +111,7 @@ friend class Server; bool _force_ssl; std::shared_ptr _ssl_ctx; - // Choose to use a certain socket: 0 TCP, 1 RDMA + // Choose to use a certain socket: 0 TCP, 1 RDMA, 2 GDR SocketMode _socket_mode; // Acceptor belongs to this tag diff --git a/src/brpc/channel.cpp b/src/brpc/channel.cpp index 83fc37b077..c3387484d1 100644 --- a/src/brpc/channel.cpp +++ b/src/brpc/channel.cpp @@ -139,6 +139,8 @@ static ChannelSignature ComputeChannelSignature(const ChannelOptions& opt) { } if (opt.socket_mode == SOCKET_MODE_RDMA) { buf.append("|rdma"); + } else if (opt.socket_mode == SOCKET_MODE_GDR) { + buf.append("|gdr"); } butil::MurmurHash3_x64_128_Update(&mm_ctx, buf.data(), buf.size()); buf.clear(); diff --git a/src/brpc/gdr_transport.cpp b/src/brpc/gdr_transport.cpp new file mode 100644 index 0000000000..005b191380 --- /dev/null +++ b/src/brpc/gdr_transport.cpp @@ -0,0 +1,35 @@ +// 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 "brpc/gdr_transport.h" +#include "brpc/rdma/rdma_helper.h" + +namespace brpc { + +void GdrTransport::Init(Socket *socket, const SocketOptions &options) { + DoInit(socket, options, true); +} + +int GdrTransport::GdrContextInitOrDie() { + rdma::GlobalGdrInitializeOrDie(); + return 0; +} + +} // namespace brpc +#endif diff --git a/src/brpc/gdr_transport.h b/src/brpc/gdr_transport.h new file mode 100644 index 0000000000..0f41c0c4db --- /dev/null +++ b/src/brpc/gdr_transport.h @@ -0,0 +1,32 @@ +// 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. + +#ifndef BRPC_GDR_TRANSPORT_H +#define BRPC_GDR_TRANSPORT_H + +#if BRPC_WITH_GDR +#include "brpc/rdma_transport.h" + +namespace brpc { +class GdrTransport : public RdmaTransport { +public: + void Init(Socket* socket, const SocketOptions& options) override; + static int GdrContextInitOrDie(); +}; +} // namespace brpc +#endif // BRPC_WITH_GDR +#endif //BRPC_GDR_TRANSPORT_H diff --git a/src/brpc/policy/baidu_rpc_protocol.cpp b/src/brpc/policy/baidu_rpc_protocol.cpp index 49863b2c06..6c0f2e43b7 100644 --- a/src/brpc/policy/baidu_rpc_protocol.cpp +++ b/src/brpc/policy/baidu_rpc_protocol.cpp @@ -104,6 +104,12 @@ static void SerializeRpcHeaderAndMeta( ParseResult ParseRpcMessage(butil::IOBuf* source, Socket* socket, bool /*read_eof*/, const void*) { +#if BRPC_WITH_GDR + bool is_gpu_memory = source->is_gpu_memory(); + if (is_gpu_memory) { + return ParseRpcMessageGpu(source, socket, false /* not use */, nullptr /* not use */); + } +#endif // BRPC_WITH_GDR char header_buf[12]; const size_t n = source->copy_to(header_buf, sizeof(header_buf)); if (n >= 4) { @@ -810,7 +816,15 @@ void ProcessRpcRequest(InputMessageBase* msg_base) { butil::IOBuf req_buf; int body_without_attachment_size = req_size - meta.attachment_size(); - msg->payload.cutn(&req_buf, body_without_attachment_size); +#if BRPC_WITH_GDR + bool is_gpu_memory = msg->payload.is_gpu_memory(); + if (is_gpu_memory) { + FillReqBufGpu(&req_buf, msg.get(), body_without_attachment_size); + } else +#endif // BRPC_WITH_GDR + { + msg->payload.cutn(&req_buf, body_without_attachment_size); + } if (meta.attachment_size() > 0) { cntl->request_attachment().swap(msg->payload); } @@ -989,17 +1003,25 @@ void ProcessRpcResponse(InputMessageBase* msg_base) { butil::IOBuf res_buf; const int res_size = msg->payload.length(); butil::IOBuf* res_buf_ptr = &msg->payload; - if (meta.has_attachment_size()) { - if (meta.attachment_size() > res_size) { - cntl->SetFailed( - ERESPONSE, "attachment_size=%d is larger than response_size=%d", - meta.attachment_size(), res_size); - break; +#if BRPC_WITH_GDR + bool is_gpu_memory = msg->payload.is_gpu_memory(); + if (is_gpu_memory) { + FillResBufGpu(&res_buf, msg.get(), meta, &res_buf_ptr, cntl); + } else +#endif // BRPC_WITH_GDR + { + if (meta.has_attachment_size()) { + if (meta.attachment_size() > res_size) { + cntl->SetFailed( + ERESPONSE, "attachment_size=%d is larger than response_size=%d", + meta.attachment_size(), res_size); + break; + } + int body_without_attachment_size = res_size - meta.attachment_size(); + msg->payload.cutn(&res_buf, body_without_attachment_size); + res_buf_ptr = &res_buf; + cntl->response_attachment().swap(msg->payload); } - int body_without_attachment_size = res_size - meta.attachment_size(); - msg->payload.cutn(&res_buf, body_without_attachment_size); - res_buf_ptr = &res_buf; - cntl->response_attachment().swap(msg->payload); } ContentType content_type = meta.content_type(); diff --git a/src/brpc/policy/baidu_rpc_protocol.h b/src/brpc/policy/baidu_rpc_protocol.h index 77ecc780a2..6a3c379142 100644 --- a/src/brpc/policy/baidu_rpc_protocol.h +++ b/src/brpc/policy/baidu_rpc_protocol.h @@ -19,6 +19,8 @@ #ifndef BRPC_POLICY_BRPC_PROTOCOL_H #define BRPC_POLICY_BRPC_PROTOCOL_H +#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta +#include "brpc/policy/most_common_message.h" #include "brpc/protocol.h" namespace brpc { @@ -53,6 +55,17 @@ void PackRpcRequest(butil::IOBuf* buf, // Returns the `name' of the 'content_type'. const char* ContentTypeToCStr(ContentType content_type); +#if BRPC_WITH_GDR +// Parse binary format of baidu_std +ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket *socket, bool read_eof, + const void *arg); + +void FillReqBufGpu(butil::IOBuf* req_buf, MostCommonMessage* msg, int body_without_attachment_size); + +void FillResBufGpu(butil::IOBuf* res_buf, MostCommonMessage* msg, const RpcMeta& meta, + butil::IOBuf** res_buf_ptr, Controller* cntl); + +#endif } // namespace policy } // namespace brpc diff --git a/src/brpc/policy/baidu_rpc_protocol_gpu.cpp b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp new file mode 100644 index 0000000000..392cdbbff0 --- /dev/null +++ b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp @@ -0,0 +1,228 @@ +// 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 // MethodDescriptor +#include // Message +#include +#include +#include + +#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!!!"; + } + butil::gdr::BlockPoolAllocator* host_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_cpu_allocator(); + prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size); + if (prefetch_d2h_data == NULL) { + LOG(FATAL) << "alloc host data failed!!!"; + } + + // 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 bytes we real frefetch. We set n as the meta and n will be used in ProcessRpcRequest/ProcessRpcResponse. + // This is a trick, we should keep n in another better way. + 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!!!"; + } + 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) { + void* data = msg->meta.get_first_data_ptr(); + if (data == nullptr) { + LOG(FATAL) << "illegal data!!!"; + } + req_buf->append((char*)data + meta_size, body_without_attachment_size); + msg->payload.pop_front(body_without_attachment_size); + } else { + msg->payload.cutn_from_gpu(req_buf, body_without_attachment_size); + } +} + +void FillResBufGpu(butil::IOBuf* res_buf, MostCommonMessage* msg, const RpcMeta& meta, + butil::IOBuf** res_buf_ptr, Controller* cntl) { + const int res_size = msg->payload.length(); + 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!!!"; + } + if (meta.has_attachment_size()) { + if (meta.attachment_size() > res_size) { + cntl->SetFailed( + ERESPONSE, "attachment_size=%d is larger than response_size=%d", + meta.attachment_size(), res_size); + return; + } + int body_without_attachment_size = res_size - meta.attachment_size(); + + 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) { + void* data = msg->meta.get_first_data_ptr(); + if (data == nullptr) { + LOG(FATAL) << "illegal data!!!"; + } + res_buf->append((char*)data + meta_size, body_without_attachment_size); + msg->payload.pop_front(body_without_attachment_size); + } else { + msg->payload.cutn_from_gpu(res_buf, body_without_attachment_size); + } + *res_buf_ptr = res_buf; + cntl->response_attachment().swap(msg->payload); + } else { + int64_t real_prefetch_d2h_size = msg->meta.get_first_data_meta(); + if (header_size + meta_size + res_size <= real_prefetch_d2h_size) { + void* data = msg->meta.get_first_data_ptr(); + if (data == nullptr) { + LOG(FATAL) << "illegal data!!!"; + } + res_buf->append((char*)data + meta_size, res_size); + msg->payload.pop_front(res_size); + } else { + msg->payload.cutn_from_gpu(res_buf, res_size); + } + *res_buf_ptr = res_buf; + } +} + +} // namespace policy +} // namespace brpc + +#endif diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index 6c20ea3994..874c1c899c 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -19,6 +19,9 @@ #include #include "butil/fd_utility.h" +#if BRPC_WITH_GDR +#include "butil/gpu/gpu_block_pool.h" +#endif #include "butil/logging.h" // CHECK, LOG #include "butil/sys_byteorder.h" // HostToNet,NetToHost #include "bthread/bthread.h" @@ -59,7 +62,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."); DEFINE_bool(rdma_trace_verbose, false, "Print log message verbosely"); BRPC_VALIDATE_GFLAG(rdma_trace_verbose, brpc::PassValidate); DEFINE_bool(rdma_use_polling, false, "Use polling mode for RDMA."); @@ -76,6 +79,9 @@ extern const size_t RESERVED_WR_NUM = 3; // The local recv block size, set during GlobalInitialize. uint32_t g_rdma_recv_block_size = 0; +// The local recv block size for GDR, set during GlobalInitialize. +uint32_t g_gdr_recv_block_size = 0; + // static const uint32_t MAX_INLINE_DATA = 64; static const uint8_t MAX_HOP_LIMIT = 16; static const uint8_t TIMEOUT = 14; @@ -105,8 +111,9 @@ RdmaResource::~RdmaResource() { } } -RdmaEndpoint::RdmaEndpoint(Socket* s) +RdmaEndpoint::RdmaEndpoint(Socket* s, bool use_gdr) : _socket(s) + , _use_gdr(use_gdr) , _state(UNINIT) , _handshake_version(0) , _resource(NULL) @@ -911,8 +918,15 @@ ssize_t RdmaEndpoint::HandleCompletion(ibv_wc& wc) { case IBV_WC_RECV: { // recv completion // Please note that only the first wc.byte_len bytes is valid if (wc.byte_len > 0) { - if (wc.byte_len < (uint32_t)FLAGS_rdma_zerocopy_min_size) { - zerocopy = false; +#if BRPC_WITH_GDR + if (_use_gdr) { + zerocopy = true; + } else +#endif // BRPC_WITH_GDR + { + if (wc.byte_len < (uint32_t)FLAGS_rdma_zerocopy_min_size) { + zerocopy = false; + } } CHECK_NE(_state.load(butil::memory_order_relaxed), FALLBACK_TCP); if (zerocopy) { @@ -973,26 +987,76 @@ int RdmaEndpoint::DoPostRecv(void* block, size_t block_size) { return 0; } +int RdmaEndpoint::DoPostRecvGDR(void* block, size_t block_size, uint32_t lkey) { + ibv_recv_wr wr; + memset(&wr, 0, sizeof(wr)); + ibv_sge sge; + sge.addr = (uint64_t)block; + sge.length = block_size; + sge.lkey = lkey; + wr.num_sge = 1; + wr.sg_list = &sge; + //LOG(INFO) << "POST recv: addr=0x" << std::hex << sge.addr + // << std::dec << " length=0x" << sge.length + // << " lkey=0x" << sge.lkey; + //LOG(INFO) << block << " " << _device_allocator->get_lkey(); + ibv_recv_wr* bad = NULL; + int err = ibv_post_recv(_resource->qp, &wr, &bad); + if (err != 0) { + LOG(WARNING) << "Fail to ibv_post_recv: " << berror(err); + return -1; + } + return 0; +} + int RdmaEndpoint::PostRecv(uint32_t num, bool zerocopy) { // We do the post repeatedly from the _rbuf[_rq_received]. while (num > 0) { +#if BRPC_WITH_GDR + uint32_t lkey = 0; +#endif // if BRPC_WITH_GDR if (zerocopy) { _rbuf[_rq_received].clear(); - butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received], - g_rdma_recv_block_size + IOBUF_BLOCK_HEADER_LEN); - int size = 0; - if (!os.Next(&_rbuf_data[_rq_received], &size)) { - // Memory is not enough for preparing a block - PLOG(WARNING) << "Fail to allocate rbuf"; - return -1; - } else { - CHECK_EQ(static_cast(size), g_rdma_recv_block_size); + +#if BRPC_WITH_GDR + if (_use_gdr) { + butil::gdr::BlockPoolAllocator* device_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_gpu_allocator(); + void* device_ptr = device_allocator->AllocateRaw(g_gdr_recv_block_size); + auto deleter = [device_allocator](void* data) { device_allocator->DeallocateRaw(data); }; + lkey = device_allocator->get_lkey(device_ptr); + // we keep lkey into the meta, and this is a thick. we also keep prefetch d2h size in meta too. + _rbuf[_rq_received].append_user_data_with_meta(device_ptr, g_gdr_recv_block_size, deleter, lkey); + _rbuf_data[_rq_received] = device_ptr; + } else +#endif // if BRPC_WITH_GDR + { + butil::IOBufAsZeroCopyOutputStream os(&_rbuf[_rq_received], + g_rdma_recv_block_size + IOBUF_BLOCK_HEADER_LEN); + int size = 0; + if (!os.Next(&_rbuf_data[_rq_received], &size)) { + // Memory is not enough for preparing a block + PLOG(WARNING) << "Fail to allocate rbuf"; + return -1; + } else { + CHECK_EQ(static_cast(size), g_rdma_recv_block_size); + } } } - if (DoPostRecv(_rbuf_data[_rq_received], g_rdma_recv_block_size) < 0) { - _rbuf[_rq_received].clear(); - return -1; +#if BRPC_WITH_GDR + if (_use_gdr) { + if (DoPostRecvGDR(_rbuf_data[_rq_received], g_gdr_recv_block_size, lkey) < 0) { + _rbuf[_rq_received].clear(); + return -1; + } + } else +#endif // if BRPC_WITH_GDR + { + if (DoPostRecv(_rbuf_data[_rq_received], g_rdma_recv_block_size) < 0) { + _rbuf[_rq_received].clear(); + return -1; + } } + --num; ++_rq_received; if (_rq_received == _rq_size) { @@ -1589,6 +1653,14 @@ void RdmaEndpoint::DebugInfo(std::ostream& os, butil::StringPiece connector) con << connector << "rdma_unsignaled_sq_wr=" << _sq_unsignaled; } +int RdmaEndpoint::GlobalGdrInitialize() { +#if BRPC_WITH_GDR + g_gdr_recv_block_size = butil::gdr::GetGdrBlockSize() - IOBUF_BLOCK_HEADER_LEN; + LOG(INFO) << "g_gdr_recv_block_size: " << g_gdr_recv_block_size; +#endif // BRPC_WITH_GDR + return 0; +} + int RdmaEndpoint::GlobalInitialize() { g_rdma_recv_block_size = GetRdmaBlockSize() - IOBUF_BLOCK_HEADER_LEN; if (g_rdma_recv_block_size <= 0) { diff --git a/src/brpc/rdma/rdma_endpoint.h b/src/brpc/rdma/rdma_endpoint.h index 36e22ad28d..d58936b731 100644 --- a/src/brpc/rdma/rdma_endpoint.h +++ b/src/brpc/rdma/rdma_endpoint.h @@ -33,7 +33,6 @@ #include "brpc/socket.h" #include "brpc/rdma/rdma_handshake_server.h" - namespace brpc { class Socket; namespace rdma { @@ -105,15 +104,21 @@ friend void v3_wire::FillLocalRdmaHello(const RdmaEndpoint*, RdmaHello*); friend int v3_wire::ReadAndParseV3Hello(RdmaEndpoint*, RdmaHello*); friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); public: - explicit RdmaEndpoint(Socket* s); + explicit RdmaEndpoint(Socket* s, bool use_gdr = false); ~RdmaEndpoint() override; - // Global initialization + // Global Rdma initialization // Return 0 if success, -1 if failed and errno set static int GlobalInitialize(); + // Global Gdr initialization + // Return 0 if success, -1 if failed and errno set + static int GlobalGdrInitialize(); + static void GlobalRelease(); + bool use_gdr() { return _use_gdr; } + // Reset the endpoint (for next use) void Reset(); @@ -208,6 +213,16 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // -1: failed, errno set int DoPostRecv(void* block, size_t block_size); + // Post a WR pointing to the gpu block to the local Recv Queue + // Arguments: + // block: the gpu addr to receive data (ibv_sge.addr) + // block_size: the maximum length can be received (ibv_sge.length) + // lkey: the lkey of block + // Return: + // 0: success + // -1: failed, errno set + int DoPostRecvGDR(void* block, size_t block_size, uint32_t lkey); + // Read at most len bytes from fd in _socket to data // wait for _read_butex if encounter EAGAIN // return -1 if encounter other errno (including EOF) @@ -281,6 +296,9 @@ friend int v3_wire::WriteV3Hello(RdmaEndpoint*, const RdmaHello&); // QP reached RTS (filled in BringUpQp). butil::optional _outgoing_ece; + // whether open gpu direct rdma + bool _use_gdr; + // rdma resource RdmaResource* _resource; diff --git a/src/brpc/rdma/rdma_helper.cpp b/src/brpc/rdma/rdma_helper.cpp index b0e13ad72c..d7a02e6625 100644 --- a/src/brpc/rdma/rdma_helper.cpp +++ b/src/brpc/rdma/rdma_helper.cpp @@ -25,6 +25,9 @@ #include "butil/containers/flat_map.h" // butil::FlatMap #include "butil/fd_guard.h" #include "butil/fd_utility.h" // butil::make_non_blocking +#if BRPC_WITH_GDR +#include "butil/gpu/gpu_block_pool.h" +#endif #include "butil/logging.h" #include "brpc/socket.h" #include "brpc/rdma/block_pool.h" @@ -86,6 +89,8 @@ static uint16_t g_lid; static int g_max_sge = 0; static uint8_t g_port_num = 1; +static int g_gpu_index = 0; + static int g_comp_vector_index = 0; butil::atomic g_rdma_available(false); @@ -95,7 +100,7 @@ DEFINE_string(rdma_device, "", "The name of the HCA device used " "(Empty means using the first active device)"); DEFINE_int32(rdma_port, 1, "The port number to use. For RoCE, it is always 1."); DEFINE_int32(rdma_gid_index, -1, "The GID index to use. -1 means using the last one."); - +DEFINE_int32(gpu_index, 0, "The GPU device index to use. In GDR, we suggest to use the GPU that is connected to the same PCIe switch with rdma devices"); // static const size_t SYSFS_SIZE = 4096; static ibv_device** g_devices = NULL; static ibv_context* g_context = NULL; @@ -601,7 +606,26 @@ static void GlobalRdmaInitializeOrDieImpl() { g_rdma_available.store(true, butil::memory_order_relaxed); } +static void GlobalGdrInitializeOrDieImpl() { +#if BRPC_WITH_GDR + g_gpu_index = FLAGS_gpu_index; + + if (!butil::gdr::InitGPUBlockPool(g_gpu_index, GetRdmaPd())) { + PLOG(ERROR) << "Fail to initialize RDMA GPU memory pool"; + ExitWithError(); + } + if (RdmaEndpoint::GlobalGdrInitialize() < 0) { + LOG(ERROR) << "gdr_block_size_kb incorrect " + << "(must be larger than 0)"; + ExitWithError(); + } + +#endif // if BRPC_WITH_GDR + +} + static pthread_once_t initialize_rdma_once = PTHREAD_ONCE_INIT; +static pthread_once_t initialize_gdr_once = PTHREAD_ONCE_INIT; void GlobalRdmaInitializeOrDie() { if (pthread_once(&initialize_rdma_once, @@ -611,6 +635,14 @@ void GlobalRdmaInitializeOrDie() { } } +void GlobalGdrInitializeOrDie() { + if (pthread_once(&initialize_gdr_once, + GlobalGdrInitializeOrDieImpl) != 0) { + LOG(FATAL) << "Fail to pthread_once GlobalGdrInitializeOrDie"; + exit(1); + } +} + uint32_t RegisterMemoryForRdma(void* buf, size_t len) { ibv_mr* mr = IbvRegMr(g_pd, buf, len, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_RELAXED_ORDERING); if (!mr) { @@ -704,6 +736,10 @@ uint8_t GetRdmaPortNum() { return g_port_num; } +int GetGPUIndex() { + return g_gpu_index; +} + bool IsRdmaAvailable() { return g_rdma_available.load(butil::memory_order_acquire); } diff --git a/src/brpc/rdma/rdma_helper.h b/src/brpc/rdma/rdma_helper.h index 052763325b..06cbb1f5c2 100644 --- a/src/brpc/rdma/rdma_helper.h +++ b/src/brpc/rdma/rdma_helper.h @@ -33,6 +33,10 @@ namespace rdma { // Exit if failed void GlobalRdmaInitializeOrDie(); +// Initialize GDR environment +// Exit if failed +void GlobalGdrInitializeOrDie(); + // Initialize RDMA polling mode with tag bool InitPollingModeWithTag(bthread_tag_t tag, std::function callback = nullptr, @@ -74,6 +78,9 @@ int GetRdmaCompVector(); // Return current port number used uint8_t GetRdmaPortNum(); +// Get GPU index used +int GetGPUIndex(); + // Get max_sge supported by the device int GetRdmaMaxSge(); diff --git a/src/brpc/rdma_transport.cpp b/src/brpc/rdma_transport.cpp index b172e47f34..367c086bb9 100644 --- a/src/brpc/rdma_transport.cpp +++ b/src/brpc/rdma_transport.cpp @@ -30,10 +30,13 @@ 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 == NULL); - if (options.socket_mode == SOCKET_MODE_RDMA) { - _rdma_ep = new(std::nothrow)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"; @@ -64,6 +67,10 @@ void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { _tcp_transport->Init(socket, options); } +void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { + DoInit(socket, options, false); +} + void RdmaTransport::Release() { if (_rdma_ep) { delete _rdma_ep; diff --git a/src/brpc/rdma_transport.h b/src/brpc/rdma_transport.h index 0ac8a25324..932362c7c5 100644 --- a/src/brpc/rdma_transport.h +++ b/src/brpc/rdma_transport.h @@ -31,6 +31,7 @@ friend class rdma::RdmaConnect; friend class rdma::RdmaHandshakeServerV2; friend class rdma::RdmaHandshakeServerV3; public: + void DoInit(Socket* socket, const SocketOptions& options, bool use_gdr); void Init(Socket* socket, const SocketOptions& options) override; void Release() override; int Reset(int32_t expected_nref) override; @@ -64,4 +65,4 @@ friend class rdma::RdmaHandshakeServerV3; }; } // namespace brpc #endif // BRPC_WITH_RDMA -#endif //BRPC_RDMA_TRANSPORT_H \ No newline at end of file +#endif //BRPC_RDMA_TRANSPORT_H diff --git a/src/brpc/socket_mode.h b/src/brpc/socket_mode.h index b4ac7dfbca..1d6d3dcd3d 100644 --- a/src/brpc/socket_mode.h +++ b/src/brpc/socket_mode.h @@ -21,7 +21,8 @@ namespace brpc { enum SocketMode { SOCKET_MODE_TCP = 0, SOCKET_MODE_RDMA = 1, - SOCKET_MODE_UBRING = 2 + SOCKET_MODE_GDR = 2, + SOCKET_MODE_UBRING = 3 }; } // namespace brpc -#endif //BRPC_SOCKET_MODE_H \ No newline at end of file +#endif //BRPC_SOCKET_MODE_H diff --git a/src/brpc/transport_factory.cpp b/src/brpc/transport_factory.cpp index 36fdaaed05..1b0a9c453e 100644 --- a/src/brpc/transport_factory.cpp +++ b/src/brpc/transport_factory.cpp @@ -18,6 +18,7 @@ #include "brpc/transport_factory.h" #include "brpc/tcp_transport.h" #include "brpc/rdma_transport.h" +#include "brpc/gdr_transport.h" #include "brpc/ubshm_transport.h" namespace brpc { @@ -30,6 +31,15 @@ int TransportFactory::ContextInitOrDie(SocketMode mode, bool serverOrNot, const return RdmaTransport::ContextInitOrDie(serverOrNot, _options); } #endif +#if BRPC_WITH_GDR + else if (mode == SOCKET_MODE_GDR) { + // gdr is a special case of rdma, so we should init rdma first; + if (RdmaTransport::ContextInitOrDie(serverOrNot, _options) < 0) { + return -1; + } + return GdrTransport::GdrContextInitOrDie(); + } +#endif #if BRPC_WITH_UBRING else if (mode == SOCKET_MODE_UBRING) { return UBShmTransport::ContextInitOrDie(serverOrNot, _options); @@ -50,6 +60,11 @@ std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { return std::unique_ptr(new RdmaTransport()); } #endif +#if BRPC_WITH_GDR + else if (mode == SOCKET_MODE_GDR) { + return std::unique_ptr(new GdrTransport()); + } +#endif #if BRPC_WITH_UBRING else if (mode == SOCKET_MODE_UBRING) { return std::unique_ptr(new UBShmTransport()); @@ -60,4 +75,4 @@ std::unique_ptr TransportFactory::CreateTransport(SocketMode mode) { return nullptr; } } -} // namespace brpc \ No newline at end of file +} // namespace brpc diff --git a/src/butil/gpu/gpu_block_pool.cpp b/src/butil/gpu/gpu_block_pool.cpp new file mode 100644 index 0000000000..86a307ba06 --- /dev/null +++ b/src/butil/gpu/gpu_block_pool.cpp @@ -0,0 +1,475 @@ +// 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 +#include +#include +#include +#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 int original_device = -1; + static bool first_call = true; + + int current_device; + cudaGetDevice(¤t_device); + + if (first_call) { + original_device = current_device; + first_call = false; + return true; + } + + 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::instance_ = nullptr; + +BlockPoolAllocators* BlockPoolAllocators::singleton() { + static std::mutex mutex; + if (instance_ == nullptr) { + std::lock_guard l(mutex); + if (instance_ == nullptr) { + instance_ = new BlockPoolAllocators(); + std::atomic_thread_fence(std::memory_order_release); + } + } + std::atomic_thread_fence(std::memory_order_acquire); + return instance_; +} + +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::singleton()->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 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 lock(mu_); + if (objects_.size() < kMaxObjects) { + objects_.push_back(obj); + return; + } + } + delete obj; + } + + private: + static const int kMaxObjects = 100000; + + std::mutex mu_; + std::vector 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; + + extendRegion(); +} + +BlockPoolAllocator::~BlockPoolAllocator() { +#ifdef DEBUG + printStatistics(); +#endif + + for (int i = 0; i < FLAGS_max_gdr_regions; i++) { + Region* r = &g_regions[i]; + if (!r->mr) { + return; + } + + 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(r->start))); + } else { + CHECK_CUDA(cudaFreeHost(reinterpret_cast(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) { + LOG(FATAL) << "try to alloc " << num_bytes << " bytes, its bigger than block_size " << BLOCK_SIZE; + } + + auto startTime = std::chrono::high_resolution_clock::now(); + + std::lock_guard lock(poolMutex); + + if (!freeList) { + extendRegion(); + } + + 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(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 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(totalAllocated - totalDeallocated) / + (g_region_num * (REGION_SIZE / BLOCK_SIZE)) * 100) + << "%"; +} + +void BlockPoolAllocator::extendRegion() { + if (g_region_num == FLAGS_max_gdr_regions) { + LOG(FATAL) << "Gdr Memory pool reaches max regions"; + return ; + } + + 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); + } + + 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; + auto mr = ibv_reg_mr(pd, aligned_ptr, aligned_bytes, + IBV_ACCESS_LOCAL_WRITE | + IBV_ACCESS_REMOTE_READ | + IBV_ACCESS_REMOTE_WRITE | + IBV_ACCESS_RELAXED_ORDERING); + + if (!mr) { + LOG(FATAL) << "Failed to register MR: " << strerror(errno) + << ", pd " << pd << ", aligned_ptr:" << aligned_ptr; + } 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(static_cast(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(endTime - startTime); + + LOG(INFO) << "Extended region #" << g_region_num << ": " << blockCount + << " blocks (" << (REGION_SIZE / (1024 * 1024)) << " MB)" << ", on_gpu " << on_gpu + << ", cost " << duration.count() << " ns"; +} + +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& src_list, + std::vector& 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 stream_lb_lock(d2d_lb_lock_); + d2d_cnt_.fetch_add(1); + stream_idx = d2d_cnt_ % kMaxConcurrent; + } + std::lock_guard 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(dst) + offset, src_list[i], length_list[i], + cudaMemcpyDeviceToDevice, d2d_streams_[stream_idx])); + offset += length_list[i]; + } + CHECK_CUDA(cudaStreamSynchronize(d2d_streams_[stream_idx])); +} + +void GPUStreamPool::fast_d2h(std::vector& src_list, + std::vector& length_list, + void* dst) { + if (!verify_same_context()) { + LOG(FATAL) << "Context mismatch!"; + return; + } + 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 stream_lb_lock(d2h_lb_lock_); + d2h_cnt_.fetch_add(1); + stream_idx = d2h_cnt_ % kMaxConcurrent; + } + std::lock_guard stream_lock(d2h_locks_[stream_idx]); + CHECK_CUDA(cudaStreamSynchronize(d2h_streams_[stream_idx])); + for (int i = 0; i < segs; i++) { + if (length_list[i] == 0) { + continue; + } + CHECK_CUDA(cudaMemcpyAsync(static_cast(dst) + offset, src_list[i], length_list[i], + cudaMemcpyDeviceToHost, d2h_streams_[stream_idx])); + offset += length_list[i]; + } + CHECK_CUDA(cudaStreamSynchronize(d2h_streams_[stream_idx])); +} + +} +} + +#endif // BRPC_WITH_GDR diff --git a/src/butil/gpu/gpu_block_pool.h b/src/butil/gpu/gpu_block_pool.h new file mode 100644 index 0000000000..655c487a92 --- /dev/null +++ b/src/butil/gpu/gpu_block_pool.h @@ -0,0 +1,175 @@ +// 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. +#ifndef BUTIL_GPU_GPU_BLOCK_POOL_H +#define BUTIL_GPU_GPU_BLOCK_POOL_H + +#if BRPC_WITH_GDR + +#include +#include +#include +#include +#include +#include +#include +#include "butil/containers/hash_tables.h" +#include "butil/logging.h" +#include +#include "cuda.h" + +namespace butil { +namespace gdr { + +size_t GetGdrBlockSize(); +void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size); +void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size); + +bool InitGPUBlockPool(int gpu_id, ibv_pd* pd); + +struct Region { + Region() { start = 0; aligned_start = 0;} + uintptr_t start; + uintptr_t aligned_start; + + size_t size; + size_t aligned_size; + size_t blockCount; + struct ibv_mr *mr {nullptr}; +}; + +struct BlockHeader { + BlockHeader() { addr = nullptr; next = nullptr;} + void Reset() { addr = nullptr; next = nullptr; } + void* addr; + BlockHeader* next; +}; + +class BlockPoolAllocator { + private: + int gpu_id; + bool on_gpu; + ibv_pd* pd {nullptr}; + + const size_t BLOCK_SIZE; + const size_t REGION_SIZE; + + BlockHeader* freeList; + int g_region_num {0}; + std::vector g_regions; + std::mutex poolMutex; + + // stat + size_t totalAllocated; + size_t totalDeallocated; + size_t peakUsage; + + public: + explicit BlockPoolAllocator(int gpuId, + bool onGpu, ibv_pd* ibvPd, + size_t blockSize, size_t regionSize); + + ~BlockPoolAllocator(); + + void* AllocateRaw(size_t num_bytes); + + void DeallocateRaw(void* ptr); + + void printStatistics() const; + + int64_t getCurrentUsage() const { + return totalAllocated - totalDeallocated; + } + + int64_t getTotalMemory() const { + return g_region_num * REGION_SIZE; + } + + int64_t get_block_size() const { + return BLOCK_SIZE; + } + + Region* GetRegion(const void* buf); + + uint32_t get_lkey(const void* buf); + + private: + void extendRegion(); +}; + +class GPUStreamPool { +public: + explicit GPUStreamPool(int gpu_id); + + ~GPUStreamPool(); + + GPUStreamPool(const GPUStreamPool&) = delete; + GPUStreamPool& operator=(const GPUStreamPool&) = delete; + + void fast_d2h(std::vector& src_list, std::vector& length_list, void* dst); + + void fast_d2d(std::vector& src_list, std::vector& length_list, void* dst); + + static constexpr int kMaxConcurrent = 32; +private: + int gpu_id_ {-1}; + std::atomic d2h_cnt_ {0}; + std::atomic d2d_cnt_ {0}; + std::mutex d2h_locks_[kMaxConcurrent]; + std::mutex d2d_locks_[kMaxConcurrent]; + std::mutex d2h_lb_lock_; + std::mutex d2d_lb_lock_; + std::vector d2h_streams_; + std::vector d2d_streams_; +}; + +class BlockPoolAllocators { +public: + static BlockPoolAllocators* singleton(); + BlockPoolAllocators() {} + virtual ~BlockPoolAllocators() { + CHECK_EQ(this, instance_); + instance_ = nullptr; + } + + void init(int gpu_id, ibv_pd* pd); + + BlockPoolAllocator* get_gpu_allocator() { + return gpu_mem_alloc; + } + + BlockPoolAllocator* get_cpu_allocator() { + return cpu_mem_alloc; + } + + GPUStreamPool* get_gpu_stream_pool() { + return gpu_stream_pool; + } + +public: + static BlockPoolAllocators* instance_; + +private: + BlockPoolAllocator* gpu_mem_alloc {nullptr}; + BlockPoolAllocator* cpu_mem_alloc {nullptr}; + GPUStreamPool* gpu_stream_pool {nullptr}; +}; +} +} + +#endif // BRPC_WITH_GDR + +#endif diff --git a/src/butil/iobuf.cpp b/src/butil/iobuf.cpp index be751d318c..075e6364c1 100644 --- a/src/butil/iobuf.cpp +++ b/src/butil/iobuf.cpp @@ -42,6 +42,9 @@ #include "butil/fd_guard.h" // butil::fd_guard #include "butil/iobuf.h" #include "butil/iobuf_profiler.h" +#ifdef BRPC_WITH_GDR +#include "butil/gpu/gpu_block_pool.h" +#endif namespace butil { static size_t default_block_size = 8192; @@ -768,6 +771,46 @@ size_t IOBuf::cutn(IOBuf* out, size_t n) { return saved_n; } +#if BRPC_WITH_GDR +size_t IOBuf::cutn_from_gpu(IOBuf* out, size_t n) { + if (n == 0) { + return 0; + } + + butil::gdr::BlockPoolAllocator* host_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_cpu_allocator(); + bool alloc_from_host_alloc = (n <= host_allocator->get_block_size()); + void* mem = NULL; + if (alloc_from_host_alloc) { + mem = host_allocator->AllocateRaw(n); + } else { + mem = malloc(n); + } + + if (mem == NULL) { + return 0; + } + size_t saved_n = copy_from_gpu(mem, n, 0, false); + if (saved_n > 0) { + if (alloc_from_host_alloc) { + auto deleter = [host_allocator](void* data) { host_allocator->DeallocateRaw(data); }; + out->append_user_data(mem, saved_n, deleter); + } else { + auto deleter = [](void* data) { free(data); }; + out->append_user_data(mem, saved_n, deleter); + } + pop_front(saved_n); + } else { + if (alloc_from_host_alloc) { + host_allocator->DeallocateRaw(mem); + } else { + free(mem); + } + } + + return saved_n; +} +#endif // BRPC_WITH_GDR + size_t IOBuf::cutn(void* out, size_t n) { const size_t len = length(); if (n > len) { @@ -1201,6 +1244,15 @@ uint64_t IOBuf::get_first_data_meta() { return r.block->u.data_meta; } +void* IOBuf::get_first_data_ptr() { + if (_ref_num() == 0) { + return 0; + } + IOBuf::BlockRef const& r = _ref_at(0); + return r.block->data; +} + + int IOBuf::resize(size_t n, char c) { const size_t saved_len = length(); if (n < saved_len) { @@ -1363,6 +1415,46 @@ size_t IOBuf::copy_to(void* d, size_t n, size_t pos) const { return n - m; } +#if BRPC_WITH_GDR +size_t IOBuf::copy_from_gpu(void* d, size_t n, size_t pos, bool to_gpu) const { + if (n == 0) { + return 0; + } + const size_t nref = _ref_num(); + // Skip `pos' bytes. `offset' is the starting position in starting BlockRef. + size_t offset = pos; + size_t i = 0; + for (; offset != 0 && i < nref; ++i) { + IOBuf::BlockRef const& r = _ref_at(i); + if (offset < (size_t)r.length) { + break; + } + offset -= r.length; + } + + butil::gdr::GPUStreamPool* gpu_stream_pool = butil::gdr::BlockPoolAllocators::singleton()->get_gpu_stream_pool(); + size_t m = n; + std::vector src_list; + std::vector length_list; + for (; m != 0 && i < nref; ++i) { + IOBuf::BlockRef const& r = _ref_at(i); + const size_t nc = std::min(m, (size_t)r.length - offset); + void* gpu_src = r.block->data + r.offset + offset; + src_list.push_back(gpu_src); + length_list.push_back(nc); + offset = 0; + m -= nc; + } + if (to_gpu) { + gpu_stream_pool->fast_d2d(src_list, length_list, d); + } else { + gpu_stream_pool->fast_d2h(src_list, length_list, d); + } + // If nref == 0, here returns 0 correctly + return n - m; +} +#endif // BRPC_WITH_GDR + size_t IOBuf::copy_to(std::string* s, size_t n, size_t pos) const { const size_t len = length(); if (len <= pos) { @@ -1509,6 +1601,16 @@ bool IOBuf::equals(const butil::IOBuf& other) const { return true; } +#if BRPC_WITH_GDR +// when IOBuf is used for send, data_meta is set by user; +// when IOBf is used for recv and gdr is open, data_meta is set by brpc +// and it is lkey. +bool IOBuf::is_gpu_memory() { + uint64_t data_meta = get_first_data_meta(); + return (data_meta > 0 && data_meta <= UINT_MAX); +} +#endif + ////////////////////////////// IOPortal ////////////////// IOPortal::~IOPortal() { return_cached_blocks(); } diff --git a/src/butil/iobuf.h b/src/butil/iobuf.h index b92a2e3da3..c371b67a10 100644 --- a/src/butil/iobuf.h +++ b/src/butil/iobuf.h @@ -146,6 +146,13 @@ friend class SingleIOBuf; size_t cutn(IOBuf* out, size_t n); size_t cutn(void* out, size_t n); size_t cutn(std::string* out, size_t n); + +#if BRPC_WITH_GDR + size_t cutn_from_gpu(IOBuf* out, size_t n); + size_t copy_from_gpu(void* d, size_t n, size_t pos = 0, bool to_gpu = false) const; + bool is_gpu_memory(); +#endif // BRPC_WITH_GDR + // Cut off 1 byte from the front side and set to *c // Return true on cut, false otherwise. bool cut1(void* c); @@ -265,6 +272,9 @@ friend class SingleIOBuf; // 0 means the meta is invalid. uint64_t get_first_data_meta(); + // Get the data addr of the first byte in this IOBuf. + void* get_first_data_ptr(); + // Resizes the buf to a length of n characters. // If n is smaller than the current length, all bytes after n will be // truncated. From 051d0b2199fa749001e20d7eb36a4143453d02f4 Mon Sep 17 00:00:00 2001 From: randomkang <550941794@qq.com> Date: Tue, 4 Aug 2026 01:10:54 +0800 Subject: [PATCH 2/4] review by deepseek --- src/brpc/policy/baidu_rpc_protocol_gpu.cpp | 2 ++ src/brpc/socket_mode.h | 4 ++-- src/butil/gpu/gpu_block_pool.cpp | 27 +++++++++++----------- src/butil/gpu/gpu_block_pool.h | 7 +++--- 4 files changed, 21 insertions(+), 19 deletions(-) diff --git a/src/brpc/policy/baidu_rpc_protocol_gpu.cpp b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp index 392cdbbff0..500875eb70 100644 --- a/src/brpc/policy/baidu_rpc_protocol_gpu.cpp +++ b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp @@ -83,11 +83,13 @@ ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket, 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::singleton()->get_cpu_allocator(); prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size); if (prefetch_d2h_data == NULL) { LOG(FATAL) << "alloc host data failed!!!"; + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); } // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size; diff --git a/src/brpc/socket_mode.h b/src/brpc/socket_mode.h index 1d6d3dcd3d..3b6177cc10 100644 --- a/src/brpc/socket_mode.h +++ b/src/brpc/socket_mode.h @@ -21,8 +21,8 @@ namespace brpc { enum SocketMode { SOCKET_MODE_TCP = 0, SOCKET_MODE_RDMA = 1, - SOCKET_MODE_GDR = 2, - SOCKET_MODE_UBRING = 3 + SOCKET_MODE_UBRING = 2, + SOCKET_MODE_GDR = 3 }; } // namespace brpc #endif //BRPC_SOCKET_MODE_H diff --git a/src/butil/gpu/gpu_block_pool.cpp b/src/butil/gpu/gpu_block_pool.cpp index 86a307ba06..2efcefa513 100644 --- a/src/butil/gpu/gpu_block_pool.cpp +++ b/src/butil/gpu/gpu_block_pool.cpp @@ -43,17 +43,15 @@ size_t GetGdrBlockSize() { } bool verify_same_context() { + static std::once_flag init_flag; static int original_device = -1; - static bool first_call = true; int current_device; cudaGetDevice(¤t_device); - if (first_call) { + std::call_once(init_flag, [&]() { original_device = current_device; - first_call = false; - return true; - } + }); return (current_device == original_device); } @@ -84,19 +82,20 @@ void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) { } -BlockPoolAllocators* BlockPoolAllocators::instance_ = nullptr; +std::atomic BlockPoolAllocators::instance_{nullptr}; BlockPoolAllocators* BlockPoolAllocators::singleton() { - static std::mutex mutex; - if (instance_ == nullptr) { + BlockPoolAllocators* tmp = instance_.load(std::memory_order_acquire); + if (tmp == nullptr) { + static std::mutex mutex; std::lock_guard l(mutex); - if (instance_ == nullptr) { - instance_ = new BlockPoolAllocators(); - std::atomic_thread_fence(std::memory_order_release); + tmp = instance_.load(std::memory_order_relaxed); + if (tmp == nullptr) { + tmp = new BlockPoolAllocators(); + instance_.store(tmp, std::memory_order_release); } } - std::atomic_thread_fence(std::memory_order_acquire); - return instance_; + return tmp; } void BlockPoolAllocators::init(int gpu_id, ibv_pd* pd) { @@ -191,7 +190,7 @@ BlockPoolAllocator::~BlockPoolAllocator() { for (int i = 0; i < FLAGS_max_gdr_regions; i++) { Region* r = &g_regions[i]; if (!r->mr) { - return; + continue; } LOG(INFO) << "try to free " << r->size << " bytes from gpu " << gpu_id << ", on_gpu " << on_gpu; diff --git a/src/butil/gpu/gpu_block_pool.h b/src/butil/gpu/gpu_block_pool.h index 655c487a92..73b557f3d0 100644 --- a/src/butil/gpu/gpu_block_pool.h +++ b/src/butil/gpu/gpu_block_pool.h @@ -26,6 +26,7 @@ #include #include #include +#include #include "butil/containers/hash_tables.h" #include "butil/logging.h" #include @@ -141,8 +142,8 @@ class BlockPoolAllocators { static BlockPoolAllocators* singleton(); BlockPoolAllocators() {} virtual ~BlockPoolAllocators() { - CHECK_EQ(this, instance_); - instance_ = nullptr; + CHECK_EQ(this, instance_.load(std::memory_order_relaxed)); + instance_.store(nullptr, std::memory_order_relaxed); } void init(int gpu_id, ibv_pd* pd); @@ -160,7 +161,7 @@ class BlockPoolAllocators { } public: - static BlockPoolAllocators* instance_; + static std::atomic instance_; private: BlockPoolAllocator* gpu_mem_alloc {nullptr}; From 90eef405a1de170de8a8976c3b6a1b3234b0deb0 Mon Sep 17 00:00:00 2001 From: randomkang <550941794@qq.com> Date: Tue, 4 Aug 2026 01:51:01 +0800 Subject: [PATCH 3/4] review by glm5.2 --- BUILD.bazel | 7 ++- CMakeLists.txt | 47 +++++++++++++++++++ config.h.in | 6 +++ config_brpc.sh | 32 ++++++++++++- example/rdma_performance/CMakeLists.txt | 35 ++++++++++++++ example/rdma_performance/Makefile | 3 ++ src/brpc/policy/baidu_rpc_protocol_gpu.cpp | 13 ++++-- src/brpc/rdma/rdma_endpoint.cpp | 11 ++++- src/butil/gpu/gpu_block_pool.cpp | 53 ++++++++++++++++++---- src/butil/gpu/gpu_block_pool.h | 5 +- src/butil/iobuf.cpp | 46 +++++++++++++++++-- src/butil/iobuf.h | 12 +++++ src/butil/iobuf_inl.h | 4 ++ 13 files changed, 249 insertions(+), 25 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index ef4bf402de..b03177414a 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -52,7 +52,9 @@ DEFINES = [ "//bazel/config:brpc_with_rdma": ["BRPC_WITH_RDMA=1"], "//conditions:default": [], }) + select({ - "//bazel/config:brpc_with_gdr": ["-DBRPC_WITH_GDR=1"], + # NOTE: `defines` is auto-prefixed with -D by Bazel, so do NOT add -D here. + # Compare with the rdma/glog/thrift entries above which also omit the prefix. + "//bazel/config:brpc_with_gdr": ["BRPC_WITH_GDR=1"], "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_debug_bthread_sche_safety": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1"], @@ -99,7 +101,8 @@ LINKOPTS = [ "//conditions:default": [], }) + select({ "//bazel/config:brpc_with_gdr": [ - "-lcuda -lcudart", + "-lcuda", + "-lcudart", ], "//conditions:default": [], }) + select({ diff --git a/CMakeLists.txt b/CMakeLists.txt index 915b7d2977..1d5bd88253 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,6 +27,7 @@ option(WITH_THRIFT "With thrift framed protocol supported" OFF) option(WITH_BTHREAD_TRACER "With bthread tracer supported" OFF) option(WITH_SNAPPY "With snappy" OFF) option(WITH_RDMA "With RDMA" OFF) +option(WITH_GDR "With GPU Direct RDMA (requires WITH_RDMA and CUDA)" OFF) option(WITH_UBRING "With UB" OFF) option(WITH_DEBUG_BTHREAD_SCHE_SAFETY "With debugging bthread sche safety" OFF) option(WITH_DEBUG_LOCK "With debugging lock" OFF) @@ -114,6 +115,17 @@ if(WITH_RDMA) set(WITH_RDMA_VAL "1") endif() +set(WITH_GDR_VAL "0") +if(WITH_GDR) + # GDR is a special mode of RDMA; it implies RDMA. + if(NOT WITH_RDMA) + message(STATUS "WITH_GDR=ON implies WITH_RDMA=ON") + set(WITH_RDMA ON) + set(WITH_RDMA_VAL "1") + endif() + set(WITH_GDR_VAL "1") +endif() + set(WITH_UBRING_VAL "0") if(WITH_UBRING) set(WITH_UBRING_VAL "1") @@ -153,6 +165,7 @@ endif() list(APPEND BRPC_COMMON_DEFINITIONS BRPC_WITH_GLOG=${WITH_GLOG_VAL} BRPC_WITH_RDMA=${WITH_RDMA_VAL} + BRPC_WITH_GDR=${WITH_GDR_VAL} BRPC_WITH_UBRING=${WITH_UBRING_VAL} BRPC_DEBUG_BTHREAD_SCHE_SAFETY=${WITH_DEBUG_BTHREAD_SCHE_SAFETY_VAL} BRPC_DEBUG_LOCK=${WITH_DEBUG_LOCK_VAL} @@ -309,6 +322,34 @@ if(WITH_RDMA) list(APPEND BRPC_COMMON_INCLUDE_DIRS ${RDMA_INCLUDE_PATH}) endif() +if(WITH_GDR) + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR "GDR (GPU Direct RDMA) is only supported on Linux") + endif() + message("brpc compile with gdr") + # Honor CUDA_HOME if set by the user, otherwise fall back to common defaults. + if(DEFINED ENV{CUDA_HOME} AND NOT CUDA_HOME) + set(CUDA_HOME $ENV{CUDA_HOME}) + endif() + if(CUDA_HOME) + set(CUDA_INCLUDE_PATH ${CUDA_HOME}/include) + find_library(CUDA_LIB NAMES cuda PATHS ${CUDA_HOME}/lib64 ${CUDA_HOME}/lib NO_DEFAULT_PATH) + find_library(CUDART_LIB NAMES cudart PATHS ${CUDA_HOME}/lib64 ${CUDA_HOME}/lib NO_DEFAULT_PATH) + else() + find_path(CUDA_INCLUDE_PATH NAMES cuda_runtime.h + PATHS /usr/local/cuda/include /usr/include) + find_library(CUDA_LIB NAMES cuda + PATHS /usr/local/cuda/lib64 /usr/lib64) + find_library(CUDART_LIB NAMES cudart + PATHS /usr/local/cuda/lib64 /usr/lib64) + endif() + if((NOT CUDA_INCLUDE_PATH) OR (NOT CUDA_LIB) OR (NOT CUDART_LIB)) + message(FATAL_ERROR "Fail to find CUDA (set CUDA_HOME or install CUDA toolkit). " + "Needed headers: cuda_runtime.h; libs: cuda, cudart") + endif() + list(APPEND BRPC_COMMON_INCLUDE_DIRS ${CUDA_INCLUDE_PATH}) +endif() + find_library(PROTOC_LIB NAMES protoc) if(NOT PROTOC_LIB) message(FATAL_ERROR "Fail to find protoc lib") @@ -361,6 +402,11 @@ if(WITH_RDMA) list(APPEND DYNAMIC_LIB ${RDMA_LIB}) endif() +if(WITH_GDR) + list(APPEND DYNAMIC_LIB ${CUDA_LIB}) + list(APPEND DYNAMIC_LIB ${CUDART_LIB}) +endif() + if(WITH_UBRING) message(STATUS "brpc compile with ubring") list(APPEND DYNAMIC_LIB ${UB_LIB}) @@ -538,6 +584,7 @@ set(BUTIL_SOURCES ${PROJECT_SOURCE_DIR}/src/butil/iobuf.cpp ${PROJECT_SOURCE_DIR}/src/butil/single_iobuf.cpp ${PROJECT_SOURCE_DIR}/src/butil/iobuf_profiler.cpp + ${PROJECT_SOURCE_DIR}/src/butil/gpu/gpu_block_pool.cpp ${PROJECT_SOURCE_DIR}/src/butil/binary_printer.cpp ${PROJECT_SOURCE_DIR}/src/butil/recordio.cc ${PROJECT_SOURCE_DIR}/src/butil/popen.cpp diff --git a/config.h.in b/config.h.in index 4f26e57775..81fe4aa6ad 100644 --- a/config.h.in +++ b/config.h.in @@ -21,4 +21,10 @@ #endif #cmakedefine BRPC_WITH_GLOG @WITH_GLOG_VAL@ +// Defined to 1 when brpc is built with GPU Direct RDMA support. +#ifdef BRPC_WITH_GDR +#undef BRPC_WITH_GDR +#endif +#cmakedefine BRPC_WITH_GDR @WITH_GDR_VAL@ + #endif // BUTIL_CONFIG_H diff --git a/config_brpc.sh b/config_brpc.sh index ed7d8f1872..e8a98c0a39 100755 --- a/config_brpc.sh +++ b/config_brpc.sh @@ -539,8 +539,36 @@ if [ $WITH_RDMA != 0 ]; then fi if [ $WITH_GDR != 0 ]; then - CUDA_LIB="/usr/local/cuda/lib64" - CUDA_HDR="/usr/local/cuda/include" + # GDR is a special mode of RDMA and requires RDMA to be enabled as well. + if [ $WITH_RDMA -eq 0 ]; then + print_info "WITH_GDR implies WITH_RDMA, enabling RDMA" + WITH_RDMA=1 + RDMA_LIB=$(find_dir_of_lib_or_die ibverbs) + RDMA_HDR=$(find_dir_of_header_or_die infiniband/verbs.h) + append_to_output_libs "$RDMA_LIB" + append_to_output_headers "$RDMA_HDR" + CPPFLAGS="${CPPFLAGS} -DBRPC_WITH_RDMA" + append_to_output "DYNAMIC_LINKINGS+=-libverbs" + append_to_output "WITH_RDMA=1" + fi + + # Locate the CUDA toolkit. Honor $CUDA_HOME when set, otherwise fall back to + # the well-known /usr/local/cuda installation path. + if [ -z "${CUDA_HOME}" ]; then + CUDA_HOME="/usr/local/cuda" + fi + if [ ! -d "${CUDA_HOME}" ]; then + echo "Fail to find CUDA toolkit at CUDA_HOME=${CUDA_HOME}." >&2 + echo "Please install CUDA or set CUDA_HOME to point at your toolkit." >&2 + exit 1 + fi + CUDA_LIB="${CUDA_HOME}/lib64" + CUDA_HDR="${CUDA_HOME}/include" + if [ ! -f "${CUDA_HDR}/cuda_runtime.h" ]; then + echo "Fail to find cuda_runtime.h under ${CUDA_HDR}." >&2 + echo "Please check CUDA_HOME=${CUDA_HOME}." >&2 + exit 1 + fi append_to_output_libs "$CUDA_LIB" append_to_output_headers "$CUDA_HDR" diff --git a/example/rdma_performance/CMakeLists.txt b/example/rdma_performance/CMakeLists.txt index 213cff0f24..e703ba7e49 100644 --- a/example/rdma_performance/CMakeLists.txt +++ b/example/rdma_performance/CMakeLists.txt @@ -21,6 +21,7 @@ project(rdma_performance C CXX) include(${CMAKE_CURRENT_LIST_DIR}/../cmake/BrpcExample.cmake) option(LINK_SO "Whether examples are linked dynamically" OFF) +option(WITH_GDR "Build with GPU Direct RDMA support (implies RDMA)" OFF) brpc_example_find_common_deps(DYNAMIC_LIB) @@ -33,10 +34,44 @@ if ((NOT RDMA_INCLUDE_PATH) OR (NOT RDMA_LIB)) endif() list(APPEND DYNAMIC_LIB ${RDMA_LIB}) +if(WITH_GDR) + # GDR is a special mode of RDMA and requires CUDA. + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message(FATAL_ERROR "GDR (GPU Direct RDMA) is only supported on Linux") + endif() + if(DEFINED ENV{CUDA_HOME} AND NOT CUDA_HOME) + set(CUDA_HOME $ENV{CUDA_HOME}) + endif() + if(CUDA_HOME) + find_path(CUDA_INCLUDE_PATH NAMES cuda_runtime.h + PATHS ${CUDA_HOME}/include NO_DEFAULT_PATH) + find_library(CUDA_LIB NAMES cuda PATHS ${CUDA_HOME}/lib64 ${CUDA_HOME}/lib NO_DEFAULT_PATH) + find_library(CUDART_LIB NAMES cudart PATHS ${CUDA_HOME}/lib64 ${CUDA_HOME}/lib NO_DEFAULT_PATH) + else() + find_path(CUDA_INCLUDE_PATH NAMES cuda_runtime.h + PATHS /usr/local/cuda/include /usr/include) + find_library(CUDA_LIB NAMES cuda + PATHS /usr/local/cuda/lib64 /usr/lib64) + find_library(CUDART_LIB NAMES cudart + PATHS /usr/local/cuda/lib64 /usr/lib64) + endif() + if((NOT CUDA_INCLUDE_PATH) OR (NOT CUDA_LIB) OR (NOT CUDART_LIB)) + message(FATAL_ERROR "Fail to find CUDA (set CUDA_HOME or install CUDA toolkit)") + endif() + list(APPEND DYNAMIC_LIB ${CUDA_LIB} ${CUDART_LIB}) +endif() + add_executable(client client.cpp ${PROTO_SRC} ${PROTO_HEADER}) brpc_example_configure_target(client) add_executable(server server.cpp ${PROTO_SRC} ${PROTO_HEADER}) brpc_example_configure_target(server) +if(WITH_GDR) + foreach(t client server) + target_compile_definitions(${t} PRIVATE BRPC_WITH_GDR=1 BRPC_WITH_RDMA=1) + target_include_directories(${t} PRIVATE ${CUDA_INCLUDE_PATH}) + endforeach() +endif() + target_link_libraries(client PRIVATE ${BRPC_LIB} ${DYNAMIC_LIB}) target_link_libraries(server PRIVATE ${BRPC_LIB} ${DYNAMIC_LIB}) diff --git a/example/rdma_performance/Makefile b/example/rdma_performance/Makefile index 60c59702a7..c5cc3753ed 100644 --- a/example/rdma_performance/Makefile +++ b/example/rdma_performance/Makefile @@ -60,6 +60,9 @@ endif .PHONY:all ifdef WITH_RDMA all: client server +else ifdef WITH_GDR +# GDR implies RDMA; libbrpc may have been built with only WITH_GDR recorded. +all: client server else all: endif diff --git a/src/brpc/policy/baidu_rpc_protocol_gpu.cpp b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp index 500875eb70..2352fbecc1 100644 --- a/src/brpc/policy/baidu_rpc_protocol_gpu.cpp +++ b/src/brpc/policy/baidu_rpc_protocol_gpu.cpp @@ -88,7 +88,10 @@ ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket, butil::gdr::BlockPoolAllocator* host_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_cpu_allocator(); prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size); if (prefetch_d2h_data == NULL) { - LOG(FATAL) << "alloc host data failed!!!"; + // 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); } @@ -146,8 +149,12 @@ ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket, if (header_size + meta_size <= n) { auto deleter = [host_allocator, prefetch_d2h_data](void* data) { host_allocator->DeallocateRaw(prefetch_d2h_data); }; - // n is the bytes we real frefetch. We set n as the meta and n will be used in ProcessRpcRequest/ProcessRpcResponse. - // This is a trick, we should keep n in another better way. + // `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 { diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index 874c1c899c..0e6eccbc55 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -1022,10 +1022,17 @@ int RdmaEndpoint::PostRecv(uint32_t num, bool zerocopy) { if (_use_gdr) { butil::gdr::BlockPoolAllocator* device_allocator = butil::gdr::BlockPoolAllocators::singleton()->get_gpu_allocator(); void* device_ptr = device_allocator->AllocateRaw(g_gdr_recv_block_size); + if (device_ptr == nullptr) { + // GPU pool exhausted; abort this post, the caller will retry. + PLOG(WARNING) << "Fail to allocate gpu rbuf"; + return -1; + } auto deleter = [device_allocator](void* data) { device_allocator->DeallocateRaw(data); }; lkey = device_allocator->get_lkey(device_ptr); - // we keep lkey into the meta, and this is a thick. we also keep prefetch d2h size in meta too. - _rbuf[_rq_received].append_user_data_with_meta(device_ptr, g_gdr_recv_block_size, deleter, lkey); + // Tag the block as GPU memory so that IOBuf::is_gpu_memory() + // works deterministically. The lkey is carried via `data_meta' + // so the recv path can locate the correct memory region key. + _rbuf[_rq_received].append_user_data_gpu(device_ptr, g_gdr_recv_block_size, deleter, lkey); _rbuf_data[_rq_received] = device_ptr; } else #endif // if BRPC_WITH_GDR diff --git a/src/butil/gpu/gpu_block_pool.cpp b/src/butil/gpu/gpu_block_pool.cpp index 2efcefa513..4c594d9473 100644 --- a/src/butil/gpu/gpu_block_pool.cpp +++ b/src/butil/gpu/gpu_block_pool.cpp @@ -179,7 +179,12 @@ BlockPoolAllocator::BlockPoolAllocator(int gpuId, bool onGpu, ibv_pd* ibvPd, << ", region_size=" << REGION_SIZE << ", max_gdr_regions=" << FLAGS_max_gdr_regions << ", gpu_id=" << gpu_id << ", on_gpu=" << on_gpu << ", pd=" << pd; - extendRegion(); + // 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() { @@ -243,7 +248,12 @@ void* BlockPoolAllocator::AllocateRaw(size_t num_bytes) { return nullptr; } if (num_bytes > BLOCK_SIZE) { - LOG(FATAL) << "try to alloc " << num_bytes << " bytes, its bigger than block_size " << 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(); @@ -251,7 +261,11 @@ void* BlockPoolAllocator::AllocateRaw(size_t num_bytes) { std::lock_guard lock(poolMutex); if (!freeList) { - extendRegion(); + if (extendRegion() != 0) { + // Out of regions or registration failed; let the caller handle it. + errno = ENOMEM; + return nullptr; + } } BlockHeader* block = freeList; @@ -301,10 +315,11 @@ void BlockPoolAllocator::printStatistics() const { << "%"; } -void BlockPoolAllocator::extendRegion() { +int BlockPoolAllocator::extendRegion() { if (g_region_num == FLAGS_max_gdr_regions) { - LOG(FATAL) << "Gdr Memory pool reaches max regions"; - return ; + LOG(ERROR) << "Gdr Memory pool reaches max regions (" << FLAGS_max_gdr_regions << ")"; + errno = ENOMEM; + return -1; } auto startTime = std::chrono::high_resolution_clock::now(); @@ -317,6 +332,11 @@ void BlockPoolAllocator::extendRegion() { } 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)); @@ -330,15 +350,27 @@ void BlockPoolAllocator::extendRegion() { } LOG(INFO) << "reg_mr for ptr: " << aligned_ptr << ", size:" << aligned_bytes; - auto mr = ibv_reg_mr(pd, aligned_ptr, 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 | - IBV_ACCESS_RELAXED_ORDERING); + 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(FATAL) << "Failed to register MR: " << strerror(errno) + 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(ptr)); + } else { + cudaFreeHost(reinterpret_cast(ptr)); + } + return -1; } else { LOG(INFO) << "Success to register MR: " << ", pd " << pd << ", aligned_ptr:" << aligned_ptr; @@ -378,6 +410,7 @@ void BlockPoolAllocator::extendRegion() { 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) : diff --git a/src/butil/gpu/gpu_block_pool.h b/src/butil/gpu/gpu_block_pool.h index 73b557f3d0..b263dbd720 100644 --- a/src/butil/gpu/gpu_block_pool.h +++ b/src/butil/gpu/gpu_block_pool.h @@ -108,7 +108,10 @@ class BlockPoolAllocator { uint32_t get_lkey(const void* buf); private: - void extendRegion(); + // Extend the pool by one region. + // Returns 0 on success; -1 on failure (e.g. allocation/registration error, + // or max regions reached) with errno/LOG set. Does NOT abort the process. + int extendRegion(); }; class GPUStreamPool { diff --git a/src/butil/iobuf.cpp b/src/butil/iobuf.cpp index 075e6364c1..6f70751aa0 100644 --- a/src/butil/iobuf.cpp +++ b/src/butil/iobuf.cpp @@ -1233,6 +1233,37 @@ int IOBuf::append_user_data_with_meta(void* data, return 0; } +#if BRPC_WITH_GDR +// Same as append_user_data_with_meta, but also tags the block as GPU memory so +// that is_gpu_memory() can report it deterministically. +int IOBuf::append_user_data_gpu(void* data, + size_t size, + std::function deleter, + uint64_t meta) { + if (size > 0xFFFFFFFFULL - 100) { + LOG(FATAL) << "data_size=" << size << " is too large"; + return -1; + } + if (!deleter) { + deleter = ::cudaFree; + } + if (!size) { + deleter(data); + return 0; + } + char* mem = (char*)malloc(sizeof(IOBuf::Block) + sizeof(UserDataExtension)); + if (mem == NULL) { + return -1; + } + IOBuf::Block* b = new (mem) IOBuf::Block((char*)data, size, std::move(deleter)); + b->u.data_meta = meta; + b->flags |= IOBUF_BLOCK_FLAGS_GPU_MEMORY; + const IOBuf::BlockRef r = { 0, b->cap, b }; + _move_back_ref(r); + return 0; +} +#endif // BRPC_WITH_GDR + uint64_t IOBuf::get_first_data_meta() { if (_ref_num() == 0) { return 0; @@ -1602,12 +1633,17 @@ bool IOBuf::equals(const butil::IOBuf& other) const { } #if BRPC_WITH_GDR -// when IOBuf is used for send, data_meta is set by user; -// when IOBf is used for recv and gdr is open, data_meta is set by brpc -// and it is lkey. +// Returns true iff the first block was explicitly tagged as GPU memory via +// append_user_data_gpu() (done by the GDR receive path in RdmaEndpoint). +// We deliberately check the block flag rather than guessing from `data_meta', +// because `data_meta' is legitimately used by callers/users to carry non-zero +// values (e.g. RDMA lkey, prefetch size) that must NOT be treated as a GPU hint. bool IOBuf::is_gpu_memory() { - uint64_t data_meta = get_first_data_meta(); - return (data_meta > 0 && data_meta <= UINT_MAX); + if (_ref_num() == 0) { + return false; + } + IOBuf::BlockRef const& r = _ref_at(0); + return (r.block->flags & IOBUF_BLOCK_FLAGS_GPU_MEMORY) != 0; } #endif diff --git a/src/butil/iobuf.h b/src/butil/iobuf.h index c371b67a10..7720c7b7e4 100644 --- a/src/butil/iobuf.h +++ b/src/butil/iobuf.h @@ -150,6 +150,10 @@ friend class SingleIOBuf; #if BRPC_WITH_GDR size_t cutn_from_gpu(IOBuf* out, size_t n); size_t copy_from_gpu(void* d, size_t n, size_t pos = 0, bool to_gpu = false) const; + // Returns true iff the first block of this IOBuf was tagged as GPU memory + // via append_user_data_gpu(). This is deterministic and independent of the + // value of `data_meta' (which is reserved for user/RDMA metadata such as + // the lkey or the prefetch size). bool is_gpu_memory(); #endif // BRPC_WITH_GDR @@ -267,6 +271,14 @@ friend class SingleIOBuf; // The meta is associated with this piece of user-data. int append_user_data_with_meta(void* data, size_t size, std::function deleter, uint64_t meta); +#if BRPC_WITH_GDR + // Same as append_user_data_with_meta, but additionally tags the underlying + // block with the GPU-memory flag so that is_gpu_memory() can be answered + // deterministically (instead of guessing from `meta'). Used by the GDR + // receive path to attach device pointers received via RDMA. + int append_user_data_gpu(void* data, size_t size, std::function deleter, uint64_t meta); +#endif // BRPC_WITH_GDR + // Get the data meta of the first byte in this IOBuf. // The meta is specified with append_user_data_with_meta before. // 0 means the meta is invalid. diff --git a/src/butil/iobuf_inl.h b/src/butil/iobuf_inl.h index 756cf8bf63..a75e8dad87 100644 --- a/src/butil/iobuf_inl.h +++ b/src/butil/iobuf_inl.h @@ -43,6 +43,10 @@ void SubmitIOBufSample(IOBuf::Block* block, int64_t ref); const uint16_t IOBUF_BLOCK_FLAGS_USER_DATA = 1 << 0; const uint16_t IOBUF_BLOCK_FLAGS_SAMPLED = 1 << 1; +// Marks a USER_DATA block as holding GPU memory (used by the GDR path). +// `is_gpu_memory()` checks this flag instead of guessing from `data_meta`, +// which avoids false positives when users attach their own non-zero meta. +const uint16_t IOBUF_BLOCK_FLAGS_GPU_MEMORY = 1 << 2; inline ssize_t IOBuf::cut_into_file_descriptor(int fd, size_t size_hint) { return pcut_into_file_descriptor(fd, -1, size_hint); From 3cdba904291bb6846dcb2bc635e2995c839c172b Mon Sep 17 00:00:00 2001 From: randomkang <75484924+randomkang@users.noreply.github.com> Date: Tue, 4 Aug 2026 10:53:46 +0800 Subject: [PATCH 4/4] fix build error --- src/brpc/rdma/rdma_endpoint.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/brpc/rdma/rdma_endpoint.cpp b/src/brpc/rdma/rdma_endpoint.cpp index 0e6eccbc55..235dc89002 100644 --- a/src/brpc/rdma/rdma_endpoint.cpp +++ b/src/brpc/rdma/rdma_endpoint.cpp @@ -113,9 +113,9 @@ RdmaResource::~RdmaResource() { RdmaEndpoint::RdmaEndpoint(Socket* s, bool use_gdr) : _socket(s) - , _use_gdr(use_gdr) , _state(UNINIT) , _handshake_version(0) + , _use_gdr(use_gdr) , _resource(NULL) , _send_cq_events(0) , _recv_cq_events(0)