はじめに

猫も杓子もAI agentの話をしていますね、私もそんな感じです。

 私はこれまで、MN-Core™向けコンパイラを開発してきました。最近は、AI agentがモデルをMN-Core向けにコンパイル・最適化する「AI compiler」の実現に取り組んでいます。ここでいうAI compilerは、AIモデルをコンパイルするだけでなく、AI agentをコンパイラ開発に利用するだけでもなく、AI agentをコード生成に利用する仕組みです。その一環として、PLaMo 3 31BをMN-Core L1000向けに全自動でポーティングすることに成功しました。具体的には、L1000で効率的に実行可能な計算カーネルをそろえ、現段階では実機がないので、エミュレータ上で推論結果の一致を確認するところまでを、全自動で行いました。 従来だと1人月程度を想定するくらいの規模感でしょうか。

 AI compilerといっても、まるっきり今までと違うものを開発してそちらに移行、という話ではありません。MN-Core向けコンパイラを開発してきた延長線上で、ポーティングや最適化の中心を、急速にAI agentへ移しつつあります。このポストでは我々のやったことや、今後やりたいことなどを紹介します。

きっかけ

この活動を始めたきっかけは、趣味で使うとメチャクチャ賢いAI agentが、仕事で使っていると、ちょっと信じられないくらい何もしてくれなかったことでした。MN-Coreや我々のコードベース特有の事象を知らないのはしょうがないな、と、作業に必要そうな情報を一生懸命説明しました。だいぶそれらしいことを言うようになってから、いざ作業をしてもらうと、何も成果が出ない、そんな感じでした。

 説明しても難しいということは、必要な情報を検索して補うだけでは埋めきれないギャップがあるのかもしれません。CPUやGPUについては、記事やコードが世の中に満ちあふれているため、LLMはトレーニングの段階から、それらに触れている。一方で、MN-Coreについては独自性の強い技術なこともあり、公開情報も限られています。そもそも学習時から身につけた「筋肉」がないのではないか。だとすると、トレーニング時点からMN-Coreのことを教える必要があるのでは、と。

 最近見かけた論文に、Customizing an LLM for Enterprise Software Engineering があります。グーグル社内のソフトウェア開発にあわせてGeminiをドメイン特化させたGemini for Googleを開発し、グーグル社内の開発者のAI体験を大幅に向上させた、という内容です。

 グーグルでは独自技術が大量に使われているため、ビルドルールの更新のような簡単な作業すら、一般的なLLMだと苦戦するそうです。Gemini for Googleの性能があまりに良かったため、A/Bテストのコントロール側に割りあてられたユーザが、なんとかしてGemini for Googleの方に切りかえられないかと模索した、なんていうエピソードも面白かったです。だいぶ前なのでいろいろ変わっているとは思いますが、私もグーグルで働いていたので、いかにもありそうな話だなぁ、と。

 規模の差は大きいですが、独自技術を扱っていて、自社LLMをはじめとするAI技術を持っているという点では、我々の状況にも共通しています。比較的自然な方向性を選べたのかな、と思っています。 

やりたいこと

私たちは、この問題に対して、複数の方向から同時にアプローチすることにしました。具体的には、次の4つです。

• CPT/SFT: MN-Coreに詳しいドメイン特化PLaMoを訓練し、活用する。
• AI向け知識基盤の整備: ドキュメントを整備し、AI agentの挙動を補強する。
• コンパイラ基盤の改善: MN-Coreの意味や制約を明示的に表現でき、AIが生成・検証しやすいプログラム記述DSLを定義する。
• AI agentの全面適用: 現在のAI agentでも成果を出せる領域を見極め、コンパイラ開発やカーネル最適化に利用する。

これら4つの活動は、独立したものではありません。相補的で循環的なものです。例えば、AI agentの全面適用を進める際、その成果や失敗例を、トレーニングデータとしてモデルに、ドキュメントとしてAI向け知識基盤の整備に、AI agentがひっかかるポイントの修正としてコンパイラ基盤の改善に還元することができます。AI向け知識基盤の整備の結果、コンパイラ基盤の改善をAI agentを使ってやる時の速度も以前より加速していますし、CPT/SFTのためのデータ生成にも役立てています。

CPT/SFTはデータ準備や適切なトレーニング設定の調査など、少し時間がかかる活動で、まだ循環フィードバックは始まっていませんが、人間による日々の開発活動にも、AI compilerの強化にも、活用していく目論見です。

LEIR

ここからは、私の最近の作業について書きます。まずはLEIRについて。

LEIRは、MN-CoreカーネルをPythonで記述できるDSLです。MN-Core特有の命令やデータ転送を陽に記述できますが、instruction-levelではなく、tensor-levelのopを基本単位としています。そのため、ループや個々の命令を直接記述するのではなく、テンソルに対する計算やデータ移動を組み合わせてプログラムを構成します。

以下は、attentionのスコアを計算し、softmaxの前段のmax reductionまでを記述したコードです。

k_cache = leir.tensor(
    leir.block_half,
    leir.layout.PadLayout.parse("(2,2048,16)/((2_PE:2), (16:32, 8:2, 2_PE:1, 8_MAB:1), (4_W:1, 2:1, 2:16))"),
    leir.LDRAM,
    name="k_cache",
)
q = leir.broadcast_download(q)
q = leir.distribute_download(q)
q = leir.bf(q)
qk = leir.linear(q, k_cache)
qk = leir.unicast_download(qk)
neg_inf = leir.immediate_float(value=float("-inf"), dtype=qk.dtype())
scores = leir.where(attn_mask, qk, neg_inf)
max_scores = leir.reduce_max(scores, axis=2, keepdims=1)

linear,where,reduce_maxといったnumpyやtorchにありそうなopの他に、broadcast_downloadやbfといった、MN-Core特有のopが入っていることがわかると思います。また、leir.tensorで示されているように、値はMN-Core特有のlayoutのようなメタデータを陽に保持しています。

leir.linearなどのPython関数を実行すると、入力のdtypeやlayoutやlocationが制約を満たしているかを検証し、DAGのノードを追加します。define-and-runスタイルのディープラーニングフレームワークがdtypeとshapeに対して制約チェックをするところを、MN-Core特有のlayoutやlocationに対してもチェックしていると考えてもらうと良いと思います。

このように制約を明示した言語でプログラムをAI agentに書かせることで、次のことを狙っています。

• メモリレイアウトや大域的な計算手順など、AI agentが得意な探索に集中させる
• 効率を悪化させるopをそもそも選べないようにし、探索空間を制限する
• 具体的な命令列へのloweringは、アルゴリズムやcompilerに任せる

GPUの分野では、AI agentにTritonなどの比較的抽象度の高い言語を記述させ、ハードウェアの詳細を反映したkernelを探索させる研究が増えています。例えばTritonBenchは、LLMにTriton operatorを生成させ、その正しさや実行性能を評価するbenchmarkです。また、Dr. Kernelは、コンパイル結果や実行性能をフィードバックとして利用し、SFTや強化学習によってTriton kernel生成モデルを訓練する取り組みです。

粒度や記述内容は大きく違いますが、抽象度は高いがハードウェアの情報が露出している探索空間をAI agentに探索させる、というアプローチに共通点があるように思っています。 

循環

個人的な開発スタイルも大きく変化しました。

たとえばPyTorchモデルのポーティングでは、最初は自分が書いていたものを、AI agentに手伝わせるようになり、自分はコミットごとの方針出しとコード確認だけになり、コミットごとにコードを見るようなこともやめ、と、だんだんとAI agentに比重を移していきました。AI agentに任せる範囲を広げる過程で、agentがつまずいた箇所を手がかりに、AI向け知識基盤や前章で紹介したLEIRを改善し、さらにコンパイラスタック全体の問題も修正していきました。

こうした改善の効果は、PLaMo 3 31Bそのものではなく、他のカーネル生成タスクで確認していました。10個中3個しか完了できなかったものが、5個、7個と増えていくことを繰り返していった結果、冒頭で述べたPLaMo 3 31Bのend-to-endポーティングに成功しました。

開発に当てている時間の半分くらいは、「なぜAI agentは自力でこのタスクを完了できず、私のアドバイスが必要になってしまったか」を考察し、解消する時間にあてています。具体的には、AI向け知識基盤の整備や、LEIR opの制約を精緻化したりしています。また、現時点のAI agentには難しいと判明したタスクをベンチマークセットとして集めることもしています。

いったん完了した作業をすぐにマージせず、対処によってAI agentが自律的に問題解決できるようになったかを、再実験することもあります。ゴールを達成できなかったAI agentのセッションに、完成例を見せることもあります。AI agentは「これがあるなら、私の作業は完全にムダでした」などと返してきたりします。ちょっとかわいそうですが、なぜ自律的に進められなかったのかをAI agentと協力して特定し、次の改善につなげることができます。

完了した作業を繰り返すなんていうのは、今まででは考えられなかった時間のムダです。しかし、今後、私自身の生産性が数倍になることは考えにくい一方、AI agentの生産性はまだまだ急速に増えるでしょうから、AI agentの生産性に投資することは、今の時間を補って余りある価値があると考えています。

AI agent自体の進化にも助けられていますが、AIが作業しやすいリソースや基盤を整理した結果、モデルポーティングや最適化のようなアプリケーション寄りのタスクだけでなく、コンパイラ基盤の改善などの作業も加速していて、少しずつ循環が機能しつつある、と感じています。

まとめ

ここで紹介した以外にも、LLMとコンパイラの両面から、広い範囲でアプローチを検討・実装しています。例えば、開発者がMN-Coreについて聞くことができるチャットボットの用意や、AI agentによる機械語の直接生成・コンパイラ出力の事後最適化などです。

AI compilerはまだ黎明期にあり、数年後にどのような形が主流になっているのかは、まだわかりません。個人的には、LLMに詳しいメンバーと、雑多なアイデアを気軽に話せる機会が増えたのを、とても嬉しく思っています。異なる専門性を持つ人たちが協力して新しいものを作れることは、PFNの良さの一つだと思います。こういった活動に興味を持っていただけた方は、ぜひPFNへの応募を検討してみてください。