【Swift Metal】MetalでGPGPUを実装する方法を解説

import UIKit

// Metalをインポートする
import Metal

class ViewController: UIViewController {

    // MTLDeviceの生成
    private let device = MTLCreateSystemDefaultDevice()!
    private var library: MTLLibrary!
    private var commandQueue: MTLCommandQueue!
    private var computePipelineState: MTLComputePipelineState!

    private let yCount = 10
    private let xCount = 10

    private var inputData: [Float] = []

    override func viewDidLoad() {
        super.viewDidLoad()

        for _ in 0..<yCount*xCount {
            inputData.append(Float.random(in: 0..<10))
        }

        initMetal()
        startCPU()
        startGPU()
    }

    // Metalの初期化処理
    private func initMetal() {
        // (1) MTLLibraryの生成
        guard let libUrl = Bundle.main.url(forResource: "default.metallib", withExtension: nil) else { return }
        library = try! device.makeLibrary(URL: libUrl)

        // (2) MTLFunctionの生成
        let function = library.makeFunction(name: "testFunction")!

        // (3) MTLComputePipelineStateの生成
        computePipelineState = try! device.makeComputePipelineState(function: function)

        // (4) MTLCommandQueueの生成
        commandQueue = device.makeCommandQueue()!
    }

    // CPU処理
    private func startCPU() {
        var outputData = [Float](repeating: 0, count: inputData.count)

        let start = Date().timeIntervalSince1970

        for y in 0..<yCount {
            for x in 0..<xCount {
                outputData[y * xCount + x] = inputData[y * xCount + x]
            }
        }

        let end = Date().timeIntervalSince1970
        print("CPU || outputData.first: \(outputData.first ?? 0), outputData.last: \(outputData.last ?? 0), time: " + String(format: "%.5f ms", (end - start) * 1000))
    }

    // GPU処理
    private func startGPU() {
        var outputData = [Float](repeating: 0, count: inputData.count)

        let start = Date().timeIntervalSince1970

        // (5) 入力バッファと出力バッファの生成
        let inputBuffer = device.makeBuffer(bytes: inputData, length: MemoryLayout<Float>.stride * inputData.count, options: [])
        let outputBuffer = device.makeBuffer(bytes: outputData, length: MemoryLayout<Float>.stride * outputData.count, options: [])

        // (6) MTLCommandBufferの生成
        let commandBuffer = commandQueue.makeCommandBuffer()!

        // (7) MTLComputeCommandEncoderの生成
        let computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeCommandEncoder.setComputePipelineState(computePipelineState)
        computeCommandEncoder.setBuffer(inputBuffer, offset: 0, index: 0)
        computeCommandEncoder.setBuffer(outputBuffer, offset: 0, index: 1)

        // (8) スレッドグループ数、スレッドの数の設定
        let width = computePipelineState.threadExecutionWidth
        let threadgroupsPerGrid = MTLSize(width: (outputData.count + width - 1) / width, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: width, height: 1, depth: 1)
        computeCommandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)

        // (9) エンコードの終了
        computeCommandEncoder.endEncoding()

        // (10) コマンドバッファを実行
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        // (11) 結果の取得
        let resultData = Data(bytesNoCopy: outputBuffer!.contents(), count: MemoryLayout<Float>.stride * outputData.count, deallocator: .none)
        outputData = resultData.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress!.assumingMemoryBound(to: Float.self ), count: $0.count / MemoryLayout<Float>.size)) }

        let end = Date().timeIntervalSince1970
        print("GPU || outputData.first: \(outputData.first ?? 0), outputData.last: \(outputData.last ?? 0), time: " + String(format: "%.5f ms", (end - start) * 1000))
    }
}

#include 
using namespace metal;

kernel void testFunction(const device float* inputData [[ buffer(0) ]],
                         device float* outputData      [[ buffer(1) ]],
                         uint thread_position_in_grid [[thread_position_in_grid]] // Gridにおけるthreadの位置
) {
    outputData[thread_position_in_grid] = inputData[thread_position_in_grid];
}

10 * 10
CPU || outputData.first: 3.0272727, outputData.last: 0.8784741, time: 0.27895 ms
GPU || outputData.first: 3.0272727, outputData.last: 0.8784741, time: 4.32491 ms

50 * 50
CPU || outputData.first: 9.532412, outputData.last: 7.8694916, time: 4.61698 ms
GPU || outputData.first: 9.532412, outputData.last: 7.8694916, time: 2.49100 ms

100 * 100
CPU || outputData.first: 2.0698965, outputData.last: 1.419571, time: 8.14319 ms
GPU || outputData.first: 2.0698965, outputData.last: 1.419571, time: 2.00367 ms

500 * 500
CPU || outputData.first: 2.014732, outputData.last: 5.279184, time: 96.63892 ms
GPU || outputData.first: 2.014732, outputData.last: 5.279184, time: 2.34222 ms

1000 * 1000
CPU || outputData.first: 3.9278536, outputData.last: 0.89728, time: 380.40805 ms
GPU || outputData.first: 3.9278536, outputData.last: 0.89728, time: 4.67777 ms

5000 * 5000
CPU || outputData.first: 3.1419396, outputData.last: 7.6604886, time: 9489.20012 ms
GPU || outputData.first: 3.1419396, outputData.last: 7.6604886, time: 62.89768 ms