Vulkan Kompute 高级用法详解：从自定义操作到并行计算-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00265/article/details/148890277

Vulkan Kompute 高级用法详解：从自定义操作到并行计算

kompute 项目地址: https://gitcode.com/gh_mirrors/vu/vulkan-kompute

Vulkan Kompute 是一个基于 Vulkan 的计算框架，它简化了 GPU 计算任务的开发流程。本文将深入探讨 Vulkan Kompute 的高级用法，帮助开发者充分利用 GPU 的计算能力。

自定义 Kompute 操作

在 Vulkan Kompute 中，开发者可以创建自定义操作来封装特定的计算逻辑。这种方式特别适合需要重复使用的计算模式。

实现自定义操作

自定义操作通过继承 OpAlgoDispatch 类来实现。下面是一个实现张量乘法的自定义操作示例：

class OpTensorMultiply : public OpAlgoDispatch {
public:
    OpTensorMultiply(std::vector<std::shared_ptr<Tensor>> tensors,
         std::shared_ptr<kp::Algorithm> algorithm)
      : OpAlgoBase(algorithm)
    {
         if (tensors.size() != 3) {
             throw std::runtime_error("需要3个张量，但提供了" + tensors.size());
         }

         std::vector<uint32_t> spirv = compileSource(R"(
             #version 450
             layout(set = 0, binding = 0) buffer tensorLhs {
                float valuesLhs[];
             };
             layout(set = 0, binding = 1) buffer tensorRhs {
                float valuesRhs[];
             };
             layout(set = 0, binding = 2) buffer tensorOutput {
                float valuesOutput[];
             };
             layout (local_size_x = 1) in;
             void main() {
                 uint index = gl_GlobalInvocationID.x;
                 valuesOutput[index] = valuesLhs[index] * valuesRhs[index];
             }
         )");

         algorithm->rebuild(tensors, spirv);
    }
};

使用自定义操作

创建自定义操作后，可以像内置操作一样使用：

kp::Manager mgr;
auto tensorLhs = mgr.tensor({0., 1., 2.});
auto tensorRhs = mgr.tensor({2., 4., 6.});
auto tensorOut = mgr.tensor({0., 0., 0.});

mgr.sequence()
     ->record<kp::OpTensorSyncDevice>({tensorLhs, tensorRhs, tensorOut})
     ->record<OpTensorMultiply>({tensorLhs, tensorRhs, tensorOut}, mgr.algorithm())
     ->record<kp::OpTensorSyncLocal>({tensorOut})
     ->eval();

异步操作与等待机制

Vulkan Kompute 支持异步操作，允许在 GPU 执行计算的同时进行 CPU 端的其他工作。

基本异步流程

kp::Manager mgr;
auto tensor = mgr.tensor(10, 0.0);

// 同步数据到设备
mgr.sequence()->eval<kp::OpTensorSyncDevice>({tensor});

// 定义计算着色器
std::string shader = "..."; // 计算密集型着色器代码

// 异步执行计算
auto algo = mgr.algorithm({tensor}, compileSource(shader));
auto sq = mgr.sequence();
sq->evalAsync<kp::OpAlgoDispatch>(algo);

// 可以在这里执行其他CPU工作

// 等待GPU计算完成
sq->evalAwait();

// 同步结果回主机
sq->eval<kp::OpTensorSyncLocal>({tensor});

超时控制

evalAwait 方法可以指定超时时间（纳秒）：

// 等待最多100毫秒
bool completed = sq->evalAwait(100000000);
if (!completed) {
    // 处理超时情况
}

并行计算策略

Vulkan Kompute 支持通过多队列实现真正的并行计算，但需要注意硬件限制。

硬件队列特性

不同GPU厂商的队列实现差异很大：

NVIDIA显卡通常有多个队列家族
AMD显卡可能有不同的并行能力
集成显卡的并行能力通常较弱

并行计算实现

初始化多队列管理器：

// 选择设备0，使用家族0和家族2的队列
uint32_t deviceIndex = 0;
std::vector<uint32_t> familyIndices = {0, 2};
kp::Manager mgr(deviceIndex, familyIndices);

创建专用序列：

// 为每个队列创建专用序列
auto sqGraphics = mgr.sequence(0); // 使用第一个队列(图形家族)
auto sqCompute = mgr.sequence(1);  // 使用第二个队列(计算家族)

并行执行任务：

// 准备数据
auto tensorA = mgr.tensor({...});
auto tensorB = mgr.tensor({...});

// 编译着色器
std::vector<uint32_t> spirv = compileSource("...");

// 创建算法
auto algoA = mgr.algorithm({tensorA}, spirv);
auto algoB = mgr.algorithm({tensorB}, spirv);

// 并行提交任务
sqGraphics->evalAsync<kp::OpAlgoDispatch>(algoA);
sqCompute->evalAsync<kp::OpAlgoDispatch>(algoB);

// 等待两个任务完成
sqGraphics->evalAwait();
sqCompute->evalAwait();

Vulkan 扩展支持

Vulkan Kompute 允许启用特定的 Vulkan 扩展，如原子浮点操作：

// 启用VK_EXT_shader_atomic_float扩展
kp::Manager mgr(0, {}, {"VK_EXT_shader_atomic_float"});

// 使用原子操作的着色器
std::string shader = R"(
    #version 450
    #extension GL_EXT_shader_atomic_float: enable
    layout(binding = 0) buffer b { float pb[]; };
    void main() {
        atomicAdd(pb[0], 1.0);
    }
)";

// 创建张量和算法
auto tensor = mgr.tensor({0.0});
auto algo = mgr.algorithm({tensor}, compileSource(shader));

// 执行计算
mgr.sequence()
   ->record<kp::OpTensorSyncDevice>({tensor})
   ->record<kp::OpAlgoDispatch>(algo)
   ->record<kp::OpTensorSyncLocal>({tensor})
   ->eval();