Android设备上Vulkan计算着色器16位浮点运算实现咨询

当然可以在Vulkan计算着色器中使用16位浮点类型！这是优化计算性能、降低内存带宽开销的常用方案，不过有几个关键细节需要你留意，我结合你的矩阵乘法场景逐一说明：

1. 先确认设备支持16位浮点着色器操作 #

不是所有Vulkan设备都原生支持float16的着色器运算，你需要先做兼容性检查：

• 如果目标是Vulkan 1.2及以上版本：shaderFloat16是核心特性，可通过VkPhysicalDeviceVulkan12Features结构体检查。
• 如果是Vulkan 1.2以下版本：需要启用VK_KHR_shader_float16_int8扩展，并验证对应的特性是否开启。

示例检查代码：

VkPhysicalDeviceFeatures2 features{};
features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;

VkPhysicalDeviceShaderFloat16Int8FeaturesKHR float16Features{};
float16Features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR;
features.pNext = &float16Features;

vkGetPhysicalDeviceFeatures2(physicalDevice, &features);

if (!float16Features.shaderFloat16) {
    // 设备不支持16位浮点着色器操作，需要做降级处理（比如改用float32）
}

2. 在计算着色器中使用float16类型 #

在GLSL着色器里，你需要启用对应扩展（或依赖核心版本支持），然后直接使用float16_t、vecNh（如vec2h、vec4h）、matNxMh等类型定义16位浮点变量和矩阵。

针对你的矩阵乘法场景，着色器示例如下：

#version 450
// Vulkan 1.2以下版本需要启用该扩展
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : enable

// 可替换为通过uniform传入的动态维度
const int M = 1024;
const int K = 512;
const int N = 1024;

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

// 绑定你主机端的单缓冲区，按顺序存储A、B、C矩阵
layout(set = 0, binding = 0) buffer MatrixBuffer {
    matKhAh matA; // M行K列的float16矩阵
    matNhKh matB; // K行N列的float16矩阵
    matNhAh matC; // 输出的M行N列float16矩阵
} buf;

void main() {
    ivec2 threadPos = ivec2(gl_GlobalInvocationID.xy);
    // 跳过超出矩阵范围的线程
    if (threadPos.x >= M || threadPos.y >= N) return;

    float16_t sum = 0.0h;
    for (int k = 0; k < K; ++k) {
        sum += buf.matA[threadPos.x][k] * buf.matB[k][threadPos.y];
    }
    buf.matC[threadPos.x][threadPos.y] = sum;
}

3. 主机端缓冲区处理的注意事项 #

你目前用half库将float32转float16的方式是完全可行的，补充几个细节：

• 缓冲区大小计算：确保缓冲区总大小为(M*K + K*N + M*N) * sizeof(half)，每个float16元素占2字节，别算错总字节数。
• 内存对齐：Vulkan对缓冲区偏移有对齐要求，float16的最小对齐是2字节，你代码中按元素数计算偏移（如offset = M*K），实际字节偏移为offset * sizeof(half)，完全符合对齐要求，无需额外调整。
• 缓冲区用法：创建缓冲区时，VkBufferCreateInfo的usage字段必须包含VK_BUFFER_USAGE_STORAGE_BUFFER_BIT，确保计算着色器能读写该缓冲区。

你的主机端转换代码片段是正确的，这里再补充完整一点逻辑：

half* matrixTmp;
vkMapMemory(device, bufferMemory, 0, matrixSize, 0, (void**)&matrixTmp);

int offset = 0;
// 写入矩阵A
for(int i = 0; i < M*K; i++) {
    matrixTmp[offset+i] = half_cast<half>(matrixA[i]);
}
offset += M*K;
// 写入矩阵B
for(int i = 0; i < K*N; i++) {
    matrixTmp[offset+i] = half_cast<half>(matrixB[i]);
}

// 执行计算着色器指令...

// 计算完成后读取矩阵C
offset += K*N;
for(int i = 0; i < M*N; i++) {
    matrixC[i] = static_cast<float>(matrixTmp[offset+i]);
}

vkUnmapMemory(device, bufferMemory);

4. 精度与性能的权衡 #

使用float16能大幅减少内存带宽占用、提升计算效率，但它的精度只有约3位十进制有效数字。如果你的矩阵乘法对精度要求极高，建议先评估是否接受这种精度损失；如果需要折中，可以考虑中间计算用float32，最终结果再转成float16存储。

内容的提问来源于stack exchange，提问作者Tsesar Rizqi