1# mypy: allow-untyped-defs 2import math 3 4import torch 5from torch.utils import benchmark 6from torch.utils.benchmark import FuzzedParameter, FuzzedTensor, ParameterAlias 7 8 9__all__ = ['SpectralOpFuzzer'] 10 11MIN_DIM_SIZE = 16 12MAX_DIM_SIZE = 16 * 1024 13 14def power_range(upper_bound, base): 15 return (base ** i for i in range(int(math.log(upper_bound, base)) + 1)) 16 17# List of regular numbers from MIN_DIM_SIZE to MAX_DIM_SIZE 18# These numbers factorize into multiples of prime factors 2, 3, and 5 only 19# and are usually the fastest in FFT implementations. 20REGULAR_SIZES = [] 21for i in power_range(MAX_DIM_SIZE, 2): 22 for j in power_range(MAX_DIM_SIZE // i, 3): 23 ij = i * j 24 for k in power_range(MAX_DIM_SIZE // ij, 5): 25 ijk = ij * k 26 if ijk > MIN_DIM_SIZE: 27 REGULAR_SIZES.append(ijk) 28REGULAR_SIZES.sort() 29 30class SpectralOpFuzzer(benchmark.Fuzzer): 31 def __init__(self, *, seed: int, dtype=torch.float64, 32 cuda: bool = False, probability_regular: float = 1.0): 33 super().__init__( 34 parameters=[ 35 # Dimensionality of x. (e.g. 1D, 2D, or 3D.) 36 FuzzedParameter("ndim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), 37 38 # Shapes for `x`. 39 # It is important to test all shapes, however 40 # regular sizes are especially important to the FFT and therefore 41 # warrant special attention. This is done by generating 42 # both a value drawn from all integers between the min and 43 # max allowed values, and another from only the regular numbers 44 # (both distributions are loguniform) and then randomly 45 # selecting between the two. 46 [ 47 FuzzedParameter( 48 name=f"k_any_{i}", 49 minval=MIN_DIM_SIZE, 50 maxval=MAX_DIM_SIZE, 51 distribution="loguniform", 52 ) for i in range(3) 53 ], 54 [ 55 FuzzedParameter( 56 name=f"k_regular_{i}", 57 distribution={size: 1. / len(REGULAR_SIZES) for size in REGULAR_SIZES} 58 ) for i in range(3) 59 ], 60 [ 61 FuzzedParameter( 62 name=f"k{i}", 63 distribution={ 64 ParameterAlias(f"k_regular_{i}"): probability_regular, 65 ParameterAlias(f"k_any_{i}"): 1 - probability_regular, 66 }, 67 strict=True, 68 ) for i in range(3) 69 ], 70 71 # Steps for `x`. (Benchmarks strided memory access.) 72 [ 73 FuzzedParameter( 74 name=f"step_{i}", 75 distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, 76 ) for i in range(3) 77 ], 78 ], 79 tensors=[ 80 FuzzedTensor( 81 name="x", 82 size=("k0", "k1", "k2"), 83 steps=("step_0", "step_1", "step_2"), 84 probability_contiguous=0.75, 85 min_elements=4 * 1024, 86 max_elements=32 * 1024 ** 2, 87 max_allocation_bytes=2 * 1024**3, # 2 GB 88 dim_parameter="ndim", 89 dtype=dtype, 90 cuda=cuda, 91 ), 92 ], 93 seed=seed, 94 ) 95