gpt4 book ai didi

ios - 计算 Metal 内核中的平均值

转载 作者:行者123 更新时间:2023-12-01 23:10:51 25 4
gpt4 key购买 nike

有人知道用 Metal 内核中的随机 float 计算缓冲区平均值的正确方法吗?

在计算命令编码器上分派(dispatch)工作:

threadsPerGroup = MTLSizeMake(1, 1, inputTexture.arrayLength);
numThreadGroups = MTLSizeMake(1, 1, inputTexture.arrayLength / threadsPerGroup.depth);

[commandEncoder dispatchThreadgroups:numThreadGroups
threadsPerThreadgroup:threadsPerGroup];

内核代码:

kernel void mean(texture2d_array<float, access::read> inTex [[ texture(0) ]],
device float *means [[ buffer(1) ]],
uint3 id [[ thread_position_in_grid ]]) {

if (id.x == 0 && id.y == 0) {
float mean = 0.0;
for (uint i = 0; i < inTex.get_width(); ++i) {
for (uint j = 0; j < inTex.get_height(); ++j) {
mean += inTex.read(uint2(i, j), id.z)[0];
}
}

float textureArea = inTex.get_width() * inTex.get_height();
mean /= textureArea;
out[id.z] = mean;
}
}

缓冲区以 R32Float 像素格式的texture2d_array 类型的纹理表示。

最佳答案

如果您可以使用 uint 数组(而不是 float )作为数据源,我建议使用“原子获取和修改函数”(如 Metal 着色语言 spec 中所述)以原子方式写入缓冲。

这是一个内核函数的示例,它采用输入缓冲区(数据:Float 数组)并将缓冲区的总和写入原子缓冲区(sum,指向 uint 的指针):

kernel void sum(device uint *data [[ buffer(0) ]],
volatile device atomic_uint *sum [[ buffer(1) ]],
uint gid [[ thread_position_in_grid ]])
{
atomic_fetch_add_explicit(sum, data[gid], memory_order_relaxed);
}

在您的 swift 文件中,您将设置缓冲区:

...
let data: [UInt] = [1, 2, 3, 4]
let dataBuffer = device.makeBuffer(bytes: &data, length: (data.count * MemoryLayout<UInt>.size), options: [])
commandEncoder.setBuffer(dataBuffer, offset: 0, at: 0)

var sum:UInt = 0
let sumBuffer = device!.makeBuffer(bytes: &sum, length: MemoryLayout<UInt>.size, options: [])
commandEncoder.setBuffer(sumBuffer, offset: 0, at: 1)
commandEncoder.endEncoding()

提交,等待,然后从 GPU 获取数据:

commandBuffer.commit()
commandBuffer.waitUntilCompleted()

let nsData = NSData(bytesNoCopy: sumBuffer.contents(),
length: sumBuffer.length,
freeWhenDone: false)
nsData.getBytes(&sum, length:sumBuffer.length)

let mean = Float(sum/data.count)
print(mean)

或者,如果您的初始数据源必须是 float 组,您可以使用 vDSP_meanv Accelerate 框架的方法对于此类计算来说非常快。

希望有帮助,干杯!

关于ios - 计算 Metal 内核中的平均值,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/40347083/

25 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com