"""End-to-end example to test a PR for regressions:
$ python -m examples.end_to_end --pr 39850
$ python -m examples.end_to_end --pr 39967
$ python -m examples.end_to_end --pr 39744
NOTE:
This example assumes that you have and environment prefixed with
`ref_`, and another prefixed with `pr_` for the PR
in question. (e.g. `ref_39850` and `pr_39850`).
A helper script (examples/prepare_e2e.sh) is provided to build
the required environments with the correct configuration.
"""
import argparse
import itertools as it
import multiprocessing
import multiprocessing.dummy
import os
import pickle
import queue
import subprocess
import tempfile
import textwrap
import numpy as np
import torch
from torch.utils.benchmark.op_fuzzers import unary
from torch.utils.benchmark import Timer, Measurement
from typing import Dict, Tuple, List
_MAIN, _SUBPROCESS = "main", "subprocess"
_PR_ENV_TEMPLATE = "pr_{pr}"
_REF_ENV_TEMPLATE = "ref_{pr}"
_PR_LIST = (
# Optimize topk performance for tensor with a large dimension size
"39850",
# Migrate `var` & `std` to ATen
"39967",
# Introducing (Const)StridedRandomAccessor + CompositeRandomAccessor + migrate `sort` to ATen (CPU)
"39744",
)
_CPU, _GPU = "cpu", "gpu"
_MIN_RUN_SEC = 1
_REPLICATES = {
_CPU: 5, # CPU has a higher variance.
_GPU: 1,
}
_RUNS_PER_LOOP = 3
_NUM_LOOPS = {
_CPU: 32,
_GPU: 64,
}
_DEVICES_TO_TEST = {
"39850": {_CPU: False, _GPU: True},
"39967": {_CPU: True, _GPU: True},
"39744": {_CPU: True, _GPU: True},
}
_AVAILABLE_GPUS = queue.Queue[int]()
_DTYPES_TO_TEST = {
"39850": ("int8", "float32", "float64"),
"39967": ("float32", "float64"),
"39744": ("int8", "float32", "float64"),
}
_DTYPE_STR_TO_DTYPE = {
"float64": torch.float64,
"float32": torch.float32,
"int8": torch.int8,
}
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--pr", type=str, default=_PR_LIST[0], choices=_PR_LIST)
parser.add_argument("--num-gpus", "--num_gpus", type=int, default=None)
parser.add_argument("--test-variance", "--test_variance", action="store_true")
# (Implementation details)
parser.add_argument("--DETAIL-context", "--DETAIL_context", type=str, choices=(_MAIN, _SUBPROCESS), default=_MAIN)
parser.add_argument("--DETAIL-device", "--DETAIL_device", type=str, choices=(_CPU, _GPU), default=None)
parser.add_argument("--DETAIL-env", "--DETAIL_env", type=str, default=None)
parser.add_argument("--DETAIL-result-file", "--DETAIL_result_file", type=str, default=None)
parser.add_argument("--DETAIL-seed", "--DETAIL_seed", type=int, default=None)
args = parser.parse_args()
if args.num_gpus is None:
args.num_gpus = torch.cuda.device_count()
return args
_SUBPROCESS_CMD_TEMPLATE = (
"source activate {source_env} && python -m examples.end_to_end "
"--pr {pr} "
"--DETAIL-context subprocess "
"--DETAIL-device {device} "
"--DETAIL-env {env} "
"--DETAIL-result-file {result_file} "
"--DETAIL-seed {seed}"
)
def construct_stmt_and_label(pr, params):
if pr == "39850":
k0, k1, k2, dim = (params[i] for i in ["k0", "k1", "k2", "dim"])
state = np.random.RandomState(params["random_value"])
topk_dim = state.randint(low=0, high=dim)
dim_size = [k0, k1, k2][topk_dim]
k = max(int(np.floor(2 ** state.uniform(low=0, high=np.log2(dim_size)))), 1)
return f"torch.topk(x, dim={topk_dim}, k={k})", "topk"
if pr == "39967":
return "torch.std(x)", "std"
if pr == "39744":
state = np.random.RandomState(params["random_value"])
sort_dim = state.randint(low=0, high=params["dim"])
return f"torch.sort(x, dim={sort_dim})", "sort"
raise ValueError("Unknown PR")
def subprocess_main(args):
seed = args.DETAIL_seed
cuda = (args.DETAIL_device == _GPU)
with open(args.DETAIL_result_file, "ab") as f:
for dtype_str in _DTYPES_TO_TEST[args.pr]:
dtype = _DTYPE_STR_TO_DTYPE[dtype_str]
iterator = unary.UnaryOpFuzzer(
seed=seed, dtype=dtype, cuda=cuda).take(_RUNS_PER_LOOP)
for i, (tensors, tensor_parameters, params) in enumerate(iterator):
params["dtype_str"] = dtype_str
stmt, label = construct_stmt_and_label(args.pr, params)
timer = Timer(
stmt=stmt,
globals=tensors,
label=label,
description=f"[{i}, seed={seed}] ({dtype_str}), stmt = {stmt}",
env=args.DETAIL_env,
)
measurement = timer.blocked_autorange(min_run_time=_MIN_RUN_SEC)
measurement.metadata = {
"tensor_parameters": tensor_parameters,
"params": params,
}
print(measurement)
pickle.dump(measurement, f)
def _main(args):
pools, map_iters, finished_counts = {}, {}, {}
pr = args.pr
envs = (_REF_ENV_TEMPLATE.format(pr=pr), _PR_ENV_TEMPLATE.format(pr=pr))
# We initialize both pools at the start so that they run simultaneously
# if applicable
if _DEVICES_TO_TEST[args.pr][_GPU]:
finished_counts[_GPU] = 0
for i in range(args.num_gpus):
_AVAILABLE_GPUS.put(i)
pools[_GPU] = multiprocessing.dummy.Pool(args.num_gpus)
trials = [
(seed, envs, pr, True, finished_counts, args.test_variance)
for seed in range(_NUM_LOOPS[_GPU])] * _REPLICATES[_GPU]
map_iters[_GPU] = pools[_GPU].imap(map_fn, trials)
if _DEVICES_TO_TEST[args.pr][_CPU]:
finished_counts[_CPU] = 0
cpu_workers = int(multiprocessing.cpu_count() / 3)
pools[_CPU] = multiprocessing.dummy.Pool(cpu_workers)
trials = [
(seed, envs, pr, False, finished_counts, args.test_variance)
for seed in range(_NUM_LOOPS[_CPU])] * _REPLICATES[_CPU]
map_iters[_CPU] = pools[_CPU].imap(map_fn, trials)
results = []
for map_iter in map_iters.values():
for r in map_iter:
results.append(r)
progress = [
f"{k}: {v} / {_NUM_LOOPS[k] * _REPLICATES[k]}"
for k, v in finished_counts.items()]
print(f"\r{(' ' * 10).join(progress)}", end="")
print()
for pool in pools.values():
pool.close()
process_results(results, args.test_variance)
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# == Data processing and string formatting ====================================
# /////////////////////////////////////////////////////////////////////////////
def merge(measurements):
if not measurements:
return None
states = [m.__getstate__() for m in measurements]
for k in states[0].keys():
if k in ("number_per_run", "times", "metadata"):
continue
assert all(s[k] == states[0][k] for s in states)
numbers_per_run = {m.number_per_run for m in measurements}
n = numbers_per_run.pop() if len(numbers_per_run) == 1 else 1
merged_state = states[0]
times = [[t / m.number_per_run * n for t in m.times] for m in measurements]
merged_state["times"] = list(it.chain(*times))
merged_state["number_per_run"] = n
merged_state["metadata"] = states[0]["metadata"]
return Measurement(**merged_state)
def process_results(results, test_variance):
paired_results: Dict[Tuple[str, str, int, bool, int], List] = {}
for (seed, use_gpu), result_batch in results:
for r in result_batch:
key = (r.label, r.description, r.num_threads, use_gpu, seed)
paired_results.setdefault(key, [[], []])
index = 0 if r.env.startswith("ref") else 1
paired_results[key][index].append(r)
paired_results = {
key: [merge(r_ref_list), merge(r_pr_list)]
for key, (r_ref_list, r_pr_list) in paired_results.items()
}
flagged_for_removal = set()
for key, (r_ref, r_pr) in paired_results.items():
if any(r is None or r.has_warnings for r in (r_ref, r_pr)):
flagged_for_removal.add(key)
paired_results = {
k: v for k, v in paired_results.items()
if k not in flagged_for_removal
}
print(f"{len(flagged_for_removal)} samples were culled, {len(paired_results)} remain")
gpu_results = [(k, v) for k, v in paired_results.items() if k[3]]
cpu_results = [(k, v) for k, v in paired_results.items() if not k[3]]
if cpu_results:
construct_table(cpu_results, "CPU", test_variance)
if gpu_results:
construct_table(gpu_results, "GPU", test_variance)
def construct_table(results, device_str, test_variance):
device_str = f"== {device_str} {' (Variance Test)' if test_variance else ''} ".ljust(40, "=")
print(f"{'=' * 40}\n{device_str}\n{'=' * 40}\n")
results = sorted((
(key, (r_ref, r_pr), r_pr.median / r_ref.median - 1)
for key, (r_ref, r_pr) in results
), key=lambda i: i[2])
n = len(results)
n_regressed = len([i for i in results if i[2] > 0.05])
n_improved = len([i for i in results if i[2] < -0.05])
n_unchanged = n - n_improved - n_regressed
legends = ["Improved (>5%):", "Regressed (>5%):", "Within 5%:"]
for legend, count in zip(legends, [n_improved, n_regressed, n_unchanged]):
print(f"{legend:<17} {count:>6} ({count / len(results) * 100:>3.0f}%)")
keys_to_print = (
{i[0] for i in results[20:30]} |
{i[0] for i in results[int(n // 2 - 5):int(n // 2 + 5)]} |
{i[0] for i in results[-30:-20]}
)
ellipsis_after = {results[29][0], results[int(n // 2 + 4)][0]}
column_labels = (
f"Relative Δ Absolute Δ | numel{'':>8}dtype{'':>14}"
f"shape{'':>10}steps{'':>10}layout{'':>7}task specific\n{'=' * 126}"
)
_, result_log_file = tempfile.mkstemp(suffix=".log")
with open(result_log_file, "w") as f:
f.write(f"{device_str}\n\n{column_labels}\n")
print(f"\n{column_labels}\n[First twenty omitted (these tend to be noisy) ]")
for key, (r_ref, r_pr), rel_diff in results:
row = row_str(rel_diff, r_pr.median - r_ref.median, r_ref)
f.write(f"{row}\n")
if key in keys_to_print:
print(row)
if key in ellipsis_after:
print("...")
print("[Last twenty omitted (these tend to be noisy) ]")
print(textwrap.dedent("""
steps:
Indicates that `x` is sliced from a larger Tensor. For instance, if
shape is [12, 4] and steps are [2, 1], then a larger Tensor of size
[24, 4] was created, and then x = base_tensor[::2, ::1]. Omitted if
all elements are ones.
layout:
Indicates that `x` is not contiguous due to permutation. Invoking
`x.permute(layout)` (e.g. x.permute((2, 0, 1)) if layout = [2, 0, 1])
would produce a Tensor with physical memory layout matching logical
memory layout. (Though still not contiguous if `steps` contains
non-one elements.)
"""))
print(f"\nComplete results in: {result_log_file}")
def row_str(rel_diff, diff_seconds, measurement):
params = measurement.metadata["params"]
tensor_parameters = measurement.metadata["tensor_parameters"]
dim = params["dim"]
x_numel = tensor_parameters["x"]["numel"]
steps = [params[f"x_step_{i}"] for i in range(dim)]
order = tensor_parameters['x']["order"]
order = str("" if all(i == j for i, j in zip(order, range(dim))) else order)
task_specific = ""
if measurement.stmt.startswith("torch.topk"):
dim_str, k_str = measurement.stmt[:-1].replace("torch.topk(x, ", "").split(", ")
task_specific = f"{dim_str}, {k_str:<8}"
elif measurement.stmt.startswith("torch.std"):
pass
elif measurement.stmt.startswith("torch.sort"):
task_specific = measurement.stmt[:-1].replace("torch.sort(x, ", "")
return (
f"{rel_diff * 100:>5.0f}% {abs(diff_seconds) * 1e6:>11.1f} us{'':>6}|"
f"{x_numel:>12} {params['dtype_str']:>10} "
f"{str([params[f'k{i}'] for i in range(dim)]):>17} "
f"{str(steps) if not all(i == 1 for i in steps) else '':>12} {order:>12}"
f"{'':>8}{task_specific}"
)
# \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
# == Subprocess and environment management ====================================
# /////////////////////////////////////////////////////////////////////////////
def read_results(result_file: str):
output = []
with open(result_file, "rb") as f:
while True:
try:
output.append(pickle.load(f))
except EOFError:
break
return output
def run(cmd, cuda_visible_devices=""):
return subprocess.run(
cmd,
env={
"CUDA_VISIBLE_DEVICES": str(cuda_visible_devices),
"PATH": os.getenv("PATH", ""),
},
stdout=subprocess.PIPE,
shell=True
)
def test_source(envs):
"""Ensure that subprocess"""
for env in envs:
result = run(f"source activate {env}")
if result.returncode != 0:
raise ValueError(f"Failed to source environment `{env}`")
def map_fn(args):
seed, envs, pr, use_gpu, finished_counts, test_variance = args
gpu = _AVAILABLE_GPUS.get() if use_gpu else None
try:
_, result_file = tempfile.mkstemp(suffix=".pkl")
for env in envs:
cmd = _SUBPROCESS_CMD_TEMPLATE.format(
source_env=envs[0] if test_variance else env,
env=env, pr=pr, device=_GPU if use_gpu else _CPU,
result_file=result_file, seed=seed,
)
run(cmd=cmd, cuda_visible_devices=gpu if use_gpu else "")
finished_counts[_GPU if use_gpu else _CPU] += 1
return (seed, use_gpu), read_results(result_file)
except KeyboardInterrupt:
pass # Handle ctrl-c gracefully.
finally:
if gpu is not None:
_AVAILABLE_GPUS.put(gpu)
if os.path.exists(result_file):
os.remove(result_file)
def main(args):
test_source([
_REF_ENV_TEMPLATE.format(pr=args.pr),
_PR_ENV_TEMPLATE.format(pr=args.pr),
])
_main(args)
if __name__ == "__main__":
args = parse_args()
if args.DETAIL_context == "main":
main(args)
if args.DETAIL_context == "subprocess":
try:
subprocess_main(args)
except KeyboardInterrupt:
pass # Handle ctrl-c gracefully.