Define a problem config (#158)

Author: Flpha0830

compiler_opt/rl/regalloc_priority/__init__.py

@@ -0,0 +1,18 @@
+
+import gin
+
+from compiler_opt.rl import problem_configuration
+from compiler_opt.rl.regalloc import config
+from compiler_opt.rl.regalloc_priority import regalloc_priority_runner
+
+
+@gin.register(module='configs')
+class RegallocPriorityConfig(problem_configuration.ProblemConfiguration):
+    def get_runner(self, *args, **kwargs):
+        return regalloc_priority_runner.RegAllocPriorityRunner(*args, **kwargs)
+
+    def get_signature_spec(self):
+        return config.get_regalloc_signature_spec()
+
+    def get_preprocessing_layer_creator(self):
+        return config.get_observation_processing_layer_creator()

compiler_opt/rl/regalloc_priority/config.py

@@ -0,0 +1,38 @@
+import gin
+import tensorflow as tf
+from tf_agents.specs import tensor_spec
+from tf_agents.trajectories import time_step
+from compiler_opt.rl import feature_ops
+
+
+@gin.configurable()
+def get_regalloc_signature_spec():
+    observation_spec = dict(
+        (key, tf.TensorSpec(dtype=tf.int64, shape=(), name=key))
+        for key in ('li_size', 'stage'))
+    observation_spec['weight'] = tf.TensorSpec(dtype=tf.float32, shape=(), name='weight')
+
+    reward_spec = tf.TensorSpec(dtype=tf.float32, shape=(), name='reward')
+    time_step_spec = time_step.time_step_spec(observation_spec, reward_spec)
+
+    action_spec = tensor_spec.TensorSpec(
+        dtype=tf.float32,
+        shape=(),
+        name='priority'
+    )
+
+    return time_step_spec, action_spec
+
+
+@gin.configurable
+def get_observation_processing_layer_creator():
+    def observation_processing_layer(obs_spec):
+        """Creates the layer to process observation given obs_spec."""
+
+        if obs_spec.name in ('li_size', 'stage', 'weight'):
+            return tf.keras.layers.Lambda(feature_ops.identity_fn)
+
+        # Make sure all features have a preprocessing function.
+        raise KeyError('Missing preprocessing function for some feature.')
+
+    return observation_processing_layer

compiler_opt/rl/regalloc_priority/gin_configs/common.gin

@@ -0,0 +1,7 @@
+config_registry.get_configuration.implementation=@configs.RegallocPriorityConfig
+
+launcher_path=None
+clang_path=None
+
+runners.RegAllocPriorityRunner.clang_path=%clang_path
+runners.RegAllocPriorityRunner.launcher_path=%launcher_path

compiler_opt/rl/regalloc_priority/gin_configs/ppo_nn_agent.gin

@@ -0,0 +1,76 @@
+import gin.tf.external_configurables
+import compiler_opt.rl.constant
+import compiler_opt.rl.constant_value_network
+import compiler_opt.rl.gin_external_configurables
+import compiler_opt.rl.regalloc_priority.config
+import compiler_opt.rl.regalloc_network
+import tf_agents.agents.ppo.ppo_agent
+import tf_agents.networks.actor_distribution_network
+
+include 'compiler_opt/rl/regalloc_priority/gin_configs/common.gin'
+
+train_eval.agent_name=%constant.AgentName.PPO
+train_eval.warmstart_policy_dir=''
+train_eval.num_policy_iterations=3000
+train_eval.num_iterations=200
+train_eval.batch_size=256
+train_eval.train_sequence_length=16
+train_eval.deploy_policy_name='saved_collect_policy'
+train_eval.moving_average_decay_rate=0.8
+train_eval.use_random_network_distillation=False
+
+#######################################
+# Turn on if using train_with_rpc.py
+# train_eval.additional_compilation_flags=()
+#######################################
+
+# RandomNetworkDistillation configs, off if train_eval.use_random_network_distillation=False.
+RandomNetworkDistillation.encoding_network = @regalloc_network.RegAllocRNDEncodingNetwork
+RandomNetworkDistillation.learning_rate = 1e-4
+RandomNetworkDistillation.update_frequency = 2
+RandomNetworkDistillation.fc_layer_params = [32, 128]
+RandomNetworkDistillation.initial_intrinsic_reward_scale = 1.0
+RandomNetworkDistillation.half_decay_steps = 10000
+
+create_agent.policy_network =  @actor_distribution_network.ActorDistributionNetwork
+
+ActorDistributionNetwork.preprocessing_combiner=@tf.keras.layers.Concatenate()
+ActorDistributionNetwork.fc_layer_params=(40, 40, 20)
+ActorDistributionNetwork.dropout_layer_params=None
+ActorDistributionNetwork.activation_fn=@tf.keras.activations.relu
+
+NormalProjectionNetwork.mean_transform=None
+
+ConstantValueNetwork.constant_output_val=0
+
+tf.train.AdamOptimizer.learning_rate = 0.0003
+tf.train.AdamOptimizer.epsilon = 0.0003125
+
+PPOAgent.optimizer = @tf.train.AdamOptimizer()
+PPOAgent.importance_ratio_clipping = 0.2
+PPOAgent.lambda_value = 0.0
+PPOAgent.discount_factor = 0.0
+PPOAgent.entropy_regularization = 0.005
+PPOAgent.policy_l2_reg = 0.00001
+PPOAgent.value_function_l2_reg = 0.0
+PPOAgent.shared_vars_l2_reg = 0.0
+PPOAgent.value_pred_loss_coef = 0.0
+PPOAgent.num_epochs = 1
+PPOAgent.use_gae = False
+PPOAgent.use_td_lambda_return = False
+PPOAgent.normalize_rewards = False
+PPOAgent.reward_norm_clipping = 10.0
+PPOAgent.normalize_observations = False
+PPOAgent.log_prob_clipping = 0.0
+PPOAgent.kl_cutoff_factor = 2.0
+PPOAgent.kl_cutoff_coef = 1000.0
+PPOAgent.initial_adaptive_kl_beta = 1.0
+PPOAgent.adaptive_kl_target = 0.01
+PPOAgent.adaptive_kl_tolerance = 0.3
+PPOAgent.gradient_clipping = None
+PPOAgent.value_clipping = None
+PPOAgent.check_numerics = False
+PPOAgent.compute_value_and_advantage_in_train = True
+PPOAgent.update_normalizers_in_train=True
+PPOAgent.debug_summaries = True
+PPOAgent.summarize_grads_and_vars = True

compiler_opt/rl/regalloc_priority/regalloc_priority_runner.py

@@ -0,0 +1,70 @@
+import gin
+import tensorflow as tf
+
+import base64
+import io
+import os
+import tempfile
+from typing import Dict, Optional, Tuple
+from absl import logging
+
+from google.protobuf import struct_pb2
+from compiler_opt.rl import compilation_runner
+
+
+@gin.configurable(module='runners')
+class RegAllocPriorityRunner(compilation_runner.CompilationRunner):
+    def _compile_fn(
+            self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
+            cancellation_manager: Optional[
+                compilation_runner.WorkerCancellationManager]
+    ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
+
+        file_paths = file_paths[0].replace('.bc', '')
+        working_dir = tempfile.mkdtemp()
+
+        log_path = os.path.join(working_dir, 'log')
+        output_native_path = os.path.join(working_dir, 'native')
+
+        result = {}
+        try:
+            command_line = []
+            if self._launcher_path:
+                command_line.append(self._launcher_path)
+            command_line.extend([self._clang_path] + [
+                                    '-c', file_paths, '-O3',
+                                    '-mllvm', '-regalloc-priority-training-log=' + log_path,
+                                    '-mllvm', '-regalloc-enable-priority-advisor=development',
+                                    '-o', output_native_path
+                                ])
+
+            if tf_policy_path:
+                command_line.extend(['-mllvm', '-regalloc-priority-model=' + tf_policy_path])
+            compilation_runner.start_cancellable_process(command_line,
+                                                         self._compilation_timeout,
+                                                         cancellation_manager)
+
+            sequence_example = struct_pb2.Struct()
+
+            with io.open(log_path, 'rb') as f:
+                sequence_example.ParseFromString(f.read())
+
+            for key, value in sequence_example.fields.items():
+                e = tf.train.SequenceExample()
+                e.ParseFromString(base64.b64decode(value.string_value))
+                print(e)
+                if not e.HasField('feature_lists'):
+                    continue
+                r = (
+                    e.feature_lists.feature_list['reward'].feature[-1].float_list
+                        .value[0])
+                if reward_only:
+                    result[key] = (None, r)
+                else:
+                    del e.feature_lists.feature_list['reward']
+                    result[key] = (e, r)
+
+        finally:
+            tf.io.gfile.rmtree(working_dir)
+
+        return result