diff --git a/compiler_opt/rl/imitation_learning/generate_bc_trajectories.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories.py index 76b069ed..c333642e 100644 --- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories.py +++ b/compiler_opt/rl/imitation_learning/generate_bc_trajectories.py @@ -21,19 +21,19 @@ import gin from compiler_opt.rl.imitation_learning import generate_bc_trajectories_lib -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 +_GIN_FILES = flags.DEFINE_multi_string( + 'gin_files', [], 'List of paths to gin configuration files.') +_GIN_BINDINGS = flags.DEFINE_multi_string( + 'gin_bindings', [], + 'Gin bindings to override the values set in the config files.') def main(_): gin.parse_config_files_and_bindings( - FLAGS.gin_files, bindings=FLAGS.gin_bindings, skip_unknown=True) + _GIN_FILES.value, bindings=_GIN_BINDINGS.value, skip_unknown=True) logging.info(gin.config_str()) generate_bc_trajectories_lib.gen_trajectories() 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..fbd3cdd8 100644 --- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py +++ b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py @@ -20,7 +20,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Generator, Union import json -from absl import flags +# from absl import flags from absl import logging import bisect import dataclasses @@ -46,13 +46,6 @@ from compiler_opt.distributed import buffered_scheduler from compiler_opt.distributed.local import local_worker_manager -from compiler_opt.tools import generate_test_model # pylint:disable=unused-import - -flags.FLAGS['gin_files'].allow_override = True -flags.FLAGS['gin_bindings'].allow_override = True - -FLAGS = flags.FLAGS - ProfilingDictValueType = Dict[str, Union[str, float, int]] @@ -348,6 +341,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 +695,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 +772,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 +851,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 +876,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 +898,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 +939,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/imitation_learning/generate_bc_trajectories_test.py b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py index c3037833..f87551e1 100644 --- a/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py +++ b/compiler_opt/rl/imitation_learning/generate_bc_trajectories_test.py @@ -16,7 +16,6 @@ import functools from absl import app -from absl import flags import gin import json from typing import List @@ -36,8 +35,8 @@ from compiler_opt.rl import env from compiler_opt.rl import env_test -flags.FLAGS['gin_files'].allow_override = True -flags.FLAGS['gin_bindings'].allow_override = True +# flags.FLAGS['gin_files'].allow_override = True +# flags.FLAGS['gin_bindings'].allow_override = True _eps = 1e-5 @@ -648,17 +647,17 @@ def select_best_exploration(self, mock_popen, loaded_module_spec): class GenTrajectoriesTest(tf.test.TestCase): def setUp(self): - with gin.unlock_config(): - gin.parse_config_files_and_bindings( - config_files=['compiler_opt/rl/inlining/gin_configs/common.gin'], - bindings=[ - ('generate_bc_trajectories_test.' - 'MockModuleWorker.clang_path="/test/clang/path"'), - ('generate_bc_trajectories_test.' - 'MockModuleWorker.exploration_frac=1.0'), - ('generate_bc_trajectories_test.' - 'MockModuleWorker.reward_key="default"'), - ]) + with gin.config_scope('gen_trajectories_test'): + with gin.unlock_config(): + gin.bind_parameter( + 'generate_bc_trajectories_test.MockModuleWorker.clang_path', + '/test/clang/path') + gin.bind_parameter( + 'generate_bc_trajectories_test.MockModuleWorker.exploration_frac', + 1.0) + gin.bind_parameter( + 'generate_bc_trajectories_test.MockModuleWorker.reward_key', + 'default') return super().setUp() def test_gen_trajectories(self): 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..9b69ffdd --- /dev/null +++ b/compiler_opt/rl/inlining/gin_configs/imitation_learning.gin @@ -0,0 +1,30 @@ +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='' + +generate_bc_trajectories_lib.ModuleWorker.clang_path='' +generate_bc_trajectories_lib.ModuleWorker.mlgo_task_type=@env.InliningForSizeTask +generate_bc_trajectories_lib.ModuleWorker.policy_paths=[''] +generate_bc_trajectories_lib.ModuleWorker.exploration_policy_paths=[] +generate_bc_trajectories_lib.ModuleWorker.explore_on_features=None +generate_bc_trajectories_lib.ModuleWorker.base_path='' +generate_bc_trajectories_lib.ModuleWorker.partitions=[ + 285.0, 376.0, 452.0, 512.0, 571.0, 627.5, 720.0, 809.5, 1304.0, 1832.0, + 2467.0, 3344.0, 4545.0, 6459.0, 9845.0, 17953.0, 29430.5, 85533.5, + 124361.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='' +generate_bc_trajectories_lib.gen_trajectories.delete_flags=('-split-dwarf-file', '-split-dwarf-output') +generate_bc_trajectories_lib.gen_trajectories.output_file_name='' +generate_bc_trajectories_lib.gen_trajectories.output_path='' +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='' +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..9d5bbfb1 --- /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 tf_agents.system import system_multiprocessing as multiprocessing + +_GIN_FILES = flags.DEFINE_multi_string( + 'gin_files', [], 'List of paths to gin configuration files.') +_GIN_BINDINGS = flags.DEFINE_multi_string( + 'gin_bindings', [], + 'Gin bindings to override the values set in the config files.') + + +def main(_): + gin.parse_config_files_and_bindings( + _GIN_FILES.value, bindings=_GIN_BINDINGS.value, 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))