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_debug_bthread_sche_safety": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=1"],
      "//conditions:default": ["BRPC_DEBUG_BTHREAD_SCHE_SAFETY=0"],
@@ -94,6 +99,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": [],
@@ -237,6 +248,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

@@ -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

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

@@ -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"],
-)
+)

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

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,46 @@ if [ $WITH_RDMA != 0 ]; then
     append_to_output "WITH_RDMA=1"
 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
@@ -667,6 +709,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

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； 

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.

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",
     ],
-)
+)

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