# 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** 字节码(一种可移植的二进制格式)。 ### 向量加法 ```glsl // 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 ```glsl // 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),共享内存是必不可少的。 ### 并行归约(求和) ```glsl // 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。 ### 使用分块的矩阵乘法 ```glsl // 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++ 样板代码。以下是完整管线的精简版本: ```cpp // vulkan_compute.cpp — 一个最小但完整的 Vulkan 计算示例 // 编译命令: g++ -O3 -o vulkan_compute vulkan_compute.cpp -lvulkan // 要求: 已安装 Vulkan SDK,已从 add.comp 编译 add.spv #include #include #include #include #include // 辅助函数:读取 SPIR-V 文件 std::vector readSPIRV(const std::string& filename) { std::ifstream file(filename, std::ios::ate | std::ios::binary); size_t fileSize = file.tellg(); std::vector buffer(fileSize / sizeof(uint32_t)); file.seekg(0); file.read(reinterpret_cast(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 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 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 计算的样板代码。同样的向量加法变成: ```cpp #include 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({tensorA, tensorB, tensorC}) ->record(algorithm) ->record({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: ```wgsl // add.wgsl — WebGPU 计算着色器 @group(0) @binding(0) var a: array; @group(0) @binding(1) var b: array; @group(0) @binding(2) var c: array; @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) id: vec3) { let i = id.x; c[i] = a[i] + b[i]; } ``` - **JavaScript 设置**(精简版): ```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(包括最大值和求和归约步骤)。用已知值进行测试。 ```glsl // 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; } } ``` 3. 修改 C++ 宿主代码以对计算着色器进行基准测试:使用 Vulkan 时间戳查询或 CPU 端栅栏对调度(排除设置阶段)计时,并计算以 GB/s 为单位的实际带宽。