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 BinaryOpFuzzer(Fuzzer): 17 def __init__(self, seed, dtype=torch.float32, cuda=False): 18 super().__init__( 19 parameters=[ 20 # Dimensionality of x and y. (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` and `y`. 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 # Moreover, `y` will occasionally have singleton 32 # dimensions in order to test broadcasting. 33 [ 34 FuzzedParameter( 35 name=f"k_any_{i}", 36 minval=_MIN_DIM_SIZE, 37 maxval=_MAX_DIM_SIZE, 38 distribution="loguniform", 39 ) for i in range(3) 40 ], 41 [ 42 FuzzedParameter( 43 name=f"k_pow2_{i}", 44 distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} 45 ) for i in range(3) 46 ], 47 [ 48 FuzzedParameter( 49 name=f"k{i}", 50 distribution={ 51 ParameterAlias(f"k_any_{i}"): 0.8, 52 ParameterAlias(f"k_pow2_{i}"): 0.2, 53 }, 54 strict=True, 55 ) for i in range(3) 56 ], 57 58 [ 59 FuzzedParameter( 60 name=f"y_k{i}", 61 distribution={ 62 ParameterAlias(f"k{i}"): 0.8, 63 1: 0.2, 64 }, 65 strict=True, 66 ) for i in range(3) 67 ], 68 69 # Steps for `x` and `y`. (Benchmarks strided memory access.) 70 [ 71 FuzzedParameter( 72 name=f"{name}_step_{i}", 73 distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, 74 ) 75 for i in range(3) 76 for name in ("x", "y") 77 ], 78 79 # Repeatable entropy for downstream applications. 80 FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), 81 ], 82 tensors=[ 83 FuzzedTensor( 84 name="x", 85 size=("k0", "k1", "k2"), 86 steps=("x_step_0", "x_step_1", "x_step_2"), 87 probability_contiguous=0.75, 88 min_elements=4 * 1024, 89 max_elements=32 * 1024 ** 2, 90 max_allocation_bytes=2 * 1024**3, # 2 GB 91 dim_parameter="dim", 92 dtype=dtype, 93 cuda=cuda, 94 ), 95 FuzzedTensor( 96 name="y", 97 size=("y_k0", "y_k1", "y_k2"), 98 steps=("x_step_0", "x_step_1", "x_step_2"), 99 probability_contiguous=0.75, 100 max_allocation_bytes=2 * 1024**3, # 2 GB 101 dim_parameter="dim", 102 dtype=dtype, 103 cuda=cuda, 104 ), 105 ], 106 seed=seed, 107 ) 108