Zig GPU 计算:Vulkan Compute 与着色器绑定

Vulkan 不只是图形 API,它的 compute queue 能直接做通用计算,且跨 NVIDIA/AMD/Intel/移动 GPU。本文用 Zig 走完 instance/device 初始化、设备内存分配、描述符集绑定、SPIR-V 管线与命令提交,并给出向量加法的完整代码与性能调优要点。

1. 为什么用 Vulkan 而不是 CUDA

CUDA 生态成熟、工具链完善,但只跑在 NVIDIA 显卡上。Vulkan 的 compute queue 提供了同等的通用计算能力,且:

  • 跨厂商:NVIDIA、AMD、Intel、Apple(MoltenVK)、Android 都能跑。
  • 显式控制:内存类型、同步原语、队列调度全部由应用决定,没有隐式同步带来的性能意外。
  • 无运行时开销:驱动层极薄,适合嵌入到已有应用里。

代价是样板代码量巨大:一个向量加法要写 300 行初始化。Zig 在这里的优势是 @cImport 直接吃 vulkan.h,不需要手写绑定,同时保留对内存布局的完全控制。

如果你只是想理解 GPU 的硬件执行模型(SIMT、warp、occupancy),先看 GPU 架构与并行计算模型 ;本文聚焦 Vulkan 的 API 层面。若你的目标是推理而非通用计算,/zig-ml-inference/ 里讨论了推理框架的选型。

2. 初始化:instance、物理设备与队列

2.1 创建 instance

const std = @import("std");
const vk = @cImport({
    @cInclude("vulkan/vulkan.h");
});

const app_info = vk.VkApplicationInfo{
    .sType = vk.VK_STRUCTURE_TYPE_APPLICATION_INFO,
    .pApplicationName = "zig-compute",
    .applicationVersion = vk.VK_MAKE_VERSION(1, 0, 0),
    .pEngineName = "none",
    .engineVersion = vk.VK_MAKE_VERSION(1, 0, 0),
    .apiVersion = vk.VK_API_VERSION_1_2,
};

var inst: vk.VkInstance = undefined;
const ci = vk.VkInstanceCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
    .pApplicationInfo = &app_info,
    .enabledLayerCount = 0,
    .ppEnabledLayerNames = null,
    .enabledExtensionCount = 0,
    .ppEnabledExtensionNames = null,
};

if (vk.vkCreateInstance(&ci, null, &inst) != vk.VK_SUCCESS) {
    return error.InstanceCreationFailed;
}
defer vk.vkDestroyInstance(inst, null);

Zig 的 @cImport 会把 C 的 VkApplicationInfo 转成 extern struct,sType 字段必须显式赋值——Vulkan 靠它做结构体版本校验,漏填会返回 VK_ERROR_INITIALIZATION_FAILED。

2.2 选择支持计算的物理设备

关键判据是**队列族(Queue Family)**里有没有 COMPUTE 位:

pub fn pickDevice(inst: vk.VkInstance) !struct { dev: vk.VkPhysicalDevice, family: u32 } {
    var count: u32 = 0;
    _ = vk.vkEnumeratePhysicalDevices(inst, &count, null);
    if (count == 0) return error.NoVulkanDevice;

    const devs = try std.heap.page_allocator.alloc(vk.VkPhysicalDevice, count);
    defer std.heap.page_allocator.free(devs);
    _ = vk.vkEnumeratePhysicalDevices(inst, &count, devs.ptr);

    for (devs) |dev| {
        var fam_count: u32 = 0;
        vk.vkGetPhysicalDeviceQueueFamilyProperties(dev, &fam_count, null);
        const fams = try std.heap.page_allocator.alloc(vk.VkQueueFamilyProperties, fam_count);
        defer std.heap.page_allocator.free(fams);
        vk.vkGetPhysicalDeviceQueueFamilyProperties(dev, &fam_count, fams.ptr);

        for (fams, 0..) |f, i| {
            if (f.queueFlags & vk.VK_QUEUE_COMPUTE_BIT != 0) {
                return .{ .dev = dev, .family = @intCast(i) };
            }
        }
    }
    return error.NoComputeQueue;
}

优先选只有 COMPUTE 而没有 GRAPHICS 的队列族(如果存在)——这类纯计算队列不会被图形任务抢占,延迟更稳定。很多独显有独立的 compute-only 队列。

2.3 逻辑设备与队列

const qp = vk.VkDeviceQueueCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
    .queueFamilyIndex = family,
    .queueCount = 1,
    .pQueuePriorities = &[_]f32{1.0},
};
const dci = vk.VkDeviceCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
    .queueCreateInfoCount = 1,
    .pQueueCreateInfos = &qp,
    .enabledExtensionCount = 0,
};

var device: vk.VkDevice = undefined;
if (vk.vkCreateDevice(phys, &dci, null, &device) != vk.VK_SUCCESS) {
    return error.DeviceCreationFailed;
}
var queue: vk.VkQueue = undefined;
vk.vkGetDeviceQueue(device, family, 0, &queue);

3. 设备内存:Vulkan 最复杂的部分

Vulkan 不让你直接 malloc 显存。要先选内存类型(Memory Type),再分配 VkDeviceMemory,最后把它绑到 VkBuffer 上。

3.1 内存类型与堆

vkGetPhysicalDeviceMemoryProperties 返回一个内存类型数组与堆数组。每个类型有属性位与所属堆:

pub fn findMemoryType(
    phys: vk.VkPhysicalDevice,
    type_filter: u32,
    props: vk.VkMemoryPropertyFlags,
) !u32 {
    var mp: vk.VkPhysicalDeviceMemoryProperties = undefined;
    vk.vkGetPhysicalDeviceMemoryProperties(phys, &mp);

    for (0..mp.memoryTypeCount) |i| {
        const supported = (type_filter & (@as(u32, 1) << @intCast(i))) != 0;
        const has = (mp.memoryTypes[i].propertyFlags & props) == props;
        if (supported and has) return @intCast(i);
    }
    return error.NoSuitableMemoryType;
}

属性位里最常用的两个:

属性含义用途
DEVICE_LOCAL显存,GPU 访问最快计算输入输出缓冲区
HOST_VISIBLECPU 可映射上传/回读数据
HOST_COHERENTCPU 写无需 flush简化上传路径
HOST_CACHEDCPU 读走缓存频繁回读结果

3.2 缓冲区创建

pub fn createBuffer(
    device: vk.VkDevice,
    phys: vk.VkPhysicalDevice,
    size: u64,
    usage: vk.VkBufferUsageFlags,
    props: vk.VkMemoryPropertyFlags,
) !struct { buf: vk.VkBuffer, mem: vk.VkDeviceMemory } {
    const bci = vk.VkBufferCreateInfo{
        .sType = vk.VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        .size = size,
        .usage = usage,
        .sharingMode = vk.VK_SHARING_MODE_EXCLUSIVE,
    };
    var buf: vk.VkBuffer = undefined;
    if (vk.vkCreateBuffer(device, &bci, null, &buf) != vk.VK_SUCCESS) {
        return error.BufferCreationFailed;
    }

    var req: vk.VkMemoryRequirements = undefined;
    vk.vkGetBufferMemoryRequirements(device, buf, &req);

    const alloc = vk.VkMemoryAllocateInfo{
        .sType = vk.VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        .allocationSize = req.size,
        .memoryTypeIndex = try findMemoryType(phys, req.memoryTypeBits, props),
    };
    var mem: vk.VkDeviceMemory = undefined;
    if (vk.vkAllocateMemory(device, &alloc, null, &mem) != vk.VK_SUCCESS) {
        return error.AllocationFailed;
    }
    _ = vk.vkBindBufferMemory(device, buf, mem, 0);
    return .{ .buf = buf, .mem = mem };
}

注意 req.size 可能大于你请求的 size(对齐填充),分配时必须用 req.size,否则 vkBindBufferMemory 报 VK_ERROR_OUT_OF_DEVICE_MEMORY。

3.3 暂存缓冲区上传

GPU 不能直接高效读 HOST_VISIBLE 内存。标准做法是两段式:写一块 host-visible 的暂存缓冲区,再用 vkCmdCopyBuffer 拷到 device-local 缓冲区。

pub fn upload(
    device: vk.VkDevice,
    queue: vk.VkQueue,
    pool: vk.VkCommandPool,
    dst: vk.VkBuffer,
    data: []const u8,
    staging: vk.VkBuffer,
    staging_mem: vk.VkDeviceMemory,
) !void {
    var mapped: ?*anyopaque = null;
    _ = vk.vkMapMemory(device, staging_mem, 0, data.len, 0, &mapped);
    @memcpy(@as([*]u8, @ptrCast(mapped.?))[0..data.len], data);
    vk.vkUnmapMemory(device, staging_mem);

    // 记录一条拷贝命令并提交
    var cmd: vk.VkCommandBuffer = undefined;
    const ai = vk.VkCommandBufferAllocateInfo{
        .sType = vk.VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        .commandPool = pool,
        .level = vk.VK_COMMAND_BUFFER_LEVEL_PRIMARY,
        .commandBufferCount = 1,
    };
    _ = vk.vkAllocateCommandBuffers(device, &ai, &cmd);

    const bi = vk.VkCommandBufferBeginInfo{
        .sType = vk.VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = vk.VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
    };
    _ = vk.vkBeginCommandBuffer(cmd, &bi);
    const region = vk.VkBufferCopy{ .srcOffset = 0, .dstOffset = 0, .size = data.len };
    vk.vkCmdCopyBuffer(cmd, staging, dst, 1, &region);
    _ = vk.vkEndCommandBuffer(cmd);

    var fence: vk.VkFence = undefined;
    _ = vk.vkCreateFence(device, &vk.VkFenceCreateInfo{
        .sType = vk.VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
    }, null, &fence);
    defer vk.vkDestroyFence(device, fence, null);

    const si = vk.VkSubmitInfo{
        .sType = vk.VK_STRUCTURE_TYPE_SUBMIT_INFO,
        .commandBufferCount = 1,
        .pCommandBuffers = &cmd,
    };
    _ = vk.vkQueueSubmit(queue, 1, &si, fence);
    _ = vk.vkWaitForFences(device, 1, &fence, vk.VK_TRUE, std.math.maxInt(u64));
}

若设备支持 VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | HOST_VISIBLE(AMD 的 SAM、部分集显),可以省掉暂存,直接映射 device-local 内存——但可移植性差,要做能力检测。

4. 描述符集:把缓冲区绑定到着色器

着色器不能直接看到 VkBuffer。它通过 描述符(Descriptor) 访问:描述符集布局(Descriptor Set Layout)声明「有哪些绑定槽、每个槽是什么类型」,描述符集(Descriptor Set)填入具体缓冲区,着色器用 binding = N 引用。

4.1 布局与池

const bindings = [_]vk.VkDescriptorSetLayoutBinding{
    .{
        .binding = 0,
        .descriptorType = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
        .stageFlags = vk.VK_SHADER_STAGE_COMPUTE_BIT,
    },
    .{
        .binding = 1,
        .descriptorType = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
        .stageFlags = vk.VK_SHADER_STAGE_COMPUTE_BIT,
    },
};

var set_layout: vk.VkDescriptorSetLayout = undefined;
const slci = vk.VkDescriptorSetLayoutCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
    .bindingCount = bindings.len,
    .pBindings = &bindings,
};
_ = vk.vkCreateDescriptorSetLayout(device, &slci, null, &set_layout);

const pool_sizes = [_]vk.VkDescriptorPoolSize{.{
    .type = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
    .descriptorCount = 2,
}};
var dpool: vk.VkDescriptorPool = undefined;
_ = vk.vkCreateDescriptorPool(device, &vk.VkDescriptorPoolCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .maxSets = 1,
    .poolSizeCount = pool_sizes.len,
    .pPoolSizes = &pool_sizes,
}, null, &dpool);

描述符池的容量要精确预算。常见错误是 maxSets = 100 但 poolSizeCount 只给了 1 个 storage buffer,第二次分配就 VK_ERROR_OUT_OF_POOL_MEMORY。池不可扩容,只能重建。

4.2 写入描述符

var dset: vk.VkDescriptorSet = undefined;
_ = vk.vkAllocateDescriptorSets(device, &vk.VkDescriptorSetAllocateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
    .descriptorPool = dpool,
    .descriptorSetCount = 1,
    .pSetLayouts = &set_layout,
}, &dset);

const info_a = vk.VkDescriptorBufferInfo{ .buffer = buf_a, .offset = 0, .range = vk.VK_WHOLE_SIZE };
const info_b = vk.VkDescriptorBufferInfo{ .buffer = buf_b, .offset = 0, .range = vk.VK_WHOLE_SIZE };

const writes = [_]vk.VkWriteDescriptorSet{
    .{
        .sType = vk.VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = dset,
        .dstBinding = 0,
        .descriptorCount = 1,
        .descriptorType = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .pBufferInfo = &info_a,
    },
    .{
        .sType = vk.VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = dset,
        .dstBinding = 1,
        .descriptorCount = 1,
        .descriptorType = vk.VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .pBufferInfo = &info_b,
    },
};
vk.vkUpdateDescriptorSets(device, writes.len, &writes, 0, null);

range = VK_WHOLE_SIZE 表示整块缓冲区。若只绑定一部分,range 必须按 minStorageBufferOffsetAlignment 对齐——这个值在 VkPhysicalDeviceLimits 里,常见是 16 或 256 字节。描述符绑定的细节与常见陷阱,参见 Vulkan 描述符与着色器绑定 。

5. 计算管线与 SPIR-V

Vulkan 只接受 SPIR-V 字节码。从 Zig 侧有几条路:

方式工具优点缺点
GLSL → SPIR-VglslangValidator -V生态成熟多一步构建
HLSL → SPIR-Vdxc -spirv熟悉 HLSL 的人友好同上
Zig → SPIR-V手写或用 spirv 库无外部工具门槛高
预编译嵌入@embedFile运行时零依赖改动需重编

最实用的是构建期编译 + @embedFile 嵌入,产出的可执行文件不带任何外部资源:

const shader_spv = @embedFile("shaders/vector_add.spv");

var module: vk.VkShaderModule = undefined;
const smci = vk.VkShaderModuleCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
    .codeSize = shader_spv.len,
    .pCode = @ptrCast(@alignCast(shader_spv.ptr)),
};
_ = vk.vkCreateShaderModule(device, &smci, null, &module);

pCode 必须是 4 字节对齐的 u32 指针。@embedFile 返回的是 []const u8,@alignCast 在 Debug 构建下会做对齐检查,若编译器没把数组按 4 字节对齐,这里会 panic——可以在声明处加 align(4) 的包装。

5.1 着色器代码

#version 450
layout(local_size_x = 256) in;

layout(binding = 0) readonly buffer A { float a[]; };
layout(binding = 1) readonly buffer B { float b[]; };
layout(binding = 2) writeonly buffer C { float c[]; };

void main() {
    uint i = gl_GlobalInvocationID.x;
    if (i < a.length()) {
        c[i] = a[i] + b[i];
    }
}

local_size_x = 256 是工作组大小。经验取值:

GPU 类型建议工作组大小说明
NVIDIA128~256需为 32(warp)的倍数
AMD64~256wavefront 64,需为其倍数
Intel64~256SIMD8/16/32
移动 GPU64~128寄存器压力大,别开太大

5.2 管线创建

const stage = vk.VkPipelineShaderStageCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
    .stage = vk.VK_SHADER_STAGE_COMPUTE_BIT,
    .module = module,
    .pName = "main",
};

var layout: vk.VkPipelineLayout = undefined;
const plci = vk.VkPipelineLayoutCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
    .setLayoutCount = 1,
    .pSetLayouts = &set_layout,
    // push constants 用于传 N 等标量,避免额外的 UBO
    .pushConstantRangeCount = 0,
};
_ = vk.vkCreatePipelineLayout(device, &plci, null, &layout);

var pipeline: vk.VkPipeline = undefined;
const pci = vk.VkComputePipelineCreateInfo{
    .sType = vk.VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
    .stage = stage,
    .layout = layout,
};
if (vk.vkCreateComputePipelines(device, vk.VK_NULL_HANDLE, 1, &pci, null, &pipeline) != vk.VK_SUCCESS) {
    return error.PipelineCreationFailed;
}

6. 提交与同步

vk.vkCmdBindPipeline(cmd, vk.VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
vk.vkCmdBindDescriptorSets(
    cmd,
    vk.VK_PIPELINE_BIND_POINT_COMPUTE,
    layout,
    0,
    1,
    &dset,
    0,
    null,
);

// 工作组数量 = ceil(N / 256)
const groups = (n + 255) / 256;
vk.vkCmdDispatch(cmd, @intCast(groups), 1, 1);

_ = vk.vkEndCommandBuffer(cmd);

// 提交前若该 fence 处于已触发状态,必须先重置
_ = vk.vkResetFences(device, 1, &fence);
const si = vk.VkSubmitInfo{
    .sType = vk.VK_STRUCTURE_TYPE_SUBMIT_INFO,
    .commandBufferCount = 1,
    .pCommandBuffers = &cmd,
};
_ = vk.vkQueueSubmit(queue, 1, &si, fence);
_ = vk.vkWaitForFences(device, 1, &fence, vk.VK_TRUE, std.math.maxInt(u64));

同步原语的选用:

原语方向典型用途
VkFenceGPU → CPU等待整批命令完成
VkSemaphoreGPU → GPU队列间/批次间依赖
vkQueueWaitIdle全阻塞仅调试期用

生产代码里 vkQueueWaitIdle 是性能杀手——它阻塞整个队列。要用 fence 或 semaphore 做细粒度同步。Vulkan 的同步模型还涉及管线屏障(pipeline barrier)与可用性/可见性两套内存语义,写复杂多阶段计算前务必先吃透这部分规则。

7. 性能要点

  • 缓冲区复用:每帧重新分配 VkBuffer 会触发驱动内部锁。用池化管理,按尺寸档位(4K/64K/1M/16M)预分配。
  • 减少提交次数:把多个 dispatch 录进同一个 command buffer 一次提交,比多次 vkQueueSubmit 快得多。
  • push constants 传标量:N、stride 这类小参数用 push constants(最多 128 字节),免去额外的 uniform buffer 与描述符更新。
  • 避免 GPU→CPU 回读:回读要走 staging buffer + fence,延迟以毫秒计。能用 GPU 侧归约就不要拉回 CPU。
  • 局部大小与寄存器压力:local_size_x 越大,每线程可用寄存器越少,可能反而降 occupancy。用厂商工具(Nsight Compute、Radeon GPU Profiler)实测。

7.1 计算与 SIMD 的关系

GPU 的 SIMT 模型和 CPU 的 SIMD 解决的是不同粒度的问题。小规模(几千元素)的向量运算在 CPU 上用 /zig-simd-vectorization/ 往往比启动 Vulkan 更快——因为 Vulkan 的初始化与提交开销是毫秒级的。经验阈值是百万元素以上再考虑 GPU。

小结

用 Zig 写 Vulkan 计算的核心链条是:

  1. vkCreateInstance → 选物理设备 → 建逻辑设备与 compute 队列。
  2. 按内存类型分配 VkDeviceMemory,绑到 VkBuffer;跨设备传输走 staging。
  3. 描述符集布局声明绑定槽,描述符集填入具体缓冲区。
  4. 着色器编译成 SPIR-V,@embedFile 嵌入,构建 compute pipeline。
  5. 录制 command buffer,vkCmdDispatch 提交,fence 等待。

最容易出错的三个点:sType 漏填、描述符池容量预算不足、local_size_x 与厂商 warp 大小不匹配。图形与计算共享 device/queue 时,还要注意队列族的 GRAPHICS | COMPUTE 组合位与独立计算队列的调度差异。

Zig 在这条链路上的价值是零绑定开销:@cImport 直接映射 C API,extern struct 保证内存布局,@embedFile 把着色器编进二进制——没有反射、没有运行时类型信息,每个字节都在你的掌控中。

继续阅读

探索更多技术文章

浏览归档,发现更多关于系统设计、工具链和工程实践的内容。

全部文章 返回首页

「系统编程」更多文章

  1. Zig 终端 TUI 开发:终端控制、布局与交互
  2. Zig 打包与分发:容器镜像、系统包与 Homebrew
  3. Zig WASI 与组件模型:沙箱运行时与宿主嵌入