[BC] Inlining for size configs. (#410)

* Adding configs for collecting imitation learning tarjectories for
inlining for size.

* Removing values from imitation_learning.gin

* Removed unused flags for gin bindings and gin configs.

Author: tvmarino

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()

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]]
 
 
@@ -350,6 +343,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:
@@ -703,7 +697,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))
 
 
@@ -782,6 +776,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
 
@@ -860,7 +855,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,
@@ -878,6 +880,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,
@@ -898,6 +902,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.
@@ -937,6 +943,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:
 

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):

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 = []

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

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]

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))