Add self-supervised training code. This is only for staging - DO NOT MERGE.
diff --git a/compiler_opt/rl/attention.py b/compiler_opt/rl/attention.py new file mode 100644 index 0000000..30915b6 --- /dev/null +++ b/compiler_opt/rl/attention.py
@@ -0,0 +1,183 @@ +"""attention.py + +Defines various building-blocks for attention in neural networks. Namely, the +transformer encoder, which is a sequence-to-sequence model which uses +self-attention to model relationships within the sequence. +""" + +import tensorflow as tf +import numpy as np + +from typing import Optional, List + + +def positional_encoding(length, depth): + """Build a positional encoding tensor. + + Taken from https://www.tensorflow.org/text/tutorials/transformer. + + Args: + length: the number of sin/cos samples to generate. + depth: the depth of the embedding which the encoding will be summed with. + + Returns: + A tensor of shape (length, depth) representing the positional encoding. + """ + depth = depth / 2 + + positions = np.arange(length)[:, np.newaxis] # (seq, 1) + depths = np.arange(depth)[np.newaxis, :] / depth # (1, depth) + + angle_rates = 1 / (10000**depths) # (1, depth) + angle_rads = positions * angle_rates # (pos, depth) + + pos_encoding = np.concatenate( + [np.sin(angle_rads), np.cos(angle_rads)], axis=-1 + ) + + return tf.cast(pos_encoding, dtype=tf.float32) + + +class PositionalEmbedding(tf.keras.layers.Layer): + """A positional embedding layer. + + A "positional embedding" is a sum of a token embedding with a positional + encoding, which is used as the initial layer of a transformer encoder + network. + + Taken from https://www.tensorflow.org/text/tutorials/transformer. + """ + + def __init__(self, vocab_size, d_model): + """Initialize the positional embedding. + + Args: + vocab_size: the size of the vocab, which should be one more than the + maximum token value which will be seen during training/inference. + d_model: the dimension of the model (size of the embedding vector) + """ + super().__init__() + self.d_model = d_model + self.embedding = tf.keras.layers.Embedding( + vocab_size, d_model, mask_zero=True + ) + self.pos_encoding = tf.constant( + positional_encoding(length=2048, depth=d_model), dtype=tf.float32 + ) + + def compute_mask(self, *args, **kwargs): + """Returns a mask for the given input.""" + return self.embedding.compute_mask(*args, **kwargs) + + def call(self, x): + """Perform the positional embedding.""" + length = tf.shape(x)[-1] + x = self.embedding(x) + # This factor sets the relative scale of the embedding and positonal_encoding. + x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32)) + x = x + self.pos_encoding[tf.newaxis, :length, :] + return x + + +class TransformerEncoderLayer(tf.keras.layers.Layer): + """Transformer Encoder. + + See https://arxiv.org/abs/1706.03762 for more details. + """ + + def __init__( + self, + num_heads: int, + model_dim: int, + fcn_dim: int, + attention_axes: Optional[List] = None, + ): + """Initialize the transformer encoder. + + Args: + num_heads: number of distinct attention heads within the layer. + model_dim: dimension of the model, also the dimension of the embedding. + fcn_dim: dimension of the fully-connected layers between attention layers. + attention_axes: which axes in the input tensors perform attention across. + """ + super().__init__() + self._mha = tf.keras.layers.MultiHeadAttention( + num_heads=num_heads, key_dim=model_dim, attention_axes=attention_axes + ) + self._mha_norm = tf.keras.layers.LayerNormalization() + self._mha_add = tf.keras.layers.Add() + + self._fcn = tf.keras.Sequential([ + tf.keras.layers.Dense(fcn_dim), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(model_dim), + tf.keras.layers.Dropout(0.1), + ]) + self._fcn_norm = tf.keras.layers.LayerNormalization() + self._fcn_add = tf.keras.layers.Add() + + def call(self, x, attention_mask=None): + """Call the transformer encoder.""" + x_attended = self._mha( + query=x, value=x, key=x, attention_mask=attention_mask + ) + x = self._mha_add([x_attended, x]) + x = self._mha_norm(x) + + x_fcn = self._fcn(x) + x = self._fcn_add([x_fcn, x]) + x = self._fcn_norm(x) + + return x + + +class TransformerClassifier(tf.keras.layers.Layer): + + def __init__( + self, + *, + num_tokens: int, + num_layers: int, + num_heads: int, + model_dim: int, + fcn_dim: int, + num_extra_features: int = 0, + ): + super().__init__() + + self._model_dim = model_dim + self._ctx_token = tf.constant(num_tokens, dtype=tf.int64) + self._embed_layer = PositionalEmbedding( + num_tokens + 1, model_dim - num_extra_features + ) + self._transformer_layers = [ + TransformerEncoderLayer( + num_heads=num_heads, + model_dim=model_dim, + fcn_dim=fcn_dim, + attention_axes=(2,), + ) + for _ in range(num_layers) + ] + + def __call__(self, x, extra_hidden_state): + # [B, 33, 1 + I] --> [B, 33, 1 + I, E] + mask = self._embed_layer.compute_mask(x) + x = self._embed_layer(x) + + # Append the extra hidden state + x = tf.concat([x, extra_hidden_state], axis=-1) + + mask1 = mask[:, :, :, tf.newaxis] + mask2 = mask[:, :, tf.newaxis, :] + attn_mask = tf.cast(mask1, dtype=tf.int64) + tf.cast(mask2, dtype=tf.int64) + attn_mask = attn_mask > 0 + + for transformer_layer in self._transformer_layers: + x = transformer_layer(x, attention_mask=attn_mask) + + mask_reduce = tf.cast(mask[:, :, :, tf.newaxis], dtype=tf.float32) + x_reduced = (tf.reduce_sum(mask_reduce * x, axis=-2)) / ( + tf.reduce_sum(mask_reduce, axis=-2) + 1e-3 + ) + return x, x_reduced
diff --git a/compiler_opt/rl/regalloc/lr_encoder/BUILD b/compiler_opt/rl/regalloc/lr_encoder/BUILD new file mode 100644 index 0000000..7fde6aa --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/BUILD
@@ -0,0 +1,96 @@ +# Inlining configurations. +load("//devtools/python/blaze:strict.bzl", "py_strict_library") +load("//devtools/python/blaze:pytype.bzl", "pytype_strict_binary", "pytype_strict_library") + +licenses(["notice"]) + +package( + default_applicable_licenses = ["//third_party/ml_compiler_opt:license"], + default_visibility = [ + "//third_party/ml_compiler_opt:default_visibility", + ], +) + +filegroup( + name = "gin_files", + srcs = glob(["gin_configs/**"]), +) + +py_strict_library( + name = "config", + srcs = ["config.py"], + deps = [ + "//third_party/ml_compiler_opt/compiler_opt/rl:feature_ops", + "//third_party/ml_compiler_opt/compiler_opt/rl/regalloc:config", + "//third_party/py/gin", + "//third_party/py/tensorflow:tensorflow_no_contrib", + "//third_party/py/tf_agents", + "//third_party/py/tf_agents/specs:tensor_spec", + "//third_party/py/tf_agents/trajectories:time_step", + ], +) + +pytype_strict_library( + name = "model", + srcs = ["model.py"], + deps = [ + ":config", + "//third_party/ml_compiler_opt/compiler_opt/rl:attention", + "//third_party/py/gin", + "//third_party/py/tensorflow:tensorflow_no_contrib", + "//third_party/py/tf_agents/utils:nest_utils", + ], +) + +pytype_strict_library( + name = "dataset_ops", + srcs = ["dataset_ops.py"], + deps = [ + ":config", + "//third_party/py/tensorflow:tensorflow_no_contrib", + "//third_party/py/tensorflow_text", + ], +) + +pytype_strict_library( + name = "lr_encoder_runner", + srcs = ["lr_encoder_runner.py"], + deps = [ + "//third_party/ml_compiler_opt/compiler_opt/rl:compilation_runner", + "//third_party/ml_compiler_opt/compiler_opt/rl:corpus", + "//third_party/ml_compiler_opt/compiler_opt/rl:log_reader", + "//third_party/py/gin", + "//third_party/py/tensorflow:tensorflow_no_contrib", + ], +) + +pytype_strict_library( + name = "lr_encoder_problem_config", + srcs = ["__init__.py"], + deps = [ + ":config", + ":lr_encoder_runner", + "//third_party/ml_compiler_opt/compiler_opt/rl:problem_configuration", + "//third_party/py/gin", + ], +) + +pytype_strict_binary( + name = "train", + srcs = ["train.py"], + data = [":gin_files"], + deps = [ + ":config", + ":dataset_ops", + ":model", + "//third_party/ml_compiler_opt/compiler_opt/rl:policy_saver", + "//third_party/ml_compiler_opt/compiler_opt/rl:registry", + "//third_party/ml_compiler_opt/compiler_opt/rl/regalloc:config", + "//third_party/py/absl:app", + "//third_party/py/absl/flags", + "//third_party/py/absl/logging", + "//third_party/py/gin", + "//third_party/py/tensorflow:tensorflow_google", # build_cleaner: keep + "//third_party/py/tensorflow:tensorflow_no_contrib", + ], +)
diff --git a/compiler_opt/rl/regalloc/lr_encoder/__init__.py b/compiler_opt/rl/regalloc/lr_encoder/__init__.py new file mode 100644 index 0000000..8100aa8 --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/__init__.py
@@ -0,0 +1,24 @@ +"""Implementation of the 'lr_encoder' problem.""" + +import gin + +from google3.third_party.ml_compiler_opt.compiler_opt.rl import problem_configuration +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import lr_encoder_runner + + +@gin.register(module='configs') +class LREncoderConfig(problem_configuration.ProblemConfiguration): + """Expose the LR encoder configuration.""" + + def get_runner_type(self): + return lr_encoder_runner.LREncoderRunner + + def get_signature_spec(self): + return config.get_lr_encoder_signature_spec() + + def get_preprocessing_layer_creator(self): + return config.get_preprocessing_layer_creator() + + def get_nonnormalized_features(self): + return config.get_nonnormalized_features()
diff --git a/compiler_opt/rl/regalloc/lr_encoder/config.py b/compiler_opt/rl/regalloc/lr_encoder/config.py new file mode 100644 index 0000000..11a458e --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/config.py
@@ -0,0 +1,219 @@ +"""LR encoder training config.""" + +import gin +import tensorflow as tf +import tf_agents +from tf_agents.specs import tensor_spec +from tf_agents.trajectories import time_step +from google3.third_party.ml_compiler_opt.compiler_opt.rl import feature_ops +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import config + +_NUM_REGISTERS = 33 +_NUM_INSTRUCTIONS = 64 +_OPCODE_VOCAB_SIZE = 20000 +_ENCODING_SIZE = 16 +_OPCODE_KEY = 'lr_use_def_opcode' + +_ENCODER_FEATURE_PREFIX = 'lr_use_def' + +_ACTION_KEY = 'action' +_STATE_KEY = 'state' +_NEXT_STATE_KEY = 'next_state' +_NEXT_ACTION_KEY = 'next_action' +_MLM_KEY = 'mlm' + + +@gin.configurable +def get_lr_encoder_signature_spec(): + observation_spec = {} + observation_spec[_OPCODE_KEY] = tensor_spec.BoundedTensorSpec( + dtype=tf.int64, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS), + name=_OPCODE_KEY, + minimum=0, + maximum=_OPCODE_VOCAB_SIZE, + ) + observation_spec['lr_use_def_read'] = tensor_spec.BoundedTensorSpec( + dtype=tf.int64, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS), + name='lr_use_def_read', + minimum=0, + maximum=1, + ) + observation_spec['lr_use_def_write'] = tensor_spec.BoundedTensorSpec( + dtype=tf.int64, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS), + name='lr_use_def_write', + minimum=0, + maximum=1, + ) + observation_spec.update( + { + name: tensor_spec.BoundedTensorSpec( + dtype=tf.int64, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS), + name=name, + minimum=0, + maximum=1, + ) + for name in [ + 'lr_use_def_is_use', + 'lr_use_def_is_def', + 'lr_use_def_is_implicit', + 'lr_use_def_is_renamable', + 'lr_use_def_is_ind_var_update', + 'lr_use_def_is_hint', + ] + } + ) + observation_spec['lr_use_def_freq'] = tf.TensorSpec( + dtype=tf.float32, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS), + name='lr_use_def_freq', + ) + + encoding_spec = tf.TensorSpec( + dtype=tf.float32, shape=(33, _ENCODING_SIZE), name='lr_encoding' + ) + + return observation_spec, encoding_spec + + +def get_input_specs(): + encoder_observation_spec, encoding_spec = get_lr_encoder_signature_spec() + del encoding_spec + + regalloc_time_step_spec, regalloc_action_spec = ( + config.get_regalloc_signature_spec() + ) + + # Ensure that there are no overlaps in feature names between the encoder and the RL problem + common_keys = ( + encoder_observation_spec.keys() + & regalloc_time_step_spec.observation.keys() + ) + assert len(common_keys) == 0 + + supervised_input_spec = { + **encoder_observation_spec, + **regalloc_time_step_spec.observation, + } + + return supervised_input_spec + + +def get_output_specs(regalloc_preprocessing_layer_creator): + regalloc_time_step_spec, regalloc_action_spec = ( + config.get_regalloc_signature_spec() + ) + del regalloc_action_spec + + random_regalloc_state = tf_agents.specs.sample_spec_nest( + regalloc_time_step_spec.observation + ) + for key in random_regalloc_state: + preprocessing_layer = regalloc_preprocessing_layer_creator( + regalloc_time_step_spec.observation[key] + ) + random_regalloc_state[key] = preprocessing_layer( + tf.expand_dims(random_regalloc_state[key], axis=0) + ) + random_regalloc_state = tf.concat( + list(random_regalloc_state.values()), axis=-1 + ) + random_regalloc_state = tf.squeeze(random_regalloc_state, axis=0) + regalloc_state_shape = random_regalloc_state.shape + + action_spec = tf.TensorSpec( + dtype=tf.float32, shape=(_NUM_REGISTERS, 1), name=_ACTION_KEY + ) + state_spec = tf.TensorSpec( + dtype=tf.float32, shape=regalloc_state_shape, name=_STATE_KEY + ) + mlm_spec = tf.TensorSpec( + dtype=tf.float32, + shape=(_NUM_REGISTERS, _NUM_INSTRUCTIONS, _OPCODE_VOCAB_SIZE), + name=_MLM_KEY, + ) + return { + _ACTION_KEY: action_spec, + _STATE_KEY: state_spec, + _NEXT_ACTION_KEY: action_spec, + _NEXT_STATE_KEY: state_spec, + _MLM_KEY: mlm_spec, + } + + +def get_loss(): + loss_fns = { + _STATE_KEY: 'mse', + _ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True + ), + _NEXT_STATE_KEY: 'mse', + _NEXT_ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True + ), + _MLM_KEY: tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), + } + loss_weights = { + _ACTION_KEY: 1.0, + _STATE_KEY: 1.0, + _NEXT_ACTION_KEY: 1.0, + _NEXT_STATE_KEY: 1.0, + _MLM_KEY: 10.0, + } + return loss_fns, loss_weights + + +def get_metrics(): + return { + _ACTION_KEY: ['accuracy'], + _STATE_KEY: [], + _NEXT_ACTION_KEY: ['accuracy'], + _NEXT_STATE_KEY: [], + _MLM_KEY: ['accuracy'], + } + + +@gin.configurable +def get_preprocessing_layer_creator( + quantile_file_dir='/cns/oz-d/home/mlcompileropt-dev/regalloc-transformer/vocab', + with_sqrt=True, + with_z_score_normalization=True, + eps=1e-8, +): + """Wrapper for observation_processing_layer.""" + quantile_map = feature_ops.build_quantile_map(quantile_file_dir) + + def preprocessing_layer_creator(obs_spec): + """Creates the layer to process observation given obs_spec.""" + if obs_spec.name == 'lr_use_def_freq': + quantile = quantile_map[obs_spec.name] + first_non_zero = 0 + for x in quantile: + if x > 0: + first_non_zero = x + break + + normalize_fn = feature_ops.get_normalize_fn( + quantile, with_sqrt, with_z_score_normalization, eps + ) + return tf.keras.layers.Lambda(normalize_fn) + return tf.keras.layers.Lambda(feature_ops.identity_fn) + + return preprocessing_layer_creator + + +def get_nonnormalized_features(): + return [ + 'lr_use_def_opcode', + 'lr_use_def_read', + 'lr_use_def_write', + 'lr_use_def_is_use', + 'lr_use_def_is_def', + 'lr_use_def_is_implicit', + 'lr_use_def_is_renamable', + 'lr_use_def_is_ind_var_update', + 'lr_use_def_is_hint', + ]
diff --git a/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py b/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py new file mode 100644 index 0000000..9967f16 --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/dataset_ops.py
@@ -0,0 +1,158 @@ +import tensorflow as tf +import tensorflow_text as text + +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config + +_MAX_PREDICTIONS_PER_BATCH = 128 + +_PAD_TOKEN = 0 +_MLM_IGNORE_TOKEN = -1 +_MLM_MASK_TOKEN = 18000 - 1 + + +def _get_masked_language_fn(*, selection_rate=0.33): + random_selector = text.RandomItemSelector( + max_selections_per_batch=_MAX_PREDICTIONS_PER_BATCH, + selection_rate=0.2, + unselectable_ids=[_PAD_TOKEN], + ) + mask_values_chooser = text.MaskValuesChooser( + config._OPCODE_VOCAB_SIZE, _MLM_MASK_TOKEN, 0.8 + ) + + def fn(opcodes): + masked_token_ids, masked_pos, masked_lm_ids = text.mask_language_model( + tf.RaggedTensor.from_tensor(opcodes, padding=_PAD_TOKEN), + item_selector=random_selector, + mask_values_chooser=mask_values_chooser, + ) + + masked_pos = masked_pos.to_tensor( + default_value=-1, + shape=(config._NUM_REGISTERS, _MAX_PREDICTIONS_PER_BATCH), + ) + + ii = tf.tile( + tf.range(config._NUM_REGISTERS, dtype=tf.int64)[:, tf.newaxis], + [1, _MAX_PREDICTIONS_PER_BATCH], + ) + masked_pos = tf.stack([ii, masked_pos], axis=-1) + scatter_values = tf.where(masked_pos[:, :, 1] < 0, 0.0, 1.0) + masked_pos = tf.where( + masked_pos < 0, tf.constant(_PAD_TOKEN, dtype=tf.int64), masked_pos + ) + weights = tf.scatter_nd( + masked_pos, + scatter_values[:, :, tf.newaxis], + (config._NUM_REGISTERS, config._NUM_INSTRUCTIONS, 1), + ) + return ( + masked_token_ids.to_tensor( + default_value=0, + shape=(config._NUM_REGISTERS, config._NUM_INSTRUCTIONS), + ), + opcodes, + weights, + ) + + return fn + + +def _roll_experience(seq_ex, *, shift=-1): + def _roll(atom): + return tf.roll(atom, shift=shift, axis=0) + + def _cutoff(atom): + return atom[:shift] + + seq_ex_roll = tf.nest.map_structure(_roll, seq_ex) + seq_ex_roll = tf.nest.map_structure(_cutoff, seq_ex_roll) + seq_ex = tf.nest.map_structure(_cutoff, seq_ex) + return seq_ex, seq_ex_roll + + +def _split_sequence_example(seq_ex, seq_ex_roll): + action_name = 'index_to_evict' + obs = {k: seq_ex[k] for k in seq_ex if k != action_name} + action = seq_ex[action_name] + obs_roll = {k: seq_ex_roll[k] for k in seq_ex_roll if k != action_name} + action_roll = seq_ex_roll[action_name] + return { + 'obs': obs, + 'action': action, + 'obs_roll': obs_roll, + 'action_roll': action_roll, + } + + +def _get_state_preprocessing_layer( + regalloc_input_spec, preprocessing_layer_creator +): + preprocessing_layers = {} + for name, spec in regalloc_input_spec.items(): + preprocessing_layers[name] = preprocessing_layer_creator(spec) + + def _preprocessing_layer(seq_ex): + pp = [] + for layer_name, layer in preprocessing_layers.items(): + pp.append(layer(seq_ex[layer_name])) + return tf.concat(pp, axis=-1) + + def _layer(obs_dict): + obs_dict['obs'] = obs_dict['obs'] + obs_dict['obs_cur'] = _preprocessing_layer(obs_dict['obs']) + obs_dict['obs_roll'] = _preprocessing_layer(obs_dict['obs_roll']) + return obs_dict + + return _layer + + +def _get_to_inputs_and_labels_fn(): + masked_language_fn = _get_masked_language_fn() + + def fn(obs_dict): + inputs = {'obs': obs_dict['obs'], 'action': obs_dict['action']} + mlm_input, mlm_label, mlm_weight = masked_language_fn( + obs_dict['obs']['lr_use_def_opcode'][:, : config._NUM_INSTRUCTIONS] + ) + inputs['lr_use_def_opcode'] = mlm_input + + labels = { + config._STATE_KEY: obs_dict['obs_cur'], + config._ACTION_KEY: tf.expand_dims(obs_dict['action'], axis=-1), + config._NEXT_STATE_KEY: obs_dict['obs_roll'], + config._NEXT_ACTION_KEY: tf.expand_dims( + obs_dict['action_roll'], axis=-1 + ), + config._MLM_KEY: mlm_label, + } + mask = obs_dict['obs']['mask'] + sample_weights = { + config._STATE_KEY: mask, + config._ACTION_KEY: mask, + config._NEXT_STATE_KEY: None, + config._NEXT_ACTION_KEY: None, + config._MLM_KEY: mlm_weight, + } + return (inputs, labels, sample_weights) + + return fn + + +def process_dataset( + dataset, regalloc_input_spec, regalloc_preprocessing_layer_creator +): + shuffle_buffer_size = 128 + num_map_threads = 128 + state_preprocessing_layer = _get_state_preprocessing_layer( + regalloc_input_spec, regalloc_preprocessing_layer_creator + ) + to_inputs_and_labels_fn = _get_to_inputs_and_labels_fn() + return ( + dataset.map(_roll_experience, num_parallel_calls=num_map_threads) + .map(_split_sequence_example, num_parallel_calls=num_map_threads) + .map(state_preprocessing_layer, num_parallel_calls=num_map_threads) + .unbatch() + .shuffle(shuffle_buffer_size) + .map(to_inputs_and_labels_fn, num_parallel_calls=num_map_threads) + )
diff --git a/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin b/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin new file mode 100644 index 0000000..c0fcb35 --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/gin_configs/common.gin
@@ -0,0 +1,14 @@ +config_registry.get_configuration.implementation=@configs.LREncoderConfig + +launcher_path=None +clang_path=None + +runners.LREncoderRunner.clang_path=%clang_path +runners.LREncoderRunner.launcher_path=%launcher_path + +problem_config.flags_to_add.add_flags=() +problem_config.flags_to_delete.delete_flags=('-split-dwarf-file','-split-dwarf-output',) +# For AFDO profile reinjection set: +# problem_config.flags_to_replace.replace_flags={'-fprofile-sample-use':'/path/to/gwp.afdo','-fprofile-remapping-file':'/path/to/prof_remap.txt'} +problem_config.flags_to_replace.replace_flags={'-fprofile-instrument-use-path': '/tmp/corpus_1_21/MergedCS.profdata'} +#problem_config.flags_to_replace.replace_flags={'-fprofile-instrument-use-path': 'MergedCS_profdata_jacobhegna/MergedCS_profdata'}
diff --git a/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py b/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py new file mode 100644 index 0000000..3c90edc --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/lr_encoder_runner.py
@@ -0,0 +1,94 @@ +"""Module for collect data of the LR encoder.""" + +import os +import tempfile +from typing import Dict, Tuple + +import gin +import tensorflow as tf + +from google3.third_party.ml_compiler_opt.compiler_opt.rl import compilation_runner +from google3.third_party.ml_compiler_opt.compiler_opt.rl import corpus +from google3.third_party.ml_compiler_opt.compiler_opt.rl import log_reader + + +@gin.configurable(module='runners') +class LREncoderRunner(compilation_runner.CompilationRunner): + """Class for collecting data for the LR encoder.""" + + def compile_fn( + self, + command_line: corpus.FullyQualifiedCmdLine, + tf_policy_path: str, + reward_only: bool, + ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: + """Run compilation for the given IR file under the given policy. + + Args: + command_line: the fully qualified command line. + tf_policy_path: path to TF policy direcoty on local disk. + reward_only: whether only return reward. + + Returns: + A dict mapping from example identifier to tuple containing: + sequence_example: A tf.SequenceExample proto describing compilation + trace, None if reward_only == True. + reward: reward of register allocation. + + Raises: + subprocess.CalledProcessError: if process fails. + compilation_runner.ProcessKilledError: (which it must pass through) on + cancelled work. + RuntimeError: if llvm-size produces unexpected output. + """ + assert not tf_policy_path + + working_dir = tempfile.mkdtemp() + + log_path = os.path.join(working_dir, 'log') + output_native_path = os.path.join(working_dir, 'native') + + result = {} + try: + cmdline = [] + if self._launcher_path: + cmdline.append(self._launcher_path) + cmdline.extend( + [self._clang_path] + + list(command_line) + + [ + '-mllvm', + '-regalloc-enable-advisor=development', + '-mllvm', + '-regalloc-lr-encoder-training-log=' + log_path, + '-mllvm', + '-regalloc-training-log=/dev/null', + '-o', + output_native_path, + ] + ) + + compilation_runner.start_cancellable_process( + cmdline, self._compilation_timeout, self._cancellation_manager + ) + + if not os.path.exists(log_path): + return {} + + # TODO(#202) + log_result = log_reader.read_log_as_sequence_examples(log_path) + + for fct_name, trajectory in log_result.items(): + if not trajectory.HasField('feature_lists'): + continue + # score = ( + # trajectory.feature_lists.feature_list['reward'] + # .feature[-1] + # .float_list.value[0] + # ) + result[fct_name] = (trajectory, 1.0) + + finally: + tf.io.gfile.rmtree(working_dir) + + return result
diff --git a/compiler_opt/rl/regalloc/lr_encoder/model.py b/compiler_opt/rl/regalloc/lr_encoder/model.py new file mode 100644 index 0000000..a953aa4 --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/model.py
@@ -0,0 +1,114 @@ +import gin +import tensorflow as tf +from tf_agents.utils import nest_utils +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config +from google3.third_party.ml_compiler_opt.compiler_opt.rl import attention + + +class MultiHeadEncoderModel(tf.keras.Model): + + def __init__(self, encoder, heads): + super().__init__(name='MultiHeadEncoderModel') + + self._encoder = encoder + self._heads = heads + + def get_encoder(self): + return self._encoder + + def call(self, inputs, *args): + observation = inputs['obs'] + action = tf.one_hot(inputs['action'], depth=config._NUM_REGISTERS)[ + :, :, tf.newaxis + ] + + use_def_obs = { + k: v + for k, v in observation.items() + if k.startswith(config._ENCODER_FEATURE_PREFIX) + } + encoded_state_per_token, encoded_state = self._encoder(use_def_obs) + encoded_state_with_action = tf.concat([encoded_state, action], axis=-1) + + def get_input(head_name): + if head_name == 'mlm': + return encoded_state_per_token + if head_name.startswith('next_'): + return encoded_state_with_action + return encoded_state + + outputs = {} + for name in self._heads: + head = self._heads[name] + head_input = get_input(name) + outputs[name] = head(head_input) + + return outputs + + +@gin.configurable +class LiveRangeEncoder(tf.keras.Model): + + def __init__( + self, + input_specs, + preprocessing_layer_creator, + *, + num_layers=2, + num_heads=2, + model_dim=32, + fcn_dim=128, + num_extra_features=10, + ): + super().__init__(name='LiveRangeEncoder') + self._encoder = attention.TransformerClassifier( + num_tokens=config._OPCODE_VOCAB_SIZE, + num_layers=num_layers, + num_heads=num_heads, + model_dim=model_dim, + fcn_dim=fcn_dim, + num_extra_features=num_extra_features, + ) + self._linear_reshape = tf.keras.layers.Dense(config._ENCODING_SIZE) + + self._preprocessing_layers = {} + for key in input_specs: + self._preprocessing_layers[key] = preprocessing_layer_creator(key) + + def preprocessor(observations): + pp_inputs = [] + for key in input_specs: + pp_inputs.append(self._preprocessing_layers[key](observations[key])) + return tf.concat(pp_inputs, axis=-1) + + self._preprocessor = preprocessor + + def call(self, observations): + extra_hidden_state = self._preprocessor(observations) + tokens = observations[config._OPCODE_KEY] + encoded_state_per_token, encoded_state = self._encoder( + tokens, extra_hidden_state + ) + encoded_state = self._linear_reshape(encoded_state) + return encoded_state_per_token, encoded_state + + +def create_model( + input_specs: dict, + output_specs: dict, + preprocessing_layer_creator: dict, + *, + mlp_width: int = 128, +): + encoder = LiveRangeEncoder(input_specs, preprocessing_layer_creator) + + output_heads = {} + for name, spec in output_specs.items(): + size = spec.shape[-1] + output_heads[name] = tf.keras.Sequential([ + tf.keras.layers.Dense(mlp_width), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(size), + ]) + + return MultiHeadEncoderModel(encoder=encoder, heads=output_heads)
diff --git a/compiler_opt/rl/regalloc/lr_encoder/train.py b/compiler_opt/rl/regalloc/lr_encoder/train.py new file mode 100644 index 0000000..ca0b628 --- /dev/null +++ b/compiler_opt/rl/regalloc/lr_encoder/train.py
@@ -0,0 +1,183 @@ +import os + +from absl import app +from absl import flags +from absl import logging + +import gin +import tensorflow as tf + +from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import config as regalloc_config +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import dataset_ops +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import model +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc.lr_encoder import config as encoder_config +from google3.third_party.ml_compiler_opt.compiler_opt.rl import registry # pylint: disable=unused-import + + +flags.DEFINE_multi_string( + 'gin_files', [], 'List of paths to gin configuration files.' +) + +flags.DEFINE_string('trace', None, '') +flags.DEFINE_string('root_dir', None, '') + +flags.DEFINE_string('tpu', None, 'BNS address for the TPU') + +FLAGS = flags.FLAGS + + +def _create_parser_fn(time_step_spec, action_spec): + context_features = {} + sequence_features = { + tensor_spec.name: tf.io.FixedLenSequenceFeature( + shape=tensor_spec.shape, dtype=tensor_spec.dtype + ) + for tensor_spec in time_step_spec.values() + } + sequence_features[action_spec.name] = tf.io.FixedLenSequenceFeature( + shape=action_spec.shape, dtype=action_spec.dtype + ) + + def _parser_fn(serialized_proto): + with tf.name_scope('parse'): + _, parsed_sequence = tf.io.parse_single_sequence_example( + serialized_proto, + context_features=context_features, + sequence_features=sequence_features, + ) + return parsed_sequence + + return _parser_fn + + +def get_dataset_creator_fn( + batch_size, + input_spec, + action_spec, + regalloc_input_spec, + regalloc_preprocessing_layer_creator, +): + files_buffer_size = 64 + num_map_threads = 128 + num_readers = 64 + parser_fn = _create_parser_fn(input_spec, action_spec) + + def _file_dataset_fn(data_path): + dataset = ( + tf.data.Dataset.list_files(data_path) + .shuffle(files_buffer_size) + .interleave( + tf.data.TFRecordDataset, + cycle_length=num_readers, + block_length=1, + ) + .filter(lambda string: tf.strings.length(string) > 0) + .map(parser_fn, num_parallel_calls=num_map_threads) + ) + dataset = dataset_ops.process_dataset( + dataset, regalloc_input_spec, regalloc_preprocessing_layer_creator + ) + return dataset.batch(batch_size, drop_remainder=True).prefetch( + tf.data.experimental.AUTOTUNE + ) + + return _file_dataset_fn + + +def get_strategy(): + if FLAGS.tpu: + logging.info('Using TPU strategy.') + resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu) + tf.config.experimental_connect_to_cluster(resolver) + tf.tpu.experimental.initialize_tpu_system(resolver) + return tf.distribute.TPUStrategy(resolver) + logging.info('Using CPU strategy.') + return tf.distribute.get_strategy() + + +class SaveEncoderCallback(tf.keras.callbacks.Callback): + + def __init__(self, encoder, path): + self._encoder = encoder + self._path = path + + def on_epoch_end(self, epoch, logs=None): + sm_path = os.path.join(self._path, f'epoch{epoch}') + tflite_path = os.path.join(sm_path, policy_saver.TFLITE_MODEL_NAME) + self._encoder.save(sm_path) + policy_saver.convert_saved_model(sm_path, tflite_path) + + +def main(_): + gin.parse_config_files_and_bindings( + FLAGS.gin_files, bindings=None, skip_unknown=False + ) + logging.info(gin.config_str()) + + lr_input_specs, lr_encoding_spec = ( + encoder_config.get_lr_encoder_signature_spec() + ) + + regalloc_time_step_spec, regalloc_action_spec = ( + regalloc_config.get_regalloc_signature_spec() + ) + regalloc_preprocessing_creator = ( + regalloc_config.get_observation_processing_layer_creator() + ) + + dataset_fn = get_dataset_creator_fn( + batch_size=64, + input_spec=encoder_config.get_input_specs(), + action_spec=regalloc_action_spec, + regalloc_input_spec=regalloc_time_step_spec.observation, + regalloc_preprocessing_layer_creator=regalloc_preprocessing_creator, + ) + + strategy = get_strategy() + with strategy.scope(): + logging.info('Creating model.') + lr_model = model.create_model( + lr_input_specs, + encoder_config.get_output_specs(regalloc_preprocessing_creator), + encoder_config.get_preprocessing_layer_creator(), + ) + + logging.info('Compiling model.') + loss, loss_weights = encoder_config.get_loss() + lr_model.compile( + optimizer=tf.keras.optimizers.Adam(global_clipnorm=1.0), + loss=loss, + loss_weights=loss_weights, + metrics=encoder_config.get_metrics(), + ) + + logging.info('Creating dataset.') + dataset = dataset_fn(FLAGS.trace) + tb = tf.keras.callbacks.TensorBoard( + log_dir=FLAGS.root_dir, + histogram_freq=0, + write_graph=True, + write_steps_per_second=True, + update_freq='batch', + ) + + policy_dir = os.path.join(FLAGS.root_dir, 'policy') + checkpoint = tf.keras.callbacks.ModelCheckpoint( + filepath=policy_dir, save_weights_only=False, save_freq=1024 + ) + + encoder_dir = os.path.join(FLAGS.root_dir, 'encoder') + logging.info('Saving the encoder to %s', encoder_dir) + encoder_saver = SaveEncoderCallback( + encoder=lr_model.get_encoder(), path=encoder_dir + ) + + logging.info('Starting training.') + lr_model.fit(dataset, epochs=8, callbacks=[tb, checkpoint, encoder_saver]) + + logging.info('Training complete.') + + +if __name__ == '__main__': + app.run(main)
diff --git a/compiler_opt/rl/train_semisupervised.py b/compiler_opt/rl/train_semisupervised.py new file mode 100644 index 0000000..209f2f9 --- /dev/null +++ b/compiler_opt/rl/train_semisupervised.py
@@ -0,0 +1,490 @@ +from absl import app +import tensorflow as tf +import tensorflow_text as text + +import numpy as np + +import os +from typing import Optional, List + +from google3.third_party.ml_compiler_opt.compiler_opt.rl import feature_ops +from google3.third_party.ml_compiler_opt.compiler_opt.rl import registry +from google3.third_party.ml_compiler_opt.compiler_opt.rl import attention +from google3.third_party.ml_compiler_opt.compiler_opt.rl.regalloc import regalloc_network +from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver +from absl import flags +from absl import logging +import gin +# from google3.third_party.ml_compiler_opt.compiler_opt.rl import policy_saver + + +# Have one class for the Encoder +# Have one class for the full semisupervised model +# only save the Encoder +# Output heads: +# - reward +# - next state +# - + +flags.DEFINE_multi_string( + 'gin_files', [], 'List of paths to gin configuration files.' +) + +flags.DEFINE_string('trace', None, '') +flags.DEFINE_string('root_dir', None, '') + +flags.DEFINE_string('tpu', None, 'BNS address for the TPU') + + +FLAGS = flags.FLAGS + +_ACTION_KEY = 'action' +_CUR_STATE_KEY = 'cur_state' +_NEXT_STATE_KEY = 'next_state' +_NEXT_ACTION_KEY = 'next_action' +_MLM_KEY = 'mlm' + +_INSTR_COUNT = 64 + +_MLM_IGNORE_TOKEN = -1 +_MLM_MASK_TOKEN = 18000 - 1 + + +class Model(tf.keras.Model): + + def __init__( + self, + encoder_network, + state_head, + action_head, + next_state_head, + next_action_head, + mlm_head, + ): + super().__init__(name='Model') + self._encoder_network = encoder_network + self._state_head = state_head + self._action_head = action_head + self._next_state_head = next_state_head + self._next_action_head = next_action_head + self._mlm_head = mlm_head + + def call(self, inputs): + # TODO: + # 1) export the non-reduced tensor because need it for the MLM outputs + # 2) when reducing, mask the tensors that don't matter + # for masking, I think I should set the relevant outputs to zero and also zero out the labels. + # 3) figure out the correct masks when training, might need custom training loop or custom loss functions. + observation = inputs['obs'] + action = tf.one_hot(inputs['action'], depth=33)[:, :, tf.newaxis] + + use_def_obs = { + k: v for k, v in observation.items() if k.startswith('use_def_') + } + encoded_state_per_token, encoded_state = self._encoder_network(use_def_obs) + + encoded_state_with_action = tf.concat([encoded_state, action], axis=-1) + + state = self._state_head(encoded_state) + action = self._action_head(encoded_state) + next_state = self._next_state_head(encoded_state_with_action) + next_action = self._next_action_head(encoded_state_with_action) + mlm = self._mlm_head(encoded_state_per_token) + return { + _CUR_STATE_KEY: state, + _ACTION_KEY: action, + _NEXT_STATE_KEY: next_state, + _NEXT_ACTION_KEY: next_action, + _MLM_KEY: mlm, + } + + +def action_loss(label, pred): + # Use the current mask + print(label) + pass + + +def state_loss(label, pred): + # Decide whether to use the intersection of the masks or just the current mask + # Probably use the current mask + print(label) + pass + + +def masked_language_loss(label, pred): + mask = label != 0 + loss_object = tf.keras.losses.SparseCategoricalCrossEntropy( + from_logits=True, reduction='none' + ) + loss = loss_object(label, pred) + + mask = tf.cast(mask, dtype=loss.dtype) + loss *= mask + + loss = tf.reduce_sum(loss) / tf.reduce_sum(mask) + return loss + + +def create_model(encoder_network): + action_head = tf.keras.Sequential([ + tf.keras.layers.Dense(128), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(1), + ]) + state_head = tf.keras.Sequential([ + tf.keras.layers.Dense(128), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(65), + ]) + next_state_head = tf.keras.Sequential([ + tf.keras.layers.Dense(128), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(65), + ]) + next_action_head = tf.keras.Sequential([ + tf.keras.layers.Dense(128), + tf.keras.layers.ReLU(), + tf.keras.layers.Dense(1), + ]) + mlm = tf.keras.Sequential([ + tf.keras.layers.Dense(18000), + ]) + return Model( + encoder_network=encoder_network, + action_head=action_head, + state_head=state_head, + next_state_head=next_state_head, + next_action_head=next_action_head, + mlm_head=mlm, + ) + + +def _create_parser_fn(time_step_spec, action_spec): + context_features = {} + sequence_features = { + tensor_spec.name: tf.io.FixedLenSequenceFeature( + shape=tensor_spec.shape, dtype=tensor_spec.dtype + ) + for tensor_spec in time_step_spec.observation.values() + } + sequence_features[action_spec.name] = tf.io.FixedLenSequenceFeature( + shape=action_spec.shape, dtype=action_spec.dtype + ) + + def _parser_fn(serialized_proto): + with tf.name_scope('parse'): + _, parsed_sequence = tf.io.parse_single_sequence_example( + serialized_proto, + context_features=context_features, + sequence_features=sequence_features, + ) + return parsed_sequence + + return _parser_fn + + +def _get_masked_input_and_labels(encoded_texts): + # https://keras.io/examples/nlp/masked_language_modeling/ + # 15% BERT masking + encoded_texts_shape = tf.shape(encoded_texts) + inp_mask = tf.random.uniform(encoded_texts_shape) < 0.15 + inp_mask[encoded_texts <= 0] = False + labels = _MLM_IGNORE_TOKEN * tf.ones(encoded_texts_shape, dtype=fp.int64) + # Set labels for masked tokens + labels[inp_mask] = encoded_texts[inp_mask] + + # Prepare input + encoded_texts_masked = tf.identity(encoded_texts) + # Set input to [MASK] which is the last token for the 90% of tokens + # This means leaving 10% unchanged + inp_mask_2mask = inp_mask & (tf.random.uniform(encoded_texts_shape) < 0.90) + encoded_texts_masked[inp_mask_2mask] = _MLM_MASK_TOKEN + + # Set 10% to a random token + inp_mask_2random = inp_mask_2mask & ( + tf.random.uniform(encoded_texts_shape) < 1 / 9 + ) + encoded_texts_masked[inp_mask_2random] = tf.random.uniform( + encoded_texts_shape, 3, mask_token_id, dtype=tf.int64 + ) + + # Prepare sample_weights to pass to .fit() method + sample_weights = tf.ones(encoded_texts_shape) + sample_weights[labels == -1] = 0 + + # y_labels would be same as encoded_texts i.e input tokens + y_labels = tf.identity(encoded_texts) + + return encoded_texts_masked, y_labels, sample_weights + + +_MAX_PREDICTIONS_PER_BATCH = 32 +random_selector = text.RandomItemSelector( + max_selections_per_batch=_MAX_PREDICTIONS_PER_BATCH, + selection_rate=0.2, + unselectable_ids=[0], +) +mask_values_chooser = text.MaskValuesChooser(18000, _MLM_MASK_TOKEN, 0.8) + + +def get_masked_input_and_labels(encoded_texts): + masked_token_ids, masked_pos, masked_lm_ids = text.mask_language_model( + tf.RaggedTensor.from_tensor(encoded_texts, padding=0), + item_selector=random_selector, + mask_values_chooser=mask_values_chooser, + ) + # NEed to fix this + # Produce tensor 0 to 32, tile it and produce tensr of indices, then flatten + # Then can use it in the scatter like I want + masked_pos = masked_pos.to_tensor( + default_value=-1, shape=(33, _MAX_PREDICTIONS_PER_BATCH) + ) + + ii = tf.tile( + tf.range(33, dtype=tf.int64)[:, tf.newaxis], + [1, _MAX_PREDICTIONS_PER_BATCH], + ) + masked_pos = tf.stack([ii, masked_pos], axis=-1) + scatter_values = tf.where(masked_pos[:, :, 1] < 0, 0.0, 1.0) + masked_pos = tf.where( + masked_pos < 0, tf.constant(0, dtype=tf.int64), masked_pos + ) + weights = tf.scatter_nd( + masked_pos, scatter_values[:, :, tf.newaxis], (33, _INSTR_COUNT, 1) + ) + return ( + masked_token_ids.to_tensor(default_value=0, shape=(33, _INSTR_COUNT)), + encoded_texts, + weights, + ) + + +def create_dataset_fn( + time_step_spec, action_spec, batch_size, shift, preprocessing_layer_creator +): + assert shift < 0 + files_buffer_size = 100 + num_readers = 10 + num_map_threads = 8 + shuffle_buffer_size = 256 + + parser_fn = _create_parser_fn(time_step_spec, action_spec) + + def _roll_experience(seq_ex): + # Use tf agents nest map fn here + def _roll(atom): + return tf.roll(atom, shift=shift, axis=0) + + def _cutoff(atom): + return atom[:shift] + + seq_ex_roll = tf.nest.map_structure(_roll, seq_ex) + seq_ex_roll = tf.nest.map_structure(_cutoff, seq_ex_roll) + seq_ex = tf.nest.map_structure(_cutoff, seq_ex) + return seq_ex, seq_ex_roll + + preprocessing_layers = tf.nest.map_structure( + preprocessing_layer_creator, time_step_spec.observation + ) + + def split_experience(seq_ex, seq_ex_roll): + obs = {k: seq_ex[k] for k in seq_ex if k != action_spec.name} + action = seq_ex[action_spec.name] + obs_roll = {k: seq_ex_roll[k] for k in seq_ex_roll if k != action_spec.name} + action_roll = seq_ex_roll[action_spec.name] + return { + 'obs': obs, + 'action': action, + 'obs_roll': obs_roll, + 'action_roll': action_roll, + } + + def _preprocessing_layer(seq_ex): + for layer_name, layer in preprocessing_layers.items(): + seq_ex[layer_name] = layer(seq_ex[layer_name]) + return seq_ex + + def preprocess_experience(obs_dict): + obs_dict['obs_cur'] = _preprocessing_layer(obs_dict['obs'].copy()) + obs_dict['obs_cur'] = tf.concat( + [ + v + for k, v in obs_dict['obs_cur'].items() + if not k.startswith('use_def_') + ], + axis=-1, + ) + obs_dict['obs_roll'] = _preprocessing_layer(obs_dict['obs_roll']) + obs_dict['obs_roll'] = tf.concat( + [ + v + for k, v in obs_dict['obs_roll'].items() + if not k.startswith('use_def_') + ], + axis=-1, + ) + return obs_dict + + def to_inputs_and_labels(obs_dict): + inputs = {'obs': obs_dict['obs'], 'action': obs_dict['action']} + mlm_input, mlm_label, mlm_weight = get_masked_input_and_labels( + obs_dict['obs']['use_def_opcode'][:, :_INSTR_COUNT] + ) + inputs['use_def_opcode'] = mlm_input + + labels = { + _CUR_STATE_KEY: obs_dict['obs_cur'], + _ACTION_KEY: tf.expand_dims(obs_dict['action'], axis=-1), + _NEXT_STATE_KEY: obs_dict['obs_roll'], + _NEXT_ACTION_KEY: tf.expand_dims(obs_dict['action_roll'], axis=-1), + _MLM_KEY: mlm_label, + } + mask = obs_dict['obs']['mask'] + sample_weights = { + _CUR_STATE_KEY: mask, + _ACTION_KEY: mask, + _NEXT_STATE_KEY: None, + _NEXT_ACTION_KEY: None, + _MLM_KEY: mlm_weight, + } + return (inputs, labels, sample_weights) + + def _file_dataset_fn(data_path): + return ( + tf.data.Dataset.list_files(data_path) + .shuffle(files_buffer_size) + .interleave( + tf.data.TFRecordDataset, + cycle_length=num_readers, + block_length=1, + ) + .filter(lambda string: tf.strings.length(string) > 0) + .map(parser_fn, num_parallel_calls=num_map_threads) + .map(_roll_experience, num_parallel_calls=num_map_threads) + .map(split_experience, num_parallel_calls=num_map_threads) + .map(preprocess_experience, num_parallel_calls=num_map_threads) + .unbatch() + .shuffle(shuffle_buffer_size) + .map(to_inputs_and_labels, num_parallel_calls=num_map_threads) + .batch(batch_size, drop_remainder=True) + .prefetch(tf.data.experimental.AUTOTUNE) + ) + + return _file_dataset_fn + + +def get_strategy(): + if FLAGS.tpu: + logging.info('Using TPU strategy.') + resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=FLAGS.tpu) + tf.config.experimental_connect_to_cluster(resolver) + tf.tpu.experimental.initialize_tpu_system(resolver) + return tf.distribute.TPUStrategy(resolver) + logging.info('Using CPU strategy.') + return tf.distribute.get_strategy() + + +class SaveEncoderCallback(tf.keras.callbacks.Callback): + + def __init__(self, encoder, path): + self._encoder = encoder + self._path = path + + def on_epoch_end(self, epoch, logs=None): + sm_path = os.path.join(self._path, f'epoch{epoch}') + tflite_path = os.path.join(sm_path, policy_saver.TFLITE_MODEL_NAME) + self._encoder.save(sm_path) + policy_saver.convert_saved_model(sm_path, tflite_path) + + +def main(_): + gin.parse_config_files_and_bindings( + FLAGS.gin_files, bindings=None, skip_unknown=False + ) + logging.info(gin.config_str()) + + problem_config = registry.get_configuration() + time_step_spec, action_spec = problem_config.get_signature_spec() + preprocessing_layer_creator = problem_config.get_preprocessing_layer_creator() + dataset_fn = create_dataset_fn( + time_step_spec, + action_spec, + batch_size=256, + shift=-1, + preprocessing_layer_creator=preprocessing_layer_creator, + ) + + # inputs, labels, weights = next(iter(dataset)) + # import sys + # sys.exit() + + strategy = get_strategy() + with strategy.scope(): + logging.info('Creating model.') + encoder_network = regalloc_network.InstructionEncoderNetwork( + preprocessing_layer_creator + ) + model = create_model(encoder_network) + + logging.info('Compiling model.') + + opt = tf.keras.optimizers.Adam(global_clipnorm=1.0) + + model.compile( + optimizer=opt, + loss={ + _CUR_STATE_KEY: 'mse', + _ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True + ), + _NEXT_STATE_KEY: 'mse', + _NEXT_ACTION_KEY: tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True + ), + _MLM_KEY: tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True + ), + }, + loss_weights={ + _ACTION_KEY: 1, + _NEXT_STATE_KEY: 1, + _NEXT_ACTION_KEY: 1, + _MLM_KEY: 10, + }, + metrics={ + _CUR_STATE_KEY: [], + _ACTION_KEY: ['accuracy'], + _NEXT_STATE_KEY: [], + _NEXT_ACTION_KEY: ['accuracy'], + _MLM_KEY: ['accuracy'], + }, + ) + + logging.info('Creating dataset.') + dataset = dataset_fn(FLAGS.trace) + tb = tf.keras.callbacks.TensorBoard( + log_dir=FLAGS.root_dir, + histogram_freq=0, + write_graph=True, + write_steps_per_second=True, + update_freq='batch', + ) + + policy_dir = os.path.join(FLAGS.root_dir, 'policy') + checkpoint = tf.keras.callbacks.ModelCheckpoint( + filepath=policy_dir, save_weights_only=False, save_freq=1024 + ) + + encoder_dir = os.path.join(FLAGS.root_dir, 'encoder') + logging.info('Saving the encoder to %s', encoder_dir) + encoder_saver = SaveEncoderCallback(encoder=encoder_network, path=encoder_dir) + + logging.info('Starting training.') + model.fit(dataset, epochs=8, callbacks=[tb, checkpoint, encoder_saver]) + + logging.info('Training complete.') + + +if __name__ == '__main__': + app.run(main)