from pathlib import Path from typing import Iterable, Optional from compiler_gym.datasets import Benchmark, BenchmarkSource, Dataset from compiler_gym.datasets.benchmark import BenchmarkWithSource from compiler_gym.service.proto import BenchmarkDynamicConfig, Command from compiler_gym.util.decorators import memoized_property from compiler_gym.util.runfiles_path import runfiles_path, site_data_path from compiler_gym.util.shell_format import plural logger = logging.getLogger(__name__) # The maximum value for the --seed argument to matmul. UINT_MAX = (2**32) - 1 _matmul_bin = runfiles_path( "compiler_gym/third_party/matmul/matmul/bin/matmul") _matmul_includes = runfiles_path( "compiler_gym/third_party/matmul/matmul/include/matmul-2.3.0") matmul_sizes = [(64, 64, 64)] class MatmulBenchmark(BenchmarkWithSource): """A matmul benchmark.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._src = None self.proto.dynamic_config.MergeFrom( BenchmarkDynamicConfig( build_cmd=Command( argument=["$CC", "$IN"],
# LICENSE file in the root directory of this source tree. """This module defines a utility function for computing LLVM observations.""" import subprocess import sys from pathlib import Path from typing import List import google.protobuf.text_format from compiler_gym.service.proto import Observation from compiler_gym.util.gym_type_hints import ObservationType from compiler_gym.util.runfiles_path import runfiles_path from compiler_gym.util.shell_format import plural from compiler_gym.views.observation_space_spec import ObservationSpaceSpec _COMPUTE_OBSERVATION_BIN = runfiles_path( "compiler_gym/envs/llvm/service/compute_observation") def pascal_case_to_enum(pascal_case: str) -> str: """Convert PascalCase to ENUM_CASE.""" word_arrays: List[List[str]] = [[]] for c in pascal_case: if c.isupper() and word_arrays[-1]: word_arrays.append([c]) else: word_arrays[-1].append(c.upper()) return "_".join(["".join(word) for word in word_arrays])
def test_deep_dataflow_file_count(subset: str): bitcode_dir = runfiles_path( "compiler_gym/third_party/DeepDataFlow") / subset num_files = len( [f for f in bitcode_dir.iterdir() if f.name.endswith(".bc")]) assert num_files == EXPECTED_NUMBER_OF_BITCODE_FILES[subset]
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Register the MLIR environments.""" from compiler_gym.envs.mlir.mlir_env import MlirEnv from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path __all__ = [ "MLIR_SERVICE_BINARY", "MlirEnv", ] MLIR_SERVICE_BINARY = runfiles_path( "compiler_gym/envs/mlir/service/compiler_gym-mlir-service") def _register_mlir_gym_service(): """Register an environment for each combination of MLIR observation/reward/benchmark.""" register( id="mlir-v0", entry_point="compiler_gym.envs.mlir:MlirEnv", kwargs={ "service": MLIR_SERVICE_BINARY, }, )
def __init__( self, local_service_binary: Path, port_init_max_seconds: float, rpc_init_max_seconds: float, process_exit_max_seconds: float, script_args: List[str], script_env: Dict[str, str], ): """Constructor. :param local_service_binary: The path of the service binary. :raises TimeoutError: If fails to establish connection within a specified time limit. """ self.process_exit_max_seconds = process_exit_max_seconds if not Path(local_service_binary).is_file(): raise FileNotFoundError(f"File not found: {local_service_binary}") self.cache = ServiceCache() # The command that will be executed. The working directory of this # command will be set to the local_service_binary's parent, so we can # use the relpath for a neater `ps aux` view. cmd = [ f"./{local_service_binary.name}", f"--working_dir={self.cache.path}", ] # Add any custom arguments cmd += script_args # Set the root of the runfiles directory. env = os.environ.copy() env["COMPILER_GYM_RUNFILES"] = str(runfiles_path(".")) env["COMPILER_GYM_SITE_DATA"] = str(site_data_path(".")) # Set the pythonpath so that executable python scripts can use absolute # import paths like `from compiler_gym.envs.foo import bar`. if "PYTHONPATH" in env: env["PYTHONPATH"] = f'{env["PYTHONPATH"]}:{env["COMPILER_GYM_RUNFILES"]}' else: env["PYTHONPATH"] = env["COMPILER_GYM_RUNFILES"] # Set the verbosity of the service. The logging level of the service is # the debug level - 1, so that COMPILER_GYM_DEBUG=3 will cause VLOG(2) # and lower to be logged to stdout. debug_level = max( get_debug_level(), logging_level_to_debug_level(logger.getEffectiveLevel())) if debug_level > 0: cmd.append("--alsologtostderr") cmd.append(f"-v={debug_level - 1}") # If we are debugging the backend, set the logbuflevel to a low # value to disable buffering of logging messages. This removes any # buffering between `LOG(INFO) << "..."` and the message being # emited to stderr. cmd.append("--logbuflevel=-1") else: # Silence the gRPC logs as we will do our own error reporting, but # don't override any existing value so that the user may debug the # gRPC backend by setting GRPC_VERBOSITY to ERROR, INFO, or DEBUG. if not os.environ.get("GRPC_VERBOSITY"): env["GRPC_VERBOSITY"] = "NONE" # Set environment variable COMPILER_GYM_SERVICE_ARGS to pass # additional arguments to the service. args = os.environ.get("COMPILER_GYM_SERVICE_ARGS", "") if args: cmd.append(args) # Add any custom environment variables env.update(script_env) logger.debug( "Exec `%s%s`", " ".join(f"{k}={v}" for k, v in script_env.items()) + " " if script_env else "", join_cmd(cmd), ) self.process = subprocess.Popen( cmd, env=env, cwd=local_service_binary.parent, ) self._process_returncode_exception_raised = False # Read the port from a file generated by the service. wait_secs = 0.1 port_path = self.cache / "port.txt" end_time = time() + port_init_max_seconds while time() < end_time: returncode = self.process.poll() if returncode is not None: try: # Try and decode the name of a signal. Signal returncodes # are negative. returncode = f"{returncode} ({Signals(abs(returncode)).name})" except ValueError: pass msg = f"Service terminated with returncode: {returncode}" # Attach any logs from the service if available. logs = truncate_lines(self.loglines(), max_line_len=100, max_lines=25, tail=True) if logs: msg = f"{msg}\nService logs:\n{logs}" self.cache.close() raise ServiceError(msg) if port_path.is_file(): try: with open(port_path) as f: self.port = int(f.read().rstrip()) break except ValueError: # ValueError is raised by int(...) on invalid input. In that # case, wait for longer. pass sleep(wait_secs) wait_secs *= 1.2 else: # kill() was added in Python 3.7. if sys.version_info >= (3, 7, 0): self.process.kill() else: self.process.terminate() self.process.communicate(timeout=rpc_init_max_seconds) self.cache.close() raise TimeoutError("Service failed to produce port file after " f"{port_init_max_seconds:.1f} seconds") url = f"localhost:{self.port}" wait_secs = 0.1 attempts = 0 end_time = time() + rpc_init_max_seconds while time() < end_time: try: channel = grpc.insecure_channel( url, options=GRPC_CHANNEL_OPTIONS, ) channel_ready = grpc.channel_ready_future(channel) attempts += 1 channel_ready.result(timeout=wait_secs) break except (grpc.FutureTimeoutError, grpc.RpcError) as e: logger.debug("Connection attempt %d = %s %s", attempts, type(e).__name__, str(e)) wait_secs *= 1.2 else: # kill() was added in Python 3.7. if sys.version_info >= (3, 7, 0): self.process.kill() else: self.process.terminate() self.process.communicate(timeout=process_exit_max_seconds) # Include the last few lines of logs generated by the compiler # service, if any. logs = truncate_lines(self.loglines(), max_line_len=100, max_lines=25, tail=True) logs_message = f" Service logs:\n{logs}" if logs else "" self.cache.close() raise TimeoutError( "Failed to connect to RPC service after " f"{rpc_init_max_seconds:.1f} seconds.{logs_message}") super().__init__(channel, url)
from compiler_gym.datasets import Benchmark from compiler_gym.envs import LlvmEnv, llvm from compiler_gym.errors import BenchmarkInitError from compiler_gym.service.proto import Benchmark as BenchmarkProto from compiler_gym.service.proto import File from compiler_gym.third_party import llvm as llvm_paths from compiler_gym.util.runfiles_path import runfiles_path from compiler_gym.util.temporary_working_directory import temporary_working_directory from tests.pytest_plugins.common import bazel_only from tests.test_main import main pytest_plugins = ["tests.pytest_plugins.llvm"] # The path of an IR file that assembles but does not compile. INVALID_IR_PATH = runfiles_path("tests/llvm/invalid_ir.ll") EXAMPLE_BITCODE_FILE = runfiles_path( "compiler_gym/third_party/cbench/cbench-v1/crc32.bc") EXAMPLE_BITCODE_IR_INSTRUCTION_COUNT = 242 def test_reset_invalid_benchmark(env: LlvmEnv): invalid_benchmark = "an invalid benchmark" with pytest.raises( LookupError, match=f"Dataset not found: benchmark://{invalid_benchmark}"): env.reset(benchmark=invalid_benchmark) def test_invalid_benchmark_data(env: LlvmEnv): benchmark = Benchmark.from_file_contents("benchmark://new",
# # $ make_cBench_llvm_module.py <in_dir> <outpath> [<cflag>...] # # This compiles the code from <in_dir> and generates an LLVM bitcode module at # the given <outpath>, using any additional <cflags> as clang arguments. import subprocess import sys import tempfile from pathlib import Path from typing import List from compiler_gym.util.runfiles_path import runfiles_path # Path of the LLVM binaries. CLANG = Path(runfiles_path("llvm/10.0.0/clang")) LLVM_LINK = Path(runfiles_path("llvm/10.0.0/llvm-link")) def make_cbench_llvm_module(benchmark_dir: Path, cflags: List[str], output_path: Path) -> str: """Compile a cBench benchmark into an unoptimized LLVM bitcode file.""" cflags = cflags or [] src_dir = benchmark_dir / "src" if not src_dir.is_dir(): src_dir = benchmark_dir assert src_dir.is_dir(), f"Source directory not found: {src_dir}" clang_command = [ str(CLANG),
import sys import tarfile import tempfile from collections import Counter, defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from typing import Callable, Dict, List, NamedTuple, Optional import fasteners from compiler_gym.datasets.dataset import Dataset from compiler_gym.util.download import download from compiler_gym.util.runfiles_path import cache_path, runfiles_path, site_data_path from compiler_gym.util.timer import Timer _CLANG = runfiles_path("CompilerGym/compiler_gym/third_party/llvm/clang") _CBENCH_DATA = site_data_path("llvm/cBench-v0-runtime-data/runtime_data") _CBENCH_DATA_URL = ( "https://dl.fbaipublicfiles.com/compiler_gym/cBench-v0-runtime-data.tar.bz2" ) _CBENCH_DATA_SHA256 = "a1b5b5d6b115e5809ccaefc2134434494271d184da67e2ee43d7f84d07329055" if sys.platform == "darwin": _COMPILE_ARGS = [ "-L", "/Library/Developer/CommandLineTools/SDKs/MacOSX.sdk/usr/lib", ] else: _COMPILE_ARGS = []
def test_expected_version(): """Test that embedded compiler gym version matches VERSION file.""" with open(runfiles_path("VERSION")) as f: version = f.read().strip() assert version == compiler_gym.__version__
# # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Tests for LlvmEnv.fork().""" import subprocess import sys import pytest from compiler_gym.envs import LlvmEnv from compiler_gym.util.runfiles_path import runfiles_path from tests.test_main import main pytest_plugins = ["tests.pytest_plugins.llvm"] EXAMPLE_BITCODE_FILE = runfiles_path( "compiler_gym/third_party/cbench/cbench-v1/crc32.bc") EXAMPLE_BITCODE_IR_INSTRUCTION_COUNT = 196 def test_with_statement(env: LlvmEnv): """Test that the `with` statement context manager works on forks.""" env.reset("cbench-v1/crc32") env.step(0) with env.fork() as fkd: assert fkd.in_episode assert fkd.actions == [0] assert not fkd.in_episode assert env.in_episode def test_fork_child_process_is_not_orphaned(env: LlvmEnv):
# $ make_cBench_llvm_module.py <in_dir> <outpath> [<cflag>...] # # This compiles the code from <in_dir> and generates an LLVM bitcode module at # the given <outpath>, using any additional <cflags> as clang arguments. import subprocess import sys import tempfile from pathlib import Path from typing import List from compiler_gym.envs.llvm.benchmarks import get_system_includes from compiler_gym.util.runfiles_path import runfiles_path # Path of the LLVM binaries. CLANG = Path(runfiles_path("compiler_gym/third_party/llvm/clang")) LLVM_LINK = Path(runfiles_path("compiler_gym/third_party/llvm/llvm-link")) def make_cbench_llvm_module( benchmark_dir: Path, cflags: List[str], output_path: Path ) -> str: """Compile a cBench benchmark into an unoptimized LLVM bitcode file.""" cflags = cflags or [] src_dir = benchmark_dir / "src" if not src_dir.is_dir(): src_dir = benchmark_dir assert src_dir.is_dir(), f"Source directory not found: {src_dir}" clang_command = [
# This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """This module defines and registers the example gym environments.""" import subprocess from pathlib import Path from typing import Iterable from compiler_gym.datasets import Benchmark, Dataset from compiler_gym.envs.llvm.llvm_benchmark import get_system_includes from compiler_gym.spaces import Reward from compiler_gym.third_party import llvm from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path, site_data_path UNROLLING_PY_SERVICE_BINARY: Path = runfiles_path( "examples/example_unrolling_service/service_py/example-unrolling-service-py" ) BENCHMARKS_PATH: Path = runfiles_path( "examples/example_unrolling_service/benchmarks") NEURO_VECTORIZER_HEADER: Path = runfiles_path( "compiler_gym/third_party/neuro-vectorizer/header.h") class RuntimeReward(Reward): """An example reward that uses changes in the "runtime" observation value to compute incremental reward. """ def __init__(self): super().__init__(
"""Pytest fixtures for the LLVM CompilerGym environments.""" import os from pathlib import Path from typing import Iterable, List import gym import pytest from compiler_gym.datasets import Dataset from compiler_gym.envs.llvm import LlvmEnv from compiler_gym.envs.llvm.datasets.cbench import VALIDATORS from compiler_gym.third_party import llvm from compiler_gym.util.runfiles_path import runfiles_path BENCHMARKS_LIST = Path( runfiles_path("compiler_gym/third_party/cbench/benchmarks.txt")) def _read_list_file(path: Path) -> Iterable[str]: with open(str(path)) as f: for action in f: if action.strip(): yield action.strip() BENCHMARK_NAMES = list(_read_list_file(BENCHMARKS_LIST)) # Skip ghostscript on CI as it is just too heavy. if bool(os.environ.get("CI")): BENCHMARK_NAMES = [ b for b in BENCHMARK_NAMES if b != "benchmark://cbench-v1/ghostscript"
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """This module demonstrates how to """ from pathlib import Path from compiler_gym.envs.gcc.gcc import Gcc, GccSpec, Option from compiler_gym.envs.gcc.gcc_env import DEFAULT_GCC, GccEnv from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path GCC_SERVICE_BINARY: Path = runfiles_path( "compiler_gym/envs/gcc/service/compiler_gym-gcc-service") register( id="gcc-v0", entry_point="compiler_gym.envs.gcc:GccEnv", kwargs={"service": GCC_SERVICE_BINARY}, ) __all__ = ["GccEnv", "GccSpec", "Gcc", "Option", "DEFAULT_GCC"]
# # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Tests for LLVM benchmark handling.""" import pytest from compiler_gym.datasets.benchmark import BenchmarkInitError from compiler_gym.envs import llvm from compiler_gym.envs.llvm import LlvmEnv from compiler_gym.util.runfiles_path import runfiles_path from tests.pytest_plugins.common import bazel_only from tests.test_main import main pytest_plugins = ["tests.pytest_plugins.llvm"] INVALID_IR_PATH = runfiles_path("tests/llvm/invalid_ir.ll") @bazel_only # invalid_ir.ll not installed def test_reset_invalid_ir(env: LlvmEnv): """Test that setting the $CXX to an invalid binary raises an error.""" benchmark = llvm.make_benchmark(INVALID_IR_PATH) with pytest.raises(BenchmarkInitError, match="Failed to compute .text size cost"): env.reset(benchmark=benchmark) if __name__ == "__main__": main()
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from itertools import product from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path _EXAMPLE_SERVICE_BINARY = runfiles_path( "CompilerGym/examples/example_compiler_gym_service/service/service") def _register_example_gym_service(): """Register environments for the example service.""" register( id="example-v0", entry_point="compiler_gym.envs:CompilerEnv", kwargs={ "service": _EXAMPLE_SERVICE_BINARY, }, ) # Register rewards for all combinations of eager observation and # reward spaces. observation_spaces = ["ir", "features"] reward_spaces = ["codesize"] configurations = product(observation_spaces, reward_spaces) for observation_space, reward_space in configurations: env_id = f"example-{observation_space}-{reward_space}-v0" register(
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """This module defines and registers the example gym environments.""" from pathlib import Path from typing import Iterable from compiler_gym.datasets import Benchmark, Dataset from compiler_gym.spaces import Reward from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path EXAMPLE_CC_SERVICE_BINARY: Path = runfiles_path( "examples/example_compiler_gym_service/service_cc/compiler_gym-example-service-cc" ) EXAMPLE_PY_SERVICE_BINARY: Path = runfiles_path( "examples/example_compiler_gym_service/service_py/compiler_gym-example-service-py" ) class RuntimeReward(Reward): """An example reward that uses changes in the "runtime" observation value to compute incremental reward. """ def __init__(self): super().__init__( id="runtime", observation_spaces=["runtime"], default_value=0,
# This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """This module registers the Loop Optimizations CompilerGym environment """ import subprocess from pathlib import Path from typing import Iterable from compiler_gym.datasets import Benchmark, Dataset from compiler_gym.envs.llvm.llvm_benchmark import get_system_includes from compiler_gym.spaces import Reward from compiler_gym.third_party import llvm from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path LOOPS_OPT_PY_SERVICE_BINARY: Path = runfiles_path( "examples/loop_optimizations_service/service_py/loops-opt-service-py" ) BENCHMARKS_PATH: Path = runfiles_path("examples/loop_optimizations_service/benchmarks") NEURO_VECTORIZER_HEADER: Path = runfiles_path( "compiler_gym/third_party/neuro-vectorizer/header.h" ) class RuntimeReward(Reward): """An example reward that uses changes in the "runtime" observation value to compute incremental reward. """ def __init__(self):
from pathlib import Path from typing import Iterable, List, Optional, Union, cast import numpy as np from gym.spaces import Dict as DictSpace from compiler_gym.envs.compiler_env import CompilerEnv, step_t from compiler_gym.envs.llvm.benchmarks import make_benchmark from compiler_gym.envs.llvm.datasets import LLVM_DATASETS from compiler_gym.spaces import Commandline, CommandlineFlag, Scalar, Sequence from compiler_gym.third_party.autophase import AUTOPHASE_FEATURE_NAMES from compiler_gym.third_party.inst2vec import Inst2vecEncoder from compiler_gym.util.runfiles_path import runfiles_path, site_data_path _ACTIONS_LIST = Path( runfiles_path("compiler_gym/envs/llvm/service/passes/actions_list.txt")) _FLAGS_LIST = Path( runfiles_path("compiler_gym/envs/llvm/service/passes/actions_flags.txt")) _DESCRIPTIONS_LIST = Path( runfiles_path( "compiler_gym/envs/llvm/service/passes/actions_descriptions.txt")) def _read_list_file(path: Path) -> Iterable[str]: with open(str(path)) as f: for action in f: if action.strip(): yield action.strip()
from compiler_gym.datasets import Benchmark, BenchmarkSource, Dataset from compiler_gym.datasets.benchmark import BenchmarkWithSource from compiler_gym.datasets.uri import BenchmarkUri from compiler_gym.envs.gcc.gcc import Gcc from compiler_gym.util.commands import Popen from compiler_gym.util.decorators import memoized_property from compiler_gym.util.runfiles_path import runfiles_path from compiler_gym.util.shell_format import plural from compiler_gym.util.truncate import truncate logger = logging.getLogger(__name__) # The maximum value for the --seed argument to csmith. UINT_MAX = (2 ** 32) - 1 _CSMITH_BIN = runfiles_path("compiler_gym/third_party/csmith/csmith/bin/csmith") _CSMITH_INCLUDES = runfiles_path( "compiler_gym/third_party/csmith/csmith/include/csmith-2.3.0" ) _CSMITH_INSTALL_LOCK = Lock() # TODO(github.com/facebookresearch/CompilerGym/issues/325): This can be merged # with the LLVM implementation. class CsmithBenchmark(BenchmarkWithSource): """A CSmith benchmark.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._src = None
# This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Tests for LLVM benchmark handling.""" import tempfile from pathlib import Path import pytest from compiler_gym.envs import CompilerEnv from compiler_gym.service.proto import Benchmark, File from compiler_gym.util.runfiles_path import runfiles_path from tests.test_main import main pytest_plugins = ["tests.envs.llvm.fixtures"] EXAMPLE_BITCODE_FILE = runfiles_path( "CompilerGym/third_party/cBench-v0/crc32.bc") def test_reset_invalid_benchmark(env: CompilerEnv): invalid_benchmark = "an invalid benchmark" with pytest.raises(ValueError) as ctx: env.reset(benchmark=invalid_benchmark) assert str(ctx.value) == f'Unknown benchmark "{invalid_benchmark}"' def test_invalid_benchmark_data(env: CompilerEnv): benchmark = Benchmark( uri="benchmark://new", program=File(contents="Invalid bitcode".encode("utf-8")))
# LICENSE file in the root directory of this source tree. """This module defines a utility function for constructing LLVM benchmarks.""" import multiprocessing import os import random import subprocess import sys import tempfile from datetime import datetime from pathlib import Path from typing import Iterable, List, Optional, Union from compiler_gym.service.proto import Benchmark, File from compiler_gym.util.runfiles_path import cache_path, runfiles_path CLANG = runfiles_path("CompilerGym/compiler_gym/third_party/llvm/clang") LLVM_LINK = runfiles_path("CompilerGym/compiler_gym/third_party/llvm/llvm-link") def _communicate(process, input=None, timeout=None): """subprocess.communicate() which kills subprocess on timeout.""" try: return process.communicate(input=input, timeout=timeout) except subprocess.TimeoutExpired as e: # kill() was added in Python 3.7. if sys.version_info >= (3, 7, 0): process.kill() else: process.terminate() raise
# # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Pytest fixtures for the LLVM CompilerGym environments.""" from pathlib import Path from typing import Iterable, List import gym import pytest from compiler_gym.envs import CompilerEnv, LlvmEnv from compiler_gym.service import CompilerGymServiceConnection from compiler_gym.util.runfiles_path import runfiles_path ACTIONS_LIST = Path( runfiles_path("CompilerGym/compiler_gym/envs/llvm/service/passes/actions_list.txt") ) BENCHMARKS_LIST = Path( runfiles_path("CompilerGym/compiler_gym/third_party/cBench/benchmarks.txt") ) SERVICE_BIN = Path(runfiles_path("CompilerGym/compiler_gym/envs/llvm/service/service")) def _read_list_file(path: Path) -> Iterable[str]: with open(str(path)) as f: for action in f: if action.strip(): yield action.strip()
# # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Pytest fixtures for the LLVM CompilerGym environments.""" from pathlib import Path from typing import Iterable, List import gym import pytest from compiler_gym.envs import CompilerEnv, LlvmEnv from compiler_gym.service import CompilerGymServiceConnection from compiler_gym.util.runfiles_path import runfiles_path ACTIONS_LIST = Path( runfiles_path( "CompilerGym/compiler_gym/envs/llvm/service/passes/actions_flags.txt")) BENCHMARKS_LIST = Path( runfiles_path( "CompilerGym/compiler_gym/third_party/cBench/benchmarks.txt")) def _read_list_file(path: Path) -> Iterable[str]: with open(str(path)) as f: for action in f: if action.strip(): yield action.strip() ACTION_NAMES = list(_read_list_file(ACTIONS_LIST)) BENCHMARK_NAMES = list(_read_list_file(BENCHMARKS_LIST))
from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path __all__ = [ "BenchmarkFromCommandLine", "ClangInvocation", "compute_observation", "get_system_library_flags", "LLVM_SERVICE_BINARY", "LlvmEnv", "make_benchmark", "observation_spaces", "reward_spaces", ] LLVM_SERVICE_BINARY = runfiles_path( "compiler_gym/envs/llvm/service/compiler_gym-llvm-service") def _register_llvm_gym_service(): """Register an environment for each combination of LLVM observation/reward/benchmark.""" observation_spaces = {"autophase": "Autophase", "ir": "Ir"} reward_spaces = { "ic": "IrInstructionCountOz", "codesize": "ObjectTextSizeOz" } register( id="llvm-v0", entry_point="compiler_gym.envs.llvm:LlvmEnv", kwargs={
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """This module demonstrates how to """ from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path # Register the example service on module import. After importing this module, # the example-v0 environment will be available to gym.make(...). register( id="example-v0", entry_point="compiler_gym.envs:CompilerEnv", kwargs={ "service": runfiles_path( "examples/example_compiler_gym_service/service/compiler_gym-example-service" ), }, )
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Register the loop_tool environment and reward.""" from pathlib import Path from typing import Iterable from compiler_gym.datasets import Benchmark, Dataset, benchmark from compiler_gym.datasets.uri import BenchmarkUri from compiler_gym.spaces import Reward from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path LOOP_TOOL_SERVICE_BINARY: Path = runfiles_path( "compiler_gym/envs/loop_tool/service/compiler_gym-loop_tool-service") class FLOPSReward(Reward): """ `loop_tool` uses "floating point operations per second" as its default reward space. """ def __init__(self): super().__init__( name="flops", observation_spaces=["flops"], default_value=0, default_negates_returns=True, deterministic=False, platform_dependent=True,
def __init__( self, local_service_binary: Path, port_init_max_seconds: float, rpc_init_max_seconds: float, process_exit_max_seconds: float, ): """Constructor. :param local_service_binary: The path of the service binary. :raises TimeoutError: If fails to establish connection within a specified time limit. """ self.process_exit_max_seconds = process_exit_max_seconds if not Path(local_service_binary).is_file(): raise FileNotFoundError(f"File not found: {local_service_binary}") self.working_dir = make_working_dir() # Set environment variable COMPILER_GYM_SERVICE_ARGS to pass # additional arguments to the service. args = os.environ.get("COMPILER_GYM_SERVICE_ARGS", "") cmd = [ str(local_service_binary), f"--port=0", f"--log_dir={self.working_dir}/logs", f"--working_dir={self.working_dir}", args, ] # Set the root of the runfiles directory. env = os.environ.copy() env["COMPILER_GYM_RUNFILES"] = str(runfiles_path("CompilerGym")) # Set environment variable COMPILER_GYM_SERVICE_DEBUG=1 to pipe # local service output to stderr. Set COMPILER_GYM_SERVICE_LOG_LEVEL=val # to set log level to <val>, where large values are increasingly # verbose. if os.environ.get("COMPILER_GYM_SERVICE_DEBUG", "0") != "0": cmd.append("--alsologtostderr") log_level = os.environ.get("COMPILER_GYM_SERVICE_LOG_LEVEL", "0") if log_level: cmd.append(f"-v={log_level}") logging.debug(f"$ {' '.join(cmd)}") self.process = subprocess.Popen(cmd, env=env) # Read the port from a file generated by the service. wait_secs = 0.1 port_path = self.working_dir / "port.txt" end_time = time() + port_init_max_seconds while time() < end_time: returncode = self.process.poll() if returncode is not None: shutil.rmtree(self.working_dir) raise ServiceError( f"Service terminated with returncode: {returncode}") if port_path.is_file(): with open(port_path) as f: self.port = int(f.read().rstrip()) break sleep(wait_secs) wait_secs *= 1.2 else: self.process.kill() self.process.communicate(timeout=rpc_init_max_seconds) shutil.rmtree(self.working_dir) raise TimeoutError("Service failed to produce port file after " f"{port_init_max_seconds:.1f} seconds") url = f"localhost:{self.port}" wait_secs = 0.1 attempts = 0 end_time = time() + rpc_init_max_seconds while time() < end_time: try: channel = grpc.insecure_channel( url, options=GRPC_CHANNEL_OPTIONS, ) channel_ready = grpc.channel_ready_future(channel) attempts += 1 channel_ready.result(timeout=wait_secs) break except (grpc.FutureTimeoutError, grpc.RpcError) as e: logging.debug(f"Connection attempt {attempts} = {e}") wait_secs *= 1.2 else: self.process.kill() self.process.communicate(timeout=process_exit_max_seconds) shutil.rmtree(self.working_dir) raise TimeoutError("Failed to connect to RPC service after " f"{process_exit_max_seconds:.1f} seconds" f"({attempts} attempts made)") super().__init__(channel, url)
"""This module defines an API for processing LLVM-IR with inst2vec.""" import pickle from typing import List import numpy as np from compiler_gym.third_party.inst2vec import inst2vec_preprocess from compiler_gym.util.runfiles_path import runfiles_path _PICKLED_VOCABULARY = runfiles_path( "compiler_gym/third_party/inst2vec/dictionary.pickle" ) _PICKLED_EMBEDDINGS = runfiles_path( "compiler_gym/third_party/inst2vec/embeddings.pickle" ) class Inst2vecEncoder(object): """An LLVM encoder for inst2vec.""" def __init__(self): with open(str(_PICKLED_VOCABULARY), "rb") as f: self.vocab = pickle.load(f) with open(str(_PICKLED_EMBEDDINGS), "rb") as f: self.embeddings = pickle.load(f) self.unknown_vocab_element = self.vocab["!UNK"] def preprocess(self, ir: str) -> List[str]: """Produce a list of pre-processed statements from an IR."""
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """Register the LLVM environments.""" from itertools import product from compiler_gym.envs.llvm.llvm_env import LlvmEnv from compiler_gym.util.registration import register from compiler_gym.util.runfiles_path import runfiles_path __all__ = ["LlvmEnv"] _LLVM_SERVICE_BINARY = runfiles_path( "CompilerGym/compiler_gym/envs/llvm/service/service") def _register_llvm_gym_service(): """Register an environment for each combination of LLVM observation/reward/benchmark.""" observation_spaces = {"autophase": "Autophase", "ir": "Ir"} reward_spaces = {"ic": "IrInstructionCountOz"} register( id="llvm-v0", entry_point="compiler_gym.envs.llvm:LlvmEnv", kwargs={ "service": _LLVM_SERVICE_BINARY, }, )