From 2982d394f9ecc10b420ecae1514fd99f9fa1f453 Mon Sep 17 00:00:00 2001 From: Aiden Grossman Date: Thu, 22 Aug 2024 16:30:37 +0000 Subject: [PATCH] Hopefully helpful modifications --- compiler_opt/es/es_trainer_lib.py | 12 ++--- compiler_opt/rl/trace_data_collector.py | 67 +++++++++++++++---------- 2 files changed, 47 insertions(+), 32 deletions(-) diff --git a/compiler_opt/es/es_trainer_lib.py b/compiler_opt/es/es_trainer_lib.py index 7ec14657..95d52675 100644 --- a/compiler_opt/es/es_trainer_lib.py +++ b/compiler_opt/es/es_trainer_lib.py @@ -76,10 +76,10 @@ def __init__(self, *, all_gin): self._template_dir = tempfile.mkdtemp() saver.save(self._template_dir) - self._corpus_dir = '/usr/local/google/home/aidengrossman/programming/test_traces/corpus_subset' - self._clang_path = '/usr/local/google/home/aidengrossman/programming/test_traces/clang' - self._trace_path = '/usr/local/google/home/aidengrossman/programming/test_traces/execution_trace.pb' - self._bb_trace_model_path = '/usr/local/google/home/aidengrossman/programming/test_traces/basic_block_trace_model' + 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/' def es_compile(self, params: list[float], baseline_score: float) -> float: @@ -98,7 +98,7 @@ def es_compile(self, params: list[float], baseline_score: float) -> float: os.path.join(tflitedir, policy_saver.OUTPUT_SIGNATURE)) trace_data_collector.compile_corpus(self._corpus_dir, tempdir, - self._clang_path, tflitedir) + self._clang_path, tflitedir, single_threaded=True) score = trace_data_collector.evaluate_compiled_corpus( tempdir, self._trace_path, self._bb_trace_model_path) @@ -106,7 +106,7 @@ def es_compile(self, params: list[float], baseline_score: float) -> float: print(reward) output_path = os.path.join(self._models_for_test_path, "model" + str(reward)) - if reward > 0.2 and not os.path.exists(output_path): + if reward > 0 and not os.path.exists(output_path): shutil.copytree(tflitedir, output_path) return compilation_runner._calculate_reward(score, baseline_score) diff --git a/compiler_opt/rl/trace_data_collector.py b/compiler_opt/rl/trace_data_collector.py index 6b3185d0..c69bffcb 100644 --- a/compiler_opt/rl/trace_data_collector.py +++ b/compiler_opt/rl/trace_data_collector.py @@ -17,42 +17,56 @@ from absl import logging import os import subprocess +import functools import json import pathlib import shutil +import multiprocessing -def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None): +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' + pathlib.Path(os.path.dirname(module_full_output_path)).mkdir( + parents=True, exist_ok=True) + with open(module_command_full_path) as module_command_handle: + module_command_line = tuple(module_command_handle.read().replace( + r'{', r'{{').replace(r'}', r'}}').split('\0')) + + command_vector = [clang_path] + command_vector.extend(module_command_line) + 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}' + ) + + +def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None, single_threaded=False): with open( os.path.join(corpus_path, 'corpus_description.json'), encoding='utf-8') as corpus_description_handle: corpus_description = json.load(corpus_description_handle) + to_compile = [] + for module_index, module_path in enumerate(corpus_description['modules']): # Compile each module. - 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' - pathlib.Path(os.path.dirname(module_full_output_path)).mkdir( - parents=True, exist_ok=True) - - module_command_full_path = os.path.join(corpus_path, module_path) + '.cmd' - with open(module_command_full_path) as module_command_handle: - module_command_line = tuple(module_command_handle.read().replace( - r'{', r'{{').replace(r'}', r'}}').split('\0')) - - command_vector = [clang_path] - command_vector.extend(module_command_line) - 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} ({module_index + 1}/{len(corpus_description["modules"])})' - ) + 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) shutil.copy( os.path.join(corpus_path, 'corpus_description.json'), @@ -61,16 +75,17 @@ def compile_corpus(corpus_path, output_path, clang_path, tflite_dir=None): 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') command_vector = [ bb_trace_model_path, '--bb_trace_path=' + trace_path, - '--corpus_path=' + compiled_corpus_path + '--corpus_path=' + corpus_description_path, '--cpu_name=skylake-avx512' ] process_return = subprocess.run( command_vector, - check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + output = process_return.stdout.decode('utf-8') return int(output)