xref: /aosp_15_r20/external/pytorch/torch/utils/benchmark/utils/sparse_fuzzer.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2from typing import Optional, Tuple, Union
3from numbers import Number
4import torch
5from torch.utils.benchmark import FuzzedTensor
6import math
7
8class FuzzedSparseTensor(FuzzedTensor):
9    def __init__(
10        self,
11        name: str,
12        size: Tuple[Union[str, int], ...],
13        min_elements: Optional[int] = None,
14        max_elements: Optional[int] = None,
15        dim_parameter: Optional[str] = None,
16        sparse_dim: Optional[str] = None,
17        nnz: Optional[str] = None,
18        density: Optional[str] = None,
19        coalesced: Optional[str] = None,
20        dtype=torch.float32,
21        cuda=False
22    ):
23        """
24        Args:
25            name:
26                A string identifier for the generated Tensor.
27            size:
28                A tuple of integers or strings specifying the size of the generated
29                Tensor. String values will replaced with a concrete int during the
30                generation process, while ints are simply passed as literals.
31            min_elements:
32                The minimum number of parameters that this Tensor must have for a
33                set of parameters to be valid. (Otherwise they are resampled.)
34            max_elements:
35                Like `min_elements`, but setting an upper bound.
36            dim_parameter:
37                The length of `size` will be truncated to this value.
38                This allows Tensors of varying dimensions to be generated by the
39                Fuzzer.
40            sparse_dim:
41                The number of sparse dimensions in a sparse tensor.
42            density:
43                This value allows tensors of varying sparsities to be generated by the Fuzzer.
44            coalesced:
45                The sparse tensor format permits uncoalesced sparse tensors,
46                where there may be duplicate coordinates in the indices.
47            dtype:
48                The PyTorch dtype of the generated Tensor.
49            cuda:
50                Whether to place the Tensor on a GPU.
51        """
52        super().__init__(name=name, size=size, min_elements=min_elements,
53                         max_elements=max_elements, dim_parameter=dim_parameter, dtype=dtype, cuda=cuda)
54        self._density = density
55        self._coalesced = coalesced
56        self._sparse_dim = sparse_dim
57
58    @staticmethod
59    def sparse_tensor_constructor(size, dtype, sparse_dim, nnz, is_coalesced):
60        """sparse_tensor_constructor creates a sparse tensor with coo format.
61
62        Note that when `is_coalesced` is False, the number of elements is doubled but the number of indices
63        represents the same amount of number of non zeros `nnz`, i.e, this is virtually the same tensor
64        with the same sparsity pattern. Moreover, most of the sparse operation will use coalesce() method
65        and what we want here is to get a sparse tensor with the same `nnz` even if this is coalesced or not.
66
67        In the other hand when `is_coalesced` is True the number of elements is reduced in the coalescing process
68        by an unclear amount however the probability to generate duplicates indices are low for most of the cases.
69        This decision was taken on purpose to maintain the construction cost as low as possible.
70        """
71        if isinstance(size, Number):
72            size = [size] * sparse_dim
73        assert all(size[d] > 0 for d in range(sparse_dim)) or nnz == 0, 'invalid arguments'
74        v_size = [nnz] + list(size[sparse_dim:])
75        if dtype.is_floating_point:
76            v = torch.rand(size=v_size, dtype=dtype, device="cpu")
77        else:
78            v = torch.randint(1, 127, size=v_size, dtype=dtype, device="cpu")
79
80        i = torch.rand(sparse_dim, nnz, device="cpu")
81        i.mul_(torch.tensor(size[:sparse_dim]).unsqueeze(1).to(i))
82        i = i.to(torch.long)
83
84        if not is_coalesced:
85            v = torch.cat([v, torch.randn_like(v)], 0)
86            i = torch.cat([i, i], 1)
87
88        x = torch.sparse_coo_tensor(i, v, torch.Size(size))
89        if is_coalesced:
90            x = x.coalesce()
91        return x
92
93    def _make_tensor(self, params, state):
94        size, _, _ = self._get_size_and_steps(params)
95        density = params['density']
96        nnz = math.ceil(sum(size) * density)
97        assert nnz <= sum(size)
98
99        is_coalesced = params['coalesced']
100        sparse_dim = params['sparse_dim'] if self._sparse_dim else len(size)
101        sparse_dim = min(sparse_dim, len(size))
102        tensor = self.sparse_tensor_constructor(size, self._dtype, sparse_dim, nnz, is_coalesced)
103
104        if self._cuda:
105            tensor = tensor.cuda()
106        sparse_dim = tensor.sparse_dim()
107        dense_dim = tensor.dense_dim()
108        is_hybrid = len(size[sparse_dim:]) > 0
109
110        properties = {
111            "numel": int(tensor.numel()),
112            "shape": tensor.size(),
113            "is_coalesced": tensor.is_coalesced(),
114            "density": density,
115            "sparsity": 1.0 - density,
116            "sparse_dim": sparse_dim,
117            "dense_dim": dense_dim,
118            "is_hybrid": is_hybrid,
119            "dtype": str(self._dtype),
120        }
121        return tensor, properties
122