要約

現代のAIワークロードは、学習および推論の両方において最適化されたコンピューティングカーネルに大きく依存している。これらのAIカーネルは、DRAMとSRAM間でタイルを移動し、そのタイル上で一連の計算を実行するといった明確なデータフロー・パターンに従う。しかし、これらのパターンの明確さにもかかわらず、高性能なカーネルを実装することは依然として複雑である。ピーク性能を達成するためには、現代のアクセラレータを最大限に活用するための、細心の配慮を要するハードウェア中心の最適化が不可欠である。ドメイン特化型コンパイラは、高性能カーネルの記述負荷を軽減しようと試みているが、 usability（使いやすさ）と表現力のギャップにより、しばしば限界に直面している。本論文では、より効率的なAIカーネル開発を実現するための汎用的なタイル化プログラミングモデル「TileLang」を提案する。TileLangは、スケジューリング空間（スレッドバインディング、レイアウト、テンソル化、パイプライン化）をデータフローから分離し、カスタマイズ可能なアノテーションとプリミティブのセットとして封裝している。このアプローチにより、ユーザーはカーネルのデータフローそのものに集中できる一方で、他の多くの最適化処理はコンパイラに任せることができる。多数の一般的に使用されるデバイス上で包括的な実験を実施した結果、評価においてTileLangは主要なカーネルにおいて最先端の性能を達成することが示された。これにより、統一されたブロック・スレッドパラダイムと透過的なスケジューリング機能が、現代のAIシステム開発に求められる強力かつ柔軟な能力を両立していることが確認された。