Kompute 面向跨厂商显卡（AMD、Qualcomm、NVIDIA 等）的通用 GPU 计算框架

极速、支持移动端、异步，并针对高级 GPU 加速用例进行了优化。

💬 加入 Discord 社区与定期会议 🔋 文档 💻 博客文章 ⌨ 示例 💾

---

Kompute 得到 Linux 基金会的支持，作为 LF AI & Data 基金会的 托管项目。

核心原则与特性

• 具有 C++ SDK 的灵活 Python 模块，用于优化
• 通过 GPU 家族队列支持 异步与并行处理
• 支持 移动端，提供基于 Android NDK 的多架构示例
• BYOV：采用 自带 Vulkan 设计，可与现有 Vulkan 应用程序良好兼容
• 明确定义 GPU 和主机之间的 内存所有权与内存管理关系
• 强健的代码库，单元测试覆盖率达 90%
• 适用于 机器学习 🤖、移动开发 📱 和 游戏开发 🎮 等高级场景。
• 活跃的社区，包括 每月例会、Discord 聊天等

使用 Kompute 的项目 ❤️ 🤖

• GPT4ALL - 一个开源的边缘大型语言模型生态系统，可在本地 CPU 和几乎所有 GPU 上运行。
• llama.cpp - Facebook LLaMA 模型的 C/C++ 移植版本（现已停止维护）。
• tpoisonooo/how-to-optimize-gemm - 行主序矩阵乘法优化。
• vkJAX - 面向 Vulkan 的 JAX 解释器。

快速上手

下方提供了使用 C++ 和 Python Kompute 接口的 GPU 乘法示例。

如有疑问或讨论，您可 加入 Discord，打开 GitHub 问题，或阅读 文档。

你的第一个 Kompute (C++)

C++ 接口提供了对 Kompute 原生组件的低级访问，从而实现高级优化以及组件扩展。

void kompute(const std::string& shader) {

    // 1. 使用默认设置创建 Kompute 管理器（设备 0、第一个队列，无扩展）
    kp::Manager mgr; 

    // 2. 通过管理器创建并初始化 Kompute 张量

    // 默认张量构造函数简化了浮点值的创建
    auto tensorInA = mgr.tensor({ 2., 2., 2. });
    auto tensorInB = mgr.tensor({ 1., 2., 3. });
    // 显式类型构造函数支持 uint32、int32、double、float 和 bool
    auto tensorOutA = mgr.tensorT<uint32_t>({ 0, 0, 0 });
    auto tensorOutB = mgr.tensorT<uint32_t>({ 0, 0, 0 });

    std::vector<std::shared_ptr<kp::Memory>> 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::OpSyncDevice>(params)
        ->record<kp::OpAlgoDispatch>(algorithm) // 绑定默认推送常量
        ->eval() // 执行记录的两个操作
        ->record<kp::OpAlgoDispatch>(algorithm, pushConstsB) // 覆盖推送常量
        ->eval(); // 仅执行最后一条记录的操作

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

    // ... 在 GPU 完成计算的同时异步执行其他工作

    sq->evalAwait();

    // 打印第一个输出：{ 4, 8, 12 }
    for (const float& elem : tensorOutA->vector()) std::cout << elem << "  ";
    // 打印第二个输出：{ 10, 10, 10 }
    for (const float& elem : tensorOutB->vector()) std::cout << elem << "  ";

} // 管理并释放所有 CPU 和 GPU 内存资源

int main() {

    // 定义原始字符串着色器（或使用 Kompute 工具将其编译为 SPIRV / C++ 头文件）。该着色器展示了包括常量、缓冲区等在内的主要组件。
    std::string shader = (R"(
        #version 450

        layout (local_size_x = 1) in;

        // 输入张量的绑定索引相对于传入参数的索引
        layout(set = 0, binding = 0) buffer buf_in_a { float in_a[]; };
        layout(set = 0, binding = 1) buffer buf_in_b { float in_b[]; };
        layout(set = 0, binding = 2) buffer buf_out_a { uint out_a[]; };
        layout(set = 0, binding = 3) buffer buf_out_b { uint out_b[]; };

        // Kompute 支持在调度时更新的推送常量
        layout(push_constant) uniform PushConstants {
            float val;
        } push_const;

        // Kompute 还支持在初始化时使用的特殊常量
        layout(constant_id = 0) const float const_one = 0;

        void main() {
            uint index = gl_GlobalInvocationID.x;
            out_a[index] += uint( in_a[index] * in_b[index] );
            out_b[index] += uint( const_one * push_const.val );
        }
    )");

    // 使用上述定义的函数和原始字符串着色器运行程序
    kompute(shader);
}

你的第一个 Kompute（Python）

Python 软件包 提供了一个高级交互式接口，它既支持实验性操作，又能确保高性能和快速的开发流程。

from .utils import compile_source # 使用 python/test/utils 中的工具函数

def kompute(shader):
    # 1. 使用默认设置创建 Kompute 管理器（设备 0、第一个队列且无扩展）
    mgr = kp.Manager()

    # 2. 通过管理器创建并初始化 Kompute 张量

    # 默认张量构造函数简化了浮点值的创建
    tensor_in_a = mgr.tensor([2, 2, 2])
    tensor_in_b = mgr.tensor([1, 2, 3])
    # 显式类型构造函数支持 uint32、int32、double、float 和 bool
    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))
    assert(t_data.data_type() == kp.DataTypes.uint)

    params = [tensor_in_a, tensor_in_b, tensor_out_a, tensor_out_b]

    # 3. 根据着色器创建算法（支持缓冲区及推送/特殊常量）
    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)

    # 4. 使用序列同步运行操作
    (mgr.sequence()
        .record(kp.OpTensorSyncDevice(params))
        .record(kp.OpAlgoDispatch(algo)) # 绑定提供的默认推送常量
        .eval() # 执行记录的两个操作
        .record(kp.OpAlgoDispatch(algo, push_consts_b)) # 覆盖推送常量
        .eval()) # 只执行最后一条记录的操作

    # 5. 异步将结果从 GPU 同步回主机
    sq = mgr.sequence()
    sq.eval_async(kp.OpTensorSyncLocal(params))

    # ... 在 GPU 完成工作的同时异步进行其他任务

    sq.eval_await()

    # 打印第一个输出：{ 4, 8, 12 }
    print(tensor_out_a)
    # 打印第二个输出：{ 10, 10, 10 }
    print(tensor_out_b)

if __name__ == "__main__":

    # 定义一个原始字符串着色器（或使用 Kompute 工具将其编译为 SPIRV / C++ 头文件）。此着色器展示了包括常量、缓冲区等在内的主要组件
    shader = """
        #version 450

        layout (local_size_x = 1) in;

        // 输入张量的绑定索引相对于传入参数的索引
        layout(set = 0, binding = 0) buffer buf_in_a { float in_a[]; };
        layout(set = 0, binding = 1) buffer buf_in_b { float in_b[]; };
        layout(set = 0, binding = 2) buffer buf_out_a { uint out_a[]; };
        layout(set = 0, binding = 3) buffer buf_out_b { uint out_b[]; };

        // Kompute 支持在调度时更新的推送常量
        layout(push_constant) uniform PushConstants {
            float val;
        } push_const;

        // Kompute 还支持在初始化时使用的特殊常量
        layout(constant_id = 0) const float const_one = 0;

        void main() {
            uint index = gl_GlobalInvocationID.x;
            out_a[index] += uint( in_a[index] * in_b[index] );
            out_b[index] += uint( const_one * push_const.val );
        }
    """

    kompute(shader)

交互式笔记本与动手视频

你可以尝试使用免费 GPU 的 Colab 交互式笔记本。可用的示例包括以下 Python 和 C++ 示例：

尝试来自 博客文章 的交互式 C++ Colab	尝试来自 博客文章 的交互式 Python Colab

你还可以观看在 FOSDEM 2021 大会上发表的两场演讲。

这两段视频都带有时间戳，方便你跳转到最相关的内容——两者的介绍和动机部分几乎相同，因此可以直接跳到更具体的内容。

观看面向 C++ 爱好者 的视频	观看面向 Python 和机器学习爱好者 的视频

架构概览

Kompute 的核心架构包括以下组件：

• Kompute Manager - 基础编排器，用于创建和管理设备及子组件
• Kompute Sequence - 可作为批处理发送到 GPU 的操作容器
• Kompute Operation (Base) - 所有操作均继承的基类
• Kompute Tensor - 在 GPU 操作中使用的张量结构化数据
• Kompute Algorithm - 在 GPU 中执行的（着色器）逻辑抽象

如需完整分解，可进一步阅读 C++ 类参考文档。

完整架构	简化后的 Kompute 组件
（非常小，请查看 文档中的完整参考图以获取详细信息）

异步与并行操作

Kompute 通过 vk::Fences 提供了以异步方式运行操作的灵活性。此外，Kompute 还支持显式分配队列，从而允许跨队列族并行执行操作。

下图直观展示了如何将 Kompute Sequence 分配到不同队列，以根据硬件实现并行执行。您可以在 动手示例 和 详细文档页面 中了解其工作原理，其中以 NVIDIA 1650 为例进行了说明。

移动端支持

Kompute 已针对移动环境进行了优化。构建系统支持在 Android 环境中动态加载 Vulkan 共享库，并提供适用于 CPP 头文件的 Android NDK 封装。

如需深入了解，可阅读博客文章“利用 Android NDK、Vulkan 和 Kompute 为您的移动应用注入 GPU 加速机器学习动力”。 您还可以访问仓库中的 端到端示例代码,该代码可通过 Android Studio 运行。

更多示例

简单示例

• 简单乘法示例
• 使用 Kompute Sequence 记录批处理命令
• 运行异步操作
• 在多个 GPU 队列上并行运行操作
• 创建自定义 Kompute 操作
• 从零开始实现逻辑回归

端到端示例

• GPU 上的机器学习逻辑回归实现
• 通过多队列操作并行化 GPU 密集型工作负载
• Android NDK 移动 Kompute ML 应用程序
• Godot 引擎中的游戏开发 Kompute ML

Python 包

除了 C++ 核心 SDK 外，您还可以使用 Kompute 的 Python 包，它提供了相同的核心功能，并支持与 Python 对象（如列表、NumPy 数组等）的互操作。

唯一的依赖项是 Python 3.5+ 和 CMake 3.4.1+。您可以通过以下命令从 Python pypi 包 安装 Kompute：

pip install kp

您也可以从主分支安装：

pip install git+git://github.com/KomputeProject/kompute.git@master

有关更多详情，请参阅 Python 包文档 或 Python 类参考文档。

C++ 构建概览

提供的构建系统使用 cmake，支持跨平台构建。

顶级 Makefile 提供了针对开发以及 Docker 镜像构建的优化配置，但您也可以通过以下命令开始构建：

   cmake -Bbuild

您还可以使用 add_subdirectory 将 Kompute 添加到您的项目中——Android 示例 CMakeLists.txt 文件展示了具体做法。

如需更深入地了解构建配置，请参阅 构建系统深度解析 文档。

Kompute 开发

我们欢迎 PR 和问题提交。如果您想贡献代码，可以查看“Good first issue”标签；即使只是使用 Kompute 并报告问题，也是一种很好的贡献！

贡献

开发依赖

• 测试
  • GTest
• 文档
  • Doxygen（带 Dot）
  • Sphinx

开发规范

• 遵循 Mozilla C++ 代码风格指南：https://www-archive.mozilla.org/hacking/mozilla-style-guide.html
  • 使用提交后钩子运行代码检查工具；可以配置为在提交前自动执行代码检查。
  • 所有依赖项均在 vcpkg.json 文件中定义。
• 使用 CMake 作为构建系统，并提供顶层 Makefile，其中包含推荐的构建命令。
• 使用 xxd（或 Windows 64 位版本 xxd.exe）将着色器 SPIR-V 文件转换为头文件。
• 使用 Doxygen 和 Sphinx 生成文档及自动文档。
• 使用 vcpkg 管理依赖库，这是推荐的获取第三方库的方式。

若需启用调试层，可通过设置 KOMPUTE_ENV_DEBUG_LAYERS 参数来添加，例如：

export KOMPUTE_ENV_DEBUG_LAYERS="VK_LAYER_LUNARG_api_dump"

更新文档

要更新文档，您需要：

• 在构建系统中运行 gendoxygen 目标；
• 在构建系统中运行 gensphynx 目标；
• 使用 make push_docs_to_ghpages 命令将文档推送到 GitHub Pages。

运行测试

对于贡献者而言，运行单元测试已大大简化。

测试在 CPU 上执行，可通过 ACT 命令行工具（https://github.com/nektos/act）触发。安装该工具并启动 Docker 守护进程后，只需输入以下命令即可：

$ act

[Python Tests/python-tests] 🚀  Start image=axsauze/kompute-builder:0.2
[C++ Tests/cpp-tests      ] 🚀  Start image=axsauze/kompute-builder:0.2
[C++ Tests/cpp-tests      ]   🐳  docker run image=axsauze/kompute-builder:0.2 entrypoint=["/usr/bin/tail" "-f" "/dev/null"] cmd=[]
[Python Tests/python-tests]   🐳  docker run image=axsauze/kompute-builder:0.2 entrypoint=["/usr/bin/tail" "-f" "/dev/null"] cmd=[]
...

仓库包含 C++ 和 Python 代码的单元测试，分别位于 test/ 和 python/test 目录下。

目前，这些测试通过 GitHub Actions 的 CI 系统运行，并使用 docker-builders/ 中的镜像。

为降低硬件要求，测试可以在没有 GPU 的情况下直接在 CPU 上运行，借助 SwiftShader 实现。

有关 CI 和测试的具体配置，请参阅文档中的 CI、Docker 和测试章节。

动机

本项目源于观察到许多新兴且知名的机器学习和深度学习框架，如 PyTorch、TensorFlow、阿里巴巴 DNN、腾讯 NCNN 等，均已集成或正计划集成 Vulkan SDK，以支持移动端及其他厂商的 GPU 加速计算功能。

Vulkan SDK 提供了一个优秀的底层接口，能够实现高度定制化的优化；然而，其代码编写较为冗长，仅开始编写应用代码就需要 500 到 2000 行代码。这导致各个项目不得不重复实现相同的抽象层，用于屏蔽 Vulkan SDK 中与计算无关的功能。这种大量非标准化的样板代码不仅限制了知识共享，还增加了引入特定框架独有缺陷的风险等。

我们当前开发 Kompute 的目的，并非隐藏设计精良的 Vulkan SDK 接口，而是围绕 Vulkan SDK 的 GPU 计算能力进行扩展和完善。这篇文章 对 Kompute 的开发动机进行了高层次的概述，并结合实际示例介绍了 GPU 计算以及 Kompute 的核心架构。

Kompute 快速上手指南

Kompute 是一个通用的 GPU 计算框架，支持跨厂商显卡（AMD、Qualcomm、NVIDIA 等），基于 Vulkan API 构建。它提供高性能的异步并行处理能力，适用于机器学习、移动开发及游戏开发等场景。

环境准备

系统要求

• 操作系统: Linux, Windows, macOS (需 MoltenVK), Android
• GPU: 支持 Vulkan 的显卡 (AMD, NVIDIA, Intel, Qualcomm Adreno 等)
• 编译器: 支持 C++14 至 C++20 的编译器 (GCC, Clang, MSVC)
• Python: 3.7 - 3.9 (如需使用 Python 接口)

前置依赖

确保系统已安装以下工具：

• CMake: 3.15 或更高版本
• Vulkan SDK: 必须安装并配置 VULKAN_SDK 环境变量
• Git: 用于克隆代码库

提示: 国内用户可通过 华为云镜像 或 清华大学镜像站 加速下载 Vulkan SDK 及相关依赖。

安装步骤

方式一：从源码构建 (C++ SDK)

# 1. 克隆仓库
git clone https://github.com/KomputeProject/kompute.git
cd kompute

# 2. 创建构建目录
mkdir build && cd build

# 3. 配置 CMake (可选：指定 Python 路径或安装前缀)
cmake .. -DKOMPUTE_BUILD_PYTHON=OFF

# 4. 编译并安装
cmake --build . --config Release
sudo cmake --install .

方式二：安装 Python 包

推荐使用 pip 直接安装预编译包（需确保系统已安装 Vulkan 运行时）：

pip install kp

若需从源码构建 Python 绑定：

pip install .[dev]

基本使用

C++ 示例：GPU 矩阵乘法

以下示例演示如何创建 Manager、初始化张量、定义 Shader 并执行异步计算。

#include "kompute/Kompute.hpp"
#include <iostream>
#include <vector>

void kompute(const std::string& shader) {
    // 1. 创建 Kompute Manager (默认使用设备 0)
    kp::Manager mgr; 

    // 2. 创建并初始化张量
    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::Memory>> params = {tensorInA, tensorInB, tensorOutA, tensorOutB};

    // 3. 基于 Shader 创建算法
    kp::Workgroup workgroup({3, 1, 1});
    std::vector<float> specConsts({ 2 });
    std::vector<float> pushConstsA({ 2.0 });
    std::vector<float> pushConstsB({ 3.0 });

    // 假设 compileSource 函数已将 GLSL 源码编译为 SPIR-V
    auto algorithm = mgr.algorithm(params, compileSource(shader), workgroup, specConsts, pushConstsA);

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

    // 5. 异步同步结果回主机
    auto sq = mgr.sequence();
    sq->evalAsync<kp::OpSyncLocal>(params);

    // ... 此处可执行其他异步任务

    sq->evalAwait();

    // 输出结果
    for (const float& elem : tensorOutA->vector()) std::cout << elem << "  "; // 输出: 4  8  12
    std::cout << std::endl;
    for (const float& elem : tensorOutB->vector()) std::cout << elem << "  "; // 输出: 10  10  10
}

int main() {
    // 定义 GLSL 计算着色器
    std::string shader = (R"(
        #version 450
        layout (local_size_x = 1) in;
        layout(set = 0, binding = 0) buffer buf_in_a { float in_a[]; };
        layout(set = 0, binding = 1) buffer buf_in_b { float in_b[]; };
        layout(set = 0, binding = 2) buffer buf_out_a { uint out_a[]; };
        layout(set = 0, binding = 3) buffer buf_out_b { uint out_b[]; };
        layout(push_constant) uniform PushConstants { float val; } push_const;
        layout(constant_id = 0) const float const_one = 0;

        void main() {
            uint index = gl_GlobalInvocationID.x;
            out_a[index] += uint( in_a[index] * in_b[index] );
            out_b[index] += uint( const_one * push_const.val );
        }
    )");

    kompute(shader);
    return 0;
}

Python 示例：交互式 GPU 计算

Python 接口提供了更简洁的高级 API，适合快速原型开发。

import numpy as np
import kp
from utils import compile_source  # 需自行实现或使用测试工具中的 compile_source

def kompute(shader):
    # 1. 创建 Manager
    mgr = kp.Manager()

    # 2. 创建张量
    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]

    # 3. 创建算法
    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)

    # 4. 执行序列
    (mgr.sequence()
        .record(kp.OpTensorSyncDevice(params))
        .record(kp.OpAlgoDispatch(algo))
        .eval()
        .record(kp.OpAlgoDispatch(algo, push_consts_b))
        .eval())

    # 5. 异步获取结果
    sq = mgr.sequence()
    sq.eval_async(kp.OpTensorSyncLocal(params))
    
    # ... 执行其他任务
    
    sq.eval_await()

    print(tensor_out_a)  # 输出: [4 8 12]
    print(tensor_out_b)  # 输出: [10 10 10]

if __name__ == "__main__":
    shader = """
        #version 450
        layout (local_size_x = 1) in;
        layout(set = 0, binding = 0) buffer buf_in_a { float in_a[]; };
        layout(set = 0, binding = 1) buffer buf_in_b { float in_b[]; };
        layout(set = 0, binding = 2) buffer buf_out_a { uint out_a[]; };
        layout(set = 0, binding = 3) buffer buf_out_b { uint out_b[]; };
        layout(push_constant) uniform PushConstants { float val; } push_const;
        layout(constant_id = 0) const float const_one = 0;

        void main() {
            uint index = gl_GlobalInvocationID.x;
            out_a[index] += uint( in_a[index] * in_b[index] );
            out_b[index] += uint( const_one * push_const.val );
        }
    """
    kompute(shader)