基于“嵩山”超級計算機的UCX 庫分析與優化

劉 康，萬 偉，劉 波，李俊宏，李 柱

（鄭州大學 計算機與人工智能學院，鄭州 450001）

0 概述

“嵩山”超級計算機部署于國家超級計算鄭州中心，是我國國產自研的新一代超高性能計算平臺，它采用32 核CPU+國產加速器的異構計算架構、InfiniBand（直譯為“無限帶寬”技術，縮寫為IB）［1-2］超高速網絡以及高性能分布式并行存儲系統，理論峰值算力可達100 PFlops，整機實測性能達到65 PFlops。InfiniBand 是一個用于高性能計算平臺的計算機網絡通信標準，具有極高的吞吐量和極低的延遲，用于計算機與計算機之間的數據互連，也用作服務器與存儲的互連以及存儲系統之間的互連。InfiniBand搭載的Mellanox ConnectX-6 網卡提供了最高200 Gb/s帶寬的數據傳輸性能。

“嵩山”計算平臺所搭載的主要通信框架為UCX（Unified Communication X）［3-4］。UCX 作為底層InfiniBand 網絡和上層并行編程模型的通信中間件，定義了一組統一的標準化通信編程接口，以滿足主流的并行編程模型，如MPI（Message Passing Interface）［5］、UPC（Unified Parallel C）［6］、PGAS（Partitioned Global Address Space）［7］等的需求，同時又可以在各種高性能平臺上實現，以便在互聯網絡上更好地滿足高性能、可移植、可伸縮等并行應用的開發需求［8-10］。

但是，由于RDMA（Remote Direct Memory Access）系統具有復雜性，因此存在很多未知的問題［11-12］，UCX 作為“嵩山”RDMA 系統中的通信框架，在“嵩山”特色互聯網絡架構上還有一定的優化空間。在存在復雜通信環境的集合通信中，通信有時會成為瓶頸而拖累了整體計算速度。UCX 的通信性能直接影響了上層并行編程模型的數據傳輸與計算性能。因此，在“嵩山”超級計算平臺上對UCX 進行研究與優化具有重要的工程意義。

本文基于“嵩山”超級計算平臺，以MPI 為例，使用osu_benchmark 測試工具［13］在不同的傳輸下進行多種集合通信測試，獲得各種情形下的延遲與帶寬數據，以發現節點設備存在的瓶頸。同時，對UCX的代碼進行優化，以解決節點內通信占用網卡資源的問題。在此基礎上，實現UCX 在“嵩山”超級計算平臺上的最優傳輸選擇，以提升平臺的集合通信能力以及整體的計算性能。

1 “嵩山”超級計算機互聯網絡

1.1 InfiniBand 互聯網絡

InfiniBand 是一種高速互聯網絡，用于連接大型集群和超級計算機，它是目前應用最廣泛的高速互連網絡之一，2016 年在Top500 互連網絡中就已經達到37.5%的份額［14］。InfiniBand 網絡為通信提供了雙邊（發送-接收）和單邊（RDMA）語義。InfiniBand上的通信使用隊列對（QP）模型，其中，發送和接收隊列分別用于發送和接收消息，工作請求被提交到這些隊列中，硬件可以在隊列中讀取工作請求以執行通信。此外，將隊列與每個QP 相關聯，用于通知通信完成。在通過InfiniBand 進行通信時，需要注意注冊硬件訪問的所有內存區域。為了減少內存注冊的開銷，短消息（Short）可以內聯到工作請求中，而較大的消息可以利用零拷貝（ZCopy）協議。這種策略意味著工作請求只獲取內存緩沖區的描述信息，然后直接從緩沖區讀取數據，而不需要CPU 的參與。

當前InfiniBand 結構實現了各種傳輸機制［15-16］，最常見的是RC（面向連接的可靠連接）和UD（無連接的不可靠數據報），后者只實現了雙邊通信語義。此外，UD 一次只能傳輸一個MTU 的數據（通常是4 KB），而RC 通常能提供比UD 更高的帶寬和更低的延遲，代價是RC 對資源有很高的要求：要完全連接N個進程；要求每個進程有O（N2）個連接和O（N）個隊列對。而UD 是無連接的，因此，每個進程只需要一個UD 隊列對。為了減少RC 的內存消耗，InfiniBand 規范引入了共享接受隊列以及擴展的RC傳輸。Mellanox 后續又推出了動態連接（DC）傳輸服務［17］，該服務動態地創建和銷毀連接，將內存消耗限制在接近UD 的級別，同時提供與RC 類似的內存語義。然而，DC 的可擴展性設計是以性能為代價的，主要是因為存在連接事務的開銷［18］。

InfiniBand 的用戶 空間接口是Verbs API［19］，它是一個帶有OFED 堆棧的用戶級別的庫，位于內核級Verbs API 之上。內核API 與特定供應商的InfiniBand 驅動程序和驅動程序庫協作，以實現InfiniBand 硬件訪問。InfiniBand 軟件棧示意圖如圖1 所示。

圖1 InfiniBand 軟件棧和UCXFig.1 InfiniBand software stack and UCX

1.2 Socket Direct 技術

隨著數據量的指數級增長，使用者對服務器和計算資源提出了更高的性能要求，以便對海量數據進行實時分析。“嵩山”超級計算平臺采用非均勻內存訪問（NUMA）架構［20］，每個處理器擁有4 個CPU Die，在每個Die 內集成了8 個物理核心，共計32 個物理核心。

在一般情況下，數據主要依靠GMI 總線來進行跨Die 傳輸，如圖2（a）所示，此時Die3/Die4 中的數據如果要傳輸至網卡，則網絡流量需要使用GMI 總線經過Die0/Die1 然后再流入網卡。而“嵩山”平臺的網絡架構支持Socket Direct 技術，該技術可以使節點中的每個Die（NUMA node）都可以通過其專用的PCIe 接口直接連接到網絡，使得網絡流量無須遍歷內部總線（GMI）和其他Die，如圖2（b）所示。Socket Direct 不僅降低了CPU 的利用率、增加了網絡吞吐量，還顯著降低了開銷與延遲，從而提高了服務器的性能。

圖2 跨Die 傳輸的對比Fig.2 Comparison of cross-Die transmission

1.3 UCX 設計

隨著DPU 的普及以及各類DSA 芯片的廣泛使用［21］，如何在這之上抽象出統一的內存訪問語義和統一的通信方式成為一個值得研究的問題，因此，UCX 應運而生。UCX 可以在通信方面實現低級別的軟件開銷，并且提供接近原生級別的性能。UCX旨在提供一個統一的抽象通信接口，能夠適配任何通信設備，并支持各種應用的需求，從而滿足當前高性能、可移植且穩定可靠的并行應用開發需求，同時還能通過持續的迭代更新來適應未來的高速互聯網絡。

從圖1 可以看出UCX 軟件堆棧是如何放置在InfiniBand 之上的，UCX 由下層UCT 和上層UCP 這2 層組成。下文將介紹UCX 框架，討論UCP 和UCT這2 個層之間的主要區別以及UCX 內部最重要的語義。

1.3.1 UCX 框架

UCX 利用高速網絡進行節點間通信，并利用共享內存機制進行有效的節點內通信。UCX 總體采用分層結構公開一組抽象通信原語，這些原語充分利用了可用的硬件資源和負載，其中包括RDMA［22-23］（InfiniBand 和RoCE）、TCP、共享內存和網絡原子操作。圖3 顯示了UCX 的軟件棧結構。

圖3 UCX 軟件棧結構Fig.3 UCX software stack structure

UCX 通過提供高級API 促進快速開發，屏蔽低層細節，同時保持高性能和可伸縮性，其框架主要由3 個組件組成，即UCS（UC-Services）、UCT（UCTransports）和UCP（UC-Protocols）。每一 個組件都導出一個公共API，可以作為一個獨立的庫使用。底層的UCT 適配各種通信設備，上層的UCP 則是在UCT 不同設備的基礎上封裝更抽象的通信接口，以方便使用。

UCT 是傳輸層，它抽象了各種硬件架構之間的差異，并提供了一個支持通信協議實現的低級API，從單機的共享內存到常用的TCP Socket 以及“嵩山”超算底層的InfiniBand 協議，都有很好的支持。該層的主要目標是提供對硬件網絡功能的直接有效訪問，為此，UCT 依賴供應商提供的低級驅動程序，如InfiniBand Verbs、Cray 的uGNI 等。此外，該層還提供用于通信上下文管理（基于線程和應用程序級別）以及分配和管理的構造。在通信API 方面，UCT 定義了立即（short）、緩沖復制、發送（BCopy）和零拷貝（ZCopy）等通信操作的接口。

UCP 是協議層，通過使用UCT 層公開的較低級別功能來實現上層高級編程模型（如MPI、UPC、PGAS）所使用的較高級別協議。UCP 提供的功能是能夠為通信選擇不同的傳輸、消息分段、多軌通信以及初始化和完成庫。目前，API 具有的接口類別包括初始化、遠程內存訪問（RMA）通信、原子內存操作（AMO）、活動消息（Active Message）、標簽匹配（Tag-Matching）和集合（Collectives）。

1.3.2 UCX 語義

UCX 提供的最主要語義包括通信上下文、通信原語、通信實體和連接建立。這4 種語義詳細敘述如下：

1）通信上下文。UCP 和UCT 的最主要區別在于通信上下文。UCT 被設計成一個位于單個通信設備和傳輸層之上的通信層，而UCP 可以讓用戶操作不同的設備和傳輸層。因此，UCT 在設備（如InfiniBand、共享內存SM）上定義了一個內存域，用來分配和注冊進行通信的內存以及特定設備（如InfiniBand 上的UD 和RC）上的特定傳輸接口。內存域和接口都有一組它們自己的屬性，這些屬性來自于硬件功能。內存域屬性包括內存分配限制和內存訪問的憑據，接口屬性包括傳輸機制的通信和連接能力以及協議切換的閾值。UCP 將這些多個UCT內存域和接口封裝在單個通信上下文中，并根據硬件屬性和性能指標選擇適合通信操作的接口。

3）通信實體。Worker 是UCX、UCP 和UCT 的核心通信實體。Worker 的主要特征是有自己的進度引擎（progress engine），進度引擎會在所有打開的接口上強制執行當前的進度。在Worker 要啟用與另一個進程的通信時，每個進程都會創建一個端點（endpoint），并將其連接到遠程進程的endpoint。UCT endpoint 與特定接口（如UD、RC）綁定，即每個使用的接口對應一個UCT endpoint，而UCP endpoint擁有多個UCT endpoint。因此，在UCP 中，endpoint始終連接著2 個Worker。在內部，UCP 負責從可用于執行通信操作的接口/UCT endpoint 中選擇最佳的接口/UCT endpoint。

4）連接建立。當UCP 的Worker 創建endpoint時，UCP 層為每種類型的操作選擇一個或多個接口，并且在每個接口上創建并對應一個UCT endpoint，所有的這些UCT endpoint 都與父UCP endpoint 相關聯。如果一個接口對應無連接的傳輸，那么它可以立即連接到遠程接口，這也就是UCP 中發生的情況，即UCT endpoint 通過無連接傳輸立即建立連接。但是，如果接口對應P2P 的傳輸，UCP 將創建一個stub endpoint。Wireup UCT endpoint 始終是無連接的，通過立即發送Wireup 請求然后通過P2P 傳輸以實現所有UCT endpoint 的連接。當父UCP endpoint 的所有UCT endpoint 都已連接時，stub endpoint 即被銷 毀。

2 UCX 在“嵩山”中的優化

2.1 參數調優

UCX 可以適配多種設備、系統、架構等，因而具有繁雜的參數設置，調整各項參數可以使UCX 更加適配“嵩山”平臺。

“嵩山”平臺的高速網絡使用Socket Direct技術劃分CPU 為4 個NUMA nodes 并分別連接至4 塊網卡設備，實現各個Die 與網卡的直連。在UCX 中設置UCX_MAX_RNDV_LANES 為4，為Rendezvous 協 議開啟4 端口的多軌傳輸，使用多塊網卡同時進行傳輸，從而提升數據的傳輸效率。

“嵩山”平臺CPU 使用的NUMA 架構，在PCIe總線傳輸中更適合采用寬松排序（relaxed order）的事務排序方法，即允許PCIe 交換開關，將軟件確認過的事務重排在其他事務之前發送，這樣既提升了PCIe 總線效率，又能保證程序如期執行。本文使用UCX_IB_PCI_RELAXED_ORDERING=on 在UCX 中開啟寬松排序，使得所有使用UCX 通信庫的程序采用寬松排序，從而獲得更高的性能。

2.2 網卡占用優化

在使用UCX 進行節點內部通信時，進程間通信不僅會使用共享內存傳輸，還會調用網卡設備共同完成數據傳輸，這是由于ITIGIN［4］為UCX 添加了實現，即進行進程間通信時rc 會輔助共享通信，從而共同完成通信。但是在“嵩山”平臺上，實測IB 網卡對多進程的通信支持相較于共享內存并不友好，多進程會平分網卡帶寬，導致整體性能下降。而在進行大規模節點運算時，網卡是節點間通信的主力設備，應該盡可能地保證網卡用于跨節點傳輸。因此，需要對UCX 的傳輸邏輯進行修改，使其在進行節點間通信時不使用網卡。

以MPI 為例，在其進行通信時，UCX 會調用遠程內存訪問（RMA）［24］和活動消息（AM）等操作來實現快速的節點間通信。在涉及進程間通信時，UCX同樣也會選擇網卡來調用這些操作進行傳輸。

在“嵩山”超級計算平臺中，MPI 節點內通信使用的設備有memory 和mlx5。memory 會調用am、am_bw 和 rma_bw 操 作；mlx5 網卡會調用am_bw 和rma_bw 操作。因此，mlx5 網卡在節點內通信時所調用的操作完全可以被memory 所取代。程序調用am_bw 和rma_bw 操作之 前，UCT 會執行ucp_wireup_add_bw_lanes 函數，選出合適的傳輸，以此建立支持相應功能的endpoint。對此函數進行分析可以發現，函數調用ucp_wireup_select_transport，根據bitmap 選出支持am_bw 或rma_bw 的傳輸，隨后函數將傳輸放入ep 的配置文件中，等待endpoint 的創建，這些操作都發生在連接的準備階段。

在實驗確定的最佳色譜條件下，選取1#果酒樣品，分別加入10，50，100 mg/L標準混合溶液，平行進行6次實驗，實驗結果見表3。回收率為81.6%～102.8%，相對平均偏差不大于4.4%，說明方法精密度高，準確度好。

本文在ucp_wireup_add_bw_lanes 函數中添加判斷。在函數循環搜尋可用傳輸并將其添加到設備的dev_bitmap 后，讀取出Worker 的上下文信息，判斷此時的設備是否為共享內存傳輸：如果是共享內存傳輸，則break 跳出循環，不再搜索額外的傳輸；如果不是共享內存傳輸，則正常循環，搜尋可用傳輸。修改后的程序流程如圖4 所示，可以這樣做的原因是搜索過程中的設備次序memory 排在mlx5 網卡之前，當選擇出共享內存傳輸時，函數退出，便不會檢索到mlx5 網卡，從而在進行節點內的進程間通信時只使用共享內存通信而不是網卡傳輸。

圖4 優化后的程序流程Fig.4 Optimized program procedure

使用osu_benchmark 以4 MB 包長測得單節點內2～32 進程的alltoall 集合通信延遲數據，如圖5 所示，其中，before 是優化前同時使用網卡rc 傳輸和共享內存傳輸的通信數據，after 是僅使用共享內存傳輸的通信數據。從圖5 可以看出，在節點內通信時，隨著PPN（Processes Per Node）的增加，2 種傳輸方式的延遲差距愈加明顯，優化后的通信延遲相較于優化前最多降低了70%。

圖5 節點內alltoall 測試結果Fig.5 Intra-node alltoall test results

測試不同PPN 下的點對點通信帶寬數據，如圖6所示。從中可以明顯看出，在PPN 大于8 時，僅使用共享內存傳輸的通信性能優于使用IB 網卡的通信性能，此外，如果節點內的進程間通信使用了網卡傳輸，在絕大多數情況下，網卡通信還會對本來的進程間通信造成負面影響，降低整體的通信帶寬性能。

圖6 節點內點對點帶寬測試結果Fig.6 Intra-node p2p bandwidth test results

2.3 共享內存通信選擇優化

“嵩山”平臺的MPI 庫中存在不同的節點內通信機制，圖7 主要展示了其中的2 種。傳統的雙副本拷貝的共享內存實現，其傳輸數據涉及一個共享的緩沖區空間，由本地進程來交換消息，如圖7（a）所示。但是，這種方式僅適用于小消息，對于較大的消息，雙副本拷貝會給CPU 帶來額外的負擔，導致緩存污染和帶寬的浪費。圖7（b）展示了支持內核輔助的共享內存傳輸實現［25］，傳輸實現依靠內核的援助，內核模塊可以為節點內通信提供單拷貝機制，在傳輸較大消息時會大幅提升傳輸效率。

圖7 2 種共享內存通信機制Fig.7 Two shared memory communication mechanisms

對于第2 種通信方式，“嵩山”平臺的MPI 庫支持CMA 和KNEM［26］這2 種內核模塊。CMA 引入了2 個系統調用，分別是process_vm_readv 和process_vm_writev，它們根據進程的PID 和遠程虛擬地址直接讀寫另一個進程的內存［27］。對于使用KNEM 內核的通信，發送進程在KNEM 驅動中聲明一個發送緩沖區（不管是否連續），并將相應的標識符cookie傳遞給接收進程，接收進程接收到cookie 并請求KNEM 驅動從cookie 緩沖區復制到它的本地緩沖區（連續或非連續）［26］。

本文使用osu_benchmark 分別指定CMA 和KNEM 傳輸，測得2 種共享內存通信的帶寬與延遲如表1 所示。在“嵩山”超算平臺下對輸出的UCX 日志進行分析發現，節點在進行共享內存傳輸時，無論何種情況都只會選擇CMA 進行通信，并不會選擇帶寬更高、延遲更低的KNEM。

表1 CMA 與KNEM 的性能參數Table 1 Performance parameters of the CMA and KNEM

進程間在進行共享內存通信時，通過rma_bw 操作來進行高速的遠程內存訪問。在建立連接前，UCT 會根據UCX 提供的一套公式來計算傳輸評分，選出rma_bw 中評分最高的傳輸，添加到連接通道（lanes）中。

本文對rma_bw 操作傳輸選擇的評分機制進行分析。在UCT 中，計算評分時以256 KB 的消息大小為基準，調用rma_bw 操作的傳輸的評分為時間開銷的倒數，如式（1）所示：

設mcost為內存域注冊開銷，注冊開銷近似為一個線性函數，如式（2）所示：

其中：omd為固定開銷；ggrowth為增長系數；ssize為數據大小（256 KB）。設bl和br分別為本地和遠程的帶寬大小，因此，總開銷為256 KB 消息的傳輸時延、內存注冊開銷mcost與傳輸接口間延遲llr的累加，如下：

對于節點中的每個進程，其帶寬b在UCX 中的計算方式如式（4）所示：

其中：bdedicated為專用帶寬；bshared為共享帶寬。

對平臺的UCX 源代碼進行分析，可以發現：在UCX 1.9.0 的帶寬設置下，CMA 擁有11 145.00 MB的dedicated 帶 寬，而KNEM 是13 862 MB 的shared帶寬。根據srma的計算公式，在PPN 不為1 時，UCX在計算KNEM 和CMA 的rma_bw 評分時帶寬會存在巨大差異，從而導致永遠不會選擇KNEM 傳輸。這是因為早期KNEM 對多進程支持不如CMA，單進程時KNEM 會有更高的帶寬，但是存在多進程通信時，KNEM 性能將不如CMA，因而將KNEM 帶寬值設置為shared。但是，“嵩山”超級計算平臺所具有的優化KNEM 對多進程的支持極好，同時支持高性能單拷貝消息傳輸。因此，本文將KNEM 的帶寬從shared 改為dedicated，使KNEM 獲得了更合理的評分，從而在進行集合通信時共享內存方面的傳輸會更多地選擇帶寬更高、延遲更低的KNEM。

在大部分通信中，KNEM 和CMA 兩者差異較小，但是在涉及節點內進程間gather 通信時，KNEM內核相對CMA 內核有較為明顯的性能提升，并且隨著PPN 的增加提升效果愈加明顯，如圖8 所示。

圖8 節點內gather 測試結果Fig.8 Intra-node gather test results

3 數據測試與驗證

3.1 實驗環境

在“嵩山”超級計算機的固化節點上進行實驗，單個節點配置為32 核2.0 GHz CPU，網卡采用Mellanox ConnectX-6，以HDR 模式（200 Gb/s）工作。操作系統為CentOS 7.6，內核版本為3.10.0-957.el7.x86_64。

在本次測試中，使用的MPI 版本為Open MPI v4.0.4rc3，它在平臺的共享內存通信時支持KNEM和CMA 內核。使用的HPCX 版本為2.7.4，UCX 版本為UCX 1.9。對于點到點和集合通信測試，使用osu_benchmark v5.5 來測試并記錄通信性能數據。在性能測試對比數據中，before 的通信底層是目前在用的由ITIGIN 優化后的UCX 正式版本，after 采用本文優化后的UCX 庫。

3.2 節點內集合通信測試

首先測試優化后的UCX 庫在單節點內的通信表現。圖9～圖11 展示了單節點內不同PPN 下的4 MB 包長的MPI 集合通信測試數據，橫坐標為使用核心數（總進程數），縱坐標為平均通信延遲，每個核心綁定一個進程進行通信。從中可以看出，使用優化后UCX 的MPI 集合通信能力有了明顯提升，alltoall 的通信性能提升尤為明顯，延遲最多降至優化前的30%（圖9），gather 的通信延遲最多約降至優化前的55%（圖10），allreduce 的通信延遲最多約降至優化前的69%（圖11）。

圖9 優化前后節點內alltoall 測試結果Fig.9 Intra-node alltoall test results before and after optimization

圖10 優化前后節點內gather 測試結果Fig.10 Intra-node gather test results before and after optimization

圖11 優化前后節點內allreduce 測試結果Fig.11 Intra-node allreduce test results before and after optimization

3.3 節點間集合通信測試

對于節點間的集合通信，本文對2 個規模（32 節點和100 節點）進行測試。對于32 節點的規模，選取分屬lka 2 個交換機的32 個節點，每個交換機16 個節點，測試消息包長為1 MB。經過測試發現，在節點間集合通信時，其他集合通信測試效果與優化前一致，allgather通信產生了較為明顯的差異，如圖12 所示。

圖12 32 節點allgather 測試結果Fig.12 32 nodes allgather test results

本文在8 箱的節點中隨機選擇100 個節點，進行100 個節點間的集合通信測試，獲得節點間的2～18 PPN 下1 MB 包長的集合通信延遲數據。從圖13 可以看出，優化后的UCX 在allgather 集合通信中取得了極為明顯的優化效果，延遲最多可降至原來的20%，并且隨著進程的增多差距逐漸變大。其他的集合通信測試優化前后數據基本保持一致。

圖13 100 節點allgather 測試結果Fig.13 100 nodes allgather test results

4 結束語

“嵩山”超級計算平臺支持多種并行編程模型，對高速互聯網絡進行優化有助于提升平臺的整體通信性能，為平臺的并行編程模型提供良好的底層通信支持。本文對“嵩山”超級計算平臺上的節點進行測試，獲得了節點間與節點內的通信性能數據，并且發現IB 網卡在節點內多PPN 通信中存在的局限性。然后，對平臺的主要通信框架UCX 進行分析與優化，解決了節點內進程間通信占用網卡的問題，同時改善了UCX 對共享內存傳輸的選擇機制。優化后的UCX 對大PPN 下的節點間allgather 集合通信以及節點內的進程間集合通信性能提升效果明顯。

由于RDMA 具有復雜性，很多因素都可能影響RDMA 系統的整體通信性能。下一步將找出其他制約節點間通信速度的因素，對算法進行改進，使得節點間的其他集合通信能力得到加強。此外，UCX 根據PPN 來預測帶寬，依此帶寬來選擇傳輸，這種帶寬計算方法還不夠準確，未來將對此進行改進，從而改善UCX 的傳輸選擇評分機制。