ディープラーニングコンパイラについて知っておくべきこと

[[409589]]

ディープラーニングはここ 10 年ほどで急速に発展し、業界では多くのディープラーニング アルゴリズム開発フレームワークが登場しました。同時に、ディープラーニングには幅広い応用シナリオがあり、膨大な計算能力が求められるため、さまざまな種類の CPU、GPU、TPU、NPU など、さまざまな汎用ハードウェアや専用ハードウェア上でディープラーニング アルゴリズムを実行する必要があります。その結果、図 1 に示すように、フレームワークとハードウェアの組み合わせが爆発的に増加します。たとえば、TensorFlow が GPU コンピューティングをサポートする場合、TensorFlow 内のすべての演算子に対して GPU バージョンを開発する必要があります。D チップをサポートする場合は、演算子ごとに D チップ バージョンを開発する必要があります。このプロセスは間違いなく非常に時間がかかり、労力がかかります。

図1

同時に、YOLO、BERT、GPT など、多くのアルゴリズム ネットワークが存在します。これらのアルゴリズム ネットワークは、さまざまなタイプ、形状、接続関係の演算子で構成されています。最終的には、異なる種類およびモデルのハードウェア上で実行されます。その結果、各シナリオに最適なオペレーターを手動で開発して実装するのに高いコストがかかります。ここでは 2 つの例を示します。図 2 に示すように、演算子の融合は一般的なパフォーマンス最適化方法です。融合の前に、各演算子は計算の前後にメモリからキャッシュにデータを読み取り、その後キャッシュからメモリにデータを書き戻す必要があります。融合後は、演算子間のメモリの読み取りと書き込みを回避できるため、パフォーマンスが向上します。従来のアプローチでは、オペレータの接続関係に基づいて手動で融合オペレータを開発しますが、異なるネットワーク内の異なるカテゴリのオペレータ間の接続関係を完全に列挙することはほぼ不可能です。もう 1 つの例は、オペレータのチューニングです。オペレータ実装プロセスにはパフォーマンスに影響するパラメータが多数ありますが、従来の手動オペレータ開発方法では、これらのパラメータを表現して維持したり、さまざまな形状やハードウェアに最適なパフォーマンスを実現するためにパラメータを調整したりすることが困難です。

図2

ディープラーニングコンパイラは、上記の問題を解決するために作成されました。これは、フレームワークとハードウェア間の共通コンポーネントおよびブリッジとして機能します。最終的な目標は、一度開発するだけで、あらゆるデバイスに最適なコードを自動的に生成することです。たとえば、CPU 用に開発された演算子は、GPU や 3D チップでもほぼ同じように使用できるため、コストが大幅に削減されます。

図 3 に示すように、ディープラーニング コンパイラーのコンポーネントと機能について簡単に紹介します。まず、フロントエンドはさまざまなフレームワークから計算グラフを取得し、高レベル IR のデータ構造を使用してそれを表現します。次に、この段階で、定数の折りたたみ、演算子の融合、等価置換などの一連のグラフ最適化が実行されます。等価置換の例を示します。元の計算グラフは次のようになります。計算方法を変更します。結果は変わりませんが、パフォーマンスは向上する可能性があります。次に、計算グラフ内の各演算子に対して、DSL と呼ばれるドメイン固有言語を使用して演算子の計算プロセスを記述し、演算子を最適化します。たとえば、演算子は、タイリング、マルチコア、ダブルバッファを使用して最適化できます。演算子の計算プロセスは通常、複数のループを使用して実装されるため、たとえば、行列の乗算は 3 重のループになります。ディープラーニング コンパイラーは、ループ上でさまざまな変換を簡単に実行し、これらの変換のパラメーターを調整して、さまざまな形状やハードウェアに最適な演算子の実装を取得できます。最後に、低レベルの IR に基づいて、さまざまなハードウェア用の特定のコードが生成されます。

図3

最後に、業界における既存のコンパイラ プロジェクトを紹介します。現在、TVM は最もエコロジカルに完成されたオープンソースかつフレームワークに依存しないプロジェクトであり、多くの企業に採用されています。 TVM のプロセスを図 3a に示します。TVM は、TensorFlow pb、onnx、TorchScript などのさまざまなフレームワークのモデルをインポートし、TVM の Relay と呼ばれる高レベル IR で統一的に表現できます。 IR の各演算子は、Tensor 式 DSL を使用して計算を記述およびスケジュールします。この DSL は、アインシュタインの表記法を使用して演算子計算を記述します。演算子計算は通常、複数の for ループとして表現されます。次に、Halide のアイデアに基づいて、スケジュールを使用して、この複数の for ループに対して、ループのマージ、分割、シーケンス変換などのさまざまな変換を実行します。最後に、低レベルの IR に下げて特定のデバイス コードを生成し、推論を実行します。

ここでは、TVM が最適な演算子コードを生成する方法について簡単に紹介します。前述のように、演算子は compute によって記述する必要があり、次に compute に対応する複数の for ループをスケジュールして変換、つまりスケジュールする必要があります。 TVM のオペレータ生成とチューニングは、3 世代にわたる開発を経てきました。 TVM/AutoTVM の第 1 世代では、ユーザーがオペレータ計算とオペレータ スケジュールを記述する必要があります。AutoTVM と TVM の違いは、スケジュールでいくつかの可変パラメータを定義でき、遺伝的アルゴリズムを使用してパラメータを調整できることです。たとえば、ループを 2 つのセグメントに分割する場合、どこで分割するかを最適化できます。第 2 世代の AutoScheduler (Ansor) では、ユーザーはオペレータを開発するだけで、Ansor がいくつかのルールに基づいてスケジュールの変更を自動的に実行します。スケジュール開発は、TVM の表現メカニズムと基礎となるハードウェア原理の両方に精通する必要があるため、非常に難しいことがよくあります。そのため、Ansor は開発者の作業負荷と開発の難易度を大幅に軽減できます。欠点は、Ansor のチューニングに時間がかかり、1 つのオペレータのチューニングに 1 時間かかることが多いことです。畳み込みネットワークを例にとると、Ansor はいくつかのシナリオで TensorFlow 演算子のパフォーマンスを上回ることができますが、TensorRT 実装とはまだ一定のギャップがあります。第 3 世代の Meta Schedule (AutoTensorIR) はまだ初期段階であり、チューニング速度とパフォーマンスを最適化することが期待されています。まだ利用できないため、しばらく様子を見ます。

TVM の実装には、TVM に基づく D チップのコード生成サポートを追加する Huawei の D チップ TBE オペレーター開発ツールが含まれています。 TVM は Halide コンピューティング + スケジューリング ルートを採用しています。Tensor Comprehensions、Tiramisu、Huawei が独自に開発した AKG など、多面体アルゴリズム ルートを採用するコンパイラーも存在します。この方法は、ユーザーがスケジュールを開発する必要がなく、オペレーターの開発のみを必要とする点で Ansor に似ており、よりユーザーフレンドリーです。その中で、AKG は MindSpore のグラフコンピューティング融合に使用されています。その他のディープラーニングコンパイラーには、TensorFlow の XLA、TensorRT などがあり、使用したことがあるかもしれません。

要約すると、ディープラーニングコンパイラには多くの利点があります。たとえば、新しいハードウェアを簡単にサポートし、開発の重複を回避し、手動の最適化の代わりに一連の自動最適化を使用して究極のコスト効率を実現します。現在、ディープラーニングコンパイラーには依然としていくつかの欠点があり、急速な発展段階にあります。例えば、チューニングに時間がかかる、複雑な演算子を効果的に生成できない、モデル内でディープラーニングコンパイラーによって生成される演算子がライブラリから呼び出される演算子を上回る割合が低いなど、まだまだ投資と最適化を継続する必要があります。