mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
8db9dfa2d7
See #127836 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/127846 Approved by: https://github.com/ezyang ghstack dependencies: #127842, #127843, #127844, #127845
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
# mypy: allow-untyped-defs
|
|
import math
|
|
|
|
import torch
|
|
from torch.utils import benchmark
|
|
from torch.utils.benchmark import FuzzedParameter, FuzzedTensor, ParameterAlias
|
|
|
|
|
|
__all__ = ['SpectralOpFuzzer']
|
|
|
|
MIN_DIM_SIZE = 16
|
|
MAX_DIM_SIZE = 16 * 1024
|
|
|
|
def power_range(upper_bound, base):
|
|
return (base ** i for i in range(int(math.log(upper_bound, base)) + 1))
|
|
|
|
# List of regular numbers from MIN_DIM_SIZE to MAX_DIM_SIZE
|
|
# These numbers factorize into multiples of prime factors 2, 3, and 5 only
|
|
# and are usually the fastest in FFT implementations.
|
|
REGULAR_SIZES = []
|
|
for i in power_range(MAX_DIM_SIZE, 2):
|
|
for j in power_range(MAX_DIM_SIZE // i, 3):
|
|
ij = i * j
|
|
for k in power_range(MAX_DIM_SIZE // ij, 5):
|
|
ijk = ij * k
|
|
if ijk > MIN_DIM_SIZE:
|
|
REGULAR_SIZES.append(ijk)
|
|
REGULAR_SIZES.sort()
|
|
|
|
class SpectralOpFuzzer(benchmark.Fuzzer):
|
|
def __init__(self, *, seed: int, dtype=torch.float64,
|
|
cuda: bool = False, probability_regular: float = 1.0):
|
|
super().__init__(
|
|
parameters=[
|
|
# Dimensionality of x. (e.g. 1D, 2D, or 3D.)
|
|
FuzzedParameter("ndim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
|
|
|
|
# Shapes for `x`.
|
|
# It is important to test all shapes, however
|
|
# regular sizes are especially important to the FFT and therefore
|
|
# warrant special attention. This is done by generating
|
|
# both a value drawn from all integers between the min and
|
|
# max allowed values, and another from only the regular numbers
|
|
# (both distributions are loguniform) and then randomly
|
|
# selecting between the two.
|
|
[
|
|
FuzzedParameter(
|
|
name=f"k_any_{i}",
|
|
minval=MIN_DIM_SIZE,
|
|
maxval=MAX_DIM_SIZE,
|
|
distribution="loguniform",
|
|
) for i in range(3)
|
|
],
|
|
[
|
|
FuzzedParameter(
|
|
name=f"k_regular_{i}",
|
|
distribution={size: 1. / len(REGULAR_SIZES) for size in REGULAR_SIZES}
|
|
) for i in range(3)
|
|
],
|
|
[
|
|
FuzzedParameter(
|
|
name=f"k{i}",
|
|
distribution={
|
|
ParameterAlias(f"k_regular_{i}"): probability_regular,
|
|
ParameterAlias(f"k_any_{i}"): 1 - probability_regular,
|
|
},
|
|
strict=True,
|
|
) for i in range(3)
|
|
],
|
|
|
|
# Steps for `x`. (Benchmarks strided memory access.)
|
|
[
|
|
FuzzedParameter(
|
|
name=f"step_{i}",
|
|
distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04},
|
|
) for i in range(3)
|
|
],
|
|
],
|
|
tensors=[
|
|
FuzzedTensor(
|
|
name="x",
|
|
size=("k0", "k1", "k2"),
|
|
steps=("step_0", "step_1", "step_2"),
|
|
probability_contiguous=0.75,
|
|
min_elements=4 * 1024,
|
|
max_elements=32 * 1024 ** 2,
|
|
max_allocation_bytes=2 * 1024**3, # 2 GB
|
|
dim_parameter="ndim",
|
|
dtype=dtype,
|
|
cuda=cuda,
|
|
),
|
|
],
|
|
seed=seed,
|
|
)
|