公開日：2020/03/20
更新日：2020/08/03

【Swift Metal】atomic変数の基本的な使い方を解説

Swift
Metal

スレッド間で、データのやりとりがしたい
atomic変数の使い方が知りたい

atomic変数とは、スレッド間でのデータをやり取りをスレッドセーフに行うことができる変数です。

Metalには、atomic_int、atomic_uint、atomic_bool、atomicという4つのatomic変数が定義されております。

【公式ドキュメント】： 2.5 Atomic Data Types | Metal Shading Language Specification

本記事では、Metalでのatomic変数の基本的な使い方を解説致します。

１．atomic変数の基本的な使い方
２．まとめ

atomic変数の基本的な使い方

保存(store)と取得(load)

atomic変数に対し、値の保存を行うには、atomic_store_explicitを使用します。

また、atomic変数から値を取得するには、atomic_load_explicitを使用します。

※ サンプルコードのinputDataには、[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]のデータが格納されています。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    // inputDataをatomicDataに保存
    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // atomicDataの値を取得し、outputDataに保存
    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, 4, 5, 6, 7, 8, 9

atomic_store_explicitを使用し、inputDataの値をatomicDataに保存しております。そして、atomic_load_explicitでatomicDataから値を取得し、outputDataに保存しております。

書き換え(exchange)

atomic変数の値の書き換えを行うには、atomic_store_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の書き換え
    atomic_exchange_explicit( &atomicData[5], 100, memory_order_relaxed );

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, 4, 100, 6, 7, 8, 9

atomic_store_explicitを使用し、atomicData[5]の値(5)を100で書き換えております。

加算(add)

atomic変数に対し、値の加算を行うには、atomic_fetch_add_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の加算
    if (index == 3) {
        atomic_fetch_add_explicit( &atomicData[3], 10, memory_order_relaxed );
    }

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 13, 4, 5, 6, 7, 8, 9

atomic_fetch_add_explicitを使用し、atomicData[3]の値(3)に10を加算しております。

減算(sub)

atomic変数に対し、値の減算を行うには、atomic_fetch_sub_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の減算
    if (index == 4) {
        atomic_fetch_sub_explicit( &atomicData[4], 10, memory_order_relaxed );
    }

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, -6, 5, 6, 7, 8, 9

atomic_fetch_sub_explicitを使用し、atomicData[4]の値(4)から6を減算しております。

値の比較(max)

atomic変数の値と、指定した値を比較し、大きい方を選択するには、atomic_fetch_max_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の比較(大きい方を選択)
    atomic_fetch_max_explicit( &atomicData[index], 4, memory_order_relaxed );

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：4, 4, 4, 4, 4, 5, 6, 7, 8, 9

atomic_fetch_max_explicitでinputDataの値を4と比較し、4よりも小さい値を4に書き換えております。

値の比較(min)

atomic変数の値と、指定した値を比較し、小さい方を選択するには、atomic_fetch_min_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の比較(小さい方を選択)
    atomic_fetch_min_explicit( &atomicData[index], 4, memory_order_relaxed );

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, 4, 4, 4, 4, 4, 4

atomic_fetch_min_explicitでinputDataの値を4と比較し、4よりも大きい値を4に書き換えております。

OR演算(or)

atomic変数の値に対し、指定した値とOR演算をするには、atomic_fetch_or_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // OR演算
    atomic_fetch_or_explicit( &atomicData[8], 2, memory_order_relaxed ); // 00001000 | 00000010 = 00001010 -> 10
    atomic_fetch_or_explicit( &atomicData[7], 2, memory_order_relaxed ); // 00000111 | 00000010 = 00000111 -> 7

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, 4, 5, 6, 7, 10, 9

atomic_fetch_or_explicitを使用し、atomicData[8]の値(8)と2、atomicData[7]の値(7)と2をOR演算しております。

XOR演算(xor)

atomic変数の値に対し、指定した値とXOR演算をするには、atomic_fetch_xor_explicitを使用します。

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // XOR演算
    if (index == 0) {
        atomic_fetch_xor_explicit( &atomicData[8], 2, memory_order_relaxed ); // 00001000 ^ 00000010 = 00001010 -> 10
        atomic_fetch_xor_explicit( &atomicData[7], 2, memory_order_relaxed ); // 00000111 ^ 00000010 = 00000101 -> 5
    }

    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}
// outputDataの出力結果：0, 1, 2, 3, 4, 5, 6, 5, 10, 9

atomic_fetch_or_explicitを使用し、atomicData[8]の値(8)と2、atomicData[7]の値(7)と2をXOR演算しております。

まとめ

最後に、全体のサンプルコードを載せておきます。

ViewController.swift
Shader.metal

import UIKit
import Metal

class ViewController: UIViewController {

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

    private var commandBuffer: MTLCommandBuffer!
    private var computeCommandEncoder: MTLComputeCommandEncoder!

    private var count = 10

    private var inputData: [Int32] = []

    override func viewDidLoad() {
        super.viewDidLoad()

        initMetal()

        for i in 0..<count {
            inputData.append(Int32(i))
        }
    }

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

        commandQueue = device.makeCommandQueue()!
    }

    @IBAction func test1(_ sender: UIButton) {
        let function = library.makeFunction(name: "test")!
        computePipelineState = try! device.makeComputePipelineState(function: function)

        autoreleasepool {
            commandBuffer = commandQueue.makeCommandBuffer()!
            computeCommandEncoder = commandBuffer.makeComputeCommandEncoder()!
            computeCommandEncoder.setComputePipelineState(computePipelineState)
        }

        let outputData = [Int32](repeating: 0, count: inputData.count)

        let inputBuffer  = device.makeBuffer(bytes: inputData,  length: MemoryLayout<Int32>.stride * inputData.count, options: [])
        let outputBuffer = device.makeBuffer(bytes: outputData, length: MemoryLayout<Int32>.stride * outputData.count, options: [])

        computeCommandEncoder.setBuffer(inputBuffer,  offset: 0, index: 0)
        computeCommandEncoder.setBuffer(outputBuffer, offset: 0, index: 1)

        let width = computePipelineState.threadExecutionWidth
        let threadgroupsPerGrid = MTLSize(width: (count + width - 1) / width, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: width, height: 1, depth: 1)

        computeCommandEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)

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

        let result = Data(bytesNoCopy: outputBuffer!.contents(), count: MemoryLayout<Int32>.stride * count, deallocator: .none)
        let res = result.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress!.assumingMemoryBound(to: Int32.self ), count: $0.count / MemoryLayout<nt32>.size)) }

        for i in 0..<res.count {
            print(res[i])
        }
    }
}

kernel void test(device int* inputData  [[ buffer(0) ]],
                 device int* outputData [[ buffer(1) ]],

                 uint2 thread_position_in_grid [[ thread_position_in_grid ]],
                 uint2 threads_per_grid        [[ threads_per_grid ]])
{
    int index = thread_position_in_grid.y * threads_per_grid.x + thread_position_in_grid.x;

    threadgroup atomic_int atomicData[10];

    // inputDataをatomicDataに保存
    atomic_store_explicit( &atomicData[index], inputData[index], memory_order_relaxed );

    // 値の書き換え
//    atomic_exchange_explicit( &atomicData[5], 100, memory_order_relaxed );

    // 値の加算
//    if (index == 3) {
//        atomic_fetch_add_explicit( &atomicData[3], 10, memory_order_relaxed );
//    }

    // 値の減算
//    if (index == 4) {
//        atomic_fetch_sub_explicit( &atomicData[4], 10, memory_order_relaxed );
//    }

    // 値の比較(大きい方を選択)
//    atomic_fetch_max_explicit( &atomicData[index], 4, memory_order_relaxed );

    // 値の比較(小さい方を選択)
//    atomic_fetch_min_explicit( &atomicData[index], 4, memory_order_relaxed );

    // OR演算
//    atomic_fetch_or_explicit( &atomicData[8], 2, memory_order_relaxed ); // 00001000 | 00000010 = 00001010 -> 10
//    atomic_fetch_or_explicit( &atomicData[7], 2, memory_order_relaxed ); // 00000111 | 00000010 = 00000111 -> 7

    // XOR演算
//    if (index == 0) {
//        atomic_fetch_xor_explicit( &atomicData[8], 2, memory_order_relaxed ); // 00001000 ^ 00000010 = 00001010 -> 10
//        atomic_fetch_xor_explicit( &atomicData[7], 2, memory_order_relaxed ); // 00000111 ^ 00000010 = 00000101 -> 5
//    }

    // atomicDataの値を取得し、outputDataに保存
    outputData[index] = atomic_load_explicit( &atomicData[index], memory_order_relaxed );
}

値を保存するには、atomic_store_explicitを使用する
値を取得するには、atomic_load_explicitを使用する
値を書き換えるには、atomic_store_explicitを使用する
値を加算するには、atomic_fetch_add_explicitを使用する
値を減算するには、atomic_fetch_sub_explicitを使用する
値を比較し、大きい方を選択するには、atomic_fetch_max_explicitを使用する
値を比較し、小さい方を選択するには、atomic_fetch_min_explicitを使用する
値をOR演算するには、atomic_fetch_or_explicitを使用する
値をXOR演算するには、atomic_fetch_xor_explicitを使用する