Kompute: 通用GPU计算框架助力跨厂商高性能计算

Kompute: 通用GPU计算框架助力跨厂商高性能计算

Kompute是一个基于Vulkan构建的通用GPU计算框架,旨在支持包括AMD、Qualcomm、NVIDIA等多家厂商的数千种显卡。作为一个由Linux基金会支持的开源项目,Kompute具有高速、移动友好、异步和针对高级GPU数据处理场景优化等特点,为开发者提供了一个强大而灵活的GPU计算工具。

核心特性与优势

Kompute的主要特性和优势包括:

1. 跨厂商支持: 支持AMD、Qualcomm、NVIDIA等多家厂商的显卡,提供广泛的硬件兼容性。
2. 高性能: 基于Vulkan构建,具有极快的执行速度,适用于要求苛刻的计算任务。
3. 移动友好: 针对移动环境进行了优化,支持Android NDK跨多种架构的应用开发。
4. 异步处理: 通过GPU家族队列支持异步和并行处理,提高计算效率。
5. 灵活性: 采用"自带Vulkan"(BYOV)设计,可以与现有的Vulkan应用程序无缝集成。
6. 内存管理: 对GPU和主机内存所有权和管理提供清晰的关系定义。
7. 代码质量: 拥有90%的单元测试覆盖率,确保代码的稳定性和可靠性。
8. 多语言支持: 提供C++核心SDK和优化的Python绑定,满足不同开发者的需求。
9. 社区支持: 活跃的社区提供月度会议、Discord聊天等交流渠道。

架构概览

Kompute的核心架构包括以下几个主要组件:

1. Kompute Manager: 基础编排器,负责创建和管理设备及子组件。
2. Kompute Sequence: 操作容器,可将操作批量发送到GPU执行。
3. Kompute Operation: 所有操作的基类。
4. Kompute Tensor: 用于GPU操作的张量结构化数据。
5. Kompute Algorithm: GPU中执行的(着色器)逻辑抽象。

这种设计使得Kompute能够高效地管理GPU资源和计算任务,为开发者提供了一个强大而灵活的GPU计算框架。

异步和并行操作

Kompute通过vk::Fences提供了异步运行操作的灵活性。此外,Kompute还支持显式分配队列,允许跨队列族并行执行操作。这种设计使得开发者可以充分利用GPU的并行计算能力,显著提高计算效率。

上图展示了Kompute序列如何分配到不同队列以实现基于硬件的并行执行。以NVIDIA 1650为例,它支持跨家族队列的并行化,但不支持同一家族内的并行化。了解硬件的这些特性对于充分发挥Kompute的性能至关重要。

移动端支持

Kompute已针对移动环境进行了优化。其构建系统支持在Android环境中动态加载Vulkan共享库,并为C++头文件提供了可用的Android NDK包装器。这使得开发者可以轻松地将Kompute集成到移动应用中,充分利用移动设备的GPU计算能力。

使用示例

以下是一个简单的GPU乘法示例,展示了如何使用Kompute的C++接口:

void kompute(const std::string& shader) {
    // 1. 创建Kompute Manager (默认设置:设备0,第一个队列,无扩展)
    kp::Manager mgr; 

    // 2. 通过manager创建并初始化Kompute Tensors
    auto tensorInA = mgr.tensor({ 2., 2., 2. });
    auto tensorInB = mgr.tensor({ 1., 2., 3. });
    auto tensorOutA = mgr.tensorT<uint32_t>({ 0, 0, 0 });
    auto tensorOutB = mgr.tensorT<uint32_t>({ 0, 0, 0 });

    std::vector<std::shared_ptr<kp::Tensor>> params = {tensorInA, tensorInB, tensorOutA, tensorOutB};

    // 3. 基于着色器创建算法(支持缓冲区和推送/规范常量)
    kp::Workgroup workgroup({3, 1, 1});
    std::vector<float> specConsts({ 2 });
    std::vector<float> pushConstsA({ 2.0 });
    std::vector<float> pushConstsB({ 3.0 });

    auto algorithm = mgr.algorithm(params, compileSource(shader), workgroup, specConsts, pushConstsA);

    // 4. 使用序列同步运行操作
    mgr.sequence()
        ->record<kp::OpTensorSyncDevice>(params)
        ->record<kp::OpAlgoDispatch>(algorithm)
        ->eval()
        ->record<kp::OpAlgoDispatch>(algorithm, pushConstsB)
        ->eval();

    // 5. 异步从GPU同步结果
    auto sq = mgr.sequence();
    sq->evalAsync<kp::OpTensorSyncLocal>(params);

    // ... 在GPU完成时异步执行其他工作

    sq->evalAwait();

    // 打印输出
    for (const float& elem : tensorOutA->vector()) std::cout << elem << "  ";
    for (const float& elem : tensorOutB->vector()) std::cout << elem << "  ";
}

这个例子展示了Kompute的基本使用流程,包括创建管理器、张量、算法,以及执行同步和异步操作。

Python支持

除了C++ SDK,Kompute还提供了Python包,暴露了相同的核心功能,并支持与Python对象(如列表、Numpy数组等)的互操作性。以下是一个简单的Python示例:

def kompute(shader):
    mgr = kp.Manager()

    tensor_in_a = mgr.tensor([2, 2, 2])
    tensor_in_b = mgr.tensor([1, 2, 3])
    tensor_out_a = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32))
    tensor_out_b = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32))

    params = [tensor_in_a, tensor_in_b, tensor_out_a, tensor_out_b]

    workgroup = (3, 1, 1)
    spec_consts = [2]
    push_consts_a = [2]
    push_consts_b = [3]

    spirv = compile_source(shader)

    algo = mgr.algorithm(params, spirv, workgroup, spec_consts, push_consts_a)

    (mgr.sequence()
        .record(kp.OpTensorSyncDevice(params))
        .record(kp.OpAlgoDispatch(algo))
        .eval()
        .record(kp.OpAlgoDispatch(algo, push_consts_b))
        .eval())

    sq = mgr.sequence()
    sq.eval_async(kp.OpTensorSyncLocal(params))

    sq.eval_await()

    print(tensor_out_a)
    print(tensor_out_b)

这个Python示例展示了与C++版本相似的功能,但使用了更加Python化的语法和接口。

项目动机

Kompute项目的诞生源于对现有机器学习和深度学习项目(如PyTorch、TensorFlow、Alibaba DNN、Tencent NCNN等)的观察。这些项目要么已经集成了Vulkan SDK,要么正在考虑集成,以添加移动(和跨厂商)GPU支持。

Vulkan SDK提供了出色的低级接口,可以实现高度专门化的优化。然而,这也带来了代码冗长的代价,通常需要500-2000行代码才能开始编写应用程序代码。这导致每个项目都必须实现相同的基线来抽象Vulkan SDK的非计算相关功能。这种大量非标准化的样板代码可能会导致知识转移受限,增加引入独特框架实现错误的机会等问题。

Kompute的开发并不是为了隐藏Vulkan SDK接口(因为它设计得非常出色),而是为了增强它,直接关注Vulkan SDK的GPU计算能力。Kompute旨在提供一个更高级别的抽象,简化GPU计算的开发过程,同时保留Vulkan的强大功能和灵活性。

总结

Kompute作为一个通用的GPU计算框架,为开发者提供了一个强大而灵活的工具,用于在各种硬件平台上进行高性能计算。它的跨厂商支持、高性能、移动友好性以及异步处理能力,使其成为GPU计算领域的一个重要选择。无论是在机器学习、移动开发还是游戏开发等领域,Kompute都展现出了巨大的潜力。

随着GPU计算在各个领域的应用日益广泛,Kompute这样的框架将在推动技术创新和提高计算效率方面发挥越来越重要的作用。对于希望充分利用GPU计算能力的开发者来说,Kompute无疑是一个值得关注和尝试的项目。