1# mypy: allow-untyped-defs 2import numpy as np 3import torch 4 5from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedSparseTensor 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 BinaryOpSparseFuzzer(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_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), 22 FuzzedParameter( 23 name="sparse_dim", 24 distribution={1: 0.4, 2: 0.4, 3: 0.2}, 25 strict=True 26 ), 27 # Shapes for `x` and `y`. 28 # It is important to test all shapes, however 29 # powers of two are especially important and therefore 30 # warrant special attention. This is done by generating 31 # both a value drawn from all integers between the min and 32 # max allowed values, and another from only the powers of two 33 # (both distributions are loguniform) and then randomly 34 # selecting between the two. 35 # Moreover, `y` will occasionally have singleton 36 # dimensions in order to test broadcasting. 37 [ 38 FuzzedParameter( 39 name=f"k_any_{i}", 40 minval=_MIN_DIM_SIZE, 41 maxval=_MAX_DIM_SIZE, 42 distribution="loguniform", 43 ) for i in range(3) 44 ], 45 [ 46 FuzzedParameter( 47 name=f"k_pow2_{i}", 48 distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} 49 ) for i in range(3) 50 ], 51 [ 52 FuzzedParameter( 53 name=f"k{i}", 54 distribution={ 55 ParameterAlias(f"k_any_{i}"): 0.8, 56 ParameterAlias(f"k_pow2_{i}"): 0.2, 57 }, 58 strict=True, 59 ) for i in range(3) 60 ], 61 [ 62 FuzzedParameter( 63 name=f"y_k{i}", 64 distribution={ 65 ParameterAlias(f"k{i}"): 1.0}, 66 strict=True, 67 ) for i in range(3) 68 ], 69 FuzzedParameter( 70 name="density", 71 distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3}, 72 ), 73 FuzzedParameter( 74 name="coalesced", 75 distribution={True: 0.5, False: 0.5}, 76 ), 77 # Repeatable entropy for downstream applications. 78 FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), 79 ], 80 tensors=[ 81 FuzzedSparseTensor( 82 name="x", 83 size=("k0", "k1", "k2"), 84 dim_parameter="dim_parameter", 85 sparse_dim="sparse_dim", 86 density="density", 87 coalesced="coalesced", 88 min_elements=4 * 1024, 89 max_elements=32 * 1024 ** 2, 90 dtype=dtype, 91 cuda=cuda, 92 ), 93 FuzzedSparseTensor( 94 name="y", 95 size=("y_k0", "y_k1", "y_k2"), 96 dim_parameter="dim_parameter", 97 sparse_dim="sparse_dim", 98 density="density", 99 coalesced="coalesced", 100 min_elements=4 * 1024, 101 max_elements=32 * 1024 ** 2, 102 dtype=dtype, 103 cuda=cuda, 104 ), 105 ], 106 seed=seed, 107 ) 108