Adding configs for collecting imitation learning tarjectories for

inlining for size.

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:
 

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

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