1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
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,
)
|
