random_sample.py 6.99 KB
Newer Older
xgqdut2016's avatar
xgqdut2016 committed
1
import torch
PanZezhongQY's avatar
PanZezhongQY committed
2
import ctypes
xgqdut2016's avatar
xgqdut2016 committed
3
4
from ctypes import POINTER, Structure, c_int32, c_size_t, c_uint64, c_void_p, c_float
from libinfiniop import (
PanZezhongQY's avatar
PanZezhongQY committed
5
6
    infiniopHandle_t,
    infiniopTensorDescriptor_t,
xgqdut2016's avatar
xgqdut2016 committed
7
8
9
    open_lib,
    to_tensor,
    get_test_devices,
PanZezhongQY's avatar
PanZezhongQY committed
10
    check_error,
xgqdut2016's avatar
xgqdut2016 committed
11
    rearrange_if_needed,
PanZezhongQY's avatar
PanZezhongQY committed
12
    create_workspace,
xgqdut2016's avatar
xgqdut2016 committed
13
14
    test_operator,
    get_args,
xgqdut2016's avatar
xgqdut2016 committed
15
    debug_all,
xgqdut2016's avatar
xgqdut2016 committed
16
17
    get_tolerance,
    profile_operation,
xgqdut2016's avatar
xgqdut2016 committed
18
    synchronize_device,
PanZezhongQY's avatar
PanZezhongQY committed
19
20
)

xgqdut2016's avatar
xgqdut2016 committed
21
22
23
24
25
26
# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
_TEST_CASES = [
    # voc, random_val, topp, topk, temperature
xgqdut2016's avatar
xgqdut2016 committed
27
28
29
30
31
32
33
34
35
36
    (512, 0.8, 0.8, 3, 0.5),
    (4096, 0.05, 0.9, 5, 1.0),
    (16384, 0.15, 0.85, 10, 2.0),
    (512, 0.08, 0, 3, 0.5),
    (4096, 0.5, 0.9, 1, 1.0),
    (16384, 0.15, 0, 1, 2.0),
    (16384, 0.15, 0, 1, 2.0),
    (32000, 0.08, 0.8, 50, 1.0),
    (32000, 0.08, 1.0, 25, 1.0),
    # (119696, 0.01, 1.0, 100, 1.0),
xgqdut2016's avatar
xgqdut2016 committed
37
38
39
]

# Data types used for testing
xgqdut2016's avatar
xgqdut2016 committed
40
41
42
43
44
_TENSOR_DTYPES = [torch.float16]

_TOLERANCE_MAP = {
    torch.float16: {"atol": 0, "rtol": 0},
}
xgqdut2016's avatar
xgqdut2016 committed
45
46


xgqdut2016's avatar
xgqdut2016 committed
47
DEBUG = False
xgqdut2016's avatar
xgqdut2016 committed
48
49
50
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000
PanZezhongQY's avatar
PanZezhongQY committed
51
52
53
54
55
56
57
58
59
60


class RandomSampleDescriptor(Structure):
    _fields_ = [("device", c_int32)]


infiniopRandomSampleDescriptor_t = POINTER(RandomSampleDescriptor)


def random_sample(data, random_val, topp, topk, voc, temperature, torch_device):
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
    if topp > 0 and topk > 1:
        indices = torch.zeros([topk], dtype=torch.int64)
        dataNp = data.clone().detach()
        sorted_indices = torch.arange(voc)

        for i in range(topk):
            for j in range(i + 1, voc):
                if dataNp[i] < dataNp[j]:
                    tmp = dataNp[i].clone().detach()
                    dataNp[i] = dataNp[j].clone().detach()
                    dataNp[j] = tmp

                    tmpInd = sorted_indices[i].clone().detach()
                    sorted_indices[i] = sorted_indices[j].clone().detach()
                    sorted_indices[j] = tmpInd

        # sorted_indices = torch.argsort(dataNp, descending=True)
        indices = sorted_indices[:topk]

        dataNp = dataNp[sorted_indices]

        globalM = dataNp[0]
        dataNp = (dataNp - globalM) / temperature
        dataNp = torch.softmax(dataNp.float(), dim=0)
        sum_s = 0
        for end in range(topk):
            sum_s += dataNp[end]
            if sum_s >= topp:
                break
        if end < topk - 1:
            end += 1
        else:
            end = topk

        sum_s = 0
        for i in range(end):
            sum_s += dataNp[i]
        random_val *= sum_s

        sum_s = 0
        for i in range(end):
            sum_s += dataNp[i]
            if random_val < sum_s:
                return indices[i]
PanZezhongQY's avatar
PanZezhongQY committed
105
    else:
106
        return torch.argmax(data)
PanZezhongQY's avatar
PanZezhongQY committed
107

108
109
110
111
112
113
114
115
116
117
118
119
120

def test(
    lib,
    handle,
    torch_device,
    voc,
    random_val,
    topp,
    topk,
    temperature,
    x_dtype=torch.float16,
):
    print(f"Testing RandomSample on {torch_device} with voc:{voc} dtype:{x_dtype}")
xgqdut2016's avatar
xgqdut2016 committed
121

PanZezhongQY's avatar
PanZezhongQY committed
122
123
124
    data = torch.arange(voc).float() * 0.0001
    _perm = torch.randperm(voc)
    data = data[_perm].to(x_dtype).to(torch_device)
125
126
127
128

    ans = random_sample(
        data, random_val, topp, topk, voc, temperature, torch_device
    )  # 这个函数在device速度可能会很慢,可以通过data.to("cpu")方式加快计算过程
xgqdut2016's avatar
xgqdut2016 committed
129

PanZezhongQY's avatar
PanZezhongQY committed
130
    indices = torch.zeros([1], dtype=torch.int64).to(torch_device)
xgqdut2016's avatar
xgqdut2016 committed
131
132
133

    x_tensor, indices_tensor = [to_tensor(tensor, lib) for tensor in [data, indices]]

PanZezhongQY's avatar
PanZezhongQY committed
134
135
136
137
138
    indices_tensor.descriptor.contents.dt = U64  # treat int64 as uint64

    descriptor = infiniopRandomSampleDescriptor_t()
    check_error(
        lib.infiniopCreateRandomSampleDescriptor(
139
140
141
142
            handle,
            ctypes.byref(descriptor),
            indices_tensor.descriptor,
            x_tensor.descriptor,
PanZezhongQY's avatar
PanZezhongQY committed
143
144
145
146
        )
    )

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
xgqdut2016's avatar
xgqdut2016 committed
147
148
    for tensor in [x_tensor, indices_tensor]:
        tensor.descriptor.contents.invalidate()
PanZezhongQY's avatar
PanZezhongQY committed
149
150
151
152
153
154
155

    workspace_size = c_uint64(0)
    check_error(
        lib.infiniopGetRandomSampleWorkspaceSize(
            descriptor, ctypes.byref(workspace_size)
        )
    )
156
    workspace = create_workspace(workspace_size.value, torch_device)
xgqdut2016's avatar
xgqdut2016 committed
157

xgqdut2016's avatar
xgqdut2016 committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
    def lib_random_sample():
        check_error(
            lib.infiniopRandomSample(
                descriptor,
                workspace.data_ptr() if workspace is not None else None,
                workspace_size.value,
                indices_tensor.data,
                x_tensor.data,
                random_val,
                topp,
                topk,
                temperature,
                None,
            )
PanZezhongQY's avatar
PanZezhongQY committed
172
173
        )

xgqdut2016's avatar
xgqdut2016 committed
174
175
176
177
178
179
180
181
182
183
184
185
    if torch_device == "npu":
        synchronize_device(torch_device)

    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
    if DEBUG:
        debug_all(
            (indices[0].type(ans.dtype), data[indices[0]]),
            (ans, data[ans]),
            "or",
            atol=atol,
            rtol=rtol,
        )
PanZezhongQY's avatar
PanZezhongQY committed
186
    assert indices[0].type(ans.dtype) == ans or data[ans] == data[indices[0]]
xgqdut2016's avatar
xgqdut2016 committed
187

xgqdut2016's avatar
xgqdut2016 committed
188
189
190
191
192
193
194
195
196
197
198
199
    # Profiling workflow
    if PROFILE:
        # fmt: off
        if topp > 0 and topk > 1:
            profile_operation("PyTorch", lambda: random_sample(
                data.to("cpu"), random_val, topp, topk, voc, temperature, "cpu"
            ), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        else:
            profile_operation("PyTorch", lambda: random_sample_0(data), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        
        profile_operation("    lib", lambda: lib_random_sample(), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        # fmt: on
PanZezhongQY's avatar
PanZezhongQY committed
200
201
    check_error(lib.infiniopDestroyRandomSampleDescriptor(descriptor))

202

PanZezhongQY's avatar
PanZezhongQY committed
203
204
205
if __name__ == "__main__":
    args = get_args()
    lib = open_lib()
xgqdut2016's avatar
xgqdut2016 committed
206

PanZezhongQY's avatar
PanZezhongQY committed
207
208
209
210
211
212
    lib.infiniopCreateRandomSampleDescriptor.restype = c_int32
    lib.infiniopCreateRandomSampleDescriptor.argtypes = [
        infiniopHandle_t,
        POINTER(infiniopRandomSampleDescriptor_t),
        infiniopTensorDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
213

PanZezhongQY's avatar
PanZezhongQY committed
214
215
216
217
218
    lib.infiniopGetRandomSampleWorkspaceSize.restype = c_int32
    lib.infiniopGetRandomSampleWorkspaceSize.argtypes = [
        infiniopRandomSampleDescriptor_t,
        POINTER(c_uint64),
    ]
xgqdut2016's avatar
xgqdut2016 committed
219

PanZezhongQY's avatar
PanZezhongQY committed
220
221
222
223
224
225
226
227
228
229
230
231
232
    lib.infiniopRandomSample.restype = c_int32
    lib.infiniopRandomSample.argtypes = [
        infiniopRandomSampleDescriptor_t,
        c_void_p,
        c_uint64,
        c_uint64,
        c_void_p,
        c_float,
        c_float,
        c_int32,
        c_float,
        c_void_p,
    ]
xgqdut2016's avatar
xgqdut2016 committed
233

PanZezhongQY's avatar
PanZezhongQY committed
234
235
236
237
238
    lib.infiniopDestroyRandomSampleDescriptor.restype = c_int32
    lib.infiniopDestroyRandomSampleDescriptor.argtypes = [
        infiniopRandomSampleDescriptor_t,
    ]

xgqdut2016's avatar
xgqdut2016 committed
239
    DEBUG = args.debug
xgqdut2016's avatar
xgqdut2016 committed
240
241
242
243
244
245
246
247
    PROFILE = args.profile
    NUM_PRERUN = args.num_prerun
    NUM_ITERATIONS = args.num_iterations

    # Execute tests
    for device in get_test_devices(args):
        test_operator(lib, device, test, _TEST_CASES, _TENSOR_DTYPES)

PanZezhongQY's avatar
PanZezhongQY committed
248
    print("\033[92mTest passed!\033[0m")