如何获取Metal片段着色器自动计算的thread group size？

直接API获取的可行性 #

Metal没有提供直接的API来查询片段着色器自动生成的image block size。这个尺寸是GPU驱动根据硬件特性、渲染目标的分辨率/格式、着色器运算复杂度等多个动态因素自动优化决定的，属于驱动内部的实现细节，并未对外暴露接口。

片段着色器侧的推测技巧 #

可以通过共享内存（threadgroup memory）和线程位置标记的方法间接推断image block的尺寸，具体实现如下：

实现思路 #

利用threadgroup变量记录当前线程组内的最大线程坐标，通过原子操作保证线程同步，最后由线程组内的首个线程将计算出的尺寸输出到纹理中。

示例代码 #

#include <metal_stdlib>
using namespace metal;

fragment void infer_image_block_size(
    texture2d<float, access::write> output [[texture(0)]],
    uint2 thread_pos [[thread_position_in_threadgroup]],
    uint2 grid_pos [[position]]
) {
    // 声明threadgroup变量，存储线程组内的最大x、y坐标
    threadgroup int max_thread_x = 0;
    threadgroup int max_thread_y = 0;

    // 原子更新最大线程坐标，确保所有线程参与统计
    atomic_max(&max_thread_x, static_cast<int>(thread_pos.x));
    atomic_max(&max_thread_y, static_cast<int>(thread_pos.y));

    // 等待线程组内所有线程完成坐标更新
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // 仅由线程组的(0,0)位置线程输出结果
    if (thread_pos.x == 0 && thread_pos.y == 0) {
        // 线程坐标从0开始，尺寸需+1
        float2 block_size = float2(max_thread_x + 1, max_thread_y + 1);
        output.write(float4(block_size, 0.0, 0.0), grid_pos);
    } else {
        // 其他线程输出无效值，避免干扰结果读取
        output.write(float4(0.0), grid_pos);
    }
}

使用步骤 #

1. 创建一个与渲染目标同尺寸的输出纹理（推荐格式rgba32f）。
2. 绑定上述片段着色器，渲染一个全屏四边形。
3. 读取输出纹理的像素值，第一个非零的float2数据即为当前场景的image block size（如16x16或32x8）。

注意事项 #

• 同一渲染通道内的image block size通常统一，但驱动可能在不同任务中调整尺寸，需按需重新推断。
• 确保渲染区域覆盖至少一个完整线程组，全屏渲染是最稳妥的选择。

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