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,
},
)
