如何向Metal片段着色器传递二维数组？语法与CPU端实现

问题背景 #

现有可接收int2一维数组的Metal片段着色器：

fragment half4 array_FragmentShader(Vertex_Out in [[ stage_in ]],
constant int2 *int2Array [[ buffer(1) ]]){
half4 retVal = half4(0.0,0.0,0.0,1.0);

if(int2Array[0].x==1){retVal = half4(0.6,0.0,0.0,1.0);}

return retVal;
}

CPU端通过SCNProgram调用着色器的代码如下：

let gridProgram = {
        let p = SCNProgram()
        p.vertexFunctionName = "boolVertexShader"
        p.fragmentFunctionName = "array_FragmentShader"
        p.isOpaque = false;
        return p
    }()

currentMaterial.program = gridProgram

let i2 = simd_make_int2(1,1)
var int2Array = [i2]

gridProgram.handleBinding(ofBufferNamed: "int2Array", frequency: .perFrame) { (bufferStream, node, shadable, renderer) in
    bufferStream.writeBytes(&int2Array,count:MemoryLayout<SIMD2<Int32>>.stride*int2Array.count)
}

需要解决的问题：如何向片段着色器传递二维数组？参数语法是否需要双指针？CPU端如何发送二维数组？handleBinding是否要求子数组大小统一？

Metal端参数语法：不要用双指针 #

Metal的Buffer是连续的内存块，GPU无法直接处理分散的指针地址（像CPU的二维数组双指针结构），所以不能用constant int2 **int2Array [[ buffer(1) ]]这种写法。正确的实现有两种方式：

1. 固定大小的二维数组 #

如果二维数组的行列数固定，可以直接声明成二维数组形式：

// 示例：2行3列的int2二维数组
fragment half4 array_FragmentShader(Vertex_Out in [[ stage_in ]],
constant int2 int2Array[2][3] [[ buffer(1) ]]){
    half4 retVal = half4(0.0,0.0,0.0,1.0);
    
    // 直接通过[row][col]访问元素
    if(int2Array[0][0].x == 1){
        retVal = half4(0.6,0.0,0.0,1.0);
    }
    
    return retVal;
}

2. 动态大小的二维数组 #

如果行列数需要动态调整，还是用一维指针声明，同时额外传递一个Buffer存储二维数组的尺寸（行数、列数），在着色器里通过索引计算来访问元素：

fragment half4 array_FragmentShader(Vertex_Out in [[ stage_in ]],
constant int2 *int2Array [[ buffer(1) ]],
constant int2 gridSize [[ buffer(2) ]]){ // gridSize.x=行数，gridSize.y=列数
    half4 retVal = half4(0.0,0.0,0.0,1.0);
    
    // 访问第row行第col列的元素：索引 = row * 列数 + col
    int row = 0;
    int col = 0;
    if(int2Array[row * gridSize.y + col].x == 1){
        retVal = half4(0.6,0.0,0.0,1.0);
    }
    
    return retVal;
}

CPU端传递二维数组的方法 #

CPU端需要把二维数组平铺成连续的一维内存块，再写入Buffer：

1. 平铺二维数组 #

以Swift为例，将[[SIMD2<Int32>]]类型的二维数组展开为[SIMD2<Int32>]：

// 定义2行3列的二维数组
let row1 = [simd_make_int2(1,2), simd_make_int2(3,4), simd_make_int2(5,6)]
let row2 = [simd_make_int2(7,8), simd_make_int2(9,10), simd_make_int2(11,12)]
var twoDArray = [row1, row2]

// 平铺成一维数组
var flattenedArray = twoDArray.flatMap { $0 }

2. 绑定并写入Buffer #

将平铺后的一维数组写入对应Buffer，如果是动态大小的场景，还要额外传递尺寸Buffer：

// 绑定二维数组Buffer
gridProgram.handleBinding(ofBufferNamed: "int2Array", frequency: .perFrame) { (bufferStream, node, shadable, renderer) in
    let totalBytes = MemoryLayout<SIMD2<Int32>>.stride * flattenedArray.count
    bufferStream.writeBytes(&flattenedArray, count: totalBytes)
}

// 如果是动态大小，绑定尺寸Buffer
let gridSize = simd_make_int2(2, 3) // x=行数，y=列数
gridProgram.handleBinding(ofBufferNamed: "gridSize", frequency: .perFrame) { (bufferStream, node, shadable, renderer) in
    bufferStream.writeBytes(&gridSize, count: MemoryLayout<SIMD2<Int32>>.stride)
}

子数组大小的要求 #

使用handleBinding传递平铺的二维数组时，必须要求子数组大小统一。原因是：
只有每行元素数量一致，才能通过row * 列数 + col的公式正确计算元素在一维内存中的索引。如果子数组大小不统一，GPU无法自动识别每行的边界，你需要额外传递一个记录每行元素个数的数组，在着色器里手动计算偏移，但这种方式会增加计算开销，不推荐使用。

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