More than 5 years have passed since last update.

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

Swift

Last updated at 2016-05-30Posted at 2016-05-30

先日、関西モバイルアプリ研究会で「Swiftで書かれたプログラムを1000倍速くした話」のタイトルで LT をしました¹。本投稿はその原稿をベースに、多少加筆修正したものです。

去年 Google が TensorFlow というニューラルネットワークのライブラリを公開しました。僕は Qoncept という AR や画像認識が専門の会社で働いているので、最近よく TensorFlow を使うんですが、残念ながら iOS 用にはまだビルドすることができません。そこで、 TensorFlow の内、テンソルの計算の部分を Swift でシミュレートするライブラリ TensorSwift を作りました。↓は TensorFlow の手書き文字認識のチュートリアル Deep MNIST for Experts を iOS 上で再現したデモです。

（ TensorSwift についての詳細はこちらの投稿を御覧下さい。）

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

`Tensor` とは

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

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 でアクセスします。

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

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

`Tensor` の実装

Tensor は shape と elements を持つだけのシンプルな struct です²。

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

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

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 は、単純に実装すると次のように書けます。

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 を単純なかけ算とたし算に書き変えられます。

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 のキャッシュのヒット率を上げたら速くなったということでした。

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 で書かれてる部分をポインタで書き変えてみました。

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 というライブラリがあるんですが、それを使って高速化したということでした。 iOS や Mac では、標準の Accelerate Framework で提供されています。

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)$ ³ で 50 万倍も速くなるというのはおもしろいなぁと思いました。

スライドはこちらです。 ↩
ここでは説明を単純化するために shape の型を [Int] にしていますが、実際の TensorSwift では [Dimension] をラップした Shape 型になっています。 Dimension は Int をラップした型です。ただし、 Dimension は IntegerLiteralConvertible 、 Shape は ArrayLiteralConvertible なので、本投稿の説明と同じように Tensor(shape: [2, 3], elements: [1, 2, 3, 4, 5, 6]) という書き方ができます。 ↩
ただ、いくつかのパターンで計測した限りでは cblas_sgemm は $O(N^3)$ に乗ってなくて謎です。 ↩

125

115

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up