UCC简介

发表于 更新于 分类于

Unified Collective Communication（UCC）是一个集体通信操作API和库，对于当前和新兴的编程模型和运行时来说，具有灵活、完整、功能丰富等特点。

Github地址：https://github.com/openucx/ucc

ucc_logo

设计目标、架构与发展规划

UCC目前的软件架构如下图所示，UCC使用尽可能接近底层的方式实现了集合通信，并向上层应用提供了调用的接口，在HPC、AI领域的多个框架中都有所应用。UCC对集合通信操作的实现分为两层，TL即Transport Layer，主要负责针对不同的底层通信协议实现集合通信操作最基本的功能，如UCP、NCCL、SHARED Memory等。CL层为Collective Layer，在TL层的基础上进一步丰富了其功能，如Hier组件实现了分层的集合通信操作。

ucc_components

UCC的设计目标如下：

  • 设计适用于HPC、AI/ML和I/O工作负载的高度可扩展和高性能的集体操作

  • 涵盖各种编程模型的非阻塞集体操作

  • 灵活的资源分配模式

  • Support for relaxed ordering model（不知道怎么翻译）

  • 灵活的同步模型

  • 重复的集体操作（启动一次，调用多次）

  • 硬件集合通信操作

目前的UCC发布到了1.0版本（22年4月正式发布），UCC框架前两个版本的规划路线如下图（貌似实际进度有一定放鸽子嫌疑）

image-20220816174429412

安装与使用

使用源码安装的方式编译即可，需要提前安装UCX通信库，若有NCCL、CUDA等环境可以通过configure脚本配置。

1
2
3
$ ./autogen.sh
$ ./configure --prefix=<ucc-install-path> --with-ucx=<ucx-install-path>
$ make

目前Open MPI已经支持使用UCC组件进行集合通信，使用方法很简单，在Open MPI安装时通过configure脚本配置UCC库路径，运行MPI程序时加上MCA参数--mca coll_ucc_enable 1 --mca coll_ucc_priority 100即可。

程序中直接使用UCC集合通信API

与RDMA编程类似，UCC在建立通信上下文时需要以OOB（Out-of-Band）的方式交换参数，使用UCC编程时需要提供一个oob的Allgather操作，这一步骤一般使用MPI_Iallgatehr进行。eam相当于MPI的通信域，ep相当于MPI的rank。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#define STR(x) #x
#define UCC_CHECK(_call)                                            \
    if (UCC_OK != (_call)) {                                        \
        fprintf(stderr, "*** UCC TEST FAIL: %s\n", STR(_call));     \
        MPI_Abort(MPI_COMM_WORLD, -1);                              \
    }

static ucc_status_t oob_allgather(void *sbuf, void *rbuf, size_t msglen,
                                  void *coll_info, void **req)
{
    MPI_Comm    comm = (MPI_Comm)coll_info;
    MPI_Request request;

    MPI_Iallgather(sbuf, msglen, MPI_BYTE, rbuf, msglen, MPI_BYTE, comm,
                   &request);
    *req = (void *)request;
    return UCC_OK;
}

static ucc_status_t oob_allgather_test(void *req)
{
    MPI_Request request = (MPI_Request)req;
    int         completed;

    MPI_Test(&request, &completed, MPI_STATUS_IGNORE);
    return completed ? UCC_OK : UCC_INPROGRESS;
}

static ucc_status_t oob_allgather_free(void *req)
{
    return UCC_OK;
}

/* Creates UCC team for a group of processes represented by MPI
   communicator. UCC API provides different ways to create a team,
   one of them is to use out-of-band (OOB) allgather provided by
   the calling runtime. */
static ucc_team_h create_ucc_team(MPI_Comm comm, ucc_context_h ctx)
{
    int               rank, size;
    ucc_team_h        team;
    ucc_team_params_t team_params;
    ucc_status_t      status;

    MPI_Comm_rank(comm, &rank);
    MPI_Comm_size(comm, &size);

    team_params.mask          = UCC_TEAM_PARAM_FIELD_OOB;
    team_params.oob.allgather = oob_allgather;
    team_params.oob.req_test  = oob_allgather_test;
    team_params.oob.req_free  = oob_allgather_free;
    team_params.oob.coll_info = (void*)comm;
    team_params.oob.n_oob_eps = size;
    team_params.oob.oob_ep    = rank;

    UCC_CHECK(ucc_team_create_post(&ctx, 1, &team_params, &team));
    while (UCC_INPROGRESS == (status = ucc_team_create_test(team))) {
        UCC_CHECK(ucc_context_progress(ctx));
    };
    if (UCC_OK != status) {
        fprintf(stderr, "failed to create ucc team\n");
        MPI_Abort(MPI_COMM_WORLD, status);
    }
    return team;
}

ucc集合通信上下文的初始化大概是下面这些，在这之后就可以调用接口做集合通信了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
/* Init ucc library */
   ucc_lib_params_t lib_params = {
       .mask        = UCC_LIB_PARAM_FIELD_THREAD_MODE,
       .thread_mode = UCC_THREAD_SINGLE
   };
   UCC_CHECK(ucc_lib_config_read(NULL, NULL, &lib_config));
   UCC_CHECK(ucc_init(&lib_params, lib_config, &lib));
   ucc_lib_config_release(lib_config);

   /* Init ucc context for a specified UCC_TEST_TLS */
   ucc_context_params_t ctx_params = {
       .mask             = UCC_CONTEXT_PARAM_FIELD_OOB,
       .oob.allgather    = oob_allgather,
       .oob.req_test     = oob_allgather_test,
       .oob.req_free     = oob_allgather_free,
       .oob.coll_info    = (void*)MPI_COMM_WORLD,
       .oob.n_oob_eps    = size,
       .oob.oob_ep       = rank
   };

   UCC_CHECK(ucc_context_config_read(lib, NULL, &ctx_config));
   UCC_CHECK(ucc_context_create(lib, &ctx_params, ctx_config, &ctx));
   ucc_context_config_release(ctx_config);

   team = create_ucc_team(MPI_COMM_WORLD, ctx);

进行集合通信的时候需要设置参数，类似于调用MPI集合通信函数并传递参数的过程。这里以Allgather操作为例，每个集合通信函数不一样。要注意的是带UCC接收方接收数据的个数为总个数，在MPI带All的集合通信操作的接口中rcount是从每个进程收到数据的个数，二者是不一样的。

1
2
3
4
5
6
7
8
9
10
args.mask              = 0;
   args.coll_type         = UCC_COLL_TYPE_ALLGATHER;
   args.src.info.buffer   = sbuf;
   args.src.info.datatype = UCC_DT_INT8;
   args.src.info.mem_type = UCC_MEMORY_TYPE_HOST;
   args.dst.info.buffer   = rbuf;
   args.dst.info.datatype = UCC_DT_INT8;
   args.dst.info.mem_type = UCC_MEMORY_TYPE_HOST;
   args.src.info.count    = msg_size;
   args.dst.info.count    = msg_size*size;

通信的过程是这几行，每种集合通信操作都是这样，差别在上一步设置参数那里。

1
2
3
4
5
UCC_CHECK(ucc_collective_init(&args, &req, team));
UCC_CHECK(ucc_collective_post(req));    
while (UCC_INPROGRESS == ucc_collective_test(req)) {
	UCC_CHECK(ucc_context_progress(ctx));
}

最后在通信结束后需要进行如下操作，相当于MPI的Finalize

1
2
3
4
/* Cleanup UCC */
   UCC_CHECK(ucc_team_destroy(team));
   UCC_CHECK(ucc_context_destroy(ctx));
   UCC_CHECK(ucc_finalize(lib));

结束语

本篇博客对UCC的介绍就是这些，UCC是一个比较新的框架可以参考的资料比较少，作者也是刚刚开始了解，下面放上一些有用的链接，日后如果对UCC中哪一方面的细节有更深一步的了解或理解会继续更新。

Github地址：https://github.com/openucx/ucc

Work Group相关资料的地址，有API文档，并且了解一些细节技术很有用：https://github.com/openucx/ucc/tree/gh-pages

DPU卸载的Allgatherv算法，2022RDMA竞赛提供：https://github.com/yqin/ucc/tree/topic/rdma_workshop

以及如果需要了解某个功能的具体实现，查看Issue或者Commit信息会很有用。