support gpu direct rdma

BUILD.bazel

@@ -51,6 +51,11 @@ DEFINES = [
  }) + select({
      "//bazel/config:brpc_with_rdma": ["BRPC_WITH_RDMA=1"],
      "//conditions:default": [],
+}) + select({
+    # 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_urma": ["BRPC_WITH_URMA=1"],
      "//conditions:default": [],
@@ -103,6 +108,12 @@ LINKOPTS = [
         "-libverbs",
     ],
     "//conditions:default": [],
+}) + select({
+    "//bazel/config:brpc_with_gdr": [
+        "-lcuda",
+        "-lcudart",
+    ],
+    "//conditions:default": [],
 }) + select({
         "//bazel/config:brpc_with_asan": ["-fsanitize=address"],
         "//conditions:default": [],
@@ -254,6 +265,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",

CMakeLists.txt

@@ -33,6 +33,7 @@ set_property(CACHE URMA_USE_MOCK PROPERTY STRINGS AUTO ON OFF)
 option(DOWNLOAD_URMA_HEADERS
        "Download UMDK headers when WITH_URMA is enabled and headers are absent"
        ON)
+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)
@@ -139,6 +140,17 @@ if(NOT URMA_USE_MOCK_MODE STREQUAL "AUTO" AND
         "URMA_USE_MOCK must be AUTO, ON, or OFF, got: ${URMA_USE_MOCK}")
 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")
@@ -179,6 +191,7 @@ list(APPEND BRPC_COMMON_DEFINITIONS
     BRPC_WITH_GLOG=${WITH_GLOG_VAL}
     BRPC_WITH_RDMA=${WITH_RDMA_VAL}
     BRPC_WITH_URMA=${WITH_URMA_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}
@@ -397,6 +410,34 @@ if(WITH_URMA)
     endif()
 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")
@@ -462,6 +503,11 @@ if(WITH_URMA)
     endif()
 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})
@@ -643,6 +689,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

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 \

bazel/config/BUILD.bazel

@@ -134,6 +134,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"},

config.h.in

@@ -21,6 +21,12 @@
 #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@
+
 #ifdef BUTIL_USE_CPU_FREQUENCY
 #undef BUTIL_USE_CPU_FREQUENCY
 #endif

config_brpc.sh

@@ -54,12 +54,13 @@ else
     LDD=ldd
 fi
 
-TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-urma,with-urma-mock,without-urma-mock,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"`
+TEMP=`getopt -o v: --long headers:,libs:,cc:,cxx:,with-glog,with-thrift,with-rdma,with-urma,with-urma-mock,without-urma-mock,with-gdr,with-mesalink,with-bthread-tracer,with-debug-bthread-sche-safety,with-debug-lock,with-asan,with-riscv-zvbc,with-riscv-zbc,with-cpu-frequency,nodebugsymbols,werror -n 'config_brpc' -- "$@"`
 WITH_GLOG=0
 WITH_THRIFT=0
 WITH_RDMA=0
 WITH_URMA=0
 URMA_MOCK_MODE=auto
+WITH_GDR=0
 WITH_MESALINK=0
 WITH_BTHREAD_TRACER=0
 WITH_ASAN=0
@@ -95,6 +96,7 @@ while true; do
         --with-urma) WITH_URMA=1; shift 1 ;;
         --with-urma-mock) URMA_MOCK_MODE=on; shift 1 ;;
         --without-urma-mock) URMA_MOCK_MODE=off; 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 ;;
@@ -579,6 +581,46 @@ if [ $WITH_URMA != 0 ]; then
     fi
 fi
 
+if [ $WITH_GDR != 0 ]; then
+    # 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"
+
+    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
@@ -716,6 +758,7 @@ 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_URMA -ne 0 ]; then print_info "With URMA: 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

docs/cn/gdr.md

@@ -0,0 +1,44 @@
+# 编译
+
+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要求所有接收到的数据都写入post recv 时指定的 GPU 显存 buffer 中，
+而MLX5_SCATTER_TO_CQE打开时，小消息直接内联写入了内存，破坏了这个要求，从而会引发操作失败。
+
+
+# 参数
+
+可配置参数说明：
+* gdr_block_size_kb: 使用gdr传送数据时，block的大小（单位为KB），默认为512；
+* max_gdr_regions: gdr显存池所使用Region的最大个数，每个Region大小为1GB； 

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. This is because GPUDirect RDMA requires that all received data be written to the GPU memory buffer specified in the post‑receive operation. However, when MLX5_SCATTER_TO_CQE is enabled, small messages are directly inlined into host memory, which violates this requirement and consequently causes operation failures.
+
+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.

example/BUILD.bazel

@@ -37,6 +37,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": [""],
 }) + select({
     "//bazel/config:brpc_with_ubring": ["-DBRPC_WITH_UBRING=1"],
     "//conditions:default": [""],
@@ -162,4 +165,4 @@ cc_binary(
     deps = [
         "//:brpc",
     ],
-)
+)

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})