跳转至

# Vulkan Compute 与跨平台 GPU

Vulkan 是唯一能在所有主要平台上运行的 GPU 计算 API：NVIDIA、AMD、Intel、Apple（通过 MoltenVK）、Android，甚至浏览器（通过 WebGPU）。本文涵盖 Vulkan 架构、计算管线、使用 GLSL 编写计算着色器、GPU 计算程序的完整 C++ 设置、共享内存与同步、用于浏览器的 WebGPU，以及实际的机器学习推理示例。

- CUDA 在 NVIDIA 硬件上主导着 ML 训练。但并非每个部署目标都有 NVIDIA GPU。移动应用运行在 Qualcomm Adreno 或 ARM Mali GPU 上。Web 应用运行在浏览器中。游戏引擎需要同时支持 AMD、Intel 和 NVIDIA。对于所有这些场景，Vulkan 就是答案。

- Vulkan 很冗长——一个"hello world"计算程序大约有 300 行 C++ 代码。但这种冗长是 显式控制 的代价：你需要自己管理每一个 GPU 资源（内存、管线、命令缓冲区）。这种控制带来了最大性能和可移植性，代价是开发速度。

## Vulkan 架构概述

- Vulkan 是由 Khronos Group（OpenGL 背后的同一组织）创建的低级 GPU API。与 CUDA（它隐藏了 GPU 资源管理）不同，Vulkan 要求你显式地管理：

  - 实例与设备：创建 Vulkan 实例，枚举可用 GPU，并选择一个。
  - 内存：显式分配 GPU 内存，指定内存类型（设备本地内存用于速度，主机可见内存用于 CPU 访问）。
  - 缓冲区：创建引用已分配内存的缓冲区对象。
  - 描述符集：将缓冲区绑定到着色器输入（类似于计算着色器的函数参数）。
  - 计算管线：编译着色器并创建管线对象。
  - 命令缓冲区：记录一系列 GPU 命令（绑定管线、绑定描述符、调度计算）。
  - 队列提交：将命令缓冲区提交给 GPU 执行。
  - 同步：使用栅栏和屏障确保正确的执行顺序。

- 这与 CUDA 的 `cudaMalloc` + 内核启动模型截然不同。在 CUDA 中，驱动程序在幕后处理大部分工作。在 Vulkan 中，你需要自己做这一切。

### 为什么如此冗长？

- Vulkan 的显式性存在有两方面原因：

  1. 驱动简化：OpenGL 驱动极其复杂（它们必须猜测应用程序的意图并进行相应优化）。Vulkan 将该责任转移给应用程序，使驱动更精简、更可预测，并且更容易在各厂商间正确实现。

  2. 性能：对内存布局、同步和命令批处理的显式控制使应用程序能够做出最优决策。在 CUDA 中，驱动可能会插入不必要的同步。在 Vulkan 中，你只在需要时才进行同步。

## GLSL 中的计算着色器

- 计算着色器 是在 GPU 上运行的程序，类似于 CUDA 内核。它使用 GLSL（OpenGL 着色语言）编写，并编译为 SPIR-V 字节码（一种可移植的二进制格式）。

### 向量加法

    // add.comp — 编译命令: glslangValidator -V add.comp -o add.spv
    #version 450

    // 工作组大小：每个工作组有 256 个调用（= CUDA 中每块的线程数）
    layout(local_size_x = 256) in;

    // 缓冲区绑定（类似于内核参数）
    layout(set = 0, binding = 0) buffer InputA { float a[]; };
    layout(set = 0, binding = 1) buffer InputB { float b[]; };
    layout(set = 0, binding = 2) buffer Output { float c[]; };

    // 推送常量：小的统一数据（类似于内核参数）
    layout(push_constant) uniform PushConstants {
        uint n;  // 元素数量
    };

    void main() {
        uint idx = gl_GlobalInvocationID.x;  // 全局线程索引
        if (idx < n) {
            c[idx] = a[idx] + b[idx];
        }
    }

- 与 CUDA 概念的映射：

Vulkan | CUDA | 含义
--- | --- | ---
工作组 (Workgroup) | 块 (Block) | 可以共享内存的线程组
调用 (Invocation) | 线程 (Thread) | 单个执行单元
`gl_GlobalInvocationID` | `blockIdx * blockDim + threadIdx` | 全局线程索引
`gl_LocalInvocationID` | `threadIdx` | 工作组内的线程索引
`gl_WorkGroupID` | `blockIdx` | 工作组索引
`local_size_x` | `blockDim.x` | 每工作组的线程数
存储缓冲区 | 全局内存 | 可读写的 GPU 内存
共享内存 (`shared`) | `__shared__` | 每工作组的高速内存
推送常量 | 内核参数 | 小的统一数据

### 使用共享内存的 ReLU

    // relu_shared.comp
    #version 450

    layout(local_size_x = 256) in;

    layout(set = 0, binding = 0) buffer Input  { float input_data[]; };
    layout(set = 0, binding = 1) buffer Output { float output_data[]; };

    layout(push_constant) uniform PushConstants { uint n; };

    // 共享内存（等同于 CUDA 的 __shared__）
    shared float tile[256];

    void main() {
        uint gid = gl_GlobalInvocationID.x;
        uint lid = gl_LocalInvocationID.x;

        // 加载到共享内存
        if (gid < n) {
            tile[lid] = input_data[gid];
        }

        // 屏障：等待工作组中所有调用完成加载
        barrier();  // 等同于 CUDA 的 __syncthreads()

        // 计算 ReLU
        if (gid < n) {
            output_data[gid] = max(tile[lid], 0.0);
        }
    }

- 对于 ReLU，共享内存并非严格必要（该操作是按元素进行的）。但这演示了基本模式：加载到共享内存 → 屏障 → 计算 → 存储。对于需要相邻线程数据的操作（卷积、归约、softmax），共享内存是必不可少的。

### 并行归约（求和）

    // reduce_sum.comp
    #version 450

    layout(local_size_x = 256) in;

    layout(set = 0, binding = 0) buffer Input  { float input_data[]; };
    layout(set = 0, binding = 1) buffer Output { float partial_sums[]; };

    layout(push_constant) uniform PushConstants { uint n; };

    shared float sdata[256];

    void main() {
        uint gid = gl_GlobalInvocationID.x;
        uint lid = gl_LocalInvocationID.x;
        uint wgid = gl_WorkGroupID.x;

        // 加载到共享内存
        sdata[lid] = (gid < n) ? input_data[gid] : 0.0;
        barrier();

        // 工作组内的树形归约
        for (uint stride = 128; stride > 0; stride >>= 1) {
            if (lid < stride) {
                sdata[lid] += sdata[lid + stride];
            }
            barrier();
        }

        // 线程 0 写入工作组的局部和
        if (lid == 0) {
            partial_sums[wgid] = sdata[0];
        }
    }

- 这是经典的并行归约模式（与 CUDA 相同）。每个工作组产生一个局部和。第二次调度将这些局部和归约为最终结果。树形归约每一步将活跃线程减半：256 → 128 → 64 → ... → 1。

### 使用分块的矩阵乘法

    // matmul_tiled.comp
    #version 450

    #define TILE_SIZE 16

    layout(local_size_x = TILE_SIZE, local_size_y = TILE_SIZE) in;

    layout(set = 0, binding = 0) buffer MatA { float A[]; };
    layout(set = 0, binding = 1) buffer MatB { float B[]; };
    layout(set = 0, binding = 2) buffer MatC { float C[]; };

    layout(push_constant) uniform PushConstants {
        uint M, N, K;
    };

    shared float tileA[TILE_SIZE][TILE_SIZE];
    shared float tileB[TILE_SIZE][TILE_SIZE];

    void main() {
        uint row = gl_GlobalInvocationID.y;
        uint col = gl_GlobalInvocationID.x;
        uint lr = gl_LocalInvocationID.y;
        uint lc = gl_LocalInvocationID.x;

        float sum = 0.0;

        for (uint t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
            // 将 A 和 B 的分块加载到共享内存中
            uint aCol = t * TILE_SIZE + lc;
            uint bRow = t * TILE_SIZE + lr;

            tileA[lr][lc] = (row < M && aCol < K) ? A[row * K + aCol] : 0.0;
            tileB[lr][lc] = (bRow < K && col < N) ? B[bRow * N + col] : 0.0;

            barrier();

            // 计算部分点积
            for (uint k = 0; k < TILE_SIZE; k++) {
                sum += tileA[lr][k] * tileB[k][lc];
            }

            barrier();
        }

        if (row < M && col < N) {
            C[row * N + col] = sum;
        }
    }

- 这与 CUDA 版本（文件 04）中的分块算法相同，只是用了 GLSL 语法。概念完全一样：将分块加载到共享内存，屏障，计算，屏障，重复。

## C++ Vulkan 设置

- 计算着色器是简单的部分。困难的部分是创建 Vulkan 实例、分配内存、绑定缓冲区和提交命令的 C++ 样板代码。以下是完整管线的精简版本：

    // vulkan_compute.cpp — 一个最小但完整的 Vulkan 计算示例
    // 编译命令: g++ -O3 -o vulkan_compute vulkan_compute.cpp -lvulkan
    // 要求: 已安装 Vulkan SDK，已从 add.comp 编译 add.spv

    #include <vulkan/vulkan.h>
    #include <iostream>
    #include <vector>
    #include <fstream>
    #include <cassert>

    // 辅助函数：读取 SPIR-V 文件
    std::vector<uint32_t> readSPIRV(const std::string& filename) {
        std::ifstream file(filename, std::ios::ate | std::ios::binary);
        size_t fileSize = file.tellg();
        std::vector<uint32_t> buffer(fileSize / sizeof(uint32_t));
        file.seekg(0);
        file.read(reinterpret_cast<char*>(buffer.data()), fileSize);
        return buffer;
    }

    int main() {
        const uint32_t N = 1024;
        const size_t bufferSize = N * sizeof(float);

        // ========== 1. 创建 Vulkan 实例 ==========
        VkApplicationInfo appInfo{};
        appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
        appInfo.apiVersion = VK_API_VERSION_1_2;

        VkInstanceCreateInfo instanceInfo{};
        instanceInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
        instanceInfo.pApplicationInfo = &appInfo;

        VkInstance instance;
        vkCreateInstance(&instanceInfo, nullptr, &instance);

        // ========== 2. 选择物理设备 (GPU) ==========
        uint32_t deviceCount = 0;
        vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
        std::vector<VkPhysicalDevice> devices(deviceCount);
        vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());
        VkPhysicalDevice physicalDevice = devices[0];  // 使用第一个 GPU

        // 打印 GPU 名称
        VkPhysicalDeviceProperties props;
        vkGetPhysicalDeviceProperties(physicalDevice, &props);
        std::cout << "使用的 GPU: " << props.deviceName << "\n";

        // ========== 3. 查找计算队列族 ==========
        uint32_t queueFamilyCount = 0;
        vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, nullptr);
        std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
        vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice, &queueFamilyCount, queueFamilies.data());

        uint32_t computeFamily = 0;
        for (uint32_t i = 0; i < queueFamilyCount; i++) {
            if (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT) {
                computeFamily = i;
                break;
            }
        }

        // ========== 4. 创建逻辑设备和队列 ==========
        float queuePriority = 1.0f;
        VkDeviceQueueCreateInfo queueInfo{};
        queueInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueInfo.queueFamilyIndex = computeFamily;
        queueInfo.queueCount = 1;
        queueInfo.pQueuePriorities = &queuePriority;

        VkDeviceCreateInfo deviceInfo{};
        deviceInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
        deviceInfo.queueCreateInfoCount = 1;
        deviceInfo.pQueueCreateInfos = &queueInfo;

        VkDevice device;
        vkCreateDevice(physicalDevice, &deviceInfo, nullptr, &device);

        VkQueue computeQueue;
        vkGetDeviceQueue(device, computeFamily, 0, &computeQueue);

        // ========== 5. 分配缓冲区 (A, B, C) ==========
        // 为简洁起见，这里使用主机可见内存（较慢但更简单）
        auto createBuffer = [&](VkBuffer& buffer, VkDeviceMemory& memory) {
            VkBufferCreateInfo bufInfo{};
            bufInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
            bufInfo.size = bufferSize;
            bufInfo.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
            vkCreateBuffer(device, &bufInfo, nullptr, &buffer);

            VkMemoryRequirements memReqs;
            vkGetBufferMemoryRequirements(device, buffer, &memReqs);

            // 查找主机可见的内存类型
            VkPhysicalDeviceMemoryProperties memProps;
            vkGetPhysicalDeviceMemoryProperties(physicalDevice, &memProps);
            uint32_t memType = 0;
            for (uint32_t i = 0; i < memProps.memoryTypeCount; i++) {
                if ((memReqs.memoryTypeBits & (1 << i)) &&
                    (memProps.memoryTypes[i].propertyFlags &
                     (VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))) {
                    memType = i;
                    break;
                }
            }

            VkMemoryAllocateInfo allocInfo{};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memReqs.size;
            allocInfo.memoryTypeIndex = memType;
            vkAllocateMemory(device, &allocInfo, nullptr, &memory);
            vkBindBufferMemory(device, buffer, memory, 0);
        };

        VkBuffer bufA, bufB, bufC;
        VkDeviceMemory memA, memB, memC;
        createBuffer(bufA, memA);
        createBuffer(bufB, memB);
        createBuffer(bufC, memC);

        // ========== 6. 填充输入缓冲区 ==========
        float* ptrA;
        vkMapMemory(device, memA, 0, bufferSize, 0, (void**)&ptrA);
        for (uint32_t i = 0; i < N; i++) ptrA[i] = 1.0f;
        vkUnmapMemory(device, memA);

        float* ptrB;
        vkMapMemory(device, memB, 0, bufferSize, 0, (void**)&ptrB);
        for (uint32_t i = 0; i < N; i++) ptrB[i] = 2.0f;
        vkUnmapMemory(device, memB);

        // ========== 7. 创建计算管线 ==========
        auto spirvCode = readSPIRV("add.spv");
        VkShaderModuleCreateInfo shaderInfo{};
        shaderInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
        shaderInfo.codeSize = spirvCode.size() * sizeof(uint32_t);
        shaderInfo.pCode = spirvCode.data();
        VkShaderModule shaderModule;
        vkCreateShaderModule(device, &shaderInfo, nullptr, &shaderModule);

        // 描述符集布局（告诉 Vulkan 缓冲区绑定的信息）
        VkDescriptorSetLayoutBinding bindings[3] = {};
        for (int i = 0; i < 3; i++) {
            bindings[i].binding = i;
            bindings[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            bindings[i].descriptorCount = 1;
            bindings[i].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
        }

        VkDescriptorSetLayoutCreateInfo layoutInfo{};
        layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
        layoutInfo.bindingCount = 3;
        layoutInfo.pBindings = bindings;
        VkDescriptorSetLayout descLayout;
        vkCreateDescriptorSetLayout(device, &layoutInfo, nullptr, &descLayout);

        // 推送常量范围
        VkPushConstantRange pushRange{};
        pushRange.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
        pushRange.offset = 0;
        pushRange.size = sizeof(uint32_t);

        // 管线布局
        VkPipelineLayoutCreateInfo pipeLayoutInfo{};
        pipeLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipeLayoutInfo.setLayoutCount = 1;
        pipeLayoutInfo.pSetLayouts = &descLayout;
        pipeLayoutInfo.pushConstantRangeCount = 1;
        pipeLayoutInfo.pPushConstantRanges = &pushRange;
        VkPipelineLayout pipelineLayout;
        vkCreatePipelineLayout(device, &pipeLayoutInfo, nullptr, &pipelineLayout);

        // 计算管线
        VkComputePipelineCreateInfo pipeInfo{};
        pipeInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
        pipeInfo.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        pipeInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
        pipeInfo.stage.module = shaderModule;
        pipeInfo.stage.pName = "main";
        pipeInfo.layout = pipelineLayout;
        VkPipeline pipeline;
        vkCreateComputePipelines(device, VK_NULL_HANDLE, 1, &pipeInfo, nullptr, &pipeline);

        // ========== 8. 描述符集（将缓冲区绑定到着色器） ==========
        VkDescriptorPoolSize poolSize{};
        poolSize.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
        poolSize.descriptorCount = 3;

        VkDescriptorPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.maxSets = 1;
        poolInfo.poolSizeCount = 1;
        poolInfo.pPoolSizes = &poolSize;
        VkDescriptorPool descPool;
        vkCreateDescriptorPool(device, &poolInfo, nullptr, &descPool);

        VkDescriptorSetAllocateInfo descAllocInfo{};
        descAllocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        descAllocInfo.descriptorPool = descPool;
        descAllocInfo.descriptorSetCount = 1;
        descAllocInfo.pSetLayouts = &descLayout;
        VkDescriptorSet descSet;
        vkAllocateDescriptorSets(device, &descAllocInfo, &descSet);

        // 将缓冲区引用写入描述符集
        VkDescriptorBufferInfo bufInfos[3] = {
            {bufA, 0, bufferSize}, {bufB, 0, bufferSize}, {bufC, 0, bufferSize}
        };
        VkWriteDescriptorSet writes[3] = {};
        for (int i = 0; i < 3; i++) {
            writes[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            writes[i].dstSet = descSet;
            writes[i].dstBinding = i;
            writes[i].descriptorCount = 1;
            writes[i].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            writes[i].pBufferInfo = &bufInfos[i];
        }
        vkUpdateDescriptorSets(device, 3, writes, 0, nullptr);

        // ========== 9. 记录和提交命令缓冲区 ==========
        VkCommandPoolCreateInfo cmdPoolInfo{};
        cmdPoolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
        cmdPoolInfo.queueFamilyIndex = computeFamily;
        VkCommandPool cmdPool;
        vkCreateCommandPool(device, &cmdPoolInfo, nullptr, &cmdPool);

        VkCommandBufferAllocateInfo cmdAllocInfo{};
        cmdAllocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        cmdAllocInfo.commandPool = cmdPool;
        cmdAllocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        cmdAllocInfo.commandBufferCount = 1;
        VkCommandBuffer cmdBuf;
        vkAllocateCommandBuffers(device, &cmdAllocInfo, &cmdBuf);

        VkCommandBufferBeginInfo beginInfo{};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        vkBeginCommandBuffer(cmdBuf, &beginInfo);

        vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
        vkCmdBindDescriptorSets(cmdBuf, VK_PIPELINE_BIND_POINT_COMPUTE,
                                pipelineLayout, 0, 1, &descSet, 0, nullptr);
        vkCmdPushConstants(cmdBuf, pipelineLayout, VK_SHADER_STAGE_COMPUTE_BIT,
                           0, sizeof(uint32_t), &N);
        vkCmdDispatch(cmdBuf, (N + 255) / 256, 1, 1);  // 启动工作组

        vkEndCommandBuffer(cmdBuf);

        // 提交
        VkFenceCreateInfo fenceInfo{};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        VkFence fence;
        vkCreateFence(device, &fenceInfo, nullptr, &fence);

        VkSubmitInfo submitInfo{};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &cmdBuf;
        vkQueueSubmit(computeQueue, 1, &submitInfo, fence);
        vkWaitForFences(device, 1, &fence, VK_TRUE, UINT64_MAX);

        // ========== 10. 读取结果 ==========
        float* ptrC;
        vkMapMemory(device, memC, 0, bufferSize, 0, (void**)&ptrC);
        std::cout << "结果: c[0]=" << ptrC[0] << " c[1]=" << ptrC[1]
                  << " (期望值 3.0)\n";
        bool correct = true;
        for (uint32_t i = 0; i < N; i++) {
            if (ptrC[i] != 3.0f) { correct = false; break; }
        }
        std::cout << (correct ? "全部正确" : "发现错误") << "\n";
        vkUnmapMemory(device, memC);

        // ========== 清理（简写） ==========
        vkDestroyFence(device, fence, nullptr);
        vkDestroyCommandPool(device, cmdPool, nullptr);
        vkDestroyPipeline(device, pipeline, nullptr);
        vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
        vkDestroyDescriptorPool(device, descPool, nullptr);
        vkDestroyDescriptorSetLayout(device, descLayout, nullptr);
        vkDestroyShaderModule(device, shaderModule, nullptr);
        vkDestroyBuffer(device, bufA, nullptr); vkFreeMemory(device, memA, nullptr);
        vkDestroyBuffer(device, bufB, nullptr); vkFreeMemory(device, memB, nullptr);
        vkDestroyBuffer(device, bufC, nullptr); vkFreeMemory(device, memC, nullptr);
        vkDestroyDevice(device, nullptr);
        vkDestroyInstance(instance, nullptr);

        return 0;
    }

- 是的，向量加法就需要大约 200 行代码。 相比之下 CUDA 只需要大约 30 行。这就是显式性的代价。但请注意：每一行都有其目的。没有隐藏的驱动决策，没有隐式同步，没有意外的内存分配。你控制一切。

- 在实践中，你可以将这些样板代码封装到辅助库中（或使用现有的库，如 vk-bootstrap、用于内存分配的 VMA，或专注于 ML 的 Vulkan 计算库 kompute）。

## Kompute：为 ML 简化的 Vulkan

- Kompute 是一个开源 C++ 库，封装了 Vulkan 用于 GPU 计算的样板代码。同样的向量加法变成：

    #include <kompute/Kompute.hpp>

    int main() {
        kp::Manager mgr;

        auto tensorA = mgr.tensor({1, 1, 1, 1, 1});
        auto tensorB = mgr.tensor({2, 2, 2, 2, 2});
        auto tensorC = mgr.tensor({0, 0, 0, 0, 0});

        std::string shader = R"(
            #version 450
            layout(local_size_x = 1) in;
            layout(set=0, binding=0) buffer A { float a[]; };
            layout(set=0, binding=1) buffer B { float b[]; };
            layout(set=0, binding=2) buffer C { float c[]; };
            void main() {
                uint i = gl_GlobalInvocationID.x;
                c[i] = a[i] + b[i];
            }
        )";

        auto algorithm = mgr.algorithm({tensorA, tensorB, tensorC},
                                         kompute::Shader::compile_source(shader));

        mgr.sequence()
            ->record<kp::OpTensorSyncDevice>({tensorA, tensorB, tensorC})
            ->record<kp::OpAlgoDispatch>(algorithm)
            ->record<kp::OpTensorSyncLocal>({tensorC})
            ->eval();

        // tensorC 现在包含 [3, 3, 3, 3, 3]
    }

- 可读性强多了。Kompute 处理实例创建、设备选择、内存分配、描述符集和命令缓冲区管理。你只需关注着色器和数据。

## WebGPU：浏览器中的 GPU 计算

- WebGPU 是 WebGL 的继任者，提供从 JavaScript 访问现代 GPU 的能力。它基于 Vulkan（Linux/Android）、Metal（macOS/iOS）和 DirectX 12（Windows）构建，抽象了平台差异。

- WebGPU 使用 WGSL（WebGPU 着色语言）而非 GLSL：

    // add.wgsl — WebGPU 计算着色器
    @group(0) @binding(0) var<storage, read> a: array<f32>;
    @group(0) @binding(1) var<storage, read> b: array<f32>;
    @group(0) @binding(2) var<storage, read_write> c: array<f32>;

    @compute @workgroup_size(256)
    fn main(@builtin(global_invocation_id) id: vec3<u32>) {
        let i = id.x;
        c[i] = a[i] + b[i];
    }

- JavaScript 设置（精简版）：

    const adapter = await navigator.gpu.requestAdapter();
    const device = await adapter.requestDevice();

    // 创建缓冲区
    const bufferA = device.createBuffer({ size: N * 4, usage: GPUBufferUsage.STORAGE, mappedAtCreation: true });
    new Float32Array(bufferA.getMappedRange()).fill(1.0);
    bufferA.unmap();

    // ...（B 和 C 类似）

    // 从 WGSL 着色器创建管线
    const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module: device.createShaderModule({ code: wgslSource }), entryPoint: 'main' }
    });

    // 调度
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(Math.ceil(N / 256));
    pass.end();
    device.queue.submit([encoder.finish()]);

- 为什么 WebGPU 对 ML 很重要：在浏览器中运行推理意味着没有服务器成本、没有延迟，且用户数据永远不会离开设备。像 ONNX Runtime Web 和 Transformers.js 这样的库使用 WebGPU 完全在客户端运行模型（包括小型 LLM）。

## 何时使用 Vulkan

场景 | 使用 Vulkan？ | 原因 / 替代方案
--- | --- | ---
ML 训练 | 否 | CUDA/Triton 在 NVIDIA 上更简单更快速
NVIDIA GPU 上的推理 | 否 | TensorRT 或 CUDA 更好
AMD/Intel GPU 上的推理 | 是 | 唯一跨厂商的 GPU 计算选项
移动端推理（Android） | 是 | Vulkan 是 Android 上的标准 GPU API
移动端推理（iOS） | 否 | 直接使用 Metal（MoltenVK 增加开销）
浏览器推理 | WebGPU | 基于 Vulkan/Metal/DX12
游戏引擎 + ML | 是 | 引擎已使用 Vulkan 进行渲染
跨平台库 | 是 | 一套代码支持所有 GPU 厂商
学习 GPU 编程 | 视情况而定 | CUDA 更容易上手；Vulkan 能学到更多

## 编码任务（使用 g++ -lvulkan 编译，需要 Vulkan SDK）

1. 编译并运行上面的向量加法示例。修改着色器以计算 `c[i] = a[i] * b[i] + a[i]`（融合乘加）并验证结果。

2. 编写一个计算着色器，使用共享内存对一行数据应用 softmax（包括最大值和求和归约步骤）。用已知值进行测试。

    // softmax.comp — 编译命令: glslangValidator -V softmax.comp -o softmax.spv
    #version 450

    #define WG_SIZE 256

    layout(local_size_x = WG_SIZE) in;

    layout(set = 0, binding = 0) buffer Input  { float input_data[]; };
    layout(set = 0, binding = 1) buffer Output { float output_data[]; };

    layout(push_constant) uniform PC { uint n; };

    shared float sdata[WG_SIZE];

    void main() {
        uint gid = gl_GlobalInvocationID.x;
        uint lid = gl_LocalInvocationID.x;

        // 步骤 1：找最大值（数值稳定性）
        sdata[lid] = (gid < n) ? input_data[gid] : -1e30;
        barrier();
        for (uint s = WG_SIZE / 2; s > 0; s >>= 1) {
            if (lid < s) sdata[lid] = max(sdata[lid], sdata[lid + s]);
            barrier();
        }
        float maxVal = sdata[0];
        barrier();

        // 步骤 2：计算 exp(x - max)
        float expVal = (gid < n) ? exp(input_data[gid] - maxVal) : 0.0;
        sdata[lid] = expVal;
        barrier();

        // 步骤 3：exp 值求和
        for (uint s = WG_SIZE / 2; s > 0; s >>= 1) {
            if (lid < s) sdata[lid] += sdata[lid + s];
            barrier();
        }
        float sumExp = sdata[0];

        // 步骤 4：归一化
        if (gid < n) {
            output_data[gid] = expVal / sumExp;
        }
    }

1. 修改 C++ 宿主代码以对计算着色器进行基准测试：使用 Vulkan 时间戳查询或 CPU 端栅栏对调度（排除设置阶段）计时，并计算以 GB/s 为单位的实际带宽。