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"]