在金属中填充浮动缓冲区

Filling Float buffer in Metal

问题:

我需要用一个常量值填充 MTLBufferFloats — 比如 1729.68921。我也需要它尽可能快。

因此我被禁止在 CPU 端填充缓冲区(即从 MTLBuffer 获取 UnsafeMutablePointer<Float> 并以串行方式分配)。

我的做法

理想情况下我会使用 MTLBlitCommandEncoder.fill(),但是据我所知,它只能用 UInt8 值填充缓冲区(假设 UInt8 是 1 个字节长并且 Float是 4 个字节长,我不能指定我的 Float 常量的任意值)。

到目前为止,我只能看到 2 个选项,但两个选项似乎都太过分了:

  1. 创建另一个缓冲区 B 填充常量值并通过 MTLBlitCommandEncoder
  2. 将其内容复制到我的缓冲区
  3. 创建一个 kernel 函数来填充缓冲区

问题

Float 填充 MTLBuffer 的最快方法是什么 常数值?

使用从每个线程写入多个缓冲区元素的计算着色器是我实验中最快的方法。这取决于硬件,因此您应该在您希望部署该应用程序的所有设备上进行测试。

我写了两个计算着色器:一个填充 16 个连续数组元素而不检查数组边界,另一个在检查缓冲区长度后设置单个数组元素:

kernel void fill_16_unchecked(device float *buffer  [[buffer(0)]],
                              constant float &value [[buffer(1)]],
                              uint index            [[thread_position_in_grid]])
{
    for (int i = 0; i < 16; ++i) {
        buffer[index * 16 + i] = value;
    }
}

kernel void single_fill_checked(device float *buffer         [[buffer(0)]],
                                constant float &value        [[buffer(1)]],
                                constant uint &buffer_length [[buffer(2)]],
                                uint index                   [[thread_position_in_grid]])
{
    if (index < buffer_length) {
        buffer[index] = value;
    }
}

如果您知道您的缓冲区计数将始终是线程执行宽度乘以您在循环中设置的元素数量的倍数,您可以只使用第一个函数。第二个函数是一个后备函数,用于处理可能会分派否则会溢出缓冲区的网格。

根据这些函数构建了两个管道后,您可以使用一对计算命令分派工作,如下所示:

NSInteger executionWidth = [unchecked16Pipeline threadExecutionWidth];
id<MTLComputeCommandEncoder> computeEncoder = [commandBuffer computeCommandEncoder];
[computeEncoder setBuffer:buffer offset:0 atIndex:0];
[computeEncoder setBytes:&value length:sizeof(float) atIndex:1];
if (bufferCount / (executionWidth * 16) != 0) {
    [computeEncoder setComputePipelineState:unchecked16Pipeline];
    [computeEncoder dispatchThreadgroups:MTLSizeMake(bufferCount / (executionWidth * 16), 1, 1)
                   threadsPerThreadgroup:MTLSizeMake(executionWidth, 1, 1)];
}
if (bufferCount % (executionWidth * 16) != 0) {
    int remainder = bufferCount % (executionWidth * 16);
    [computeEncoder setComputePipelineState:checkedSinglePipeline];
    [computeEncoder setBytes:&bufferCount length:sizeof(bufferCount) atIndex:2];
    [computeEncoder dispatchThreadgroups:MTLSizeMake((remainder / executionWidth) + 1, 1, 1)
                   threadsPerThreadgroup:MTLSizeMake(executionWidth, 1, 1)];
}
[computeEncoder endEncoding];

请注意,以这种方式完成工作不一定比每个线程仅写入一个元素的朴素方法更快。在我的测试中,它在 A8 上快了 40%,在 A10 上大致相当,在 A9 上慢了 2-3 倍(!)。始终使用您自己的工作量进行测试。