xref: /aosp_15_r20/external/pytorch/torch/utils/benchmark/op_fuzzers/sparse_unary.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2
3import numpy as np
4import torch
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
15class UnaryOpSparseFuzzer(Fuzzer):
16    def __init__(self, seed, dtype=torch.float32, cuda=False):
17        super().__init__(
18            parameters=[
19                # Sparse dim parameter of x. (e.g. 1D, 2D, or 3D.)
20                FuzzedParameter("dim_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True),
21                FuzzedParameter(
22                    name="sparse_dim",
23                    distribution={1: 0.4, 2: 0.4, 3: 0.2},
24                    strict=True
25                ),
26                # Shapes for `x`.
27                #   It is important to test all shapes, however
28                #   powers of two are especially important and therefore
29                #   warrant special attention. This is done by generating
30                #   both a value drawn from all integers between the min and
31                #   max allowed values, and another from only the powers of two
32                #   (both distributions are loguniform) and then randomly
33                #   selecting between the two.
34                [
35                    FuzzedParameter(
36                        name=f"k_any_{i}",
37                        minval=_MIN_DIM_SIZE,
38                        maxval=_MAX_DIM_SIZE,
39                        distribution="loguniform",
40                    ) for i in range(3)
41                ],
42                [
43                    FuzzedParameter(
44                        name=f"k_pow2_{i}",
45                        distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES}
46                    ) for i in range(3)
47                ],
48                [
49                    FuzzedParameter(
50                        name=f"k{i}",
51                        distribution={
52                            ParameterAlias(f"k_any_{i}"): 0.8,
53                            ParameterAlias(f"k_pow2_{i}"): 0.2,
54                        },
55                        strict=True,
56                    ) for i in range(3)
57                ],
58                FuzzedParameter(
59                    name="density",
60                    distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3},
61                ),
62                FuzzedParameter(
63                    name="coalesced",
64                    distribution={True: 0.5, False: 0.5},
65                ),
66                FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"),
67            ],
68            tensors=[
69                FuzzedSparseTensor(
70                    name="x",
71                    size=("k0", "k1", "k2"),
72                    dim_parameter="dim_parameter",
73                    sparse_dim="sparse_dim",
74                    min_elements=4 * 1024,
75                    max_elements=32 * 1024 ** 2,
76                    density="density",
77                    coalesced="coalesced",
78                    dtype=dtype,
79                    cuda=cuda,
80                ),
81            ],
82            seed=seed,
83        )
84