Swiftで書かれたプログラムを1000倍速くした話

先日、[関西モバイルアプリ研究会](http://kanmoba.connpass.com/event/31750/)で「Swiftで書かれたプログラムを1000倍速くした話」のタイトルで LT をしました[^1]。本投稿はその原稿をベースに、多少加筆修正したものです。

-------------------------

去年 Google が [TensorFlow](https://www.tensorflow.org/) というニューラルネットワークのライブラリを公開しました。僕は [Qoncept](http://qoncept.co.jp/) という AR や画像認識が専門の会社で働いているので、最近よく TensorFlow を使うんですが、残念ながら iOS 用にはまだビルドすることができません。そこで、 TensorFlow の内、テンソルの計算の部分を Swift でシミュレートするライブラリ [TensorSwift](https://github.com/qoncept/TensorSwift) を作りました。↓は TensorFlow の手書き文字認識のチュートリアル [Deep MNIST for Experts](https://www.tensorflow.org/versions/r0.8/tutorials/mnist/pros/index.html) を iOS 上で再現したデモです。

![deep-mnist.gif](https://raw.githubusercontent.com/qoncept/TensorSwift/master/Resources/DeepMnist.gif)

（ TensorSwift についての詳細は[こちらの投稿](http://qiita.com/koher/items/2c0bfca4d6e31cde674b)を御覧下さい。）

さて、 TensorSwift を作るには行列の積を実装する必要があったんですが、その高速化をがんばったらなんと最初のバージョンより 1000 倍も速くなりました。その過程で色々 Swift のすごさを感じたので、それについて話します。

## `Tensor` とは

TensorSwift の中心は `Tensor` （テンソル）です。 `Tensor` は N 次元配列のようなもので、例えば下記のコードの `a` のように `shape` を `[2, 3]` としたら 2 × 3 の 2 次元配列に、 `b` のように `shape` を `[2, 2, 2]` としたら 2 × 2 × 2 の 3 次元配列になります。

```swift
let a = Tensor(shape: [2, 3],
  elements: [1, 2, 3, 4, 5, 6])
// [[1, 2, 3], [4, 5, 6]]

let b = Tensor(shape: [2, 2, 2],
  elements: [1, 2, 3, 4, 5, 6, 7, 8])
// [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
```

2 次元のときを 2 階のテンソル、 3 次元のときを 3 階のテンソルと言います。

`Tensor` の要素には次のように `subscript` でアクセスします。

```swift
print(a[1, 2])    // 6
print(b[1, 1, 1]) // 8
```

2 階のテンソルは 2 次元配列なのでインデックスが二つ、 3 階のテンソルは 3 次元配列なのでインデックスが三つです。

## `Tensor` の実装

`Tensor` は `shape` と `elements` を持つだけのシンプルな `struct` です[^2]。

```swift
public struct Tensor {
  public let shape: [Int]
  public private(set) var elements: [Float]
}
```

`subsucript` は次のように実装できます。詳細はおいといて、何次元でもインデックスを計算できるように `zip` と `reduce` を使って複雑な処理をしていることに注目して下さい。

```swift
extension Tensor {
  internal func index(indices: [Int]) -> Int {
    return zip(shape, indices).reduce(0) { $0 * $1.0 + $1.1 }
  }

  public subscript(indices: Int...) -> Float {
    get {
      return elements[index(indices)]
    }
    set {
      elements[index(indices)] = newValue
    }
  }
}
```

これを使って行列の積を計算するメソッドを実装します。なんでテンソルなのに行列の積かと言うと、 2 階のテンソルを行列とみなして計算したいことがあるからです。

## 行列の積

行列の積を計算するメソッド `matmul` は、単純に実装すると次のように書けます。

```swift
extension Tensor { // Matrix
  public func matmul(tensor: Tensor) -> Tensor {
    precondition(shape.count == 2)
    precondition(tensor.shape.count == 2)
    let n = shape[1]
    precondition(n == tensor.shape[0])

    let numRows = shape[0]
    let numCols = tensor.shape[1]

    var elements: [Float] = []
    elements.reserveCapacity(numCols * numRows)
    for r in 0..<numRows {
      for c in 0..<numCols {
        var e: Float = 0.0
        for i in 0..<n {
          e += self[r, i] * tensor[i, c] // ここで `subscript` を利用
        }
        elements.append(e)
      }
    }

    return Tensor(shape: [numRows, numCols], elements: elements)
  }
}
```

三重ループの中で `subscript` が繰り返し呼ばれていますが、 `struct` のメソッドは静的ディスパッチされ、インライン展開可能なので、呼び出しコストを気にしなくて大丈夫です。再利用性や可読性を考えてコードを分解してもパフォーマンスを犠牲にしなくていいのは Swift のいいところです。

## 行列に限定

`matmul` は `self` も引数の `tensor` もどっちも行列、つまり 2 次元です。 `subscript` で何次元でもいいように複雑なインデックス計算をしていましたが、 2 次元と決まっているならもっとシンプルに計算できます。

当初の用途では特に速さを重視していなかったので、計算結果さえ正しければ良いということで一番簡単な実装をしたんですが、さすがにこれはないなと思って修正をしました。

2 次元に限定すれば `zip` と `reduce` を単純なかけ算とたし算に書き変えられます。

```swift
extension Tensor { // Matrix
  public func matmul(tensor: Tensor) -> Tensor {
    precondition(shape.count == 2)
    precondition(tensor.shape.count == 2)
    let n = shape[1]
    precondition(n == tensor.shape[0])

    let numRows = shape[0]
    let numCols = tensor.shape[1]

    var elements: [Float] = []
    elements.reserveCapacity(numCols * numRows)
    for r in 0..<numRows {
      for c in 0..<numCols {
        var e: Float = 0.0
        for i in 0..<n {
          e += self.elements[r * n + i] * tensor.elements[i * numCols + c]
            // 行列として計算
        }
        elements.append(e)
      }
    }

    return Tensor(shape: [numRows, numCols], elements: elements)
  }
}
```

これで 1.84 倍速くなりました。

## ループ順を入れ替えてキャッシュのヒット率を向上

Qoncept にはこの春から優秀な新入社員 @t-ae が入ってくれたんですが、その彼が数日後に PR を送ってきました。

彼は、三重ループの順序を入れ替えて `Array` のキャッシュのヒット率を上げたら速くなったということでした。

```swift
extension Tensor { // Matrix
  public func matmul(tensor: Tensor) -> Tensor {
    precondition(shape.count == 2)
    precondition(tensor.shape.count == 2)
    let n = shape[1]
    precondition(n == tensor.shape[0])

    let numRows = shape[0]
    let numCols = tensor.shape[1]

    var elements = [Float](count: numCols * numRows, repeatedValue: 0.0)
    for r in 0..<numRows {
      for i in 0..<n {         // ここと
        let e = self.elements[r * n + i]
        for c in 0..<numCols { // ここが入れ替わった
          elements[r * numCols + c] += e * tensor.elements[i * numCols + c]
        }
      }
    }

    return Tensor(shape: [numRows, numCols], elements: elements)
  }
}
```

これで、さらに 4.69 倍速くなりました。

## ポインタに書き変え

これは僕も負けてられないと思って、 `Array` で書かれてる部分をポインタで書き変えてみました。

```swift
extension Tensor { // Matrix
  public func matmul(tensor: Tensor) -> Tensor {
    precondition(shape.count == 2)
    precondition(tensor.shape.count == 2)
    let n = shape[1]
    precondition(n == tensor.shape[0])

    let numRows = shape[0]
    let numCols = tensor.shape[1]

    // `Array` をポインタに変換
    let leftHead = UnsafeMutablePointer<Float>(self.elements)
    let rightHead = UnsafeMutablePointer<Float>(tensor.elements)

    let elements = [Float](count: numCols * numRows, repeatedValue: 0.0)
    for r in 0..<numRows {
      for i in 0..<n {
        var pointer = UnsafeMutablePointer<Float>(elements) + r * numCols
        let left = leftHead[r * n + i]
        var rightPointer = rightHead + i * numCols
        for _ in 0..<numCols {
          // ここで `Array` のインデックスによるアドレス計算が減って高速化
          pointer.memory += left * rightPointer.memory
          pointer += 1
          rightPointer += 1
        }
      }
    }

    return Tensor(shape: [numRows, numCols], elements: elements)
  }
}
```

`Array` は可変長で色んな便利メソッドをもったリッチな型です。そんな `Array` をポインタとして扱って C や C++ と同じようなパフォーマンス上のメリットを享受できるというのはすごいことです。 `NSArray` や `ArrayList` ではこうはいきません。しかも、 Java の JNI なんかと違って自然な Swift の構文のままポインタを取り扱えます。

これで、さらに 2.99 倍速くなりました。

## BLAS の利用

しばらくすると @t-ae が再び PR を送ってきました。

行列演算では C や FORTRAN で定番の [BLAS](http://www.netlib.org/blas/) というライブラリがあるんですが、それを使って高速化したということでした。 iOS や Mac では、標準の [Accelerate Framework](https://developer.apple.com/library/tvos/documentation/Accelerate/Reference/AccelerateFWRef/index.html) で提供されています。

```swift
import Accelerate

extension Tensor { // Matrix
  public func matmul(tensor: Tensor) -> Tensor {
    precondition(shape.count == 2)
    precondition(tensor.shape.count == 2)
    precondition(shape[1] == tensor.shape[0])

    let result = Tensor(shape: [shape[0], tensor.shape[1]],
      elements: [Float](count: shape[0] * tensor.shape[1],repeatedValue: 0.0))

    let n = Int32(tensor.shape[1])
    let k = Int32(shape[1])
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, Int32(shape[0]),
      n, k, 1.0, elements, k, tensor.elements, n, 1.0,
      UnsafeMutablePointer<Float>(result.elements), n)

    return result
  }
}
```

せっかくこれまで色々がんばってきたのに、ほぼ BLAS の関数を呼び出すだけになってしまいました。

ただ、これも Swift の素晴らしいところです。僕はこれまで BLAS を使ったことがなかったんですが、なんとなくライブラリ境界で無駄な変換コストがかかりそうだと思ってました。しかし、実際にはまったくそんなことはありませんでした。こういうバリバリにチューニングされたライブラリはまだまだ Swift 以外のものを使わざるを得ないので、 Swift からそれらをシームレスに呼べるというのは素晴らしいことです。

これによって、さらに 39.1 倍速くなりました。

## 結論

全部あわせると、最初のバージョンから 1012 倍高速化されました。

行列の積だけだと最初から BLAS 使えよっていうまぬけな話ですが、実際には畳み込みのようなより複雑な計算も最初はシンプルに書いて後から BLAS を導入して高速化しました。

ただ、僕が伝えたかったのは高速化の話ではなくて、 __「Swift は Swift なプログラムが書けるように作られている！！」__ ということです。

- 標準ライブラリから `struct` をはじめとした値型中心の設計がされており、静的ディスパッチやインライン展開が効果的に働く
- `Array` は可変長でリッチな型なのにポインタとして扱える
- C 等の他言語のライブラリもオーバーヘッドなしにシームレスに呼び出せる

など、モダンでリッチな構文と API を備えながらも、　C や C++ 並の速度を実現するための仕組みが備わっている Swift は改めてすごいなぁと感じたという話でした。

## 余談

前述の速度はすべて `-O -whole-module-optimization` で最適化された状態で、　XCTest の `measureBlock` で計測しました。

当初、作業を手伝ってくれていたインターンの学生が計算が遅すぎると言うので最適化は有効になっているか確認すると、 `-Onone` でした。最適化されていない状態から最適化を有効にするだけで約 500 倍速くなったので、そこから比べると最終的に 50 万倍速くなったことになります。 Swift は最適化を有効にしないととても遅い言語なので比較対象としてはアンフェアだと思いますが、同じ計算量 $O(N^3)$ [^3] で 50 万倍も速くなるというのはおもしろいなぁと思いました。

---

[^1]: スライドは[こちら](https://speakerdeck.com/koher/swiftdeshu-karetapuroguramuwo1000bei-su-kusitahua)です。

[^2]: ここでは説明を単純化するために `shape` の型を `[Int]` にしていますが、実際の TensorSwift では `[Dimension]` をラップした `Shape` 型になっています。 `Dimension` は `Int` をラップした型です。ただし、 `Dimension` は `IntegerLiteralConvertible` 、 `Shape` は `ArrayLiteralConvertible` なので、本投稿の説明と同じように `Tensor(shape: [2, 3], elements: [1, 2, 3, 4, 5, 6])` という書き方ができます。

[^3]: ただ、いくつかのパターンで計測した限りでは `cblas_sgemm` は $O(N^3)$ に乗ってなくて謎です。