【Swift Metal】thread_position_in_grid等の属性について解説

import UIKit
import Metal

class ViewController: UIViewController {

    private let device = MTLCreateSystemDefaultDevice()!
    private var library: MTLLibrary!
    private var commandQueue: MTLCommandQueue!
    private var computePipelineState: MTLComputePipelineState!

    private let xCount = 1000000
    private var outputDataCount = 0

    let attribute = ["thread_execution_width", "thread_position_in_grid", "thread_position_in_threadgroup", "threadgroup_position_in_grid", "threadgroups_per_grid", "threads_per_grid", "threads_per_threadgroup"];

    override func viewDidLoad() {
        super.viewDidLoad()

        outputDataCount = xCount

        initMetal()
        startGPU()
    }

    private func initMetal() {
        guard let libUrl = Bundle.main.url(forResource: "default.metallib", withExtension: nil) else { return }
        library = try! device.makeLibrary(URL: libUrl)

        let function = library.makeFunction(name: "sample")!
        computePipelineState = try! device.makeComputePipelineState(function: function)

        commandQueue = device.makeCommandQueue()!
    }

    private func startGPU() {
        var outputDataList: [[Float]] = []
        for _ in 0..<attribute.count {
            outputDataList.append([Float](repeating: 1, count: outputDataCount))
        }

        var outputBufferList: [MTLBuffer] = []
        for i in 0..<attribute.count {
            outputBufferList.append(device.makeBuffer(bytes: outputDataList[i], length: MemoryLayout<Float>.stride * outputDataList[i].count, options: [])!)
        }

        let commandBuffer = commandQueue.makeCommandBuffer()!
        let computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeCommandEncoder.setComputePipelineState(computePipelineState)

        for i in 0..<attribute.count {
            computeCommandEncoder.setBuffer(outputBufferList[i], offset: 0, index: i)
        }
        computeCommandEncoder.setBytes(&outputDataCount, length: MemoryLayout<Float>.size, index: attribute.count)

        let width = computePipelineState.threadExecutionWidth
        let threadgroupsPerGrid = MTLSize(width: (xCount + width - 1) / width, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: width, height: 1, depth: 1)
        print("threadgroupsPerGrid  : \(threadgroupsPerGrid)")
        print("threadsPerThreadgroup: \(threadsPerThreadgroup)")
        computeCommandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)

        computeCommandEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        var resultDataList: [[Float]] = []
        for i in 0..<attribute.count {
            let result = Data(bytesNoCopy: outputBufferList[i].contents(), count: MemoryLayout<Float>.stride * outputDataList[i].count, deallocator: .none)
            resultDataList.append(
                result.withUnsafeBytes {
                    Array(UnsafeBufferPointer(start: $0.baseAddress!.assumingMemoryBound(to: Float.self ), count: $0.count / MemoryLayout<Float>.size))
                }
            )
        }

        for i in 0..<outputDataCount {
            var str = "index: \(i) ||"
            for j in 0..<attribute.count {
                str.append(" \(attribute[j]): \(resultDataList[j][i]), ")
            }
            print(str)
        }
    }
}

#include <metal_stdlib>
using namespace metal;

kernel void sample(device float* outputData0    [[ buffer(0) ]],
                   device float* outputData1    [[ buffer(1) ]],
                   device float* outputData2    [[ buffer(2) ]],
                   device float* outputData3    [[ buffer(3) ]],
                   device float* outputData4    [[ buffer(4) ]],
                   device float* outputData5    [[ buffer(5) ]],
                   device float* outputData6    [[ buffer(6) ]],
                   const device int& dataLength [[ buffer(7) ]],

                   uint thread_execution_width         [[thread_execution_width]],         // ThreadGroupの横幅
                   uint thread_position_in_grid        [[thread_position_in_grid]],        // GridにおけるThreadの位置
                   uint thread_position_in_threadgroup [[thread_position_in_threadgroup]], // ThreadGroupにおけるThreadの位置
                   uint threadgroup_position_in_grid   [[threadgroup_position_in_grid]],   // GridにおけるThreadGroupの位置
                   uint threadgroups_per_grid          [[threadgroups_per_grid]],          // Grid毎のThreadGroup数
                   uint threads_per_grid               [[threads_per_grid]],               // Grid毎のThread数
                   uint threads_per_threadgroup        [[threads_per_threadgroup]])        // ThreadGroup毎のThread数
{
    int index = thread_position_in_grid;

    if (index > dataLength) { return; }

    outputData0[index]  = thread_execution_width;
    outputData1[index]  = thread_position_in_grid;
    outputData2[index]  = thread_position_in_threadgroup;
    outputData3[index]  = threadgroup_position_in_grid;
    outputData4[index]  = threadgroups_per_grid;
    outputData5[index]  = threads_per_grid;
    outputData6[index]  = threads_per_threadgroup;
}

import UIKit
import Metal

class ViewController: UIViewController {

    private let device = MTLCreateSystemDefaultDevice()!
    private var library: MTLLibrary!
    private var commandQueue: MTLCommandQueue!
    private var computePipelineState: MTLComputePipelineState!

    private let xCount = 1000
    private let yCount = 1000
    private var outputDataCount = 0

    let attribute = ["thread_execution_width", "thread_position_in_grid", "thread_position_in_threadgroup", "threadgroup_position_in_grid", "threadgroups_per_grid", "threads_per_grid", "threads_per_threadgroup"];

    override func viewDidLoad() {
        super.viewDidLoad()

        outputDataCount = xCount * yCount

        initMetal()
        startGPU()
    }

    private func initMetal() {
        guard let libUrl = Bundle.main.url(forResource: "default.metallib", withExtension: nil) else { return }
        library = try! device.makeLibrary(URL: libUrl)

        let function = library.makeFunction(name: "sample")!
        computePipelineState = try! device.makeComputePipelineState(function: function)

        commandQueue = device.makeCommandQueue()!
    }

    private func startGPU() {
        var outputDataList: [[Float]] = []
        for _ in 0..<attribute.count {
            outputDataList.append([Float](repeating: 1, count: outputDataCount))
        }

        var outputBufferList: [MTLBuffer] = []
        for i in 0..<attribute.count {
            outputBufferList.append(device.makeBuffer(bytes: outputDataList[i], length: MemoryLayout<Float>.stride * outputDataList[i].count, options: [])!)
        }

        let commandBuffer = commandQueue.makeCommandBuffer()!
        let computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeCommandEncoder.setComputePipelineState(computePipelineState)

        for i in 0..<attribute.count {
            computeCommandEncoder.setBuffer(outputBufferList[i], offset: 0, index: i)
        }
        computeCommandEncoder.setBytes(&outputDataCount, length: MemoryLayout<Float>.size, index: attribute.count)

        let max    = computePipelineState.maxTotalThreadsPerThreadgroup
        let width  = computePipelineState.threadExecutionWidth
        let height = max / width
        let threadgroupsPerGrid = MTLSize(width: (xCount + width - 1) / width, height: (yCount + height - 1) / height, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: 1)
        print("threadgroupsPerGrid  : \(threadgroupsPerGrid)")
        print("threadsPerThreadgroup: \(threadsPerThreadgroup)")
        computeCommandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)

        computeCommandEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        var resultDataList: [[Float]] = []
        for i in 0..<attribute.count {
            let result = Data(bytesNoCopy: outputBufferList[i].contents(), count: MemoryLayout<Float>.stride * outputDataList[i].count, deallocator: .none)
            resultDataList.append(
                result.withUnsafeBytes {
                    Array(UnsafeBufferPointer(start: $0.baseAddress!.assumingMemoryBound(to: Float.self ), count: $0.count / MemoryLayout<Float>.size))
                }
            )
        }

        for i in 0..<outputDataCount {
            var str = "index: \(i) ||"
            for j in 0..<attribute.count {
                str.append(" \(attribute[j]): \(resultDataList[j][i]), ")
            }
            print(str)
        }
    }
}

#include <metal_stdlib>
using namespace metal;

kernel void sample(device float* outputData0    [[ buffer(0) ]],
                   device float* outputData1    [[ buffer(1) ]],
                   device float* outputData2    [[ buffer(2) ]],
                   device float* outputData3    [[ buffer(3) ]],
                   device float* outputData4    [[ buffer(4) ]],
                   device float* outputData5    [[ buffer(5) ]],
                   device float* outputData6    [[ buffer(6) ]],
                   const device int& dataLength [[ buffer(7) ]],

                   uint  thread_execution_width         [[thread_execution_width]],         // ThreadGroupの横幅
                   uint2 thread_position_in_grid        [[thread_position_in_grid]],        // GridにおけるThreadの位置
                   uint2 thread_position_in_threadgroup [[thread_position_in_threadgroup]], // ThreadGroupにおけるThreadの位置
                   uint2 threadgroup_position_in_grid   [[threadgroup_position_in_grid]],   // GridにおけるThreadGroupの位置
                   uint2 threadgroups_per_grid          [[threadgroups_per_grid]],          // Grid毎のThreadGroup数
                   uint2 threads_per_grid               [[threads_per_grid]],               // Grid毎のThread数
                   uint2 threads_per_threadgroup        [[threads_per_threadgroup]])        // ThreadGroup毎のThread数
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    if (index > dataLength) { return; }
    outputData0[index] = thread_execution_width;
    outputData1[index] = thread_position_in_grid.x;
    outputData2[index] = thread_position_in_threadgroup.x;
    outputData3[index] = threadgroup_position_in_grid.x;
    outputData4[index] = threadgroups_per_grid.x;
    outputData5[index] = threads_per_grid.x;
    outputData6[index] = threads_per_threadgroup.x;
		// Y値の出力結果表示する際は、.xの部分を .yに置き換えております。
}

import UIKit
import Metal

class ViewController: UIViewController {

    private let device = MTLCreateSystemDefaultDevice()!
    private var library: MTLLibrary!
    private var commandQueue: MTLCommandQueue!
    private var computePipelineState: MTLComputePipelineState!

    private let xCount = 100
    private let yCount = 100
    private let zCount = 100
    private var outputDataCount = 0

    let attribute = ["thread_execution_width", "thread_position_in_grid", "thread_position_in_threadgroup", "threadgroup_position_in_grid", "threadgroups_per_grid", "threads_per_grid", "threads_per_threadgroup"];

    override func viewDidLoad() {
        super.viewDidLoad()

        outputDataCount = xCount * yCount * zCount

        initMetal()

        startGPU()
    }

    private func initMetal() {
        guard let libUrl = Bundle.main.url(forResource: "default.metallib", withExtension: nil) else { return }
        library = try! device.makeLibrary(URL: libUrl)

        let function = library.makeFunction(name: "sample")!
        computePipelineState = try! device.makeComputePipelineState(function: function)

        commandQueue = device.makeCommandQueue()!
    }

    private func startGPU() {
        var outputDataList: [[Float]] = []
        for _ in 0..<attribute.count {
            outputDataList.append([Float](repeating: 1, count: outputDataCount))
        }

        var outputBufferList: [MTLBuffer] = []
        for i in 0..<attribute.count {
            outputBufferList.append(device.makeBuffer(bytes: outputDataList[i], length: MemoryLayout<Float>.stride * outputDataList[i].count, options: [])!)
        }

        let commandBuffer = commandQueue.makeCommandBuffer()!
        let computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeCommandEncoder.setComputePipelineState(computePipelineState)

        for i in 0..<attribute.count {
            computeCommandEncoder.setBuffer(outputBufferList[i], offset: 0, index: i)
        }
        computeCommandEncoder.setBytes(&outputDataCount, length: MemoryLayout<Float>.size, index: attribute.count)

        let max    = computePipelineState.maxTotalThreadsPerThreadgroup
        let depth  = 4
        let width  = computePipelineState.threadExecutionWidth
        let height = max / width / depth
        let threadgroupsPerGrid = MTLSize(width: (xCount + width - 1) / width, height: (yCount + height - 1) / height, depth: (zCount + depth - 1) / depth)
        let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth)
        print("threadgroupsPerGrid  : \(threadgroupsPerGrid)")
        print("threadsPerThreadgroup: \(threadsPerThreadgroup)")
        computeCommandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)

        computeCommandEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        var resultDataList: [[Float]] = []
        for i in 0..<attribute.count {
            let result = Data(bytesNoCopy: outputBufferList[i].contents(), count: MemoryLayout<Float>.stride * outputDataList[i].count, deallocator: .none)
            resultDataList.append(
                result.withUnsafeBytes {
                    Array(UnsafeBufferPointer(start: $0.baseAddress!.assumingMemoryBound(to: Float.self ), count: $0.count / MemoryLayout<Float>.size))
                }
            )
        }

        for i in 0..<outputDataCount {
            var str = "index: \(i) ||"
            for j in 0..<attribute.count {
                str.append(" \(attribute[j]): \(resultDataList[j][i]), ")
            }
            print(str)
        }
    }
}

#include <metal_stdlib>
using namespace metal;

kernel void sample(device float* outputData0    [[ buffer(0) ]],
                   device float* outputData1    [[ buffer(1) ]],
                   device float* outputData2    [[ buffer(2) ]],
                   device float* outputData3    [[ buffer(3) ]],
                   device float* outputData4    [[ buffer(4) ]],
                   device float* outputData5    [[ buffer(5) ]],
                   device float* outputData6    [[ buffer(6) ]],
                   const device int& dataLength [[ buffer(7) ]],

                   uint  thread_execution_width         [[thread_execution_width]],         // ThreadGroupの横幅
                   uint3 thread_position_in_grid        [[thread_position_in_grid]],        // GridにおけるThreadの位置
                   uint3 thread_position_in_threadgroup [[thread_position_in_threadgroup]], // ThreadGroupにおけるThreadの位置
                   uint3 threadgroup_position_in_grid   [[threadgroup_position_in_grid]],   // GridにおけるThreadGroupの位置
                   uint3 threadgroups_per_grid          [[threadgroups_per_grid]],          // Grid毎のThreadGroup数
                   uint3 threads_per_grid               [[threads_per_grid]],               // Grid毎のThread数
                   uint3 threads_per_threadgroup        [[threads_per_threadgroup]])        // ThreadGroup毎のThread数
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x + (thread_position_in_grid.z * threads_per_grid.y * threads_per_grid.x);

    if (index > dataLength) { return; }
    outputData0[index] = thread_execution_width;
    outputData1[index] = thread_position_in_grid.x;
    outputData2[index] = thread_position_in_threadgroup.x;
    outputData3[index] = threadgroup_position_in_grid.x;
    outputData4[index] = threadgroups_per_grid.x;
    outputData5[index] = threads_per_grid.x;
    outputData6[index] = threads_per_threadgroup.x;
		// Y値、Z値の出力結果表示する際は、.xの部分を .y、.zに置き換えております。
}