def observations(
env: str = "llvm-autophase-ic-v0",
observation_spaces: List[str] = [
"Ir",
"InstCount",
"Autophase",
"Inst2vec",
"Programl",
"IrInstructionCount",
"ObjectTextSizeBytes",
"Runtime",
],
n: int = int(1e6),
num_benchmarks: int = int(1e3),
j: int = cpu_count(),
seed: int = 0xCC,
outdir: Optional[Path] = None,
) -> List[float]:
"""Benchmark the environment observation spaces."""
executor = Executor(type="local", cpus=j)
outdir = Path(outdir or create_user_logs_dir("op_benchmarks"))
benchmarks = get_benchmarks(env_name=env,
n=min(n, num_benchmarks),
seed=seed,
outdir=outdir)
with executor.get_executor(logs_dir=outdir) as session:
_observations(
session=session,
env_name=env,
benchmarks=benchmarks,
j=j,
outdir=outdir,
observation_spaces=observation_spaces,
n=n,
)
def step(
n: int = int(1e6),
num_benchmarks: int = int(1e3),
env: str = "llvm-autophase-ic-v0",
j: int = cpu_count(),
seed: int = 0xCC,
outdir: Optional[Path] = None,
):
"""Benchmark the env.step() operator."""
executor = Executor(type="local", cpus=j)
outdir = Path(outdir or create_user_logs_dir("op_benchmarks"))
benchmarks = get_benchmarks(env_name=env,
n=min(n, num_benchmarks),
seed=seed,
outdir=outdir)
with executor.get_executor(logs_dir=outdir) as session:
_step(
session=session,
outdir=outdir,
benchmarks=benchmarks,
n=n,
j=j,
env_name=env,
seed=seed,
)
def run(
env: str = "llvm-autophase-ic-v0",
observation_spaces: List[str] = [
"Ir",
"InstCount",
"Autophase",
"Inst2vec",
"Programl",
"IrInstructionCount",
"ObjectTextSizeBytes",
"Runtime",
],
n: int = int(1e6),
num_benchmarks: int = int(1e3),
j: int = cpu_count(),
outdir: Optional[Path] = None,
seed: int = 0xCC,
):
"""Run all of the environment benchmarks."""
executor = Executor(type="local", cpus=j)
outdir = Path(outdir or create_user_logs_dir("op_benchmarks"))
benchmarks = get_benchmarks(env_name=env,
n=min(n, num_benchmarks),
seed=seed,
outdir=outdir)
with executor.get_executor(logs_dir=outdir) as session:
_init(env_name=env, session=session, j=j, n=n, outdir=outdir)
_reset(
benchmarks=benchmarks,
n=n,
outdir=outdir,
j=j,
env_name=env,
session=session,
)
_step(
n=n,
j=j,
benchmarks=benchmarks,
env_name=env,
seed=seed,
outdir=outdir,
session=session,
)
_observations(
n=n,
j=j,
benchmarks=benchmarks,
env_name=env,
outdir=outdir,
session=session,
observation_spaces=observation_spaces,
)
info([outdir])
def init(
n: int = int(1e6),
j: int = cpu_count(),
env: str = "llvm-autophase-ic-v0",
outdir: Optional[Path] = None,
):
"""Benchmark the environment startup time."""
executor = Executor(type="local", cpus=j)
outdir = Path(outdir or create_user_logs_dir("op_benchmarks"))
with executor.get_executor(logs_dir=outdir) as session:
_init(n=n, outdir=outdir, j=j, env_name=env, session=session)
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 sweep(
device: List[str] = ["cuda"],
k: List[int] = [512, 1024, 2048, 4096, 8192],
vectorize: List[int] = [1],
linear: List[bool] = [False],
logdir: Optional[Path] = None,
):
logdir = logdir or create_user_logs_dir("loop_tool_sweep")
for device_, k_, vectorize_, linear_ in product(device, k, vectorize,
linear):
run_one_sweep(device=device_,
k=k_,
vectorize=vectorize_,
linear=linear_,
logdir=logdir)
def test_create_user_logs_dir(temporary_environ, tmpdir):
tmpdir = Path(tmpdir)
temporary_environ["COMPILER_GYM_LOGS"] = str(tmpdir)
dir = create_user_logs_dir("foo")
now = datetime.now()
assert dir.parent.parent == tmpdir / "foo"
year, month, day = dir.parent.name.split("-")
assert int(year) == now.year
assert int(month) == now.month
assert int(day) == now.day
hour, minute, second = dir.name.split("-")
assert int(hour) == now.hour
assert int(minute) == now.minute
assert int(second) == now.second
def run(self):
self.dir = create_user_logs_dir("foo")
def main(argv):
del argv # Unused.
# Validate the --search values now.
for search in FLAGS.search:
if search not in _SEARCH_FUNCTIONS:
raise app.UsageError(f"Invalid --search value: {search}")
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
logdir = (Path(FLAGS.output_dir)
if FLAGS.output_dir else create_user_logs_dir("gcc_autotuning"))
logdir.mkdir(exist_ok=True, parents=True)
with open(logdir / "results.csv", "w") as f:
print(
"search",
"benchmark",
"scaled_size",
"size",
"baseline_size",
sep=",",
file=f,
)
print("Logging results to", logdir)
# Parallel execution environment. Use flag --nproc to control the number of
# worker processes.
executor = Executor(type="local",
timeout_hours=12,
cpus=FLAGS.nproc,
block=True)
with executor.get_executor(logs_dir=logdir) as session:
jobs = []
# Submit each search instance as a separate job.
grid = product(range(FLAGS.gcc_search_repetitions), FLAGS.search,
get_benchmarks())
for _, search, benchmark in grid:
if not benchmark:
raise app.UsageError("Empty benchmark name not allowed")
jobs.append(
session.submit(
run_search,
search=search,
benchmark=benchmark,
seed=FLAGS.seed + len(jobs),
))
for job in jobs:
result = job.result()
print(result.benchmark, f"{result.scaled_best:.3f}x", sep="\t")
with open(logdir / "results.csv", "a") as f:
print(
result.search,
result.benchmark,
result.scaled_best,
result.best_size,
result.baseline_size,
sep=",",
file=f,
)
# Print results aggregates.
info([logdir])
def create_logging_dir(name: str) -> Path:
"""Deprecated function to create a directory for writing logs to.
Use :code:`compiler_gym.util.runfiles_path.create_user_logs_dir()` instead.
"""
return create_user_logs_dir(name)
