README.md

I used setup.Dockerfile to create an image that I then ran tests on, swapping in the base image that I was using for the mlgo-development image in the experimental directory in google/ml-compiler-opt. I swapped between ubuntu:20.04 and ubuntu:22.04 as they have different Python versions (3.8.x and 3.10.x respectively), and are what I've worked on MLGO with. I then ran the following commands to generate some data:

Default trace:

time PYTHONPATH=$PYTHONPATH:. python3 compiler_opt/tools/generate_default_trace.py \
    --data_path=/corpus \
    --output_path=/default_trace \
    --gin_files=compiler_opt/rl/regalloc/gin_configs/common.gin \
    --gin_bindings=config_registry.get_configuration.implementation=@configs.RegallocEvictionConfig \
    --gin_bindings=clang_path="'/llvm-project/build/bin/clang'" \
    --sampling_rate=0.2

Warmstart:

time PYTHONPATH=$PYTHONPATH:. python3 compiler_opt/rl/train_bc.py \
    --root_dir=/warmstart \
    --data_path=/default_trace \
    --gin_files=compiler_opt/rl/regalloc/gin_configs/behavioral_cloning_nn_agent.gin

Making sure to apply the following patch to the ml-compiler-opt repo to only run a couple training iterations:

diff --git a/compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin b/compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin
index 917b5ed..e6d1942 100644
--- a/compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin
+++ b/compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin
@@ -12,7 +12,7 @@ include 'compiler_opt/rl/regalloc/gin_configs/network.gin'
 
 train_eval.agent_name=%constant.AgentName.PPO
 train_eval.warmstart_policy_dir=''
-train_eval.num_policy_iterations=3000
+train_eval.num_policy_iterations=10
 train_eval.num_modules=512
 train_eval.num_iterations=200
 train_eval.batch_size=256

Actual training:

time PYTHONPATH=$PYTHONPATH:. python3 compiler_opt/rl/train_locally.py \
    --root_dir=/output_model \
    --data_path=/corpus \
    --gin_bindings=clang_path="'/llvm-project/build/bin/clang'" \
    --gin_files=compiler_opt/rl/regalloc/gin_configs/ppo_nn_agent.gin \
    --gin_bindings=train_eval.warmstart_policy_dir=\"/warmstart/saved_policy\"

setup.Dockerfile

#FROM mlgo-development-ubuntu-20.04
FROM mlgo-development-ubuntu-22.04
RUN git clone https://github.com/llvm/llvm-project
RUN mkdir /llvm-project/build
WORKDIR /llvm-project/build
RUN cmake -G Ninja \
	-DCMAKE_BUILD_TYPE=Release \
	-DTENSORFLOW_AOT_PATH=$(python3 -c "import tensorflow; import os; print(os.path.dirname(tensorflow.__file__))") \
	-DLLVM_ENABLE_PROJECTS="clang;lld" \
	-DLLVM_ENABLE_RUNTIMES="compiler-rt" \
	-C /tflite/tflite.cmake \
	../llvm
RUN cmake --build .
RUN mkdir /llvm-corpus
WORKDIR /llvm-corpus
RUN cmake -G Ninja \
	-DCMAKE_BUILD_TYPE=Release \
	-DLLVM_ENABLE_PROJECTS="clang" \
	-DCMAKE_C_FLAGS="-Xclang -fembed-bitcode=all" \
	-DCMAKE_CXX_FLAGS="-Xclang -fembed-bitcode=all" \
	-DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
	-DCMAKE_C_COMPILER=/llvm-project/build/bin/clang \
	-DCMAKE_CXX_COMPILER=/llvm-project/build/bin/clang++ \
	/llvm-project/llvm
RUN cmake --build .
WORKDIR /ml-compiler-opt
ENV PYTHONPATH /ml-compiler-opt
RUN python3 compiler_opt/tools/extract_ir.py \
	--cmd_filter="^-O2|-O3$" \
	--input=/llvm-corpus/compile_commands.json \
	--input_type=json \
	--llvm_objcopy_path=/llvm-project/build/bin/llvm-objcopy \
	--output_dir=/corpus