def _init_worker(env_name: str, n: int, outfile: Path):
with open(outfile, "w") as f:
for _ in range(0, n, min(100, n)):
runtimes = get_runtimes(
lambda: compiler_gym.make(env_name).close(), min(100, n)
)
print("\n".join(f"{x:.8f}" for x in runtimes), file=f, flush=True)
def test_induced_remainder(backend):
with compiler_gym.make("loop_tool-v0") as env:
env.observation_space = "loop_tree"
# reset
env.reset(
benchmark=env.datasets.benchmark(
uri=f"benchmark://loop_tool-{backend}-v0/1024"),
action_space="simple",
)
# action toggle_mode
env.step(0)
# action up
env.step(1)
# action toggle_mode
env.step(0)
# action up
env.step(1)
# action up
o = env.step(1)
expected = f"""
for a in 341 r 1 : L0 {'cpu_parallel ' if backend=='cpu' else ''}[thread]
for a' in 3 : L1
for a'' in 1 : L2
%0[a] <- read()
for a'' in 1 : L4
%1[a] <- read()
for a'' in 1 : L6
%2[a] <- add(%0, %1)
for a'' in 1 : L8
%3[a] <- write(%2)
"""
lines = o[0].strip().split("\n")
out = "\n".join(line.rstrip() for line in lines)
assert out == expected.strip(), f"{out} \n vs \n {expected.strip()}"
def test_thread_removal(backend):
with compiler_gym.make("loop_tool-v0") as env:
env.observation_space = "loop_tree"
# reset
env.reset(
benchmark=env.datasets.benchmark(
uri=f"benchmark://loop_tool-{backend}-v0/1024"),
action_space="simple",
)
# action toggle_thread
o = env.step(3)
expected = """
for a in 1024 : L0
for a' in 1 : L1
for a'' in 1 : L2
%0[a] <- read()
for a'' in 1 : L4
%1[a] <- read()
for a'' in 1 : L6
%2[a] <- add(%0, %1)
for a'' in 1 : L8
%3[a] <- write(%2)
"""
lines = o[0].strip().split("\n")
out = "\n".join(line.rstrip() for line in lines)
assert out == expected.strip(), f"{out} \n vs \n {expected.strip()}"
def make_env(self, benchmark: Union[str, Benchmark]) -> LlvmEnv:
env: LlvmEnv = compiler_gym.make("llvm-v0")
# TODO(cummins): This does not work with custom benchmarks, as the URI
# will not be known to the new environment.
if str(benchmark).startswith("file:///"):
benchmark = env.make_benchmark(Path(benchmark[len("file:///"):]))
env.benchmark = benchmark
if self.value == OptimizationTarget.CODESIZE:
env.reward_space = "IrInstructionCountOz"
elif self.value == OptimizationTarget.BINSIZE:
env.reward_space = "ObjectTextSizeOz"
elif self.value == OptimizationTarget.RUNTIME:
env = RuntimePointEstimateReward(env,
warmup_count=0,
runtime_count=3)
else:
assert False, f"Unknown OptimizationTarget: {self.value}"
# Wrap the env to ignore errors during search.
env = JustKeepGoingEnv(env)
return env
def test_benchmarks_uris_list():
cfg = Benchmarks(uris=["benchmark://cbench-v1/qsort"])
assert cfg.uris == ["benchmark://cbench-v1/qsort"]
with compiler_gym.make("llvm-v0") as env:
assert list(cfg.benchmarks_iterator(env)) == ["benchmark://cbench-v1/qsort"]
assert isinstance(list(cfg.benchmarks_iterator(env))[0], Benchmark)
assert list(cfg.benchmark_uris_iterator(env)) == ["benchmark://cbench-v1/qsort"]
def final_reward(self, env: LlvmEnv, runtime_count: int = 30) -> float:
"""Compute the final reward of the environment.
Note that this may modify the environment state. You should call
:code:`reset()` before continuing to use the environment after this.
"""
# Reapply the environment state in a retry loop.
actions = list(env.actions)
env.reset()
for i in range(1, 5 + 1):
_, _, done, info = env.step(actions)
if not done:
break
logger.warning(
"Attempt %d to apply actions during final reward failed: %s",
i,
info.get("error_details"),
)
else:
raise ValueError("Failed to replay environment's actions")
if self.value == OptimizationTarget.CODESIZE:
return env.observation.IrInstructionCountOz() / max(
env.observation.IrInstructionCount(), 1)
if self.value == OptimizationTarget.BINSIZE:
return env.observation.ObjectTextSizeOz() / max(
env.observation.ObjectTextSizeBytes(), 1)
if self.value == OptimizationTarget.RUNTIME:
with _RUNTIME_LOCK:
with compiler_gym.make("llvm-v0",
benchmark=env.benchmark) as new_env:
new_env.reset()
new_env.runtime_observation_count = runtime_count
new_env.runtime_warmup_count = 0
new_env.apply(env.state)
final_runtimes = new_env.observation.Runtime()
assert len(final_runtimes) == runtime_count
new_env.reset()
new_env.send_param("llvm.apply_baseline_optimizations",
"-O3")
o3_runtimes = new_env.observation.Runtime()
assert len(o3_runtimes) == runtime_count
logger.debug("O3 runtimes: %s", o3_runtimes)
logger.debug("Final runtimes: %s", final_runtimes)
speedup = np.median(o3_runtimes) / max(
np.median(final_runtimes), 1e-12)
logger.debug("Speedup: %.4f", speedup)
return speedup
assert False, f"Unknown OptimizationTarget: {self.value}"
def test_testing_config():
cfg = Testing(**OmegaConf.create("""\
timeout_hours: 12
runs_per_benchmark: 6
benchmarks:
- dataset: benchmark://cbench-v1
max_benchmarks: 5
"""))
assert cfg.timeout_hours == 12
with compiler_gym.make("llvm-v0") as env:
assert len(list(cfg.benchmark_uris_iterator(env))) == 5 * 6
def test_run_random_walk_smoke_test():
FLAGS.unparse_flags()
FLAGS(["argv0"])
with capture_output() as out:
with compiler_gym.make("llvm-autophase-ic-v0") as env:
env.benchmark = "cbench-v1/crc32"
run_random_walk(env=env, step_count=5)
print(out.stdout)
# Note the ".*" before and after the step count to ignore the shell
# formatting.
assert re.search(r"Completed .*5.* steps in ", out.stdout)
def make_env(self) -> CompilerEnv:
"""Construct a compiler environment from the given config."""
env = compiler_gym.make(self.id)
if self.observation_space:
env.observation_space = self.observation_space
if self.reward_space:
env.reward_space = self.reward_space
env = TimeLimit(env, max_episode_steps=self.max_episode_steps)
for wrapper in self.wrappers:
env = wrapper.wrap(env)
# Wrap the env to ignore errors during search.
env = JustKeepGoingEnv(env)
return env
def run_one_sweep(
device: str,
k: int,
vectorize: int = 1,
linear: bool = False,
logdir: Optional[Path] = None,
):
"""Run a single sweep."""
logdir = logdir or create_user_logs_dir("loop_tool_sweep")
logfile = logdir / f"k{k}-v{vectorize}-{device}-{'linear' if linear else 'log'}.txt"
print("Logging results to", logfile)
print()
print("Device", "K", "Inner", "Vec.", "FLOPS", sep="\t")
with open(logfile, "w") as f:
print("device", "k", "inner", "vectorize", "flops", sep=",", file=f)
def log(k, inner, vectorize, flops):
print(device.upper(), k, inner, vectorize, flops, sep="\t", flush=True)
with open(logfile, "a") as f:
print(device, k, inner, vectorize, flops, sep=",", file=f)
actions = [3, 0, 1, 3, 0]
k *= 1024 # raw number of elements
with compiler_gym.make("loop_tool-v0") as env:
env.reset(
benchmark=env.datasets.benchmark(
uri=f"benchmark://loop_tool-{device}-v0/{k}"),
action_space="simple",
)
if vectorize - 1:
vs = [1] * (vectorize - 1)
actions += vs + [0, 1, 0] + vs + [0, 2, 0]
for a in actions:
wrapped_step(env, a)
if linear:
for i in range(k // (vectorize * 1024)):
step_count = 1022 if i == 0 else 1023
flops = flops_after_steps(env, step_count)
log(k, (i + 1) * 1024, vectorize, flops)
else: # linear=False (log)
inner = 1
step = 512
wrapped_step(env, [1] * (step - 1))
inner += step - 1
while inner * vectorize <= k:
flops = flops_after_steps(env, step)
inner += step
log(k, inner, vectorize, flops)
step *= 2
def get_benchmarks():
benchmarks = []
with compiler_gym.make("gcc-v0", gcc_bin=FLAGS.gcc_bin) as env:
env.reset()
if FLAGS.gcc_benchmark == ["all"]:
for dataset in env.datasets:
benchmarks += islice(dataset.benchmark_uris(), 50)
elif FLAGS.gcc_benchmark:
for uri in FLAGS.gcc_benchmark:
benchmarks.append(env.datasets.benchmark(uri).uri)
else:
benchmarks = list(
env.datasets["benchmark://chstone-v0"].benchmark_uris())
benchmarks.sort()
return benchmarks
def _reset_worker(
num_measurements_per_benchmark: int,
benchmarks: List[str],
env_name: str,
outfile: Path,
):
with compiler_gym.make(env_name) as env:
with open(outfile, "w") as f:
for benchmark in benchmarks:
env.reset(benchmark=benchmark)
runtimes = get_runtimes(
lambda: env.reset(benchmark=benchmark),
num_measurements_per_benchmark,
)
print("\n".join(f"{x:.8f} {benchmark}" for x in runtimes), file=f)
def test_rand(backend):
with compiler_gym.make("loop_tool-v0") as env:
env.observation_space = "flops"
env.reset(
benchmark=env.datasets.benchmark(
uri=f"benchmark://loop_tool-{backend}-v0/128"),
action_space="simple",
)
best = 0
for i in range(10):
a = env.action_space.sample()
o = env.step(a)
flops = o[0]
if flops > best:
best = flops
print(best)
def _step_worker(
num_measurements_per_benchmark: int,
benchmarks: List[str],
env_name: str,
seed: str,
step_outfile: Path,
batched_outfile: Path,
):
def get_step_times(env: CompilerEnv, num_steps: int, batched=False):
while batched:
# Run all actions in a single step().
steps = [env.action_space.sample() for _ in range(num_steps)]
with Timer() as timer:
_, _, done, _ = env.step(steps)
if not done:
return [timer.time / num_steps] * num_steps
env.reset()
# Run each action as a step().
runtimes = []
while len(runtimes) < num_steps:
with Timer() as timer:
_, _, done, _ = env.step(env.action_space.sample())
if done:
env.reset()
else:
runtimes.append(timer.time)
return runtimes
with compiler_gym.make(env_name) as env:
with open(step_outfile, "w") as f:
for i, benchmark in enumerate(benchmarks, start=seed):
env.reset(benchmark=benchmark)
env.seed(i)
runtimes = get_step_times(env, num_measurements_per_benchmark)
print("\n".join(f"{x:.8f} {benchmark}" for x in runtimes),
file=f)
with open(batched_outfile, "w") as f:
for i, benchmark in enumerate(benchmarks, start=seed):
env.reset(benchmark=benchmark)
env.seed(i)
runtimes = get_step_times(env,
num_measurements_per_benchmark,
batched=True)
print("\n".join(f"{x:.8f} {benchmark}" for x in runtimes),
file=f)
def get_benchmarks(env_name: str, n: int, seed: int,
outdir: Path) -> List[str]:
"""Get `n` benchmarks from all datasets.
If the dataset is smaller than `n`, benchmarks are repeated. If the dataset
is larger than `n`, `n` random unique programs are sampled.
"""
benchmarks = []
with compiler_gym.make(env_name) as env:
datasets = sorted(list(env.datasets))
benchmarks_per_dataset = int(ceil(n / len(datasets)))
for ds in datasets:
logger.info(
"Enumerating %s benchmarks from dataset from %s ...",
benchmarks_per_dataset,
ds,
)
if ds.size == 0 or ds.size > benchmarks_per_dataset:
rng = np.random.default_rng(seed)
uniq_bm_uris = set()
benchmarks_from_dataset = []
while len(benchmarks_from_dataset) < benchmarks_per_dataset:
bm = ds.random_benchmark(rng)
if bm.uri in uniq_bm_uris:
continue
uniq_bm_uris.add(bm.uri)
# Start an environment to check that the benchmark can be
# initialized.
try:
env.reset(benchmark=bm)
except (BenchmarkInitError, ValueError, TimeoutError):
continue
benchmarks_from_dataset.append(bm.uri)
benchmarks += benchmarks_from_dataset
else:
bms = list(ds.benchmark_uris())
bms *= int(ceil(benchmarks_per_dataset / len(bms)))
benchmarks += bms[:benchmarks_per_dataset]
benchmarks = sorted(benchmarks)
with open(outdir / "benchmarks.txt", "w") as f:
for bm in benchmarks:
print(bm, file=f)
return benchmarks
def test_basic(backend):
with compiler_gym.make("loop_tool-v0") as env:
env.observation_space = "flops"
env.reset(
benchmark=env.datasets.benchmark(
uri=f"benchmark://loop_tool-{backend}-v0/1024"),
action_space="simple",
)
env.step(0)
env.step(1)
env.step(0)
env.step(1)
env.step(1)
env.step(0)
env.step(1)
env.step(0)
o = env.step(1)
print(o)
def test_autotune():
with compiler_gym.make("llvm-v0", reward_space="IrInstructionCount") as env:
env.reset(benchmark="benchmark://cbench-v1/crc32")
autotuner = Autotuner(
algorithm="random",
optimization_target="codesize",
search_time_seconds=3,
)
result = autotuner(env)
print(result)
assert result.benchmark == "benchmark://cbench-v1/crc32"
assert result.walltime >= 3
assert result.commandline == env.commandline()
assert env.episode_reward >= 0
assert env.benchmark == "benchmark://cbench-v1/crc32"
assert env.reward_space == "IrInstructionCount"
def _observations_worker(
observation_space: str,
num_measurements_per_benchmark: int,
benchmarks: List[str],
env_name: str,
outfile: Path,
):
with compiler_gym.make(env_name) as env:
with open(outfile, "w") as f:
for benchmark in benchmarks:
env.reset(benchmark=benchmark)
if "llvm-" in env_name and observation_space == "Runtime":
if not env.observation.IsRunnable():
return []
env.runtime_observation_count = 1
env.runtime_warmups_count = 0
runtimes = get_runtimes(
lambda: env.observation[observation_space],
num_measurements_per_benchmark,
)
print("\n".join(f"{x:.8f}" for x in runtimes), file=f, flush=True)
def test_validation_benchmarks_uris_list():
cfg = Validation(**OmegaConf.create("""\
benchmarks:
- uris:
- benchmark://cbench-v1/qsort
- dataset: benchmark://cbench-v1
max_benchmarks: 2
"""))
with compiler_gym.make("llvm-v0") as env:
assert list(cfg.benchmarks_iterator(env)) == [
"benchmark://cbench-v1/qsort",
"benchmark://cbench-v1/adpcm",
"benchmark://cbench-v1/bitcount",
]
bm = list(cfg.benchmarks_iterator(env))[0]
print(type(bm).__name__)
assert isinstance(bm, Benchmark)
assert list(cfg.benchmark_uris_iterator(env)) == [
"benchmark://cbench-v1/qsort",
"benchmark://cbench-v1/adpcm",
"benchmark://cbench-v1/bitcount",
]
def run_search(search: str, benchmark: str, seed: int) -> SearchResult:
"""Run a search and return the search class instance."""
with GCC_ENV_CONSTRUCTOR_LOCK:
env = compiler_gym.make("gcc-v0", gcc_bin=FLAGS.gcc_bin)
try:
random.seed(seed)
np.random.seed(seed)
env.reset(benchmark=benchmark)
env.step(env.action_space["-Os"])
baseline_size = objective(env)
env.reset(benchmark=benchmark)
best_size = _SEARCH_FUNCTIONS[search](env)
finally:
env.close()
return SearchResult(
search=search,
benchmark=benchmark,
best_size=best_size,
baseline_size=baseline_size,
)
def test_compiler_gym_make():
"""Test that compiler_gym.make() is equivalent to gym.make()."""
with compiler_gym.make("llvm-v0") as env:
assert isinstance(env, LlvmEnv)
def make_env() -> compiler_gym.envs.CompilerEnv:
env = compiler_gym.make("llvm-v0", observation_space="InstCount", reward_space="IrInstructionCountOz")
env = stepWrapper(env)
#env = observationWrapper(env)
return env
def env():
with compiler_gym.make("llvm-v0") as env:
yield env
import compiler_gym
import examples.loop_optimizations_service as loop_optimizations_service # noqa Register environments.
with compiler_gym.make(
"loops-opt-py-v0",
benchmark="loops-opt-v0/add",
observation_space="AutophaseDict",
reward_space="runtime",
) as env:
compiler_gym.set_debug_level(4) # TODO: check why this has no effect
observation = env.reset()
print("observation: ", observation)
print()
observation, reward, done, info = env.step(env.action_space.sample())
print("observation: ", observation)
print("reward: ", reward)
print("done: ", done)
print("info: ", info)
env.close()
# TODO: implement write_bitcode(..) or write_ir(..)
# env.write_bitcode("/tmp/output.bc")
# 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 script demonstrates how the example services defined in this directory
can be used as gym environments. Usage:
$ bazel run -c opt //examples/example_unrolling_service:example
"""
import compiler_gym
import examples.example_unrolling_service as unrolling_service # noqa Register environments.
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)
print()
observation, reward, done, info = env.step(env.action_space.sample())
print("observation: ", observation)
print("reward: ", reward)
print("done: ", done)
print("info: ", info)
from pydantic import BaseModel
import compiler_gym
from compiler_gym.envs import LlvmEnv
from compiler_gym.util.truncate import truncate
app = Flask("compiler_gym")
CORS(app)
resource_dir: Path = (Path(__file__).parent /
"frontends/compiler_gym/build").absolute()
logger = logging.getLogger(__name__)
# A single compiler environment that is used to serve all endpoints.
env: LlvmEnv = compiler_gym.make("llvm-v0")
env_lock = Lock()
class StateToVisualize(BaseModel):
"""Encapsulates the state to visualize in the frontend."""
instcount: Dict[str, int]
autophase: Dict[str, int]
# The reward signal measures how "good" the previous action was. Over time
# the sequence of actions that produces the highest cumulative reward is the
# best:
reward: float
class StepRequest(BaseModel):
