diff --git a/BUILD.bazel b/BUILD.bazel index 727af8574a..7db434fbbb 100644 --- a/BUILD.bazel +++ b/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_ubring": ["BRPC_WITH_UBRING=1"], "//conditions:default": [], @@ -100,6 +105,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": [], @@ -251,6 +262,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/CMakeLists.txt b/CMakeLists.txt index 9419ac3a76..72828be809 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) @@ -121,6 +122,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") @@ -160,6 +172,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} @@ -318,6 +331,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") @@ -370,6 +411,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}) @@ -547,6 +593,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/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 e9d07eb73a..ffff739ee1 100644 --- a/bazel/config/BUILD.bazel +++ b/bazel/config/BUILD.bazel @@ -110,6 +110,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"}, @@ -161,4 +167,4 @@ config_setting( name = "brpc_with_ubring", define_values = {"BRPC_WITH_UBRING": "true"}, visibility = ["//visibility:public"], -) \ No newline at end of file +) diff --git a/config.h.in b/config.h.in index d8de111be9..c5a26c900b 100644 --- a/config.h.in +++ b/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 diff --git a/config_brpc.sh b/config_brpc.sh index 85692de3cb..0c247656c4 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,with-cpu-frequency,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,with-cpu-frequency,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 @@ -90,6 +91,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 ;; @@ -538,6 +540,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 @@ -674,6 +716,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..88ae687180 --- /dev/null +++ b/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； diff --git a/docs/en/gdr.md b/docs/en/gdr.md new file mode 100644 index 0000000000..a03d42bc28 --- /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. 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. diff --git a/example/BUILD.bazel b/example/BUILD.bazel index 05eda2080b..d458eb2e90 100644 --- a/example/BUILD.bazel +++ b/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", ], -) \ No newline at end of file +) 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 c112f1c330..b129a97f60 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/example/rdma_performance/client.cpp b/example/rdma_performance/client.cpp index 8c07eca236..0f484a71e9 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 2f138b30ea..4d54b0c4b2 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 UBRING, 3 GDR SocketMode _socket_mode; // Acceptor belongs to this tag diff --git a/src/brpc/channel.cpp b/src/brpc/channel.cpp index 01d8637564..219e36b516 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 b74fb1805a..8000f18ec8 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) { @@ -802,7 +808,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); } @@ -981,17 +995,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..78b18f50c7 --- /dev/null +++ b/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 // 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!!!"; + return MakeParseError(PARSE_ERROR_TRY_OTHERS); + } + butil::gdr::BlockPoolAllocator* host_allocator = butil::gdr::BlockPoolAllocators::GetInstance()->get_cpu_allocator(); + prefetch_d2h_data = host_allocator->AllocateRaw(prefetch_d2h_size); + if (prefetch_d2h_data == NULL) { + // Pool exhausted or single-block limit hit: ask the caller to retry + // instead of crashing the whole server. + LOG(ERROR) << "alloc host data for prefetch d2h failed, size=" + << prefetch_d2h_size; + return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA); + } + + // n is the bytes we real frefetch, n maybe less than prefetch_d2h_size; + n = source->copy_from_gpu(prefetch_d2h_data, prefetch_d2h_size); + size_t copy_size = n > 12 ? 12 : n; + memcpy(header_buf, prefetch_d2h_data, copy_size); + + do { + if (n >= 4) { + void* dummy = header_buf; + if (*(const uint32_t*)dummy != *(const uint32_t*)"PRPC") { + pe = PARSE_ERROR_TRY_OTHERS; + break; + } + } else { + if (memcmp(header_buf, "PRPC", n) != 0) { + pe = PARSE_ERROR_TRY_OTHERS; + break; + } + } + if (n < sizeof(header_buf)) { + pe = PARSE_ERROR_NOT_ENOUGH_DATA; + break; + } + butil::RawUnpacker(header_buf + 4).unpack32(body_size).unpack32(meta_size); + if (body_size > FLAGS_max_body_size) { + // We need this log to report the body_size to give users some clues + // which is not printed in InputMessenger. + LOG(ERROR) << "body_size=" << body_size << " from " + << socket->remote_side() << " is too large"; + pe = PARSE_ERROR_TOO_BIG_DATA; + break; + } else if (source->length() < sizeof(header_buf) + body_size) { + pe = PARSE_ERROR_NOT_ENOUGH_DATA; + break; + } + if (meta_size > body_size) { + LOG(ERROR) << "meta_size=" << meta_size << " is bigger than body_size=" + << body_size; + // Pop the message + source->pop_front(sizeof(header_buf) + body_size); + pe = PARSE_ERROR_TRY_OTHERS; + break; + } + } while (0); + + if (pe != PARSE_OK) { + host_allocator->DeallocateRaw(prefetch_d2h_data); + return MakeParseError(pe); + } + + source->pop_front(sizeof(header_buf)); + MostCommonMessage* msg = MostCommonMessage::Get(); + + if (header_size + meta_size <= n) { + auto deleter = [host_allocator, prefetch_d2h_data](void* data) { host_allocator->DeallocateRaw(prefetch_d2h_data); }; + // `n` is the number of bytes actually prefetched from GPU. It is stored + // in the meta's `data_meta' so that FillReqBufGpu/FillResBufGpu can + // decide whether the prefetched host buffer also covers the body (and + // thus avoid a second D2H copy). This block is host memory (prefetched), + // NOT GPU memory, so we deliberately use append_user_data_with_meta + // (not append_user_data_gpu) to keep is_gpu_memory() returning false. + msg->meta.append_user_data_with_meta((char*)prefetch_d2h_data + header_size, meta_size, deleter, n); + source->pop_front(meta_size); + } else { + host_allocator->DeallocateRaw(prefetch_d2h_data); + source->cutn_from_gpu(&msg->meta, meta_size); + } + source->cutn(&msg->payload, body_size - meta_size); + return MakeMessage(msg); +} + + +void FillReqBufGpu(butil::IOBuf* req_buf, MostCommonMessage* msg, int body_without_attachment_size) { + int meta_size = msg->meta.size(); + bool is_gpu_memory = msg->payload.is_gpu_memory(); + if (!is_gpu_memory) { + LOG(FATAL) << "message is not on gpu!!!"; + } + 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 fd70fb5c2b..09b1a2ac69 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" @@ -63,7 +66,7 @@ DEFINE_bool(rdma_recv_zerocopy, true, "Enable zerocopy for receive side"); DEFINE_int32(rdma_zerocopy_min_size, 512, "The minimal size for receive zerocopy"); DEFINE_int32(rdma_cqe_poll_once, 32, "The maximum of cqe number polled once."); DEFINE_int32(rdma_prepared_qp_size, 128, "SQ and RQ size for prepared QP."); -DEFINE_int32(rdma_prepared_qp_cnt, 1024, "Initial count of prepared QP."); +DEFINE_int32(rdma_prepared_qp_cnt, 256, "Initial count of prepared QP."); 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."); @@ -80,6 +83,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; @@ -109,10 +115,11 @@ RdmaResource::~RdmaResource() { } } -RdmaEndpoint::RdmaEndpoint(Socket* s) +RdmaEndpoint::RdmaEndpoint(Socket* s, bool use_gdr) : _socket(s) , _state(UNINIT) , _handshake_version(0) + , _use_gdr(use_gdr) , _resource(nullptr) , _send_cq_events(0) , _recv_cq_events(0) @@ -917,8 +924,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) { @@ -979,26 +993,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(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; + } 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) { @@ -1641,6 +1712,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 388e31d78e..0106c9c44b 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(); @@ -214,6 +219,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) @@ -287,6 +302,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 38bd58cac4..741dd94ff7 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 = nullptr; static ibv_context* g_context = nullptr; @@ -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(); + } + 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, @@ -596,6 +620,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) { @@ -689,6 +721,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 ee5151c3a5..40b0766ebe 100644 --- a/src/brpc/rdma_transport.cpp +++ b/src/brpc/rdma_transport.cpp @@ -30,10 +30,19 @@ DECLARE_bool(usercode_in_pthread); extern SocketVarsCollector *g_vars; -void RdmaTransport::Init(Socket *socket, const SocketOptions &options) { +void RdmaTransport::DoInit(Socket *socket, const SocketOptions &options, bool use_gdr) { CHECK(_rdma_ep == nullptr); - if (options.socket_mode == SOCKET_MODE_RDMA) { - _rdma_ep = new rdma::RdmaEndpoint(socket); + // gdr mode is a special mode of rdma mode. + // both rdma mode and gdr mode need init rdma::RdmaEndpoint. + if (options.socket_mode == SOCKET_MODE_RDMA || + options.socket_mode == SOCKET_MODE_GDR) { + _rdma_ep = new(std::nothrow)rdma::RdmaEndpoint(socket, use_gdr); + if (!_rdma_ep) { + const int saved_errno = errno; + PLOG(ERROR) << "Fail to create RdmaEndpoint"; + socket->SetFailed( + saved_errno, "Fail to create RdmaEndpoint: %s", berror(saved_errno)); + } _rdma_state = RDMA_UNKNOWN; } else { _rdma_state = RDMA_OFF; @@ -58,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 1d78fbb430..95889a8c0e 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..3b6177cc10 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_UBRING = 2, + SOCKET_MODE_GDR = 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..2e6f11ba62 --- /dev/null +++ b/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 +#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 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>::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 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; + + // 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(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) { + // 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 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(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) + << "%"; +} + +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(ptr)); + } else { + cudaFreeHost(reinterpret_cast(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(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"; + 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& 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) { +#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(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..c61e345b8f --- /dev/null +++ b/src/butil/gpu/gpu_block_pool.h @@ -0,0 +1,181 @@ +// 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 +#include "butil/containers/hash_tables.h" +#include "butil/logging.h" +#include "butil/memory/singleton.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: + // 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 { +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: + // Returns the process-wide leaky singleton instance. Created on first + // access via butil::Singleton and intentionally never destroyed (GPU/RDMA + // resources outlive normal static destructors). + static BlockPoolAllocators* GetInstance(); + + BlockPoolAllocators() {} + ~BlockPoolAllocators() {} + + 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; + } + +private: + // Allow DefaultSingletonTraits::New() to construct the instance. + friend struct DefaultSingletonTraits; + + 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 685a7a4982..edc28806f8 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::GetInstance()->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) { @@ -1190,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; @@ -1201,6 +1275,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 nullptr; + } + 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 +1446,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::GetInstance()->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 +1632,21 @@ bool IOBuf::equals(const butil::IOBuf& other) const { return true; } +#if BRPC_WITH_GDR +// 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() { + if (_ref_num() == 0) { + return false; + } + IOBuf::BlockRef const& r = _ref_at(0); + return (r.block->flags & IOBUF_BLOCK_FLAGS_GPU_MEMORY) != 0; +} +#endif + ////////////////////////////// IOPortal ////////////////// IOPortal::~IOPortal() { return_cached_blocks(); } diff --git a/src/butil/iobuf.h b/src/butil/iobuf.h index 724ee02d2e..8aec0437ea 100644 --- a/src/butil/iobuf.h +++ b/src/butil/iobuf.h @@ -146,6 +146,17 @@ 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; + // 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 + // Cut off 1 byte from the front side and set to *c // Return true on cut, false otherwise. bool cut1(void* c); @@ -260,11 +271,22 @@ 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. 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. diff --git a/src/butil/iobuf_inl.h b/src/butil/iobuf_inl.h index 71a449d282..cd781024d1 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);