🗃️

固定長配列のためのStaticArrays.jl

2021/12/13に公開

Julia

高速化

配列

tech

これはJulia Advent Calendar 2021の13日目の記事です。

TL;DR

固定長配列のためのパッケージStaticArrays.jlが便利！
- 3次元空間の演算など、配列の長さが決まっている状況は色々あります。
- そのような場合に高速化ができるパッケージです
- 極端なケースでは通常の配列に比べて10倍以上速くなる場合があります。
StaticArraysでよく使われる型はSArrayですが、それ以外にもMArrayやSizedArrayなどがあります。以下のように使い分ければOK
- 短い固定長でimmutableな場合 → SArray
- 短い固定長でmutableな場合 → MArray
- 長い固定長の場合 → Array

実行環境

実行環境は以下の通りです。

julia> versioninfo()
Julia Version 1.7.0
Commit 3bf9d17731 (2021-11-30 12:12 UTC)
Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: AMD Ryzen 7 2700X Eight-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-12.0.1 (ORCJIT, znver1)
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

(@v1.7) pkg> status StaticArrays
      Status `~/.julia/environments/v1.7/Project.toml`
  [90137ffa] StaticArrays v1.2.13

高速化の例

行列とベクトルの積のベンチマークを見てみましょう。

using StaticArrays  # 本記事の主役
using BenchmarkTools  # ベンチマーク

# 普通のベクトル・行列
M0 = randn(3,3)
v0 = randn(3)
@benchmark M0*v0

# StaticArrays.jlのベクトル・行列
v1 = @SVector randn(3)
M1 = @SMatrix randn(3,3)
@benchmark M1*v1

v0, M0が通常のベクトル・行列で、v1, M1がStaticArrays.jlで用意されているベクトル・行列です。

上画像はベンチマークの結果のスクリーンショットで、StaticArrays.jlを使った方が3.8倍程度高速になっていることが分かります。
メモリ使用量も減っていますね。(80 bytes -> 32 bytes)

StaticArrays.jlで定義される型

using StaticArraysすれば、抽象型StaticArrayとその部分型(具象型)が使えるようになります。

julia> subtypes(StaticArray)  # StaticArrayの部分型を調べる
6-element Vector{Any}:
 FieldArray{N} where N<:Tuple
 MArray
 SArray
 SHermitianCompact
 SizedArray
 StaticArrays.SUnitRange

これらの型は要約すると以下のようになります。^[1]

SArray immutableな固定長配列
- 要素への代入をしない場合はこれを使えばOK。
MArray mutableな固定長配列
- 要素への代入をする場合はこちら。ただしSArrayより遅い。
SizedArray mutableな固定長配列
- 使い方はMArrayにほぼ同じ、内部実装が異なる。
- 多くの場合MArrayの方が速い。
FieldArray
- structで定義した固定長配列に便利
- (本記事では詳細を扱いません。)
SHermitanCompact
- 固定長のHermite行列を扱うための型
- 対称行列のための型SSymetricは用意されていません。
- (本記事では詳細を扱いません。)
SUnitRange
- Base.UnitRangeに対応する固定長配列。
- (本記事では詳細を扱いません。)

SArray/MArray/SizedArrayは通常の配列Arrayに制約(固定長、不可変(immutable)など)が加わったものになります。以下でもう少し詳しく見ていきます。

SArray

これはstatically-sized arrayの略で、固定長の(immutableな)配列を表します。
Vector{T}がArray{T,1}のaliasになっているのと同様に、SMatrix{S1, S2, T, L}がSArray{Tuple{S1, S2}, T, 2, L}のaliasになっています。

ここではこのSMatrixを例にしてみます。
インスタンスを作るには以下のようにします。

L1 = @SMatrix [1 2 3;4 5 6]  # マクロを使えば通常の定義がそのまま使える
L2 = SMatrix{2,3}([1 2 3;4 5 6])  # 行列(Matrix)から変換して作ることもできるが、一度Matrixを作ってから変換するためパフォーマンスは良くない
L3 = SMatrix{2,3}(1,4,2,5,3,6)  # パフォーマンスを維持してコンストラクタで作る際には順番に気をつける(comumn-majorなので)
L1 === L2 === L3  # true, すべて厳密に等しい

SVector, SArrayでも同様にできます。

immutableなので要素への代入はエラーになります。

L1 .*= 2  # エラー

固定長なので要素の追加もエラーになります。

push!(L1,2)  # エラー

MArray

これはstatically-sized mutable arrayの略で、固定長のmutableな配列を表します。

SMatrixと同様に、ここでもMMatrixを例にしてみます。
インスタンスを作るには以下のようにします。

M1 = @MMatrix [1 2 3;4 5 6]  # マクロを使えば通常の定義がそのまま使える
M2 = MMatrix{2,3}([1 2 3;4 5 6])  # 行列(Matrix)から変換して作ることもできるが、一度Matrixを作ってから変換するためパフォーマンスは良くない
M3 = MMatrix{2,3}(1,4,2,5,3,6)  # パフォーマンスを維持してコンストラクタで作る際には順番に気をつける(comumn-majorなので)
M1 === M2 === M3  # false, mutableなので厳密には等しくない
M1 == M2 == M3  # true, 等号は成立する

MVector, MArrayでも同様にできます。

mutableなので要素への代入はエラーになりません。

M1 .*= 2

固定長なので要素の追加はエラーになります。

push!(M1,2)  # エラー

SizedArray

これも固定長の配列ですが、SArray, MArrayとは異なり、fieldとして配列をそのまま持っています。とりあえずインスタンスの作り方から見てみましょう。

SMatrixと同様に、ここでもMMatrixを例にしてみます。
インスタンスを作るには以下のようにします。

N1 = SizedMatrix{2, 3}([1 2 3;4 5 6])  # 通常の行列Matrixから定義
N2 = SizedMatrix{2, 3}(1,4,2,5,3,6)  # やはりこちらの方が速い (@benchmarkで確認どうぞ)
N1 === N2  # false, mutableなので厳密には等しくない
N1 == N2  # true, 等号は成立する

SizedVector, SizedArrayでも同様にできます。

前述のように、SizedArrayはfieldとして配列をそのまま持っています。フィールドの中身を確認するにはdump関数が使えます。
REPLで実行した結果がこちらです。

julia> dump(L1)
SMatrix{2, 3, Int64, 6}
  data: NTuple{6, Int64}
    1: Int64 1
    2: Int64 4
    3: Int64 2
    4: Int64 5
    5: Int64 3
    6: Int64 6

julia> dump(M1)
MMatrix{2, 3, Int64, 6}
  data: NTuple{6, Int64}
    1: Int64 1
    2: Int64 4
    3: Int64 2
    4: Int64 5
    5: Int64 3
    6: Int64 6

julia> dump(N1)
SizedMatrix{2, 3, Int64, 2, Matrix{Int64}}
  data: Array{Int64}((2, 3)) [1 2 3; 4 5 6]

SMatrix, MMatrixともに内部的にはtupleとしてデータを保持していますが、SizedMatrixは内部的にはMatrixを持っているだけです。

MMatrixとSizedMatrixは内部実装が異なるだけで、出来ることは基本的には同じです。
なのでmutableなので要素への代入はエラーになりませんし

N1 .*= 2

固定長なので要素の追加はエラーになります。

push!(N1,2)  # エラー

ベンチマーク

`Array` vs `SArray`

N = 3

冒頭の繰り返しですが、以下のようにベンチマークを取ります。

N = 3

# 普通のベクトル・行列
M0 = randn(N,N)
v0 = randn(N)
result0 = @benchmark M0*v0

# StaticArrays.jlのベクトル・行列
v1 = @SVector randn(N)
M1 = @SMatrix randn(N,N)
result1 = @benchmark M1*v1

judge(mean(result1),mean(result0))

judge結果のスクリーンショットです↓

再び計測したので冒頭と結果が異なりますが、-76.40%の高速化、つまり1/4程度の計算時間になることが分かります。

N = 50

サイズが大きくなるとどうなるでしょうか。

N = 50

# 普通のベクトル・行列
M0 = randn(N,N)
v0 = randn(N)
result0 = @benchmark M0*v0

# SArray
v1 = @SVector randn(N)
M1 = @SMatrix randn(N,N)  # JITコンパイルなので最初は時間がかかる
result1 = @benchmark M1*v1

judge(mean(result1),mean(result0))

メモリ使用量は減りましたが、実行速度は通常のベクトルに比べて僅かに悪化しました。^[2]

`SArray` vs `MArray`

N = 3

# SArray
v1 = @SVector randn(N)
M1 = @SMatrix randn(N,N)
result1 = @benchmark M1*v1

# MArray
v2 = @MVector randn(N)
M2 = @MMatrix randn(N,N)
result2 = @benchmark M2*v2

judge(mean(result2),mean(result1))

僅かですが、SArrayの方がMArrayよりも速いようです。一般に、immutableよりもmutableの方が低速なので、要素への再代入を避けられる場合はSArrayを使うのが良いですね。^[3]

`MArray` vs `SizedArray`

N = 3

# MArray
v2 = @MVector randn(N)
M2 = @MMatrix randn(N,N)
result2 = @benchmark M2*v2

# SizedArray
v3 = SizedArray(v2)
M3 = SizedArray(M2)
result3 = @benchmark M3*v3

judge(mean(result3),mean(result2))

明らかにMArrayの方がパフォーマンスが良いですね。
では、SizedArrayの方が良いパフォーマンスを出すのはどのような状況でしょうか？

# インスタンス生成用の関数
f2(M) = MMatrix{3,3}(M)
f3(M) = SizedMatrix{3,3}(M)

n = 3
M = randn(n,n)

result2_ = @benchmark f2(M)
result3_ = @benchmark f3(M)

judge(mean(result3_),mean(result2_))

上記の例ではSizedArrayの方が速い結果となりました。
SizedArrayはfieldとしてArrayをそのまま持っているので、インスタンス生成時の変換のコストがかからず、高速になっています。
ただし、インスタンス生成後の演算が結局遅いので、MArrayを使う場面の方が多いと思います。

`SizedArray` vs `Array`

N = 3

# 普通のベクトル・行列
M0 = randn(N,N)
v0 = randn(N)
result0 = @benchmark M0*v0

# SizedArray
v1 = SizedVector{N}(randn(N))
M1 = SizedMatrix{N,N}(randn(N,N))
v3 = SizedArray(v1)
M3 = SizedArray(M1)
result3 = @benchmark M3*v3

judge(mean(result3),mean(result0))

StaticArrayの中ではかなり遅い方だったSizedArrayですが、通常のArrayに比べれば十分高速ではあります。

ベンチマークまとめ

計算時間としては以下のようになります。
SArray ≤ MArray ≪ SizedArray ≪ Array

短い固定長配列かつimmutable
- SArrayを使いましょう
短い固定長配列かつmutable
- MArrayを使いましょう
長い固定長配列
- Arrayを使いましょう

おまけ

`struct`との比較

繰り返すようですが、「ベクトルの次元が決まっているとき」にしかStaticArrays.jlは使えません。ここではstructとの比較を考えてみましょう。

2次元上の点を使うのであれば

struct Point2D
    x::Float64
    y::Float64
end

のような構造体を定義すれば良いですが、この場合、(多重ディスパッチで*のmethodを追加しない限りは)Point2Dに対して行列で変換を行うことは出来ません。^[4]

ベクトルとして使える演算をすべて自前で実装するのは大変なので、Point2Dの代わりにSVector{2, Int64}を使えば良いということになります。

これによって

通常の配列(e.g. [1,2])よりは高速で
自前の構造体Point2Dよりは記述量が少ない

ような2次元ベクトルが実現できるようになります。

FieldArrayはSArray/MArray/SizedArrayと用途が異なるので説明を避けていたのですが、実はこのようなstructの定義において便利です。

struct Point2D <: FieldVector{2, Float64}
    x::Float64
    y::Float64
end

のように定義すれば、FieldArrayに定義されたmethodが使えるようになって便利です。

SArrayと基本的にパフォーマンスは変わらないはずなので、SArrayかFieldArrayを使うかは単純に可読性・変更容易性などの観点から決めれば良いと思います。

幾何的な点を作るときにも便利

数字を並べただけの $n$ 次元ベクトルではなく、幾何的な点を表す場合は、GeometryBasics.jlかMeshes.jlで定義されるPointを使うのが便利なことがあります。
これらには内部的にSVectorが使われています。

次元が決まった部分型を新しく作るときにも便利

Rotations.jlがこれの例になります。

$M\times N$ 行列に対応する固定長配列の抽象型がStaticArray{M,N}でした。^[5]

3次元回転を表す行列はRotation{3,T}として作られていますが、これはStaticMatrix{3, 3, T}の部分型になっています。

このように、固定長配列の型を新しく作りたい場合はStaticArrayの部分型にすると様々なmethodがそのまま使えるので便利になります。

脚注

この他にもScalar, SOneTo, TrivialViewなどがStaticArrays.jlで用意されているますが、これらはsubtypes(StaticArray)で確認できません。本記事ではArray/SArray/MArray/SizedArrayの比較を主に扱いたいので、深く立ち入らないことにします。 ↩︎
StaticArraysによって配列のオーバーヘッドは削減できますが、配列のサイズが大きくなるとオーバーヘッドの寄与が小さくなるのでBLASで計算する方が速くなります。 ↩︎
MArrayの方が速いベンチマーク結果が得られることもあるので、パフォーマンスの差は軽微なようです。 ↩︎
Point2D<:AbstractVectorとすれば、ベクトルと見做されますが、演算を自分で定義する必要があります。 ↩︎
要素型 $T$ の $M\times N$ 行列に対応する固定長配列の抽象型はStaticArray{M,N,T}です。 ↩︎

GitHubで編集を提案