TensorFlow.jsがWebGLで並列計算を実行する仕組みを理解する

先日、機械学習モデルの学習・デプロイのためのJavaScriptライブラリ「TensorFlow.js」がGoogleから公開されました。Googleからは「deeplearn.js」という同様のライブラリが以前から公開されていましたが、TensorFlow.jsはこのdeeplearn.jsをコアライブラリとして利用しており、モデルの構築を簡単にするhigh level APIや、TensorFlow本体・Kerasで学習したモデルをインポートするツールを備えているとのことです。

このライブラリでは、WebGLを介してGPUを利用することで計算を高速化しているのが特徴です。WebGLはJavaScriptでグラフィック描画を高速に行うための機能ですが、これをどのようにして機械学習の計算に利用しているのか調べてみました。

そもそもWebGLは何をするためのもの？

TensorFlow.jsでの計算の仕組みを理解するには、そもそもWebGLがどのようにグラフィック描画をしているかを理解する必要があります。コンピュータグラフィクスの基本的な処理に「レンダリング」というものがあります。これはプログラム内のデータ表現から、画像や映像を生成する処理です。元になるデータは、物体の形や色や質感、光源の色や位置、視点の位置などがあります。同じ物体でも光の当たり方や視点の位置によって見え方が変わるので、物体を移動・回転させたりすると、その都度、物体の各点がどのように見えるかを計算する必要があります。

この計算を高速に行うために利用されるのがGPUです。GPUではこの計算を数千単位のコアで並列して実行することができます。WebGLはこれらのGPU計算をWebブラウザから実行するためのJavaScript APIです。

TensorFlow.jsではどうやって利用しているの？

グラフィック描画ために開発されたWebGLを、機械学習のような汎用計算に利用してやろうというのがTensorFlow.jsです。GPUでの計算は「シェーダ」という機能が担当しており、特にディスプレイ上の各点に何色を表示させるかを「フラグメントシェーダ」が担当しています。WebGLでは、このフラグメントシェーダに行わせる処理を自分で記述することができます。

通常フラグメントシェーダでは、物体や光源の情報を受け取り、ディスプレイの各点に表示すべき色情報を出力するのですが、この入出力機能を「関数」ととらえ、より一般的な計算に使ってしまおうというのが、TensorFlow.jsでやっていることです。TensorFlow.jsでは、入出力テンソルをテクスチャ（色情報を定義する2次元の配列）として表現し、出力テンソルの要素1つ1つに対してシェーダを並列実行することで、目的の計算を実行しています。

TensorFlow.jsの構成要素

ここからはソースコードを見ながら追っていきます。TensorFlow.jsは大きく2つの部分から構成されています。1つは機械学習モデルの構築を簡単にするhigh level APIの部分。もう1つはWebGLを利用して計算を行うコア部分。今回はWebGLの利用方法に興味があるので後者を詳しく見ていきます。

• https://github.com/tensorflow/tfjs-layers
• https://github.com/tensorflow/tfjs-core

tfjs-coreの主な構成要素

Environemnt

このクラスはモジュールのインポート時にインスタンス化され、グローバル名前空間に"ENV"として保持されています。いわゆるシングルトンクラスで、他のクラスからも常に同一のインスタンス"ENV"が使われます。

使用できるKernelBackendの一覧を持っていて、それぞれのKernelBackendは優先度が設定されています。デフォルトではMathBackendWebGLのほうが高い優先度が設定されているので、WebGLが使える環境ではMathBackendWebGLを、使えない環境ではMathBackendCPUを使うようになっています。

Engine

EnvironmentからKernelBackendを渡されてインスタンス化されます。EngineクラスにはrunKernelというメソッドが用意されていて、後述する各種オペレータがこのrunKernelメソッドを実行すると、Engineが持つKernelBackendを通じてオペレータの計算が実行されます。

各種オペレータ定義

tf.conv2d, tf.matMul, tf.addなどの各種オペレータがスタティックメソッドとして定義されています。ユーザーはこれらのメソッドを使ってモデルを構築します。たとえば行列乗算の定義はこちら。

tfjs-core/src/ops/matmul.ts l.41

  static matMul(
      a: Tensor2D, b: Tensor2D, transposeA = false, transposeB = false):
      Tensor2D {

    // 中略

    const grad = (dy: Tensor2D) => {
      if (!transposeA && !transposeB) {
        return {
          a: () => dy.matMul(b.toFloat(), false, true),
          b: () => a.toFloat().matMul(dy, true, false)
        };
      } else if (!transposeA && transposeB) {
        // 中略
      }
    };
    return ENV.engine.runKernel(
        backend => backend.matMul(a, b, transposeA, transposeB), {a, b}, grad);
  }

Env.engine.runKernelに順伝播関数、入力テンソル、逆伝播関数を渡しているのが確認できます。

KernelBackend

KernelBackendはオペレータの計算を実行するクラスのインタフェースで、バックエンドが実装すべきメソッド群の宣言がずらずらと並んでいます。KernelBackendインタフェースを実装したクラスは2つあり、それぞれCPU用・GPU用の実装になっています。

tfjs-core/src/kernels/backedn.ts l.48

export interface KernelBackend extends TensorStorage, BackendTimer {
  // 中略
  add(a: Tensor, b: Tensor): Tensor;
  subtract(a: Tensor, b: Tensor): Tensor;
  multiply(a: Tensor, b: Tensor): Tensor;
  divide(a: Tensor, b: Tensor): Tensor;
  // 中略
}

MathBackendCPU

KernelBackendインタフェースをCPU向けに実装したのがMathBackendCPU。各メソッドをCPU上で実行するためのコードが記述されています。たとえば行列乗算の実装はこちら。

tfjs-core/src/kernels/backend_cpu.ts l.219

  matMul(a: Tensor2D, b: Tensor2D, transposeA: boolean, transposeB: boolean):
      Tensor2D {
    const sharedDim = transposeA ? a.shape[0] : a.shape[1];
    const leftDim = transposeA ? a.shape[1] : a.shape[0];
    const rightDim = transposeB ? b.shape[0] : b.shape[1];

    const aValues = a.dataSync();
    const bValues = b.dataSync();

    const [aOuterStep, aInnerStep] =
      transposeA ? [1, a.strides[0]] : [a.strides[0], 1];
    const [bOuterStep, bInnerStep] =
      transposeB ? [b.strides[0], 1] : [1, b.strides[0]];

    const aOuterEnd = leftDim * aOuterStep;
    const bOuterEnd = rightDim * bOuterStep;

    const result = new Float32Array(leftDim * rightDim);
    let resultIndex = 0;

    for (let aOuter = 0; aOuter < aOuterEnd; aOuter += aOuterStep) {
      for (let bOuter = 0; bOuter < bOuterEnd; bOuter += bOuterStep) {
        let aInner = aOuter;
        let bInner = bOuter;
        let sum = 0;
        for (let k = 0; k < sharedDim; ++k) {
          sum += aValues[aInner] * bValues[bInner];
          aInner += aInnerStep;
          bInner += bInnerStep;
        }
        result[resultIndex++] = sum;
      }
    }
    return ops.tensor2d(result, [leftDim, rightDim]);
  }

transposeというオプションがあるのでちょっと複雑に見えますが、普通に3重のfor文を回して各要素を乗算していることがわかります。

MathBackendWebGL

KernelBackendインタフェースをGPU向けに実装したのがMathBackendWebGL。各メソッドをGPU上で実行するためのコードが記述されています。たとえば行列乗算の実装はこちら。

tfjs-core/src/kernels/backend_webgl.ts l.343

  matMul(a: Tensor2D, b: Tensor2D, transposeA: boolean, transposeB: boolean):
      Tensor2D {
    const program = new MatMulProgram(a.shape, b.shape, transposeA, transposeB);
    return this.compileAndRun<Tensor2D, Tensor2D>(program, [a, b]);
  }

MatMulProgramというクラスをインスタンス化してコンパイルして実行していますが、これだけだとよくわかりません。ここが今回最も興味のあるところなので、詳しく見ていきましょう。

MatMulProgram

MatMulProgramでは、行列乗算を計算する際にフラグメントシェーダで実行すべき処理が定義されています。

tfjs-core/src/kernels/webgl/matmul_gpu.ts l.20

export class MatMulProgram implements GPGPUProgram {
  variableNames = ['matrixA', 'matrixB'];
  outputShape: number[];
  userCode: string;

  constructor(
      aShape: [number, number], bShape: [number, number], transposeA = false,
      transposeB = false) {
    const outerShapeA = transposeA ? aShape[1] : aShape[0];
    const outerShapeB = transposeB ? bShape[0] : bShape[1];
    const sharedDim = transposeA ? aShape[0] : aShape[1];
    this.outputShape = [outerShapeA, outerShapeB];

    const aSnippetFromOffset = (vec4Offset: number, indexVar: string|number) =>
        transposeA ? `${indexVar} + ${vec4Offset}, aRow` :
                     `aRow, ${indexVar} + ${vec4Offset}`;
    const bSnippetFromOffset = (vec4Offset: number, indexVar: string|number) =>
        transposeB ? `bCol, ${indexVar} + ${vec4Offset}` :
                     `${indexVar} + ${vec4Offset}, bCol`;

    const sharedDimNearestVec4 = Math.floor(sharedDim / 4) * 4;
    const sharedDimVec4Remainder = sharedDim % 4;

    this.userCode = ` float dotARowBCol(int aRow, int bCol) {
      float result = 0.0;
      for (int i = 0; i < ${sharedDimNearestVec4}; i += 4) {
        vec4 a = vec4(
          getMatrixA(${aSnippetFromOffset(0, 'i')}),
          getMatrixA(${aSnippetFromOffset(1, 'i')}),
          getMatrixA(${aSnippetFromOffset(2, 'i')}),
          getMatrixA(${aSnippetFromOffset(3, 'i')})
        );
        vec4 b = vec4(
          getMatrixB(${bSnippetFromOffset(0, 'i')}),
          getMatrixB(${bSnippetFromOffset(1, 'i')}),
          getMatrixB(${bSnippetFromOffset(2, 'i')}),
          getMatrixB(${bSnippetFromOffset(3, 'i')})
        );

        result += dot(a, b);
      }

      if (${sharedDimVec4Remainder === 1}) {
        result += getMatrixA(${aSnippetFromOffset(0, sharedDimNearestVec4)}) *
          getMatrixB(${bSnippetFromOffset(0, sharedDimNearestVec4)});
      } else if (${sharedDimVec4Remainder === 2}) {
        vec2 a = vec2(
          getMatrixA(${aSnippetFromOffset(0, sharedDimNearestVec4)}),
          getMatrixA(${aSnippetFromOffset(1, sharedDimNearestVec4)})
        );
        vec2 b = vec2(
          getMatrixB(${bSnippetFromOffset(0, sharedDimNearestVec4)}),
          getMatrixB(${bSnippetFromOffset(1, sharedDimNearestVec4)})
        );
        result += dot(a, b);
      } else if (${sharedDimVec4Remainder === 3}) {
        vec3 a = vec3(
          getMatrixA(${aSnippetFromOffset(0, sharedDimNearestVec4)}),
          getMatrixA(${aSnippetFromOffset(1, sharedDimNearestVec4)}),
          getMatrixA(${aSnippetFromOffset(2, sharedDimNearestVec4)})
        );
        vec3 b = vec3(
          getMatrixB(${bSnippetFromOffset(0, sharedDimNearestVec4)}),
          getMatrixB(${bSnippetFromOffset(1, sharedDimNearestVec4)}),
          getMatrixB(${bSnippetFromOffset(2, sharedDimNearestVec4)})
        );
        result += dot(a, b);
      }

      return result;
    }

    void main() {
      ivec2 resRC = getOutputCoords();
      setOutput(dotARowBCol(resRC.x, resRC.y));
    }
    `;
  }
}

入力テンソルの変数名、出力テンソルのサイズ、そして「userCode」という文字列をプロパティとして持っています。このuserCodeがフラグメントシェーダのコードになります。フラグメントシェーダは、シェーディング言語というC言語風の独自の言語で書く必要があります。コードの最後のほうの「dotARowBCol(resRC.x, resRC.y)」で、担当している出力要素の値を計算しています。

dotARowBColの中身を見てみると、CPUでの計算と異なりfor文が1重しかないのがわかります。フラグメントシェーダでは、担当する出力1要素について計算すれば良いので、出力全要素を走査する2つのループは必要ありません。

ちなみにdotARowBColではスカラ同士を直接乗算せずに、4つのスカラを4次元ベクトルと見立ててベクトル同士の乗算を行っています。たいていのGPUにはベクトル演算用のプロセッサが搭載されているので、こうすることでスカラ同士を1つずつ乗算するよりも高速化できるようです。

compileAndRunメソッド

上で定義したフラグメントシェーダをコンパイルして実行しているのがここ。

tfjs-core/src/kernels/backend_webgl.ts l.924

  private compileAndRun<T extends Tensor, K extends Tensor>(
      program: GPGPUProgram, inputs: T[], output?: K,
      customSetup?: (gpgpu: GPGPUContext, webGLProgram: WebGLProgram) => void):
      K {
    if (output == null) {
      output = this.makeOutputArray(program.outputShape, inputs[0].dtype);
    }
    const inputsData: Array<TensorData<T>> = inputs.map(input => {
      this.uploadToGPU(input.dataId);
      return {tensor: input, texData: this.texData.get(input.dataId)};
    });
    this.uploadToGPU(output.dataId);
    const outputData = {
      tensor: output,
      texData: this.texData.get(output.dataId)
    };
    const key = gpgpu_math.makeShaderKey(program, inputsData, outputData);
    const binary = this.getAndSaveBinary(key, () => {
      return gpgpu_math.compileProgram(
          this.gpgpu, program, inputsData, outputData);
    });
    const shouldTimeProgram = this.activeTimers != null;
    let query: WebGLQuery|CPUTimerQuery;
    if (shouldTimeProgram) {
      query = this.startTimer();
    }

    gpgpu_math.runProgram(binary, inputsData, outputData, customSetup);

    if (shouldTimeProgram) {
      query = this.endTimer(query);
      this.activeTimers.push(this.getQueryTime(query));
    }
    return output;
  }

「uploadToGPU」で入出力テンソルをテクスチャに変換し、GPUからアクセスできるようにしています。実際のコンパイル処理・実行処理は「gpgpu_math.compileProgram」や「gpgpu_math.runProgram」で行われています。

このあたりをもっと深掘っていくとWebGLのAPIが見えてきますが、長くなってしまうのでここまで。

まとめ

今回はTensorFlow.jsのコードを読み、WebGLを利用した汎用計算の実装について見てきました。

• TensorFlow.jsのAPIの裏で、CPU用・GPU用のバックエンド処理が実装されている
• GPUでの計算はフラグメントシェーダとして記述されている
• 入出力テンソルはテクスチャとしてGPUに渡している

機械学習関連のライブラリはどんどん開発されて便利になってきていますが、こういった低レイヤーの複雑な処理を隠蔽してくれていることを改めて感じました。今度はTensorFlow本体のコードも読んでみたいと思います。