CompilerGymを試してみた

CompilerGymとは

CompilerGym is a toolkit for exposing compiler optimization problems for reinforcement learning. It allows machine learning researchers to experiment with program optimization techniques without requiring any experience in compilers, and provides a framework for compiler developers to expose new optimization problems for AI.
GitHub

CompilerGymとは、コンパイラ最適化を強化学習の問題として扱うためのツールです。
コンパイラについての知識を持たない機械学習の研究者でも、プログラムの最適化手法を実験で使うことができるようになります。また、コンパイラ開発者がAIに最適化問題を扱わせるための機能も提供します。

試してみる

インストール

pipで一発です。

pip install -U compiler_gym

動作確認兼デモ

README.mdより引用

import gym
import compiler_gym
env = gym.make("llvm-autophase-ic-v0") 
env.require_dataset("npb-v0")  
env.reset()   
env.render()                       
env.step(env.action_space.sample())     

これだけだと、まだ何をしているのかよくわかりませんね。

デモを一行ずつ見ていく

1. OpenAI Gym と CompilerGym 自体をインポートします。

import gym
import compiler_gym

2. llvm-autophase-ic-v0という名前の環境を呼び出します。環境の名前は、次のような命名規則になっているようです。

   1. <compiler>-<observation>-<reward>-<version>
   2. <compiler>-<reward>-<version>
   3. <compiler>-<version>

つまりデモの場合、コンパイラがLLVM、observation typeがAutophase、reward signalがInstruction Count (-Ozに対しての命令列の長さ) となります。

Rewardとしては、IR命令の数や、コード生成後のオブジェクトサイズなどが使えるようです。これらを最小化する方向で学習が進みます。(ほかのrewardも選べるのかな?)

env = gym.make("llvm-autophase-ic-v0")

3. npb-v0という名前のデータセットを用意します。色々なベンチマークが用意されているようです。

env.require_dataset("npb-v0")

4. 環境を初期化して、ランダムなプログラムで埋めます。

env.reset()

5. プログラムのIRを出力します。

env.render()

6. ランダムな最適化を実行します。いよいよ強化学習っぽくなってきましたね。

env.step(env.action_space.sample())

このあとに、env.render()を呼ぶと、確かにプログラムが改変されていることがわかります。

終わりに

今後、コンパイラ最適化の世界にも機械学習アルゴリズムの波が来るのでしょうか。私はあまり機械学習に詳しくありませんが、CompilerGymでその一端を垣間見ることが出来ました。

公式ドキュメントがかなり親切なので、この記事に載っている以上のことを簡単に試せると思います。