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(
id=env_id,
entry_point="compiler_gym.envs:CompilerEnv",
kwargs={
"service": _EXAMPLE_SERVICE_BINARY,
"observation_space": observation_space,
"reward_space": reward_space,
},
)
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 _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={
"service": LLVM_SERVICE_BINARY,
},
)
for reward_space in reward_spaces:
register(
id=f"llvm-{reward_space}-v0",
entry_point="compiler_gym.envs.llvm:LlvmEnv",
kwargs={
"service": LLVM_SERVICE_BINARY,
"reward_space": reward_spaces[reward_space],
},
)
for observation_space, reward_space in product(observation_spaces,
reward_spaces):
register(
id=f"llvm-{observation_space}-{reward_space}-v0",
entry_point="compiler_gym.envs.llvm:LlvmEnv",
kwargs={
"service": LLVM_SERVICE_BINARY,
"observation_space": observation_spaces[observation_space],
"reward_space": reward_spaces[reward_space],
},
)
return Benchmark(proto=benchmark.BenchmarkProto(uri=str(uri)))
class LoopToolCPUDataset(Dataset):
def __init__(self, *args, **kwargs):
super().__init__(
name="benchmark://loop_tool-cpu-v0",
license="MIT",
description="loop_tool dataset",
)
def benchmark_uris(self) -> Iterable[str]:
return (f"loop_tool-cpu-v0/{i}" for i in range(1, 1024 * 1024 * 8))
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
return Benchmark(proto=benchmark.BenchmarkProto(uri=str(uri)))
register(
id="loop_tool-v0",
entry_point="compiler_gym.envs.loop_tool.loop_tool_env:LoopToolEnv",
kwargs={
"datasets": [LoopToolCPUDataset(),
LoopToolCUDADataset()],
"observation_space": "action_state",
"reward_space": "flops",
"rewards": [FLOPSReward()],
"service": LOOP_TOOL_SERVICE_BINARY,
},
)
yield from (f"benchmark://example-v0{k}"
for k in self._benchmarks.keys())
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
if uri.path in self._benchmarks:
return self._benchmarks[uri.path]
else:
raise LookupError("Unknown program name")
# Register the environment for use with gym.make(...).
register(
id="example-v0",
entry_point="compiler_gym.envs:CompilerEnv",
kwargs={
"service": EXAMPLE_PY_SERVICE_BINARY,
"rewards": [RuntimeReward()],
"datasets": [ExampleDataset()],
},
)
def main():
# Use debug verbosity to print out extra logging information.
init_logging(level=logging.DEBUG)
# Create the environment using the regular gym.make(...) interface.
with gym.make("example-v0") as env:
env.reset()
for _ in range(20):
observation, reward, done, info = env.step(
# 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"
),
},
)
]
for directory in get_system_includes():
cmd += ["-isystem", str(directory)]
return subprocess.check_output(
cmd,
timeout=300,
)
def benchmark_uris(self) -> Iterable[str]:
yield from self._benchmarks.keys()
def benchmark(self, uri: str) -> Benchmark:
if uri in self._benchmarks:
return self._benchmarks[uri]
else:
raise LookupError("Unknown program name")
# Register the unrolling example service on module import. After importing this module,
# the loops-opt-py-v0 environment will be available to gym.make(...).
register(
id="loops-opt-py-v0",
entry_point="compiler_gym.envs:CompilerEnv",
kwargs={
"service": LOOPS_OPT_PY_SERVICE_BINARY,
"rewards": [RuntimeReward(), SizeReward()],
"datasets": [LoopsDataset()],
},
)
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
if uri.path in self._benchmarks:
return self._benchmarks[uri.path]
else:
raise LookupError("Unknown program name")
# Register the unrolling example service on module import. After importing this module,
# the unrolling-py-v0 environment will be available to gym.make(...).
register(
id="unrolling-py-v0",
entry_point=
"compiler_gym.service.client_service_compiler_env:ClientServiceCompilerEnv",
kwargs={
"service": UNROLLING_PY_SERVICE_BINARY,
"rewards": [RuntimeReward(), SizeReward()],
"datasets": [UnrollingDataset()],
},
)
with compiler_gym.make(
"unrolling-py-v0",
benchmark="unrolling-v0/offsets1",
observation_space="features",
reward_space="runtime",
) as env:
compiler_gym.set_debug_level(4) # TODO: check why this has no effect
observation = env.reset()
print("observation: ", observation)
for k in self._benchmarks.keys())
def benchmark_from_parsed_uri(self, uri: BenchmarkUri) -> Benchmark:
if uri.path in self._benchmarks:
return self._benchmarks[uri.path]
else:
raise LookupError("Unknown program name")
# Register the example service on module import. After importing this module,
# the example-v0 environment will be available to gym.make(...).
register(
id="example-cc-v0",
entry_point=
"compiler_gym.service.client_service_compiler_env:ClientServiceCompilerEnv",
kwargs={
"service": EXAMPLE_CC_SERVICE_BINARY,
"rewards": [RuntimeReward()],
"datasets": [ExampleDataset()],
},
)
register(
id="example-py-v0",
entry_point=
"compiler_gym.service.client_service_compiler_env:ClientServiceCompilerEnv",
kwargs={
"service": EXAMPLE_PY_SERVICE_BINARY,
"rewards": [RuntimeReward()],
"datasets": [ExampleDataset()],
},
)
# 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"]
