Vulkan Compute 与跨平台 GPU¶
一句话总结: Vulkan 是唯一能横跨 NVIDIA/AMD/Intel/Apple/Android/浏览器的 GPU 计算 API，用显式控制换来极致可移植性，但代价是啰嗦到写"hello world"要 300 行。
Vulkan 是唯一能在所有主流平台上跑的 GPU 计算 API: NVIDIA、AMD、Intel、Apple (通过 MoltenVK)、Android，甚至浏览器 (通过 WebGPU)。本节介绍 Vulkan 架构、计算流水线、用 GLSL 写计算着色器、GPU 计算程序完整的 C++ 搭建流程、共享内存与同步、面向浏览器的 WebGPU，以及机器学习推理的实战示例。
-
统一计算设备架构 (Compute Unified Device Architecture, CUDA) 在 NVIDIA 硬件上主导了机器学习 (Machine Learning, ML) 训练。但不是每个部署目标都有 NVIDIA 图形处理器 (Graphics Processing Unit, GPU)。移动 app 跑在高通 Adreno 或 ARM Mali GPU 上。Web 应用跑在浏览器里。游戏引擎需要同时支持 AMD、Intel 与 NVIDIA。对所有这些场景，Vulkan 就是答案。
-
Vulkan 很啰嗦——一个"hello world"计算程序大概要 300 行 C++。但这份啰嗦是显式控制的代价：你自己管理每一个 GPU 资源 (内存、流水线、命令缓冲区)。这种控制换来极致性能与可移植性，代价是开发速度。
Vulkan 架构概览¶
-
Vulkan 是由 Khronos Group (OpenGL 背后的同一家组织) 打造的底层 GPU 应用程序编程接口 (Application Programming Interface, API)。不像 CUDA 那样把 GPU 资源管理藏起来，Vulkan 要求你显式管理:
- 实例与设备：创建 Vulkan 实例，枚举可用 GPU，选一个用。
- 内存：显式分配 GPU 内存，指定内存类型 (device-local 走速度，host-visible 允许 CPU 访问)。
- 缓冲区：创建引用已分配内存的缓冲区对象。
- 描述符集 (descriptor set)：把缓冲区绑到着色器输入上 (类似计算着色器的函数参数)。
- 计算流水线：编译着色器，创建流水线对象。
- 命令缓冲区：录下一串 GPU 命令 (绑流水线、绑描述符、发起 dispatch)。
- 队列提交：把命令缓冲区提交给 GPU 执行。
- 同步：用 fence 与 barrier 保证正确的顺序。
-
这跟 CUDA 的
cudaMalloc+ 核函数启动模型差别巨大。在 CUDA 里，驱动帮你把这些事都处理了。在 Vulkan 里，你自己上。
为什么这么啰嗦?¶
-
Vulkan 的显式化有两个理由:
-
让驱动更简单: OpenGL 驱动极其复杂 (它得猜应用想干什么，然后据此优化)。Vulkan 把这份责任转到应用侧，驱动因此更薄、更可预测、跨厂商实现也更容易做对。
-
性能：显式控制内存布局、同步与命令批处理，让应用能做出最优决策。在 CUDA 里驱动可能塞进不必要的同步；在 Vulkan 里你只在需要时才同步。
-
用 GLSL 写计算着色器¶
- 计算着色器 (compute shader) 是一段跑在 GPU 上的程序，类似 CUDA 核函数。它用 GLSL (OpenGL Shading Language) 写，编译成 SPIR-V 字节码 (一种可移植的二进制格式)。
向量加法¶
// add.comp —— 编译命令: glslangValidator -V add.comp -o add.spv
#version 450
// 工作组大小: 每组 256 个 invocation (= 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[]; };
// Push constant: 小块 uniform 数据 (类似核函数参数)
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 |
每个工作组的线程数 |
| Storage buffer | 全局内存 | 可读写的 GPU 内存 |
共享内存 (shared) |
__shared__ |
每个工作组独享的快速内存 |
| Push constant | 核函数参数 | 小块 uniform 数据 |
带共享内存的 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: 等本工作组所有 invocation 加载完
barrier(); // 等价于 CUDA 的 __syncthreads()
// 计算 ReLU
if (gid < n) {
output_data[gid] = max(tile[lid], 0.0);
}
}
- 修正线性单元 (Rectified Linear Unit, ReLU) 其实不需要共享内存 (它是逐元素运算)。但这段代码演示了套路：加载到共享内存 → barrier → 计算 → 写回。对于需要读邻居线程数据的算子 (卷积、归约、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 一样)。每个工作组产生一个部分和，再发起第二次 dispatch 把这些部分和归约成最终结果。树形归约每一步把活跃线程折半：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 语法。思路一模一样：加载分块到共享内存 → barrier → 计算 → barrier → 循环。
C++ 端的 Vulkan 搭建¶
- 计算着色器是简单的部分。难的是那堆 C++ 样板代码——创建 Vulkan 实例、分配内存、绑定缓冲区、提交命令。下面是完整流水线的浓缩版:
// 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 << "Using 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);
// Push constant 范围
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 << "Results: c[0]=" << ptrC[0] << " c[1]=" << ptrC[1]
<< " (expected 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 ? "ALL CORRECT" : "ERRORS FOUND") << "\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，或者面向机器学习的 kompute)。
Kompute：简化版 Vulkan for ML¶
- Kompute 是一个开源 C++ 库，把 Vulkan 计算的样板代码包起来。同样的向量加法就变成:
#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 Shading Language) 而非 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 对机器学习的意义：推理跑在浏览器里，意味着零服务器成本、零网络延迟，用户数据也不必离开设备。像 ONNX Runtime Web 和 Transformers.js 这类库就用 WebGPU 完全在客户端跑模型 (甚至包括小型大语言模型 (Large Language Model, LLM))。
什么时候用 Vulkan¶
| 场景 | 用 Vulkan? | 原因 / 替代方案 |
|---|---|---|
| 机器学习训练 | 不 | CUDA/Triton 在 NVIDIA 上更简单更快 |
| NVIDIA GPU 上的推理 | 不 | TensorRT 或 CUDA 更好 |
| AMD/Intel GPU 上的推理 | 要 | 唯一跨厂商 GPU 计算方案 |
| 移动端推理 (Android) | 要 | Vulkan 是 Android 的标准 GPU API |
| 移动端推理 (iOS) | 不 | 直接用 Metal (MoltenVK 有额外开销) |
| 浏览器推理 | WebGPU | 底层跑在 Vulkan/Metal/DX12 上 |
| 游戏引擎 + 机器学习 | 要 | 引擎渲染本来就用 Vulkan |
| 跨平台库 | 要 | 一份代码打通所有 GPU 厂商 |
| 学 GPU 编程 | 看情况 | CUDA 更容易上手；Vulkan 教得更多 |
编程任务 (用 g++ -lvulkan 编译，需装 Vulkan SDK)¶
-
编译并运行上面的向量加法示例。改一下着色器，让它计算
c[i] = a[i] * b[i] + a[i](融合乘加)，并验证结果。 -
写一个计算着色器，对一行数据做 softmax，用共享内存做归约步骤 (max 与 sum)。用已知值测试。
// 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 步: 找 max (为了数值稳定)
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;
}
}
- 修改 C++ 宿主代码，对计算着色器做基准测试：用 Vulkan 时间戳查询或 CPU 端的 fence 给 dispatch 计时 (不算搭建阶段)，并计算达到的带宽 (GB/s)。
📌 本节要点¶
- Vulkan 定位：跨 NVIDIA/AMD/Intel/Apple/Android/浏览器的底层 GPU 计算 API，由 Khronos Group 维护。
- 显式控制哲学：实例、内存、缓冲区、描述符、流水线、命令缓冲区、同步全要自己写，换来薄驱动与最优性能。
- 计算着色器：用 GLSL 编写，编译为 SPIR-V 字节码，跨平台可移植。
- CUDA 对应关系: workgroup ↔ block, invocation ↔ thread,
barrier()↔__syncthreads(),shared↔__shared__。 - 典型模式：加载到共享内存 → barrier → 计算 → barrier，归约、softmax、分块 GEMM 都是这个骨架。
- 样板代价：一个向量加法 C++ 端要 200 行，CUDA 只要 30 行，显式化的代价换来零隐藏行为。
- Kompute 简化：开源库把 Vulkan 样板包起来，ML 推理场景可读性大幅提升。
- WebGPU 浏览器：用 WGSL 写着色器，底层跑在 Vulkan/Metal/DX12 上，让模型在客户端跑，无服务器成本。
- 选型建议: NVIDIA 训练走 CUDA/Triton，移动端 Android 走 Vulkan, iOS 走 Metal，浏览器走 WebGPU，跨厂商库才必须上 Vulkan。