• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# mypy: allow-untyped-defs
2"""Example of the Timer and Sparse Fuzzer APIs:
3
4$ python -m examples.sparse.fuzzer
5"""
6
7import sys
8
9import torch.utils.benchmark as benchmark_utils
10
11def main():
12    add_fuzzer = benchmark_utils.Fuzzer(
13        parameters=[
14            [
15                benchmark_utils.FuzzedParameter(
16                    name=f"k{i}",
17                    minval=16,
18                    maxval=16 * 1024,
19                    distribution="loguniform",
20                ) for i in range(3)
21            ],
22            benchmark_utils.FuzzedParameter(
23                name="dim_parameter",
24                distribution={2: 0.6, 3: 0.4},
25            ),
26            benchmark_utils.FuzzedParameter(
27                name="sparse_dim",
28                distribution={1: 0.3, 2: 0.4, 3: 0.3},
29            ),
30            benchmark_utils.FuzzedParameter(
31                name="density",
32                distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3},
33            ),
34            benchmark_utils.FuzzedParameter(
35                name="coalesced",
36                distribution={True: 0.7, False: 0.3},
37            )
38        ],
39        tensors=[
40            [
41                benchmark_utils.FuzzedSparseTensor(
42                    name=name,
43                    size=tuple([f"k{i}" for i in range(3)]),
44                    min_elements=64 * 1024,
45                    max_elements=128 * 1024,
46                    sparse_dim="sparse_dim",
47                    density="density",
48                    dim_parameter="dim_parameter",
49                    coalesced="coalesced"
50                ) for name in ("x", "y")
51            ],
52        ],
53        seed=0,
54    )
55
56    n = 100
57    measurements = []
58
59    for i, (tensors, tensor_properties, _) in enumerate(add_fuzzer.take(n=n)):
60        x = tensors["x"]
61        y = tensors["y"]
62        shape = ", ".join(tuple(f'{i:>4}' for i in x.shape))
63        x_tensor_properties = tensor_properties["x"]
64        description = "".join([
65            f"| {shape:<20} | ",
66            f"{x_tensor_properties['sparsity']:>9.2f} | ",
67            f"{x_tensor_properties['sparse_dim']:>9d} | ",
68            f"{x_tensor_properties['dense_dim']:>9d} | ",
69            f"{('True' if x_tensor_properties['is_hybrid'] else 'False'):>9} | ",
70            f"{('True' if x.is_coalesced() else 'False'):>9} | "
71        ])
72        timer = benchmark_utils.Timer(
73            stmt="torch.sparse.sum(x) + torch.sparse.sum(y)",
74            globals=tensors,
75            description=description,
76        )
77        measurements.append(timer.blocked_autorange(min_run_time=0.1))
78        measurements[-1].metadata = {"nnz": x._nnz()}
79        print(f"\r{i + 1} / {n}", end="")
80        sys.stdout.flush()
81    print()
82
83    # More string munging to make pretty output.
84    print(f"Average attempts per valid config: {1. / (1. - add_fuzzer.rejection_rate):.1f}")
85
86    def time_fn(m):
87        return m.mean / m.metadata["nnz"]
88
89    measurements.sort(key=time_fn)
90
91    template = f"{{:>6}}{' ' * 16} Shape{' ' * 17}\
92    sparsity{' ' * 4}sparse_dim{' ' * 4}dense_dim{' ' * 4}hybrid{' ' * 4}coalesced\n{'-' * 108}"
93    print(template.format("Best:"))
94    for m in measurements[:10]:
95        print(f"{time_fn(m) * 1e9:>5.2f} ns / element     {m.description}")
96
97    print("\n" + template.format("Worst:"))
98    for m in measurements[-10:]:
99        print(f"{time_fn(m) * 1e9:>5.2f} ns / element     {m.description}")
100
101if __name__ == "__main__":
102    main()
103