diff --git a/compiler_opt/es/blackbox_learner.py b/compiler_opt/es/blackbox_learner.py index 77401036..1f5a1fb1 100644 --- a/compiler_opt/es/blackbox_learner.py +++ b/compiler_opt/es/blackbox_learner.py @@ -196,10 +196,9 @@ def _log_tf_summary(self, rewards: List[float]) -> None: with self._summary_writer.as_default(): tf.summary.scalar( 'reward/average_reward_train', np.mean(rewards), step=self._step) - + tf.summary.scalar( - 'reward/maximum_reward_train', np.max(rewards), step=self._step - ) + 'reward/maximum_reward_train', np.max(rewards), step=self._step) tf.summary.histogram('reward/reward_train', rewards, step=self._step) diff --git a/compiler_opt/es/es_trainer_lib.py b/compiler_opt/es/es_trainer_lib.py index 95d52675..5204f362 100644 --- a/compiler_opt/es/es_trainer_lib.py +++ b/compiler_opt/es/es_trainer_lib.py @@ -56,15 +56,12 @@ "pretrained_policy_path", None, "The path of the pretrained policy. If not provided, it will \ construct a new policy with randomly initialized weights.") -_CORPUS_DIR = flags.DEFINE_string("corpus_dir", None, - "The path to the corpus to use") -_CLANG_PATH = flags.DEFINE_string("clang_path", None, - "The path to the clang binary to use.") -_TRACE_PATH = flags.DEFINE_string("trace_path", None, - "The path to the BB trace to use.") -_BB_TRACE_MODEL_PATH = flags.DEFINE_string( - "bb_trace_model_path", None, - "THe path to the basic_block_trace_model binary to use.") + +_CORPUS_DIR = '/usr/local/google/home/aidengrossman/opt_mlregalloc/corpus_subset' +_CLANG_PATH = '/usr/local/google/home/aidengrossman/opt_mlregalloc/clang' +_TRACE_PATH = '/usr/local/google/home/aidengrossman/opt_mlregalloc/bb_trace.pb' +_FUNCTION_INDEX_PATH = '/usr/local/google/home/aidengrossman/opt_mlregalloc/function_index.pb' +_BB_TRACE_MODEL_PATH = '/usr/local/google/home/aidengrossman/opt_mlregalloc/basic_block_trace_model' class ESWorker(worker.Worker): @@ -76,11 +73,7 @@ def __init__(self, *, all_gin): self._template_dir = tempfile.mkdtemp() saver.save(self._template_dir) - self._corpus_dir = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/corpus' - self._clang_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/clang' - self._trace_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/execution_trace.pb' - self._bb_trace_model_path = '/usr/local/google/home/aidengrossman/programming/opt_mlregalloc/basic_block_trace_model' - self._models_for_test_path = '/usr/local/google/home/aidengrossman/programming/output_traces/' + self._models_for_test_path = '/usr/local/google/home/aidengrossman/output_models/' def es_compile(self, params: list[float], baseline_score: float) -> float: with tempfile.TemporaryDirectory() as tempdir: @@ -97,15 +90,16 @@ def es_compile(self, params: list[float], baseline_score: float) -> float: policy_saver.OUTPUT_SIGNATURE), os.path.join(tflitedir, policy_saver.OUTPUT_SIGNATURE)) - trace_data_collector.compile_corpus(self._corpus_dir, tempdir, - self._clang_path, tflitedir, single_threaded=True) + trace_data_collector.compile_corpus( + _CORPUS_DIR, tempdir, _CLANG_PATH, tflitedir, thread_count=4) score = trace_data_collector.evaluate_compiled_corpus( - tempdir, self._trace_path, self._bb_trace_model_path) - + tempdir, _TRACE_PATH, _FUNCTION_INDEX_PATH, _BB_TRACE_MODEL_PATH, 4) + reward = compilation_runner._calculate_reward(score, baseline_score) print(reward) - output_path = os.path.join(self._models_for_test_path, "model" + str(reward)) + output_path = os.path.join(self._models_for_test_path, + "model" + str(reward)) if reward > 0 and not os.path.exists(output_path): shutil.copytree(tflitedir, output_path) return compilation_runner._calculate_reward(score, baseline_score) @@ -235,10 +229,9 @@ def train(worker_class=None): # Get baseline score with tempfile.TemporaryDirectory() as tempdir: - trace_data_collector.compile_corpus(_CORPUS_DIR.value, tempdir, - _CLANG_PATH.value) + trace_data_collector.compile_corpus(_CORPUS_DIR, tempdir, _CLANG_PATH) baseline_score = trace_data_collector.evaluate_compiled_corpus( - tempdir, _TRACE_PATH.value, _BB_TRACE_MODEL_PATH.value) + tempdir, _TRACE_PATH, _FUNCTION_INDEX_PATH, _BB_TRACE_MODEL_PATH) logging.info("Initializing blackbox learner.") learner = blackbox_learner.BlackboxLearner( diff --git a/compiler_opt/es/gin_configs/blackbox_learner.gin b/compiler_opt/es/gin_configs/blackbox_learner.gin index 34b97334..bbf8f49d 100644 --- a/compiler_opt/es/gin_configs/blackbox_learner.gin +++ b/compiler_opt/es/gin_configs/blackbox_learner.gin @@ -3,8 +3,8 @@ import compiler_opt.rl.gin_external_configurables import compiler_opt.es.blackbox_optimizers # Blackbox learner config -BlackboxLearnerConfig.total_steps = 100 -BlackboxLearnerConfig.total_num_perturbations = 100 +BlackboxLearnerConfig.total_steps = 10000 +BlackboxLearnerConfig.total_num_perturbations = 25 BlackboxLearnerConfig.blackbox_optimizer = %blackbox_optimizers.Algorithm.MONTE_CARLO BlackboxLearnerConfig.est_type = %blackbox_optimizers.EstimatorType.ANTITHETIC # BlackboxLearnerConfig.est_type = %blackbox_optimizers.EstimatorType.FORWARD_FD @@ -17,5 +17,5 @@ BlackboxLearnerConfig.num_top_directions = 0 BlackboxLearnerConfig.precision_parameter = 0.5 # Try the 0.0005 step size next -BlackboxLearnerConfig.step_size = 0.005 +BlackboxLearnerConfig.step_size = 0.005 # BlackboxLearnerConfig.step_size = 0.0005 \ No newline at end of file diff --git a/compiler_opt/rl/trace_data_collector.py b/compiler_opt/rl/trace_data_collector.py index c69bffcb..7afad0cd 100644 --- a/compiler_opt/rl/trace_data_collector.py +++ b/compiler_opt/rl/trace_data_collector.py @@ -24,7 +24,8 @@ import multiprocessing -def compile_module(module_path, corpus_path, clang_path, tflite_dir, output_path): +def compile_module(module_path, corpus_path, clang_path, tflite_dir, + output_path): module_full_input_path = os.path.join(corpus_path, module_path) + '.bc' module_full_output_path = os.path.join(output_path, module_path) + '.bc.o' module_command_full_path = os.path.join(corpus_path, module_path) + '.cmd' @@ -36,20 +37,21 @@ def compile_module(module_path, corpus_path, clang_path, tflite_dir, output_path command_vector = [clang_path] command_vector.extend(module_command_line) - command_vector.extend( - [module_full_input_path, '-o', module_full_output_path]) + command_vector.extend([module_full_input_path, '-o', module_full_output_path]) if tflite_dir is not None: command_vector.extend(['-mllvm', '-regalloc-enable-advisor=development']) command_vector.extend(['-mllvm', '-regalloc-model=' + tflite_dir]) subprocess.run(command_vector, check=True) - logging.info( - f'Just finished compiling {module_full_output_path}' - ) + logging.info(f'Just finished compiling {module_full_output_path}') -def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None, single_threaded=False): +def compile_corpus(corpus_path, + output_path, + clang_path, + tflite_dir=None, + thread_count=multiprocessing.cpu_count()): with open( os.path.join(corpus_path, 'corpus_description.json'), encoding='utf-8') as corpus_description_handle: @@ -61,12 +63,14 @@ def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None, single # Compile each module. to_compile.append(module_path) - if single_threaded: - for module_to_compile in to_compile: - compile_module(module_to_compile, corpus_path, clang_path, tflite_dir, output_path) - else: - with multiprocessing.Pool() as pool: - pool.map(functools.partial(compile_module, corpus_path=corpus_path, clang_path=clang_path, tflite_dir=tflite_dir, output_path=output_path), to_compile) + with multiprocessing.Pool(thread_count) as pool: + pool.map( + functools.partial( + compile_module, + corpus_path=corpus_path, + clang_path=clang_path, + tflite_dir=tflite_dir, + output_path=output_path), to_compile) shutil.copy( os.path.join(corpus_path, 'corpus_description.json'), @@ -74,18 +78,25 @@ def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None, single def evaluate_compiled_corpus(compiled_corpus_path, trace_path, - bb_trace_model_path): - corpus_description_path = os.path.join(compiled_corpus_path, 'corpus_description.json') + function_index_path, bb_trace_model_path, thread_count=multiprocessing.cpu_count()): + corpus_description_path = os.path.join(compiled_corpus_path, + 'corpus_description.json') command_vector = [ bb_trace_model_path, '--bb_trace_path=' + trace_path, - '--corpus_path=' + corpus_description_path, '--cpu_name=skylake-avx512' + '--corpus_path=' + corpus_description_path, '--cpu_name=skylake-avx512', + '--function_index_path=' + function_index_path, '--thread_count=' + str(thread_count) ] process_return = subprocess.run( - command_vector, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT) + command_vector, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = process_return.stdout.decode('utf-8') - return int(output) + total_cost = 0.0 + + for line in output.split('\n'): + if line == '': + continue + total_cost += float(line) + + return total_cost