# Path: compiler_opt/rl/compilation_runner.py
# class RewardStat:
# class DataClassJSONEncoder(json.JSONEncoder):
# class ProcessKilledError(Exception):
# class ProcessCancellationToken:
# class WorkerCancellationManager:
# class CompilationRunner:
#   def default(self, o):
# def _postprocessing_sequence_example(
#     sequence_example: tf.train.SequenceExample, moving_average_reward: float,
#     policy_reward: float) -> tf.train.SequenceExample:
# def get_command_line_for_bundle(cmd_file: str,
#                                 ir_file: str,
#                                 thinlto: Optional[str] = None) -> List[str]:
#   def __init__(self):
#   def __init__(self):
#   def signal(self):
#   def wait(self):
# def kill_process_ignore_exceptions(p: subprocess.Popen):
#   def __init__(self):
#   def register_process(self, p: subprocess.Popen):
#   def signal(self):
#   def unregister_process(self, p: subprocess.Popen):
# def start_cancellable_process(
#     cmdline: List[str],
#     timeout: float,
#     cancellation_manager: Optional[WorkerCancellationManager],
#     want_output: bool = False) -> Optional[bytes]:
#   def init_pool():
#   def _get_pool():
#   def __init__(self,
#                clang_path: str,
#                llvm_size_path: str,
#                launcher_path: Optional[str] = None,
#                moving_average_decay_rate: float = 1):
#   def _get_cancellation_manager(
#       self, cancellation_token: Optional[ProcessCancellationToken]
#   ) -> Optional[WorkerCancellationManager]:
#     def signaler():
#   def collect_data(
#       self,
#       file_paths: Tuple[str, ...],
#       tf_policy_path: str,
#       reward_stat: Optional[Dict[str, RewardStat]],
#       cancellation_token: Optional[ProcessCancellationToken] = None
#   ) -> Tuple[List[str], Dict[str, RewardStat], List[float]]:
#   def _compile_fn(
#       self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
#   _POOL: concurrent.futures.ThreadPoolExecutor = None
# 
# Path: compiler_opt/rl/inlining/config.py
# def get_inlining_signature_spec():
# def get_observation_processing_layer_creator(quantile_file_dir=None,
#                                              with_sqrt=True,
#                                              with_z_score_normalization=True,
#                                              eps=1e-8):
#   def observation_processing_layer(obs_spec):
# 
# Path: compiler_opt/rl/inlining/inlining_runner.py
# _DEADLINE_IN_SECONDS = 60
# _DEFAULT_IDENTIFIER = 'default'
# class InliningRunner(compilation_runner.CompilationRunner):
#   def _compile_fn(
#       self, file_paths: Tuple[str, str], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[
#           compilation_runner.WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
# 
# Path: compiler_opt/rl/regalloc/config.py
# def get_num_registers():
# def get_regalloc_signature_spec():
# def get_observation_processing_layer_creator(quantile_file_dir=None,
#                                              with_sqrt=True,
#                                              with_z_score_normalization=True,
#                                              eps=1e-8):
#   def observation_processing_layer(obs_spec):
#       def use_def_density_processing_fn(obs):
#       def progress_processing_fn(obs):


from typing import Callable, Text, Tuple
from compiler_opt.rl import compilation_runner
from compiler_opt.rl.inlining import config as inlining_config
from compiler_opt.rl.inlining import inlining_runner
from compiler_opt.rl.regalloc import config as regalloc_config
import tensorflow as tf
import tf_agents as tfa
"""Configuration utility helpers for RL training."""






types = tfa.typing.types

# TODO(b/214316645): get rid of the if-else statement by defining a class for
# each problem type instead.


def get_signature_spec(
    problem_type: Text
) -> Tuple[types.NestedTensorSpec, types.NestedTensorSpec]:
  """Get the signature spec for the given problem type."""
  if problem_type == 'inlining':
    return inlining_config.get_inlining_signature_spec()
  elif problem_type == 'regalloc':
    return regalloc_config.get_regalloc_signature_spec()
  else:
    raise ValueError('Unknown problem_type: {}'.format(problem_type))


def get_preprocessing_layer_creator(
    problem_type: Text,
) -> Callable[[types.TensorSpec], tf.keras.layers.Layer]:
  """Get the observation processing layer creator for the given problem type."""
  if problem_type == 'inlining':

# Path: compiler_opt/rl/compilation_runner.py
# class RewardStat:
# class DataClassJSONEncoder(json.JSONEncoder):
# class ProcessKilledError(Exception):
# class ProcessCancellationToken:
# class WorkerCancellationManager:
# class CompilationRunner:
#   def default(self, o):
# def _postprocessing_sequence_example(
#     sequence_example: tf.train.SequenceExample, moving_average_reward: float,
#     policy_reward: float) -> tf.train.SequenceExample:
# def get_command_line_for_bundle(cmd_file: str,
#                                 ir_file: str,
#                                 thinlto: Optional[str] = None) -> List[str]:
#   def __init__(self):
#   def __init__(self):
#   def signal(self):
#   def wait(self):
# def kill_process_ignore_exceptions(p: subprocess.Popen):
#   def __init__(self):
#   def register_process(self, p: subprocess.Popen):
#   def signal(self):
#   def unregister_process(self, p: subprocess.Popen):
# def start_cancellable_process(
#     cmdline: List[str],
#     timeout: float,
#     cancellation_manager: Optional[WorkerCancellationManager],
#     want_output: bool = False) -> Optional[bytes]:
#   def init_pool():
#   def _get_pool():
#   def __init__(self,
#                clang_path: str,
#                llvm_size_path: str,
#                launcher_path: Optional[str] = None,
#                moving_average_decay_rate: float = 1):
#   def _get_cancellation_manager(
#       self, cancellation_token: Optional[ProcessCancellationToken]
#   ) -> Optional[WorkerCancellationManager]:
#     def signaler():
#   def collect_data(
#       self,
#       file_paths: Tuple[str, ...],
#       tf_policy_path: str,
#       reward_stat: Optional[Dict[str, RewardStat]],
#       cancellation_token: Optional[ProcessCancellationToken] = None
#   ) -> Tuple[List[str], Dict[str, RewardStat], List[float]]:
#   def _compile_fn(
#       self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
#   _POOL: concurrent.futures.ThreadPoolExecutor = None
# 
# Path: compiler_opt/rl/inlining/config.py
# def get_inlining_signature_spec():
# def get_observation_processing_layer_creator(quantile_file_dir=None,
#                                              with_sqrt=True,
#                                              with_z_score_normalization=True,
#                                              eps=1e-8):
#   def observation_processing_layer(obs_spec):
# 
# Path: compiler_opt/rl/inlining/inlining_runner.py
# _DEADLINE_IN_SECONDS = 60
# _DEFAULT_IDENTIFIER = 'default'
# class InliningRunner(compilation_runner.CompilationRunner):
#   def _compile_fn(
#       self, file_paths: Tuple[str, str], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[
#           compilation_runner.WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
# 
# Path: compiler_opt/rl/regalloc/config.py
# def get_num_registers():
# def get_regalloc_signature_spec():
# def get_observation_processing_layer_creator(quantile_file_dir=None,
#                                              with_sqrt=True,
#                                              with_z_score_normalization=True,
#                                              eps=1e-8):
#   def observation_processing_layer(obs_spec):
#       def use_def_density_processing_fn(obs):
#       def progress_processing_fn(obs):


# Path: compiler_opt/rl/config.py
from typing import Callable, Text, Tuple
from compiler_opt.rl import compilation_runner
from compiler_opt.rl.inlining import config as inlining_config
from compiler_opt.rl.inlining import inlining_runner
from compiler_opt.rl.regalloc import config as regalloc_config
import tensorflow as tf
import tf_agents as tfa

"""Configuration utility helpers for RL training."""






types = tfa.typing.types

# TODO(b/214316645): get rid of the if-else statement by defining a class for
# each problem type instead.


def get_signature_spec(
    problem_type: Text
) -> Tuple[types.NestedTensorSpec, types.NestedTensorSpec]:
  """Get the signature spec for the given problem type."""
  if problem_type == 'inlining':
    return inlining_config.get_inlining_signature_spec()
  elif problem_type == 'regalloc':
    return regalloc_config.get_regalloc_signature_spec()
  else:
    raise ValueError('Unknown problem_type: {}'.format(problem_type))


def get_preprocessing_layer_creator(
    problem_type: Text,
) -> Callable[[types.TensorSpec], tf.keras.layers.Layer]:
  """Get the observation processing layer creator for the given problem type."""
  if problem_type == 'inlining':

    return inlining_config.get_observation_processing_layer_creator()