causal_softmax.py 5.23 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
15
16
17
    test_operator,
    get_args,
    debug,
    get_tolerance,
    profile_operation,
PanZezhongQY's avatar
PanZezhongQY committed
18
)
19
from enum import Enum, auto
PanZezhongQY's avatar
PanZezhongQY committed
20

xgqdut2016's avatar
xgqdut2016 committed
21
22
23
24
# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
25
26
27
28
29
30
31
_TEST_CASES_ = [
    # shape, x_stride, y_stride
    ((3, 3), None, None),
    ((32, 512), None, None),
    ((32, 512), (1024, 1), (1024, 1)),
    ((32, 5, 5), None, None),
    ((32, 20, 512), None, None),
32
    ((32, 20, 512), (20480, 512, 1), None),
xgqdut2016's avatar
xgqdut2016 committed
33
34
]

xgqdut2016's avatar
xgqdut2016 committed
35
# Data types used for testing
xgqdut2016's avatar
xgqdut2016 committed
36
_TENSOR_DTYPES = [torch.float16]
xgqdut2016's avatar
xgqdut2016 committed
37
38
39

# Tolerance map for different data types
_TOLERANCE_MAP = {
xgqdut2016's avatar
xgqdut2016 committed
40
    torch.float16: {"atol": 0, "rtol": 1e-2},
xgqdut2016's avatar
xgqdut2016 committed
41
42
}

43
44
45
46
47
48
49
50

class Inplace(Enum):
    OUT_OF_PLACE = auto()
    INPLACE_X = auto()


_INPLACE = [
    Inplace.INPLACE_X,
51
    Inplace.OUT_OF_PLACE,
52
53
54
55
56
57
58
59
]

_TEST_CASES = [
    test_case + (inplace_item,)
    for test_case in _TEST_CASES_
    for inplace_item in _INPLACE
]

xgqdut2016's avatar
xgqdut2016 committed
60
61
62
63
DEBUG = False
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000
PanZezhongQY's avatar
PanZezhongQY committed
64

xgqdut2016's avatar
xgqdut2016 committed
65

PanZezhongQY's avatar
PanZezhongQY committed
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class CausalSoftmaxDescriptor(Structure):
    _fields_ = [("device", c_int32)]


infiniopCausalSoftmaxDescriptor_t = POINTER(CausalSoftmaxDescriptor)


def causal_softmax(x):
    type = x.dtype
    mask = torch.tril(torch.ones_like(x), diagonal=-1).flip(dims=[-2, -1])
    y = x.clone()
    masked = torch.where(mask == 1, -torch.inf, y.to(torch.float32))
    return torch.nn.functional.softmax(masked, dim=-1).to(type)


81
82
83
84
85
86
87
88
89
def test(
    lib,
    handle,
    torch_device,
    shape,
    x_stride=None,
    y_stride=None,
    inplace=Inplace.OUT_OF_PLACE,
    dtype=torch.float16,
90
    sync=None
91
):
PanZezhongQY's avatar
PanZezhongQY committed
92
    print(
93
        f"Testing CausalSoftmax on {torch_device} with shape:{shape} x_stride:{x_stride} y_stride:{y_stride} dtype:{dtype} inplace:{inplace}"
PanZezhongQY's avatar
PanZezhongQY committed
94
    )
xgqdut2016's avatar
xgqdut2016 committed
95

96
    x = torch.rand(shape, dtype=dtype).to(torch_device)
xgqdut2016's avatar
xgqdut2016 committed
97

PanZezhongQY's avatar
PanZezhongQY committed
98
    ans = causal_softmax(x)
xgqdut2016's avatar
xgqdut2016 committed
99
100

    x = rearrange_if_needed(x, x_stride)
xgqdut2016's avatar
xgqdut2016 committed
101

PanZezhongQY's avatar
PanZezhongQY committed
102
    x_tensor = to_tensor(x, lib)
xgqdut2016's avatar
xgqdut2016 committed
103

104
105
106
107
108
109
110
    if inplace == Inplace.INPLACE_X:
        y = x
        y_tensor = x_tensor
    else:
        y = torch.zeros(shape, dtype=dtype).to(torch_device)
        y = rearrange_if_needed(y, y_stride)
        y_tensor = to_tensor(y, lib)
111
112
113
        
    if sync is not None:
        sync()
114

PanZezhongQY's avatar
PanZezhongQY committed
115
116
117
    descriptor = infiniopCausalSoftmaxDescriptor_t()
    check_error(
        lib.infiniopCreateCausalSoftmaxDescriptor(
118
            handle, ctypes.byref(descriptor), y_tensor.descriptor, x_tensor.descriptor
PanZezhongQY's avatar
PanZezhongQY committed
119
120
        )
    )
xgqdut2016's avatar
xgqdut2016 committed
121
122

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
123
    x_tensor.destroyDesc(lib)
xgqdut2016's avatar
xgqdut2016 committed
124

PanZezhongQY's avatar
PanZezhongQY committed
125
126
127
128
129
130
131
    workspace_size = c_uint64(0)
    check_error(
        lib.infiniopGetCausalSoftmaxWorkspaceSize(
            descriptor, ctypes.byref(workspace_size)
        )
    )
    workspace = create_workspace(workspace_size.value, x.device)
xgqdut2016's avatar
xgqdut2016 committed
132

xgqdut2016's avatar
xgqdut2016 committed
133
134
135
136
137
138
    def lib_causal_softmax():
        check_error(
            lib.infiniopCausalSoftmax(
                descriptor,
                workspace.data_ptr() if workspace is not None else None,
                workspace_size.value,
139
                y_tensor.data,
xgqdut2016's avatar
xgqdut2016 committed
140
141
142
                x_tensor.data,
                None,
            )
PanZezhongQY's avatar
PanZezhongQY committed
143
        )
xgqdut2016's avatar
xgqdut2016 committed
144

xgqdut2016's avatar
xgqdut2016 committed
145
    lib_causal_softmax()
xgqdut2016's avatar
xgqdut2016 committed
146

xgqdut2016's avatar
xgqdut2016 committed
147
148
    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
    if DEBUG:
149
150
        debug(y, ans, atol=atol, rtol=rtol)
    assert torch.allclose(y, ans, atol=atol, rtol=rtol)
xgqdut2016's avatar
xgqdut2016 committed
151
152
153
154
155
156
157

    # Profiling workflow
    if PROFILE:
        # fmt: off
        profile_operation("PyTorch", lambda: causal_softmax(x), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        profile_operation("    lib", lambda: lib_causal_softmax(), torch_device, NUM_PRERUN, NUM_ITERATIONS)
        # fmt: on
PanZezhongQY's avatar
PanZezhongQY committed
158

xgqdut2016's avatar
xgqdut2016 committed
159
    check_error(lib.infiniopDestroyCausalSoftmaxDescriptor(descriptor))
160

PanZezhongQY's avatar
PanZezhongQY committed
161
162
163
164

if __name__ == "__main__":
    args = get_args()
    lib = open_lib()
xgqdut2016's avatar
xgqdut2016 committed
165

PanZezhongQY's avatar
PanZezhongQY committed
166
167
168
169
170
171
    lib.infiniopCreateCausalSoftmaxDescriptor.restype = c_int32
    lib.infiniopCreateCausalSoftmaxDescriptor.argtypes = [
        infiniopHandle_t,
        POINTER(infiniopCausalSoftmaxDescriptor_t),
        infiniopTensorDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
172

PanZezhongQY's avatar
PanZezhongQY committed
173
174
175
176
177
    lib.infiniopGetCausalSoftmaxWorkspaceSize.restype = c_int32
    lib.infiniopGetCausalSoftmaxWorkspaceSize.argtypes = [
        infiniopCausalSoftmaxDescriptor_t,
        POINTER(c_uint64),
    ]
xgqdut2016's avatar
xgqdut2016 committed
178

PanZezhongQY's avatar
PanZezhongQY committed
179
180
181
182
183
184
185
186
    lib.infiniopCausalSoftmax.restype = c_int32
    lib.infiniopCausalSoftmax.argtypes = [
        infiniopCausalSoftmaxDescriptor_t,
        c_void_p,
        c_uint64,
        c_void_p,
        c_void_p,
    ]
xgqdut2016's avatar
xgqdut2016 committed
187

PanZezhongQY's avatar
PanZezhongQY committed
188
189
190
191
    lib.infiniopDestroyCausalSoftmaxDescriptor.restype = c_int32
    lib.infiniopDestroyCausalSoftmaxDescriptor.argtypes = [
        infiniopCausalSoftmaxDescriptor_t,
    ]
xgqdut2016's avatar
xgqdut2016 committed
192

xgqdut2016's avatar
xgqdut2016 committed
193
194
195
196
197
    # Configure testing options
    DEBUG = args.debug
    PROFILE = args.profile
    NUM_PRERUN = args.num_prerun
    NUM_ITERATIONS = args.num_iterations
xgqdut2016's avatar
xgqdut2016 committed
198

xgqdut2016's avatar
xgqdut2016 committed
199
200
    for device in get_test_devices(args):
        test_operator(lib, device, test, _TEST_CASES, _TENSOR_DTYPES)
PanZezhongQY's avatar
PanZezhongQY committed
201
202

    print("\033[92mTest passed!\033[0m")