diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py index 91f693c0..8f881669 100644 --- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py +++ b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py @@ -348,6 +348,7 @@ def __init__( task_type=mlgo_task_type, obs_spec=obs_spec, action_spec=action_spec, + interactive_only=True, ) if self._env.action_spec: if self._env.action_spec.dtype != tf.int64: @@ -701,7 +702,7 @@ def _save_binary(self, base_path: str, save_path: str, binary_path: str): if not os.path.exists(save_dir): os.makedirs(save_dir, exist_ok=True) shutil.copy( - os.path.join(binary_path, 'comp_binary'), + os.path.join(binary_path, 'compiled_module'), os.path.join(save_dir, path_tail)) @@ -778,6 +779,7 @@ def __init__( self._obs_action_specs: Optional[Tuple[ time_step.TimeStep, tensor_spec.BoundedTensorSpec]] = obs_action_specs self._mw_utility = ModuleWorkerResultProcessor(base_path) + self._base_path = base_path self._partitions = partitions self._envargs = envargs @@ -856,7 +858,14 @@ def select_best_exploration( for temp_dirs in working_dir_list: for temp_dir in temp_dirs: temp_dir_head = os.path.split(temp_dir)[0] - shutil.rmtree(temp_dir_head) + try: + shutil.rmtree(temp_dir_head) + except FileNotFoundError as e: + if not self._base_path: + continue + else: + raise FileNotFoundError( + f'Compilation directory {temp_dir_head} does not exist.') from e return ( num_calls, @@ -874,6 +883,8 @@ def gen_trajectories( output_path: str = gin.REQUIRED, mlgo_task_type: Type[env.MLGOTask] = gin.REQUIRED, callable_policies: List[Optional[Callable[[Any], np.ndarray]]] = [], + explore_on_features: Optional[Dict[str, Callable[[tf.Tensor], + bool]]] = None, obs_action_spec: Optional[Tuple[time_step.TimeStep, tensor_spec.BoundedTensorSpec]] = None, num_workers: Optional[int] = None, @@ -894,6 +905,8 @@ def gen_trajectories( callable_policies: list of policies in the form of callable functions, this supplements the loaded policies from policy_paths given in ModuleWorker + explore_on_features: dict of feature names and functions which specify + when to explore on the respective feature obs_action_spec: optional observation and action spec annotating the state (TimeStep) for training a policy num_workers: number of distributed workers to process the corpus. @@ -933,6 +946,7 @@ def gen_trajectories( obs_action_specs=obs_action_spec, mlgo_task_type=mlgo_task_type, callable_policies=callable_policies, + explore_on_features=explore_on_features, gin_config_str=gin.config_str(), ) as lwm: diff --git a/compiler_opt/rl/inlining/env.py b/compiler_opt/rl/inlining/env.py index fc33a28a..5f1bd523 100644 --- a/compiler_opt/rl/inlining/env.py +++ b/compiler_opt/rl/inlining/env.py @@ -44,8 +44,8 @@ def get_cmdline(self, clang_path: str, base_args: List[str], '-enable-ml-inliner=release', '-mllvm', f'-inliner-interactive-channel-base={interactive_base_path}', - #'-mllvm', - #'-inliner-interactive-include-default', + '-mllvm', + '-inliner-interactive-include-default', ] else: interactive_args = [] diff --git a/compiler_opt/rl/inlining/gin_configs/imitation_learning.gin b/compiler_opt/rl/inlining/gin_configs/imitation_learning.gin new file mode 100644 index 00000000..c7f3bf32 --- /dev/null +++ b/compiler_opt/rl/inlining/gin_configs/imitation_learning.gin @@ -0,0 +1,27 @@ +import compiler_opt.rl.inlining.env +import compiler_opt.rl.inlining.imitation_learning_config +import compiler_opt.rl.imitation_learning.generate_bc_trajectories_lib + + +env.InliningForSizeTask.llvm_size_path='/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/inlining/runfolder.rundir/llvm-size' + +generate_bc_trajectories_lib.ModuleWorker.clang_path='/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/inlining/chrome_on_android/chromium2/src/third_party/llvm-build/tflite_build_cold/bin/clang' +generate_bc_trajectories_lib.ModuleWorker.mlgo_task_type=@env.InliningForSizeTask +generate_bc_trajectories_lib.ModuleWorker.policy_paths=['/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/inlining/chrome_on_android/models/chromium_model_06-09-23_ave'] +generate_bc_trajectories_lib.ModuleWorker.exploration_policy_paths=[] +generate_bc_trajectories_lib.ModuleWorker.explore_on_features=None +generate_bc_trajectories_lib.ModuleWorker.base_path='/tmp/test_obj' +generate_bc_trajectories_lib.ModuleWorker.partitions=[0.,] +generate_bc_trajectories_lib.ModuleWorker.reward_key='default' +# generate_bc_trajectories_lib.ModuleWorker.gin_config_str=None + +generate_bc_trajectories_lib.gen_trajectories.data_path='/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/blaze_profiles/chromium_corpus64bitppo_tlto30_rsp_24_04_2024' +generate_bc_trajectories_lib.gen_trajectories.delete_flags=('-split-dwarf-file', '-split-dwarf-output') +generate_bc_trajectories_lib.gen_trajectories.output_file_name='ppo_cold_callsites_thinlto30_6_test' +generate_bc_trajectories_lib.gen_trajectories.output_path='/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/inlining/training_data/chrome_on_android' +generate_bc_trajectories_lib.gen_trajectories.mlgo_task_type=@imitation_learning_config.get_task_type() +generate_bc_trajectories_lib.gen_trajectories.obs_action_spec=@imitation_learning_config.get_inlining_signature_spec() +generate_bc_trajectories_lib.gen_trajectories.num_workers=1 +generate_bc_trajectories_lib.gen_trajectories.num_output_files=1 +generate_bc_trajectories_lib.gen_trajectories.profiling_file_path='/usr/local/google/home/tvmarinov/Documents/mlgo_compiler_opt/inlining/profiling_folder/chrome_on_android/ppo_cold_callsites_thinlto30_6_test' +generate_bc_trajectories_lib.gen_trajectories.worker_wait_sec=100 diff --git a/compiler_opt/rl/inlining/imitation_learning_config.py b/compiler_opt/rl/inlining/imitation_learning_config.py new file mode 100644 index 00000000..fe8bcc41 --- /dev/null +++ b/compiler_opt/rl/inlining/imitation_learning_config.py @@ -0,0 +1,60 @@ +# coding=utf-8 +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module for collect data of inlining-for-size.""" + +import gin +from typing import Type + +import numpy as np +import tensorflow as tf +from tf_agents.trajectories import time_step + +from compiler_opt.rl.inlining import config +from compiler_opt.rl.inlining import env + + +@gin.register +def get_inlining_signature_spec(): + """Returns (time_step_spec, action_spec) for collecting IL trajectories.""" + time_step_spec, action_spec = config.get_inlining_signature_spec() + observation_spec = time_step_spec.observation + observation_spec.update( + dict((key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key)) for key in ( + 'is_callee_avail_external', + 'is_caller_avail_external', + # inlining_default is not used as feature in training. + 'inlining_default'))) + + time_step_spec = time_step.time_step_spec(observation_spec, + time_step_spec.reward) + + return time_step_spec, action_spec + + +@gin.register +def get_task_type() -> Type[env.InliningForSizeTask]: + """Returns the task type for the trajectory collection.""" + return env.InliningForSizeTask + + +@gin.register +def greedy_policy(state: time_step.TimeStep): + """Greedy policy playing the inlining_default action.""" + return np.array(state.observation['inlining_default']) + + +@gin.register +def explore_on_avail_external(state_observation: tf.Tensor) -> bool: + return state_observation.numpy()[0] diff --git a/compiler_opt/rl/inlining/imitation_learning_runner.py b/compiler_opt/rl/inlining/imitation_learning_runner.py new file mode 100644 index 00000000..84724331 --- /dev/null +++ b/compiler_opt/rl/inlining/imitation_learning_runner.py @@ -0,0 +1,49 @@ +# coding=utf-8 +# Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module for running compilation and collect data for behavior cloning.""" + +import functools +from absl import app +from absl import flags +from absl import logging +import gin + +from compiler_opt.rl.imitation_learning import generate_bc_trajectories_lib +from compiler_opt.rl.inlining import imitation_learning_config +from compiler_opt.tools import generate_test_model # pylint:disable=unused-import + +from tf_agents.system import system_multiprocessing as multiprocessing + +flags.FLAGS['gin_files'].allow_override = True +flags.FLAGS['gin_bindings'].allow_override = True + +FLAGS = flags.FLAGS + + +def main(_): + gin.parse_config_files_and_bindings( + FLAGS.gin_files, bindings=FLAGS.gin_bindings, skip_unknown=True) + logging.info(gin.config_str()) + + generate_bc_trajectories_lib.gen_trajectories( + callable_policies=[imitation_learning_config.greedy_policy], + explore_on_features={ + 'is_callee_avail_external': + imitation_learning_config.explore_on_avail_external + }) + + +if __name__ == '__main__': + multiprocessing.handle_main(functools.partial(app.run, main))