mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
5cb05a82b4
move benchmarking out of `torch._inductor.runtime.runtime_utils` and into `torch._inductor.runtime.benchmarking`, and prefer this path over directly accessing Triton's benchmarking Fixes #ISSUE_NUMBER Pull Request resolved: https://github.com/pytorch/pytorch/pull/132827 Approved by: https://github.com/eellison
274 lines
8.7 KiB
Python
274 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import click
|
|
import numpy as np
|
|
from operator_inp_utils import OperatorInputsLoader
|
|
|
|
import torch
|
|
from torch._dynamo.backends.cudagraphs import cudagraphs_inner
|
|
from torch._dynamo.testing import same
|
|
from torch._inductor.compile_fx import compile_fx
|
|
from torch._inductor.decomposition import decompositions
|
|
from torch._inductor.lowering import lowerings
|
|
from torch._inductor.runtime.benchmarking import benchmarker
|
|
from torch._inductor.utils import gen_gm_and_inputs
|
|
from torch.utils._pytree import tree_map_only
|
|
|
|
|
|
aten = torch.ops.aten
|
|
|
|
|
|
def compute_speedups(
|
|
operator, models, example_inputs, repeats, accuracy_checking=False, device="cuda"
|
|
):
|
|
expected = models[0](*example_inputs)
|
|
if accuracy_checking:
|
|
for model in models[1:]:
|
|
actual = model(*example_inputs)
|
|
# change to assert later
|
|
try:
|
|
same(actual, expected, cos_similarity=True, equal_nan=True)
|
|
except AssertionError as e:
|
|
print(e)
|
|
print(f"Accuracy check failed: {operator}")
|
|
print((expected[0] - actual[0]).abs().max())
|
|
|
|
timings = np.zeros((repeats, len(models)), np.float64)
|
|
for rep in range(repeats):
|
|
# interleave the runs to handle frequency scaling and load changes
|
|
for m, model in enumerate(models):
|
|
if device == "cuda":
|
|
model(*example_inputs)
|
|
|
|
# benchmarker.benchmark_gpu() clears L2 cache to hide the latency of CPU launch time
|
|
# along with cuda synchronization
|
|
timings[rep, m] = benchmarker.benchmark_gpu(
|
|
lambda: model(*example_inputs)
|
|
)
|
|
else:
|
|
from torch._inductor.utils import timed
|
|
|
|
timings[rep, m] = timed(model, example_inputs)
|
|
return np.median(timings, axis=0)
|
|
|
|
|
|
def strip_overloads(gm):
|
|
"""
|
|
Modifies the target of graph nodes in :attr:`gm` to strip overloads.
|
|
Args:
|
|
gm(fx.GraphModule): The input Fx graph module to be modified
|
|
"""
|
|
for node in gm.graph.nodes:
|
|
if isinstance(node.target, torch._ops.OpOverload):
|
|
node.target = node.target.overloadpacket
|
|
gm.recompile()
|
|
|
|
|
|
def convert_to_jit(gm, gm_args):
|
|
strip_overloads(gm)
|
|
try:
|
|
return torch.jit.script(gm)
|
|
except Exception:
|
|
pass
|
|
return torch.jit.trace(gm, gm_args)
|
|
|
|
|
|
def to_channels_last(ten):
|
|
return ten if ten.ndim != 4 else ten.to(memory_format=torch.channels_last)
|
|
|
|
|
|
def microbenchmark(
|
|
operator, args, kwargs, dtype, accuracy_checking, repeats, measure_nvfuser, device
|
|
):
|
|
gm, gm_args = gen_gm_and_inputs(operator, args, kwargs)
|
|
torch.jit._builtins._register_builtin(
|
|
torch.ops.aten.convolution_backward.default, "aten::convolution_backward"
|
|
)
|
|
if device == "cuda":
|
|
cudagraphs_eager = cudagraphs_inner(
|
|
gm, gm_args, copy_outputs=False, copy_inputs=False
|
|
)
|
|
compiled_fn = compile_fx(gm, gm_args)
|
|
cudagraphs_compiled = cudagraphs_inner(
|
|
compiled_fn, gm_args, copy_outputs=False, copy_inputs=False
|
|
)
|
|
compiled = [cudagraphs_eager, cudagraphs_compiled]
|
|
else:
|
|
compiled_fn = compile_fx(gm, gm_args)
|
|
compiled = [gm, compiled_fn]
|
|
if measure_nvfuser:
|
|
g = convert_to_jit(gm, gm_args)
|
|
cudagraphs_jit = cudagraphs_inner(
|
|
g, gm_args, copy_outputs=False, copy_inputs=False
|
|
)
|
|
compiled += [cudagraphs_jit]
|
|
if accuracy_checking:
|
|
repeats = 1
|
|
|
|
medians = compute_speedups(
|
|
operator, compiled, gm_args, repeats, accuracy_checking, device
|
|
)
|
|
return medians
|
|
|
|
|
|
def skip_operator(operator):
|
|
nyi_strings = (
|
|
"aten.gather.default",
|
|
"nll_loss",
|
|
"aten.index",
|
|
"aten.scatter_",
|
|
"masked_fill_.Scalar",
|
|
)
|
|
|
|
if any(nyi_string in str(operator) for nyi_string in nyi_strings):
|
|
# maybe disable aten.native_layer_norm.default
|
|
# TODO - inputs cannot be randomly initialized, causes cyda failures
|
|
print(f"Skipping {operator}, input generator nyi")
|
|
return True
|
|
|
|
# not covered by other non-compute operator heuristics
|
|
if operator == torch.ops.aten._unsafe_view.default:
|
|
print(f"Skipping {operator}, non compute operator")
|
|
return True
|
|
|
|
# some of inductor registered to the OpOverload, some registered to OpOverloadPacket
|
|
op_impls = [operator]
|
|
if isinstance(operator, torch._ops.OpOverload):
|
|
op_impls.append(operator.overloadpacket)
|
|
|
|
# TODO - skip benchmarking fallbacks. for some ops we have both lowerings and fallbacks
|
|
# so its not clear just from operator what will be lowered.
|
|
|
|
if all(op not in decompositions and op not in lowerings for op in op_impls):
|
|
print(f"Skipping {operator}, no inductor impl")
|
|
return True
|
|
|
|
if "convolution" in str(operator):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
@click.command()
|
|
@click.option(
|
|
"--suite",
|
|
help="suite to load inps from: options: timm, huggingface, torchbench",
|
|
default="torchbench",
|
|
)
|
|
@click.option("--op", help="operator overload to benchmark")
|
|
@click.option("--dtype", help="dtype to benchmark")
|
|
@click.option("--max-samples", help="max samples per op", default=15)
|
|
@click.option("--accuracy-checking", help="check accuracy", default=False)
|
|
@click.option(
|
|
"--repeats", help="how many times to repeat for perf measurement", default=3
|
|
)
|
|
@click.option(
|
|
"--measure-nvfuser", help="default we only measure inductor", default=False
|
|
)
|
|
@click.option("--device", help="cpu or cuda", default="cuda")
|
|
@click.option("--inp-file", help="use custom input file instead of suite", default=None)
|
|
@click.option("--start-idx", help="specify start index of samples", default=0)
|
|
@click.option(
|
|
"--channels-last", help="force inputs to channels last", is_flag=True, default=False
|
|
)
|
|
def benchmark(
|
|
suite,
|
|
op,
|
|
dtype,
|
|
max_samples,
|
|
accuracy_checking,
|
|
repeats,
|
|
measure_nvfuser,
|
|
device,
|
|
inp_file,
|
|
start_idx,
|
|
channels_last,
|
|
):
|
|
if inp_file is not None:
|
|
loader = OperatorInputsLoader(inp_file)
|
|
else:
|
|
assert suite in ("timm", "huggingface", "torchbench"), f"got {suite}"
|
|
if suite == "timm":
|
|
loader = OperatorInputsLoader.get_timm_loader()
|
|
elif suite == "huggingface":
|
|
loader = OperatorInputsLoader.get_huggingface_loader()
|
|
else:
|
|
loader = OperatorInputsLoader.get_torchbench_loader()
|
|
|
|
assert dtype in ("float16", "float32"), f"got {dtype}"
|
|
|
|
if op == "all":
|
|
filename = f"timings_{suite}_{op.replace('.', '_')}{dtype}.txt"
|
|
f = open(filename, "a")
|
|
|
|
dtype = torch.float16 if dtype == "float16" else torch.float32
|
|
|
|
if op == "all":
|
|
ops = loader.get_all_ops()
|
|
else:
|
|
ops = [eval(op)]
|
|
|
|
max_samples = max_samples + start_idx
|
|
for operator in ops:
|
|
if skip_operator(operator):
|
|
continue
|
|
|
|
print(f"Running {operator}")
|
|
inp_gen = loader.get_inputs_for_operator(operator, dtype=dtype, device=device)
|
|
timings = []
|
|
|
|
for i in range(min(max_samples, 1000000)):
|
|
try:
|
|
inps = next(inp_gen)
|
|
if inps is None:
|
|
break
|
|
if i < start_idx:
|
|
continue
|
|
print(f"Iter {i}")
|
|
args, kwargs = inps
|
|
if channels_last:
|
|
args, kwargs = tree_map_only(
|
|
torch.Tensor, to_channels_last, (args, kwargs)
|
|
)
|
|
|
|
except StopIteration:
|
|
break
|
|
try:
|
|
# aten, nvfuser, inductor
|
|
timings.append(
|
|
microbenchmark(
|
|
operator,
|
|
args,
|
|
kwargs,
|
|
dtype,
|
|
accuracy_checking,
|
|
repeats,
|
|
measure_nvfuser,
|
|
device,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
print(f"error {operator}")
|
|
print(e)
|
|
# comment out this line to avoid blocking other tests
|
|
# raise e
|
|
|
|
if not timings:
|
|
continue
|
|
|
|
timings = torch.tensor(timings).T
|
|
q = torch.tensor([0.2, 0.5, 0.8], dtype=torch.float64)
|
|
output = f"{operator}:\nInductor Speedups : {(torch.quantile(timings[0] / timings[1], q)).tolist()}\n"
|
|
if measure_nvfuser:
|
|
output += f"NVFUSER Speedups :{(torch.quantile(timings[0] / timings[2], q)).tolist()}\n"
|
|
if op == "all":
|
|
f.write(output)
|
|
print(output)
|
|
|
|
if op == "all":
|
|
f.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
benchmark()
|