1# mypy: allow-untyped-defs 2import numpy as np 3import torch 4 5from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedTensor 6 7 8_MIN_DIM_SIZE = 16 9_MAX_DIM_SIZE = 16 * 1024 ** 2 10_POW_TWO_SIZES = tuple(2 ** i for i in range( 11 int(np.log2(_MIN_DIM_SIZE)), 12 int(np.log2(_MAX_DIM_SIZE)) + 1, 13)) 14 15 16class UnaryOpFuzzer(Fuzzer): 17 def __init__(self, seed, dtype=torch.float32, cuda=False): 18 super().__init__( 19 parameters=[ 20 # Dimensionality of x. (e.g. 1D, 2D, or 3D.) 21 FuzzedParameter("dim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), 22 23 # Shapes for `x`. 24 # It is important to test all shapes, however 25 # powers of two are especially important and therefore 26 # warrant special attention. This is done by generating 27 # both a value drawn from all integers between the min and 28 # max allowed values, and another from only the powers of two 29 # (both distributions are loguniform) and then randomly 30 # selecting between the two. 31 [ 32 FuzzedParameter( 33 name=f"k_any_{i}", 34 minval=_MIN_DIM_SIZE, 35 maxval=_MAX_DIM_SIZE, 36 distribution="loguniform", 37 ) for i in range(3) 38 ], 39 [ 40 FuzzedParameter( 41 name=f"k_pow2_{i}", 42 distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} 43 ) for i in range(3) 44 ], 45 [ 46 FuzzedParameter( 47 name=f"k{i}", 48 distribution={ 49 ParameterAlias(f"k_any_{i}"): 0.8, 50 ParameterAlias(f"k_pow2_{i}"): 0.2, 51 }, 52 strict=True, 53 ) for i in range(3) 54 ], 55 56 # Steps for `x`. (Benchmarks strided memory access.) 57 [ 58 FuzzedParameter( 59 name=f"x_step_{i}", 60 distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, 61 ) for i in range(3) 62 ], 63 64 # Repeatable entropy for downstream applications. 65 FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), 66 ], 67 tensors=[ 68 FuzzedTensor( 69 name="x", 70 size=("k0", "k1", "k2"), 71 steps=("x_step_0", "x_step_1", "x_step_2"), 72 probability_contiguous=0.75, 73 min_elements=4 * 1024, 74 max_elements=32 * 1024 ** 2, 75 max_allocation_bytes=2 * 1024**3, # 2 GB 76 dim_parameter="dim", 77 dtype=dtype, 78 cuda=cuda, 79 ), 80 ], 81 seed=seed, 82 ) 83