Copilot commented on code in PR #3144:
URL: https://github.com/apache/brpc/pull/3144#discussion_r4214801288
##########
BUILD.bazel:
##########
@@ -103,6 +108,12 @@ LINKOPTS = [
"-libverbs",
],
"//conditions:default": [],
+}) + select({
+ "//bazel/config:brpc_with_gdr": [
+ "-lcuda",
+ "-lcudart",
+ ],
Review Comment:
The Bazel GDR link flags do not provide a CUDA include path or CUDA target,
while the new GDR sources include `<cuda_runtime.h>` and `cuda.h`. The
documented `--define=BRPC_WITH_GDR=true` build therefore cannot find headers
for a CUDA installation under `CUDA_HOME` (and unlike CMake/config_brpc, Bazel
never uses that location). Add a CUDA include/dependency configuration to the
Bazel GDR branch.
##########
src/brpc/policy/baidu_rpc_protocol.cpp:
##########
@@ -1017,17 +1031,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);
Review Comment:
`FillResBufGpu` returns after setting `cntl->Failed()` when `attachment_size
> res_size`, but this caller continues with `res_buf_ptr` still pointing at the
GPU payload and passes it to deserialization. That can parse an invalid GPU
buffer instead of stopping as the non-GDR path does; break out of this
response-processing block when the helper reports failure.
##########
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);
Review Comment:
Adding `SOCKET_MODE_GDR` here does not make the server run the RDMA
handshake: `ExecuteServerHandshake` still dispatches only `SOCKET_MODE_RDMA`
(`src/brpc/rdma/rdma_handshake_server.cpp:211`), so GDR connections take the
fallback handshake and cannot complete setup. Add GDR to that dispatch before
enabling this mode.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch 512B from gpu to cpu.
+// if header_size + meta_size + body_size(without attachment) is less than
512, then one
+// d2h is enough for one rpc.
+
+const int prefetch_d2h_size = 512;
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
+
+ char header_buf[12];
+ size_t n = 0;
+ uint32_t body_size;
+ uint32_t meta_size;
+ ParseError pe = PARSE_OK;
+
+ void* prefetch_d2h_data = NULL;
+ bool is_gpu_memory = source->is_gpu_memory();
+ if (!is_gpu_memory) {
+ LOG(FATAL) << "RpcMessage is not in gpu!!!";
+ return MakeParseError(PARSE_ERROR_TRY_OTHERS);
+ }
+ butil::gdr::BlockPoolAllocator* host_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator();
+ prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size);
+ if (prefetch_d2h_data == NULL) {
+ // Pool exhausted or single-block limit hit: ask the caller to retry
+ // instead of crashing the whole server.
+ LOG(ERROR) << "alloc host data for prefetch d2h failed, size="
+ << prefetch_d2h_size;
+ return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
+ }
+
+ // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size;
+ n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size);
+ size_t copy_size = n > 12 ? 12 : n;
+ memcpy(header_buf, prefetch_d2h_data, copy_size);
+
+ do {
+ if (n >= 4) {
+ void* dummy = header_buf;
+ if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") {
+ pe = PARSE_ERROR_TRY_OTHERS;
+ break;
+ }
+ } else {
+ if (memcmp(header_buf, "PRPC", n) != 0) {
+ pe = PARSE_ERROR_TRY_OTHERS;
Review Comment:
Once `source` is tagged as GPU memory, returning `PARSE_ERROR_TRY_OTHERS`
sends the same IOBuf to the next protocol parser (`InputMessengerProcessor`
does this), whose parsers use ordinary `copy_to` and cannot dereference device
memory. A bad magic or malformed GPU frame can therefore crash while protocol
detection falls through. Treat these as a terminal parse error for GDR input,
or prevent dispatch to non-GDR parsers.
##########
src/brpc/rdma/rdma_helper.cpp:
##########
@@ -586,7 +591,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();
+ }
Review Comment:
The GPU pool is constructed before the later size check, and
`BlockPoolAllocator` computes `REGION_SIZE` by dividing by `blockSize`. With
`--gdr_block_size_kb=0`, this path divides by zero; values whose byte size is
no larger than `IOBUF_BLOCK_HEADER_LEN` also underflow when the receive size is
computed. Validate the effective block size before calling `InitGPUBlockPool`
(and reject nonpositive values) rather than relying on `GlobalGdrInitialize()`,
which currently always returns 0.
##########
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
Review Comment:
CMake always defines `BRPC_WITH_GDR` as `0` or `1` in
`BRPC_COMMON_DEFINITIONS` (CMakeLists.txt:190-195), so this `#ifdef` is also
true in every non-GDR CMake build. That makes `iobuf.cpp` include CUDA/GDR code
even though `iobuf.h` hides the declarations, causing ordinary builds without
CUDA to fail (and producing definition/declaration mismatches). Use `#if
BRPC_WITH_GDR` here, consistently with the other guards in this change.
##########
src/brpc/rdma/rdma_endpoint.cpp:
##########
@@ -1099,26 +1113,83 @@ 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<uint32_t>(size), g_rdma_recv_block_size);
+
+#if BRPC_WITH_GDR
+ if (_use_gdr) {
+ butil::gdr::BlockPoolAllocator* device_allocator =
butil::gdr::BlockPoolAllocators::GetInstance()->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);
+ // 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;
Review Comment:
The return value of `append_user_data_gpu` is ignored. If allocation of the
IOBuf metadata block fails, `_rbuf` remains empty while the GPU block is still
posted and no longer has an owner; completion then cuts zero bytes from `_rbuf`
and silently loses the receive. Handle the failure, release `device_ptr`, and
only publish `_rbuf_data` after the append succeeds.
##########
src/brpc/acceptor.h:
##########
@@ -135,7 +135,7 @@ friend class Server;
bool _force_ssl;
std::shared_ptr<SocketSSLContext> _ssl_ctx;
- // Choose to use a certain socket: 0 TCP, 1 RDMA
+ // Choose to use a certain socket: 0 TCP, 1 RDMA, 2 UBRING, 3 GDR
Review Comment:
This updated comment assigns the wrong numeric values: the enum is
`SOCKET_MODE_URMA = 3` and `SOCKET_MODE_GDR = 4`, not GDR = 3. Please keep the
mode documentation synchronized with `socket_mode.h`.
##########
src/butil/gpu/gpu_block_pool.cpp:
##########
@@ -0,0 +1,498 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <iostream>
+#include <chrono>
+#include <gflags/gflags.h>
+#include <sys/stat.h>
+#include "butil/fast_rand.h"
+#include "gpu_block_pool.h"
+namespace butil {
+namespace gdr {
+DEFINE_int32(gdr_block_size_kb, 512, "gdr block size in KB");
+DEFINE_int32(max_gdr_regions, 32, "max num of gdr regions");
+
+#define CHECK_CUDA(call) \
+do { \
+ auto _sts = (call); \
+ if (_sts != cudaSuccess) { \
+ LOG(FATAL) << " cuda error:" \
+ << (cudaGetErrorString(_sts)) << std::string(" at ") \
+ << __FILE__ << ": " << __LINE__; \
+ } \
+} while (0);
+
+size_t GetGdrBlockSize() {
+ return FLAGS_gdr_block_size_kb * 1024;
+}
+
+bool verify_same_context() {
+ static std::once_flag init_flag;
+ static int original_device = -1;
+
+ int current_device;
+ CHECK_CUDA(cudaGetDevice(¤t_device));
+
+ std::call_once(init_flag, [&]() {
+ original_device = current_device;
+ });
+
+ return (current_device == original_device);
+}
+
+void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ void *d_data;
+
+ LOG(INFO) << "try to alloc " << gpu_mem_size << " bytes from gpu " <<
gpu_id;
+
+ CHECK_CUDA(cudaMalloc(&d_data, gpu_mem_size));
+ cudaDeviceSynchronize();
+ return (void *)d_data;
Review Comment:
The allocator's failure path claims to return `-1` when a region cannot be
extended, but `get_gpu_mem()` aborts the process on `cudaMalloc` failure via
`CHECK_CUDA` (and the pinned-host path does the same). VRAM/pinned-memory
exhaustion therefore kills the RPC process instead of allowing
`AllocateRaw()`/`PostRecv()` to report exhaustion and retry. Make the CUDA
allocation helpers return `nullptr` on allocation errors and let
`extendRegion()` propagate the failure.
##########
src/brpc/policy/baidu_rpc_protocol_gpu.cpp:
##########
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <google/protobuf/descriptor.h> // MethodDescriptor
+#include <google/protobuf/message.h> // Message
+#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include <google/protobuf/io/coded_stream.h>
+#include <google/protobuf/text_format.h>
+
+#include "butil/gpu/gpu_block_pool.h"
+#include "butil/iobuf.h" // butil::IOBuf
+#include "butil/logging.h" // LOG()
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/memory/scope_guard.h"
+#include "butil/raw_pack.h" // RawPacker RawUnpacker
+#include "butil/strings/string_util.h"
+
+#include "json2pb/json_to_pb.h"
+#include "json2pb/pb_to_json.h"
+#include "brpc/controller.h" // Controller
+#include "brpc/socket.h" // Socket
+#include "brpc/server.h" // Server
+#include "brpc/span.h"
+#include "brpc/compress.h" // ParseFromCompressedData
+#include "brpc/checksum.h"
+#include "brpc/stream_impl.h"
+#include "brpc/rpc_dump.h" // SampledRequest
+#include "brpc/rpc_pb_message_factory.h"
+#include "brpc/policy/baidu_rpc_meta.pb.h" // RpcRequestMeta
+#include "brpc/policy/baidu_rpc_protocol.h"
+#include "brpc/policy/most_common_message.h"
+#include "brpc/policy/streaming_rpc_protocol.h"
+#include "brpc/details/usercode_backup_pool.h"
+#include "brpc/details/controller_private_accessor.h"
+#include "brpc/details/server_private_accessor.h"
+
+namespace brpc {
+namespace policy {
+
+// Notes:
+// 1. 12-byte header [PRPC][body_size][meta_size]
+// 2. body_size and meta_size are in network byte order
+// 3. Use service->full_name() + method_name to specify the method to call
+// 4. `attachment_size' is set iff request/response has attachment
+// 5. Not supported: chunk_info
+
+// Pack header into `buf'
+
+const int header_size = 12;
+// if we recv data into gpu, the header/meta/body will be copied to cpu and
processed.
+// in to to limit the count of d2h, we will prefetch 512B from gpu to cpu.
+// if header_size + meta_size + body_size(without attachment) is less than
512, then one
+// d2h is enough for one rpc.
+
+const int prefetch_d2h_size = 512;
+
+ParseResult ParseRpcMessageGpu(butil::IOBuf* source, Socket* socket,
+ bool /*read_eof*/, const void*) {
Review Comment:
The new GPU parser introduces distinct prefetch-boundary handling and GPU
attachment splitting, but the repository's existing IOBuf and RDMA tests do not
exercise this path. Add focused coverage for partial/full 512-byte prefetches
and malformed `attachment_size` responses (using a mocked copy/allocator
boundary where hardware is unavailable) so regressions in this safety-critical
receive path are caught.
##########
src/butil/gpu/gpu_block_pool.cpp:
##########
@@ -0,0 +1,498 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <iostream>
+#include <chrono>
+#include <gflags/gflags.h>
+#include <sys/stat.h>
+#include "butil/fast_rand.h"
+#include "gpu_block_pool.h"
+namespace butil {
+namespace gdr {
+DEFINE_int32(gdr_block_size_kb, 512, "gdr block size in KB");
+DEFINE_int32(max_gdr_regions, 32, "max num of gdr regions");
+
+#define CHECK_CUDA(call) \
+do { \
+ auto _sts = (call); \
+ if (_sts != cudaSuccess) { \
+ LOG(FATAL) << " cuda error:" \
+ << (cudaGetErrorString(_sts)) << std::string(" at ") \
+ << __FILE__ << ": " << __LINE__; \
+ } \
+} while (0);
+
+size_t GetGdrBlockSize() {
+ return FLAGS_gdr_block_size_kb * 1024;
+}
+
+bool verify_same_context() {
+ static std::once_flag init_flag;
+ static int original_device = -1;
+
+ int current_device;
+ CHECK_CUDA(cudaGetDevice(¤t_device));
+
+ std::call_once(init_flag, [&]() {
+ original_device = current_device;
+ });
+
+ return (current_device == original_device);
+}
+
+void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ void *d_data;
+
+ LOG(INFO) << "try to alloc " << gpu_mem_size << " bytes from gpu " <<
gpu_id;
+
+ CHECK_CUDA(cudaMalloc(&d_data, gpu_mem_size));
+ cudaDeviceSynchronize();
+ return (void *)d_data;
+}
+
+void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+
+ LOG(INFO) << "try to alloc " << cpu_mem_size << " bytes from gpu " <<
gpu_id << "on host";
+
+ void* mem = NULL;
+
+ CHECK_CUDA(cudaMallocHost(&mem, cpu_mem_size));
+
+ cudaDeviceSynchronize();
+
+ return mem;
+}
+
+
+BlockPoolAllocators* BlockPoolAllocators::GetInstance() {
+ return ::Singleton<BlockPoolAllocators,
+ LeakySingletonTraits<BlockPoolAllocators>>::get();
+}
+
+void BlockPoolAllocators::init(int gpu_id, ibv_pd* pd) {
+ LOG(INFO) << "set GPU BlockPoolAllocator for " << gpu_id;
+ size_t region_size = 1024LL * 1024 * 1024;
+ size_t block_size = FLAGS_gdr_block_size_kb * 1024;
+ gpu_mem_alloc = new BlockPoolAllocator(gpu_id, true, pd, block_size,
region_size);
+
+ region_size = 32LL * 1024 * 1024;
+ block_size = 512;
+ cpu_mem_alloc = new BlockPoolAllocator(gpu_id, false, pd, block_size,
region_size);
+
+ gpu_stream_pool = new GPUStreamPool(gpu_id);
+}
+
+bool InitGPUBlockPool(int gpu_id, ibv_pd* pd) {
+ BlockPoolAllocators::GetInstance()->init(gpu_id, pd);
+ return true;
+}
+
+class BlockHeaderList {
+ public:
+ BlockHeaderList() {
+ objects_.reserve(kMaxObjects);
+ }
+ virtual ~BlockHeaderList() {
+ for (size_t i = 0; i < objects_.size(); i++) {
+ delete objects_[i];
+ }
+ }
+
+ BlockHeader* New() {
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (!objects_.empty()) {
+ BlockHeader* result = objects_.back();
+ objects_.pop_back();
+ return result;
+ }
+ }
+ return new BlockHeader;
+ }
+ void Release(BlockHeader* obj) {
+ obj->Reset();
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (objects_.size() < kMaxObjects) {
+ objects_.push_back(obj);
+ return;
+ }
+ }
+ delete obj;
+ }
+
+ private:
+ static const int kMaxObjects = 100000;
+
+ std::mutex mu_;
+ std::vector<BlockHeader*> objects_;
+};
+
+static BlockHeaderList* get_bh_list() {
+ static BlockHeaderList* bh_list = new BlockHeaderList();
+ return bh_list;
+}
+
+BlockPoolAllocator::BlockPoolAllocator(int gpuId, bool onGpu, ibv_pd* ibvPd,
+ size_t blockSize, size_t regionSize) :
+ gpu_id(gpuId)
+ , on_gpu(onGpu)
+ , pd(ibvPd)
+ , BLOCK_SIZE(std::max(blockSize, sizeof(BlockHeader)))
+ , REGION_SIZE((regionSize / blockSize) * blockSize) // 对齐到块大小的倍数
+ , freeList(nullptr)
+ , g_region_num(0)
+ , totalAllocated(0)
+ , totalDeallocated(0)
+ , peakUsage(0) {
+ g_regions.resize(FLAGS_max_gdr_regions);
+ LOG(INFO) << "Memory Pool initialized: block_size=" << BLOCK_SIZE
+ << ", region_size=" << REGION_SIZE << ", max_gdr_regions=" <<
FLAGS_max_gdr_regions
+ << ", gpu_id=" << gpu_id << ", on_gpu=" << on_gpu << ", pd=" << pd;
+
+ // The very first region is required at construction time; failure here is
+ // fatal because the pool would be unusable.
+ if (extendRegion() != 0) {
+ LOG(FATAL) << "Failed to initialize GDR memory pool (first region) on
gpu "
+ << gpu_id << ", on_gpu " << on_gpu;
+ }
+}
+
+BlockPoolAllocator::~BlockPoolAllocator() {
+#ifdef DEBUG
+ printStatistics();
+#endif
+
+ for (int i = 0; i < FLAGS_max_gdr_regions; i++) {
+ Region* r = &g_regions[i];
+ if (!r->mr) {
+ continue;
+ }
+
+ LOG(INFO) << "try to free " << r->size << " bytes from gpu " << gpu_id
<< ", on_gpu " << on_gpu;
+ ibv_dereg_mr(r->mr);
+ if (on_gpu) {
+ CHECK_CUDA(cudaFree(reinterpret_cast<void*>(r->start)));
+ } else {
+ CHECK_CUDA(cudaFreeHost(reinterpret_cast<void*>(r->start)));
+ }
+ }
+}
+
+Region* BlockPoolAllocator::GetRegion(const void* buf) {
+ if (!buf) {
+ errno = EINVAL;
+ return NULL;
+ }
+ Region* r = NULL;
+ uintptr_t addr = (uintptr_t)buf;
+ for (int i = 0; i < FLAGS_max_gdr_regions; ++i) {
+ if (g_regions[i].aligned_start == 0) {
+ break;
+ }
+ if (addr >= g_regions[i].aligned_start &&
+ addr < g_regions[i].aligned_start + g_regions[i].aligned_size)
{
+ r = &g_regions[i];
+ break;
+ }
+ }
+ return r;
+}
+
+uint32_t BlockPoolAllocator::get_lkey(const void* buf) {
+ Region* r = GetRegion(buf);
+ if (!r) {
+ LOG(ERROR) << "can not get a region for buf " << buf;
+ return 0;
+ }
+
+ if (!r->mr) {
+ LOG(FATAL) << "region has not been registered into rdma yet, addr:" <<
r->start;
+ return 0;
+ }
+
+ return r->mr->lkey;
+}
+
+void* BlockPoolAllocator::AllocateRaw(size_t num_bytes) {
+ if (num_bytes == 0) {
+ return nullptr;
+ }
+ if (num_bytes > BLOCK_SIZE) {
+ // A single block cannot satisfy the request; fail gracefully instead
of
+ // aborting the whole server.
+ LOG(ERROR) << "try to alloc " << num_bytes
+ << " bytes, bigger than block_size " << BLOCK_SIZE;
+ errno = ENOMEM;
+ return nullptr;
+ }
+
+ auto startTime = std::chrono::high_resolution_clock::now();
+
+ std::lock_guard<std::mutex> lock(poolMutex);
+
+ if (!freeList) {
+ if (extendRegion() != 0) {
+ // Out of regions or registration failed; let the caller handle it.
+ errno = ENOMEM;
+ return nullptr;
+ }
+ }
+
+ BlockHeader* block = freeList;
+ freeList = freeList->next;
+
+ void* addr = block->addr;
+ get_bh_list()->Release(block);
+
+ totalAllocated++;
+ peakUsage = std::max(peakUsage, totalAllocated - totalDeallocated);
+
+ auto endTime = std::chrono::high_resolution_clock::now();
+ auto duration =
std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+#ifdef DEBUG
+ if (duration.count() > 1000) {
+ LOG(INFO) << "Slow allocation: " << duration.count() << " ns";
+ }
+#endif
+
+ return addr;
+}
+
+void BlockPoolAllocator::DeallocateRaw(void* ptr) {
+ if (!ptr) return;
+
+ std::lock_guard<std::mutex> lock(poolMutex);
+
+ BlockHeader* block = get_bh_list()->New();
+ block->addr = ptr;
+ block->next = freeList;
+ freeList = block;
+
+ totalDeallocated++;
+}
+
+void BlockPoolAllocator::printStatistics() const {
+ LOG(INFO) << "=== Memory Pool Statistics ===";
+ LOG(INFO) << "Total regions: " << g_region_num
+ << ", Total blocks allocated: " << totalAllocated
+ << ", Total blocks deallocated: " << totalDeallocated
+ << ", Current usage: " << (totalAllocated - totalDeallocated) << "
blocks"
+ << ", Peak usage: " << peakUsage << " blocks"
+ << ", Memory efficiency: "
+ << (static_cast<double>(totalAllocated - totalDeallocated) /
+ (g_region_num * (REGION_SIZE / BLOCK_SIZE)) * 100)
+ << "%";
+}
+
+int BlockPoolAllocator::extendRegion() {
+ if (g_region_num == FLAGS_max_gdr_regions) {
+ LOG(ERROR) << "Gdr Memory pool reaches max regions (" <<
FLAGS_max_gdr_regions << ")";
+ errno = ENOMEM;
+ return -1;
+ }
+
+ auto startTime = std::chrono::high_resolution_clock::now();
+ void* ptr = nullptr;
+ void* aligned_ptr = nullptr;
+ int alignment = 4096;
+
+ if (on_gpu) {
+ ptr = get_gpu_mem(gpu_id, REGION_SIZE);
+ } else {
+ ptr = get_cpu_mem(gpu_id, REGION_SIZE);
+ }
+ if (ptr == nullptr) {
+ LOG(ERROR) << "Failed to allocate region memory on gpu " << gpu_id
+ << ", on_gpu " << on_gpu;
+ return -1;
+ }
+
+ aligned_ptr = (void*)(((uintptr_t)ptr + alignment - 1) & ~(alignment - 1));
+
+ int64_t aligned_bytes = REGION_SIZE;
+ if (ptr != aligned_ptr) {
+ uintptr_t region_end = uintptr_t(ptr) + REGION_SIZE;
+ uintptr_t aligned_end_ptr = region_end & ~(alignment - 1);
+ aligned_bytes = uintptr_t(aligned_end_ptr) - uintptr_t(aligned_ptr);
+ LOG(WARNING) << "addr is not aligned with 4096: " << ptr << ",
aligned_bytes: " << aligned_bytes
+ << ", region_size: " << REGION_SIZE;
+ }
+
+ LOG(INFO) << "reg_mr for ptr: " << aligned_ptr << ", size:" <<
aligned_bytes;
+ // IBV_ACCESS_RELAXED_ORDERING is only available on newer libibverbs; guard
+ // it so the build does not break on older toolchains.
+ int access_flags =
+ IBV_ACCESS_LOCAL_WRITE |
+ IBV_ACCESS_REMOTE_READ |
+ IBV_ACCESS_REMOTE_WRITE;
+#ifdef IBV_ACCESS_RELAXED_ORDERING
+ access_flags |= IBV_ACCESS_RELAXED_ORDERING;
+#endif
+ auto mr = ibv_reg_mr(pd, aligned_ptr, aligned_bytes, access_flags);
+
+ if (!mr) {
+ LOG(ERROR) << "Failed to register MR: " << strerror(errno)
+ << ", pd " << pd << ", aligned_ptr:" << aligned_ptr;
+ // Free the region memory we just allocated so it is not leaked.
+ if (on_gpu) {
+ cudaFree(reinterpret_cast<void*>(ptr));
+ } else {
+ cudaFreeHost(reinterpret_cast<void*>(ptr));
+ }
+ return -1;
+ } else {
+ LOG(INFO) << "Success to register MR: "
+ << ", pd " << pd << ", aligned_ptr:" << aligned_ptr;
+ }
+
+ LOG(INFO) << "try to init region, g_region_num:" << g_region_num;
+ size_t blockCount = aligned_bytes / BLOCK_SIZE;
+ Region* region = &g_regions[g_region_num++];
+ region->start = (uintptr_t)ptr;
+ region->aligned_start = (uintptr_t)aligned_ptr;
+ region->mr = mr;
+ region->size = REGION_SIZE;
+ region->aligned_size = aligned_bytes;
+ region->blockCount = blockCount;
+
+
+ LOG(INFO) << "try to insert list, freeList:" << freeList << ",
blockCount:" << blockCount;
+ BlockHeader* lastBlock = nullptr;
+ for (size_t i = 0; i < blockCount; ++i) {
+ BlockHeader* block = get_bh_list()->New();
+ block->addr = reinterpret_cast<void*>(static_cast<char*>(aligned_ptr)
+ i * BLOCK_SIZE);
+ if (lastBlock != nullptr) {
+ lastBlock->next = block;
+ } else {
+ freeList = block;
+ }
+ lastBlock = block;
+ }
+
+ if (lastBlock) {
+ lastBlock->next = nullptr;
+ }
+
+ auto endTime = std::chrono::high_resolution_clock::now();
+ auto duration =
std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime);
+
+ LOG(INFO) << "Extended region #" << g_region_num << ": " << blockCount
+ << " blocks (" << (REGION_SIZE / (1024 * 1024)) << " MB)" << ", on_gpu
" << on_gpu
+ << ", cost " << duration.count() << " ns";
+ return 0;
+}
+
+GPUStreamPool::GPUStreamPool(int gpu_id) :
+ gpu_id_(gpu_id) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ d2d_streams_.resize(kMaxConcurrent);
+ d2h_streams_.resize(kMaxConcurrent);
+ for (int i = 0; i < kMaxConcurrent; i++) {
+ CHECK_CUDA(cudaStreamCreate(&d2d_streams_[i]));
+ CHECK_CUDA(cudaStreamCreate(&d2h_streams_[i]));
+ }
+ CHECK_CUDA(cudaDeviceSynchronize());
+}
+
+GPUStreamPool::~GPUStreamPool() {
+ CHECK_CUDA(cudaDeviceSynchronize());
+ for (int i = 0; i < kMaxConcurrent; i++) {
+ CHECK_CUDA(cudaStreamDestroy(d2d_streams_[i]));
+ CHECK_CUDA(cudaStreamDestroy(d2h_streams_[i]));
+ }
+}
+
+void GPUStreamPool::fast_d2d(std::vector<void*>& src_list,
+ std::vector<int64_t>& length_list,
+ void* dst) {
+#ifdef DEBUG
+ if (!verify_same_context()) {
+ LOG(FATAL) << "Context mismatch!";
+ return;
+ }
+#endif
Review Comment:
`cudaSetDevice(gpu_id)` in the constructor only changes the CUDA device for
the thread that creates this singleton. These copy methods run on arbitrary
RPC/bthread worker threads and never select `gpu_id_`, so a non-default
`--gpu_index` (or a worker with another current device) can make
`cudaMemcpyAsync` use the wrong context; the DEBUG-only check will instead
abort. Select the configured device in each operation (and define how the
caller's device is restored if that matters).
##########
src/butil/gpu/gpu_block_pool.cpp:
##########
@@ -0,0 +1,498 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#if BRPC_WITH_GDR
+
+#include <iostream>
+#include <chrono>
+#include <gflags/gflags.h>
+#include <sys/stat.h>
+#include "butil/fast_rand.h"
+#include "gpu_block_pool.h"
+namespace butil {
+namespace gdr {
+DEFINE_int32(gdr_block_size_kb, 512, "gdr block size in KB");
+DEFINE_int32(max_gdr_regions, 32, "max num of gdr regions");
+
+#define CHECK_CUDA(call) \
+do { \
+ auto _sts = (call); \
+ if (_sts != cudaSuccess) { \
+ LOG(FATAL) << " cuda error:" \
+ << (cudaGetErrorString(_sts)) << std::string(" at ") \
+ << __FILE__ << ": " << __LINE__; \
+ } \
+} while (0);
+
+size_t GetGdrBlockSize() {
+ return FLAGS_gdr_block_size_kb * 1024;
+}
+
+bool verify_same_context() {
+ static std::once_flag init_flag;
+ static int original_device = -1;
+
+ int current_device;
+ CHECK_CUDA(cudaGetDevice(¤t_device));
+
+ std::call_once(init_flag, [&]() {
+ original_device = current_device;
+ });
+
+ return (current_device == original_device);
+}
+
+void* get_gpu_mem(int gpu_id, int64_t gpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+ void *d_data;
+
+ LOG(INFO) << "try to alloc " << gpu_mem_size << " bytes from gpu " <<
gpu_id;
+
+ CHECK_CUDA(cudaMalloc(&d_data, gpu_mem_size));
+ cudaDeviceSynchronize();
+ return (void *)d_data;
+}
+
+void* get_cpu_mem(int gpu_id, int64_t cpu_mem_size) {
+ CHECK_CUDA(cudaSetDevice(gpu_id));
+
+ LOG(INFO) << "try to alloc " << cpu_mem_size << " bytes from gpu " <<
gpu_id << "on host";
+
+ void* mem = NULL;
+
+ CHECK_CUDA(cudaMallocHost(&mem, cpu_mem_size));
+
+ cudaDeviceSynchronize();
+
+ return mem;
+}
+
+
+BlockPoolAllocators* BlockPoolAllocators::GetInstance() {
+ return ::Singleton<BlockPoolAllocators,
+ LeakySingletonTraits<BlockPoolAllocators>>::get();
+}
+
+void BlockPoolAllocators::init(int gpu_id, ibv_pd* pd) {
+ LOG(INFO) << "set GPU BlockPoolAllocator for " << gpu_id;
+ size_t region_size = 1024LL * 1024 * 1024;
+ size_t block_size = FLAGS_gdr_block_size_kb * 1024;
+ gpu_mem_alloc = new BlockPoolAllocator(gpu_id, true, pd, block_size,
region_size);
+
+ region_size = 32LL * 1024 * 1024;
+ block_size = 512;
+ cpu_mem_alloc = new BlockPoolAllocator(gpu_id, false, pd, block_size,
region_size);
+
+ gpu_stream_pool = new GPUStreamPool(gpu_id);
+}
+
+bool InitGPUBlockPool(int gpu_id, ibv_pd* pd) {
+ BlockPoolAllocators::GetInstance()->init(gpu_id, pd);
+ return true;
+}
+
+class BlockHeaderList {
+ public:
+ BlockHeaderList() {
+ objects_.reserve(kMaxObjects);
+ }
+ virtual ~BlockHeaderList() {
+ for (size_t i = 0; i < objects_.size(); i++) {
+ delete objects_[i];
+ }
+ }
+
+ BlockHeader* New() {
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (!objects_.empty()) {
+ BlockHeader* result = objects_.back();
+ objects_.pop_back();
+ return result;
+ }
+ }
+ return new BlockHeader;
+ }
+ void Release(BlockHeader* obj) {
+ obj->Reset();
+ {
+ std::lock_guard<std::mutex> lock(mu_);
+ if (objects_.size() < kMaxObjects) {
+ objects_.push_back(obj);
+ return;
+ }
+ }
+ delete obj;
+ }
+
+ private:
+ static const int kMaxObjects = 100000;
+
+ std::mutex mu_;
+ std::vector<BlockHeader*> objects_;
+};
+
+static BlockHeaderList* get_bh_list() {
+ static BlockHeaderList* bh_list = new BlockHeaderList();
+ return bh_list;
+}
+
+BlockPoolAllocator::BlockPoolAllocator(int gpuId, bool onGpu, ibv_pd* ibvPd,
+ size_t blockSize, size_t regionSize) :
+ gpu_id(gpuId)
+ , on_gpu(onGpu)
+ , pd(ibvPd)
+ , BLOCK_SIZE(std::max(blockSize, sizeof(BlockHeader)))
+ , REGION_SIZE((regionSize / blockSize) * blockSize) // 对齐到块大小的倍数
+ , freeList(nullptr)
+ , g_region_num(0)
+ , totalAllocated(0)
+ , totalDeallocated(0)
+ , peakUsage(0) {
+ g_regions.resize(FLAGS_max_gdr_regions);
+ LOG(INFO) << "Memory Pool initialized: block_size=" << BLOCK_SIZE
+ << ", region_size=" << REGION_SIZE << ", max_gdr_regions=" <<
FLAGS_max_gdr_regions
+ << ", gpu_id=" << gpu_id << ", on_gpu=" << on_gpu << ", pd=" << pd;
+
+ // The very first region is required at construction time; failure here is
+ // fatal because the pool would be unusable.
+ if (extendRegion() != 0) {
+ LOG(FATAL) << "Failed to initialize GDR memory pool (first region) on
gpu "
+ << gpu_id << ", on_gpu " << on_gpu;
+ }
+}
+
+BlockPoolAllocator::~BlockPoolAllocator() {
+#ifdef DEBUG
+ printStatistics();
+#endif
+
+ for (int i = 0; i < FLAGS_max_gdr_regions; i++) {
+ Region* r = &g_regions[i];
+ if (!r->mr) {
+ continue;
+ }
+
+ LOG(INFO) << "try to free " << r->size << " bytes from gpu " << gpu_id
<< ", on_gpu " << on_gpu;
+ ibv_dereg_mr(r->mr);
+ if (on_gpu) {
+ CHECK_CUDA(cudaFree(reinterpret_cast<void*>(r->start)));
+ } else {
+ CHECK_CUDA(cudaFreeHost(reinterpret_cast<void*>(r->start)));
+ }
+ }
+}
+
+Region* BlockPoolAllocator::GetRegion(const void* buf) {
+ if (!buf) {
+ errno = EINVAL;
+ return NULL;
+ }
+ Region* r = NULL;
+ uintptr_t addr = (uintptr_t)buf;
+ for (int i = 0; i < FLAGS_max_gdr_regions; ++i) {
+ if (g_regions[i].aligned_start == 0) {
+ break;
+ }
+ if (addr >= g_regions[i].aligned_start &&
+ addr < g_regions[i].aligned_start + g_regions[i].aligned_size)
{
+ r = &g_regions[i];
+ break;
+ }
+ }
+ return r;
Review Comment:
`AllocateRaw` publishes a newly extended region while holding `poolMutex`,
but `get_lkey` calls `GetRegion` after that lock is released. `GetRegion` reads
`g_regions` without synchronization while another allocator may initialize the
next entry, so it can race on these non-atomic fields or observe
`aligned_start` before `mr` and report a zero lkey/fatal. Protect region
lookup/publication with the pool lock or use a synchronized immutable
publication scheme.
##########
example/rdma_performance/client.cpp:
##########
@@ -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,
Review Comment:
The attachment is allocated on GPU 0 unconditionally, while the GDR pool and
RDMA receive buffers use `FLAGS_gpu_index`. Selecting any other GPU therefore
registers a pointer from the wrong device and can make the send or CUDA
operations fail. Use the configured RDMA GPU index here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
