diff --git a/compiler_opt/rl/compilation_runner.py b/compiler_opt/rl/compilation_runner.py index 15a6d32e..40bda622 100644 --- a/compiler_opt/rl/compilation_runner.py +++ b/compiler_opt/rl/compilation_runner.py @@ -334,20 +334,14 @@ def collect_data( """ if reward_stat is None: default_result = self.compile_fn( - module_spec, - tf_policy_path='', - reward_only=bool(tf_policy_path), - cancellation_manager=self._cancellation_manager) + module_spec, tf_policy_path='', reward_only=bool(tf_policy_path)) reward_stat = { k: RewardStat(v[1], v[1]) for (k, v) in default_result.items() } if tf_policy_path: policy_result = self.compile_fn( - module_spec, - tf_policy_path, - reward_only=False, - cancellation_manager=self._cancellation_manager) + module_spec, tf_policy_path, reward_only=False) else: policy_result = default_result @@ -383,17 +377,13 @@ def collect_data( def compile_fn( self, module_spec: corpus.ModuleSpec, tf_policy_path: str, - reward_only: bool, - cancellation_manager: Optional[WorkerCancellationManager] - ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: + reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: """Compiles for the given IR file under the given policy. Args: module_spec: a ModuleSpec. tf_policy_path: path to TF policy directory on local disk. reward_only: whether only return reward. - cancellation_manager: a WorkerCancellationManager to handle early - termination Returns: A dict mapping from example identifier to tuple containing: diff --git a/compiler_opt/rl/compilation_runner_test.py b/compiler_opt/rl/compilation_runner_test.py index 4390d51d..b7724353 100644 --- a/compiler_opt/rl/compilation_runner_test.py +++ b/compiler_opt/rl/compilation_runner_test.py @@ -76,10 +76,8 @@ def _get_sequence_example(feature_value): return text_format.Parse(sequence_example_text, tf.train.SequenceExample()) -def _mock_compile_fn(file_paths, tf_policy_path, reward_only, - cancellation_manager): # pylint: disable=unused-argument +def _mock_compile_fn(file_paths, tf_policy_path, reward_only): # pylint: disable=unused-argument del file_paths - del cancellation_manager if tf_policy_path: sequence_example = _get_sequence_example(_POLICY_FEATURE_VALUE) native_size = _POLICY_REWARD diff --git a/compiler_opt/rl/inlining/inlining_runner.py b/compiler_opt/rl/inlining/inlining_runner.py index 69730f0a..d2ea8977 100644 --- a/compiler_opt/rl/inlining/inlining_runner.py +++ b/compiler_opt/rl/inlining/inlining_runner.py @@ -17,7 +17,7 @@ import io import os import tempfile -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple import gin import tensorflow as tf @@ -47,17 +47,13 @@ def __init__(self, llvm_size_path: str, *args, **kwargs): def compile_fn( self, module_spec: corpus.ModuleSpec, tf_policy_path: str, - reward_only: bool, cancellation_manager: Optional[ - compilation_runner.WorkerCancellationManager] - ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: + reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: """Run inlining for the given IR file under the given policy. Args: module_spec: a ModuleSpec. tf_policy_path: path to TF policy direcoty on local disk. reward_only: whether only return native size. - cancellation_manager: handler for early termination by killing any running - processes Returns: A dict mapping from example identifier to tuple containing: @@ -71,6 +67,7 @@ def compile_fn( cancelled work. RuntimeError: if llvm-size produces unexpected output. """ + working_dir = tempfile.mkdtemp() log_path = os.path.join(working_dir, 'log') @@ -91,12 +88,12 @@ def compile_fn( ['-mllvm', '-ml-inliner-model-under-training=' + tf_policy_path]) compilation_runner.start_cancellable_process(command_line, self._compilation_timeout, - cancellation_manager) + self._cancellation_manager) command_line = [self._llvm_size_path, output_native_path] output_bytes = compilation_runner.start_cancellable_process( command_line, timeout=self._compilation_timeout, - cancellation_manager=cancellation_manager, + cancellation_manager=self._cancellation_manager, want_output=True) if not output_bytes: raise RuntimeError(f'Empty llvm-size output: {" ".join(command_line)}') diff --git a/compiler_opt/rl/regalloc/regalloc_runner.py b/compiler_opt/rl/regalloc/regalloc_runner.py index ef6cf0ce..2bd6622c 100644 --- a/compiler_opt/rl/regalloc/regalloc_runner.py +++ b/compiler_opt/rl/regalloc/regalloc_runner.py @@ -18,7 +18,7 @@ import io import os import tempfile -from typing import Dict, Optional, Tuple +from typing import Dict, Tuple import gin import tensorflow as tf @@ -45,17 +45,13 @@ class RegAllocRunner(compilation_runner.CompilationRunner): # construction def compile_fn( self, module_spec: corpus.ModuleSpec, tf_policy_path: str, - reward_only: bool, cancellation_manager: Optional[ - compilation_runner.WorkerCancellationManager] - ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: + reward_only: bool) -> Dict[str, Tuple[tf.train.SequenceExample, float]]: """Run inlining for the given IR file under the given policy. Args: module_spec: a ModuleSpec. tf_policy_path: path to TF policy direcoty on local disk. reward_only: whether only return reward. - cancellation_manager: handler for early termination by killing any running - processes Returns: A dict mapping from example identifier to tuple containing: @@ -69,6 +65,7 @@ def compile_fn( cancelled work. RuntimeError: if llvm-size produces unexpected output. """ + working_dir = tempfile.mkdtemp() log_path = os.path.join(working_dir, 'log') @@ -88,7 +85,7 @@ def compile_fn( command_line.extend(['-mllvm', '-regalloc-model=' + tf_policy_path]) compilation_runner.start_cancellable_process(command_line, self._compilation_timeout, - cancellation_manager) + self._cancellation_manager) sequence_example = struct_pb2.Struct()