【Swift Metal】dispatchThreadgroupsの最適化ついて解説

dispatchThreadgroupsついて

MetalでGPGPUの実装をする際、Thread及びThreadgroupの数は、MTLComputeCommandEncoderのdispatchThreadgroupsで設定します。

dispatchThreadgroupsは、以下のように定義されております。

func dispatchThreadgroups(_ threadgroupsPerGrid: MTLSize,
                            threadsPerThreadgroup: MTLSize)

【公式ドキュメント】 ： dispatchThreadgroups(_:threadsPerThreadgroup:) – MTLComputeCommandEncoder | Apple Developer Documentation

引数には、threadgroupsPerGridとthreadsPerThreadgroupを指定します。

threadgroupsPerGridは、GridにおけるThreadgroupの数のことです。MTLSize(width: Int, height: Int, depth: Int)で指定します。

threadsPerThreadgroupは、ThreadgroupにおけるThreadの数のことです。Threadgroup同様、MTLSizeで指定します。

1D、2D、3Dについて

threadgroupsPerGrid及びthreadsPerThreadgroupは、MTLSize(width: Int, height: Int, depth: Int)で指定すると紹介しました。

このMTLSizeのwidth、height、depthの内、widthしか使用しない(heightとdepthに1を指定する)場合は1D(1次元)となります。Thread及びThreadgroupが線状で配置されるイメージです。

MTLSizeのwidth、heightを使用する場合は2D(2次元)となります。Thread及びThreadgroupが面状で配置されるイメージです。

MTLSizeのwidth、height、depthを使用する場合は3D(3次元)となります。Thread及びThreadgroupが立体状で配置されるイメージです。

設定値を最適化し、パフォーマンスを上げる方法

ここからは、threadsPerThreadgroup及びthreadgroupsPerGridの設定値を最適化し、パフォーマンスを上げる方法を紹介致します。

最適化のためには、Threadの最大設定可能数と、GPUによって同時実行されるスレッド数を知る必要があります。

Threadの最大設定可能数

Threadの最大設定可能数は、ComputePipelineStateの、maxTotalThreadsPerThreadgroupで取得することができます。

【公式ドキュメント】 ： maxTotalThreadsPerThreadgroup – MTLComputePipelineState | Apple Developer Documentation

threadsPerThreadgroupとして設定するMTLSizeのwidthとheightとdepthを掛け合わせた数が、この最大設定可能数以内に収まるようにする必要があります。

GPUによって同時実行されるスレッド数

GPUによって同時実行されるスレッド数は、ComputePipelineStateの、threadExecutionWidthで取得することができます。

【公式ドキュメント】 ： threadExecutionWidth – MTLComputePipelineState | Apple Developer Documentation

threadsPerThreadgroupのwidthには、この値の倍数を指定する必要があります。

threadsPerThreadgroupの最適化

ここからは、最適化のサンプルコードを紹介していきます。
※ maxTotalThreadsPerThreadgroupは1024、threadExecutionWidthは32であることを前提として実装しております。

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = maxTotalThreadsPerThreadgroup / threadExecutionWidth * threadExecutionWidth // 1024 / 32 * 32 = 1024
let height = 1
let depth  = 1
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 1024, height: 1, depth: 1)

widthには、maxTotalThreadsPerThreadgroup / threadExecutionWidth * threadExecutionWidthの結果を設定しております。こうすることで、maxTotalThreadsPerThreadgroupに最も近い、threadExecutionWidthの倍数の値がwidthに設定されるようになります。

height及びdepthには、1Dなので1を指定しております。

threadsPerThreadgroupは、MTLSize(width: 1024, height: 1, depth: 1)となりました。width(1024) * height(1) * depth(1) の値が、maxTotalThreadsPerThreadgroup(1024)の値と等しいので、最適化されていることがわかります。

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = threadExecutionWidth // 32
let height = maxTotalThreadsPerThreadgroup / width // 1024 / 32 = 32
let depth  = 1
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 32, height: 32, depth: 1)

widthには、threadExecutionWidthの値を設定しております。

heightには、maxTotalThreadsPerThreadgroup / widthの結果を設定しております。

depthには、2Dなので1を指定しております。

threadsPerThreadgroupは、MTLSize(width: 32, height: 32, depth: 1)となりました。width(32) * height(32) * depth(1) の値が、maxTotalThreadsPerThreadgroup(1024)の値と等しいので、最適化されていることがわかります。

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = threadExecutionWidth // 32
let height = 8
let depth  = maxTotalThreadsPerThreadgroup / width / height // 1024 / 32 / 8 = 4
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 32, height: 8, depth: 4)

widthには、threadExecutionWidthの値を設定しております。

heightには、32以下で、32を割り切れる値を設定しております。

depthには、maxTotalThreadsPerThreadgroup / width / heightの結果を設定しております。

threadsPerThreadgroupは、MTLSize(width: 32, height: 32, depth: 1)となりました。width(32) * height(8) * depth(4) の値が、maxTotalThreadsPerThreadgroup(1024)の値と等しいので、最適化されていることがわかります。

threadgroupsPerGridの最適化

続いて、threadgroupsPerGridのサンプルコードを紹介致します。
※ 10000の処理を並列実行させることを前提として実装しております。

let count = 10000 // 並列処理数

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = maxTotalThreadsPerThreadgroup / threadExecutionWidth * threadExecutionWidth // 1024 / 32 * 32 = 1024
let height = 1
let depth  = 1
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 1024, height: 1, depth: 1)

// width: (10000 + 1024 - 1) / 1024 = 10
// threadsPerThreadgroup(1024 * 1 * 1) * threadgroupsPerGrid(10 * 1 * 1) = 10240
// 総Thread数(10240) > 並列処理数(10000) となる
let threadgroupsPerGrid = MTLSize(width: (count + width - 1) / width, height: 1, depth: 1)

widthには、(count + width – 1) / widthの結果を設定しております。

height及びdepthには、1Dなので1を指定しております。

threadgroupsPerGridは、MTLSize(width: 10, height: 1, depth: 1)となりました。threadsPerThreadgroup(1024 * 1 * 1) * threadgroupsPerGrid(10 * 1 * 1)の値(10240)が、並列処理数(10000)以上となるように設定されております。

本当に10が最適値なのかということについてですが、
1024 * 1 * 1 * 9 * 1 * 1 = 9216
1024 * 1 * 1 * 10 * 1 * 1 = 10240
1024 * 1 * 1 * 11 * 1 * 1 = 11264
という結果からも、10が最適値であることが確認いただけると思います。

// 並列処理数
let xCount = 100
let yCount = 100

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = threadExecutionWidth // 32
let height = maxTotalThreadsPerThreadgroup / width // 1024 / 32 = 32
let depth  = 1
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 32, height: 32, depth: 1)

// width : (100 + 32 - 1) / 32 = 4
// height: (100 + 32 - 1) / 32 = 4
// threadsPerThreadgroup(32 * 32 * 1) * threadgroupsPerGrid(4 * 4 * 1) = 16384
// 総Thread数(16384) > 並列処理数(10000) となる
let threadgroupsPerGrid = MTLSize(width: (xCount + width - 1) / width, height: (yCount + height - 1) / height, depth: 1)

widthには、(xCount + width – 1) / widthの結果を設定しております。

heightには、(yCount + height – 1) / heightの結果を設定しております。

depthには、2Dなので1を指定しております。

threadgroupsPerGridは、MTLSize(width: 4, height: 4, depth: 1)となりました。threadsPerThreadgroup(32 * 32 * 1) * threadgroupsPerGrid(4 * 4 * 1)の値(16384)が、並列処理数(10000)以上となるように設定されております。

heightを3にすれば、threadsPerThreadgroup(32 * 32 * 1) * threadgroupsPerGrid(4 * 3 * 1) = 12288で、よりいいのではないかと思うかもしれませんが、その場合、threadsPerThreadgroupのheight(32) * threadgroupsPerGridのheight(3)が96となり、yCountの100未満となるため、NGです。

// 並列処理数
let xCount = 100
let yCount = 10
let zCount = 10

let maxTotalThreadsPerThreadgroup = computePipelineState.maxTotalThreadsPerThreadgroup // 1024
let threadExecutionWidth          = computePipelineState.threadExecutionWidth // 32
let width  = threadExecutionWidth // 32
let height = 8
let depth  = maxTotalThreadsPerThreadgroup / width / height // 1024 / 32 / 8 = 4
let threadsPerThreadgroup = MTLSize(width: width, height: height, depth: depth) // MTLSize(width: 32, height: 8, depth: 4)

// width : (100 + 32 - 1) / 32 = 4
// height: (10  + 8  - 1) / 8  = 2
// depth : (10  + 4  - 1) / 4  = 3
// threadsPerThreadgroup(32 * 8 * 4) * threadgroupsPerGrid(4 * 2 * 3) = 24576
// 総Thread数(24576) > 並列処理数(10000) となる
let threadgroupsPerGrid = MTLSize(width: (xCount + width - 1) / width, height: (yCount + height - 1) / height, depth: (zCount + depth - 1) / depth)

widthには、(xCount + width – 1) / widthの結果を設定しております。

heightには、(yCount + height – 1) / heightの結果を設定しております。

depthには、(zCount + depth – 1) / depthの結果を設定しております。

threadgroupsPerGridは、MTLSize(width: 4, height: 2, depth: 3)となりました。threadsPerThreadgroup(32 * 32 * 1) * threadgroupsPerGrid(4 * 2 * 3)の値(24576)が、並列処理数(10000)以上となるように設定されております。

depthを2にすると、threadsPerThreadgroupのdepth(4) * threadgroupsPerGridのdepth(2)が8となり、zCountの10をしたまわるため、3が最適値であることがわかります。

以上、threadsPerThreadgroup及びthreadgroupsPerGridの設定値を最適化し、パフォーマンスを上げる方法でした。

まとめ

• Thread及びThreadgroupの数は、MTLComputeCommandEncoderのdispatchThreadgroupsで設定する
• Threadの最大設定可能数は、ComputePipelineStateの、maxTotalThreadsPerThreadgroupで取得できる
• GPUによって同時実行されるスレッド数は、ComputePipelineStateの、threadExecutionWidthで取得できる