gemm.py 6.36 KB
Newer Older
PanZezhong's avatar
PanZezhong committed
1
import torch
PanZezhongQY's avatar
PanZezhongQY committed
2
import ctypes
PanZezhong's avatar
PanZezhong 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,
PanZezhong's avatar
PanZezhong committed
7
8
9
    open_lib,
    to_tensor,
    get_test_devices,
PanZezhongQY's avatar
PanZezhongQY committed
10
    check_error,
PanZezhong's avatar
PanZezhong committed
11
12
13
14
15
16
17
    rearrange_if_needed,
    create_workspace,
    test_operator,
    get_args,
    debug,
    get_tolerance,
    profile_operation,
PanZezhongQY's avatar
PanZezhongQY committed
18
19
)

PanZezhong's avatar
PanZezhong committed
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# ==============================================================================
#  Configuration (Internal Use Only)
# ==============================================================================
# These are not meant to be imported from other modules
_TEST_CASES = [
    # alpha, beta, a_shape, b_shape, c_shape, a_stride, b_stride, c_stride
    (1.0, 0.0, (1, 2048), (2048, 2048), (1, 2048), None, None, None),
    (1.0, 0.0, (2, 4, 2048), (2, 2048, 2048), (2, 4, 2048), None, None, None),
    (1.0, 0.0, (1, 2048), (2048, 2048), (1, 2048), (4096, 1), (4096, 1), (4096, 1)),
    (1.0, 1.0, (6, 2048), (2048, 2560), (6, 2560), (2048, 1), (1, 2048), (2560, 1)),
    (1.0 / 8.0, 0.0, (4, 8 * 6, 64), (4, 64, 6), (4, 8 * 6, 6), None, None, None),
]

# Data types used for testing
_TENSOR_DTYPES = [torch.float16, torch.float32]

# Tolerance map for different data types
_TOLERANCE_MAP = {
    torch.float16: {"atol": 0, "rtol": 1e-2},
    torch.float32: {"atol": 0, "rtol": 1e-3},
}

DEBUG = False
PanZezhongQY's avatar
PanZezhongQY committed
43
44
45
46
PROFILE = False
NUM_PRERUN = 10
NUM_ITERATIONS = 1000

47

PanZezhong's avatar
PanZezhong committed
48
49
50
51
# ==============================================================================
#  Definitions
# ==============================================================================
class GemmDescriptor(Structure):
PanZezhongQY's avatar
PanZezhongQY committed
52
53
54
    _fields_ = [("device", c_int32)]


PanZezhong's avatar
PanZezhong committed
55
infiniopGemmDescriptor_t = POINTER(GemmDescriptor)
PanZezhongQY's avatar
PanZezhongQY committed
56
57


PanZezhong's avatar
PanZezhong committed
58
# PyTorch implementation for matrix multiplication
59
60
61
62
63
64
65
66
def gemm(d, _c, beta, _a, _b, alpha):
    if _c.ndim == 2:
        torch.addmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
    elif _c.ndim == 3:
        torch.baddbmm(_c, _a, _b, beta=beta, alpha=alpha, out=d)
    else:
        torch.matmul(_a, _b, out=d)
        d.mul_(alpha).add_(_c, alpha=beta)
PanZezhongQY's avatar
PanZezhongQY committed
67
68


PanZezhong's avatar
PanZezhong committed
69
70
# The argument list should be (lib, handle, torch_device, <param list>, dtype)
# The <param list> should keep the same order as the one specified in _TEST_CASES
PanZezhongQY's avatar
PanZezhongQY committed
71
72
73
74
75
76
77
78
79
80
81
82
83
def test(
    lib,
    handle,
    torch_device,
    alpha,
    beta,
    a_shape,
    b_shape,
    c_shape,
    a_stride=None,
    b_stride=None,
    c_stride=None,
    dtype=torch.float16,
84
    sync=None
PanZezhongQY's avatar
PanZezhongQY committed
85
86
):
    print(
PanZezhong's avatar
PanZezhong committed
87
88
89
        f"Testing Gemm on {torch_device} with alpha:{alpha}, beta:{beta},"
        f" a_shape:{a_shape}, b_shape:{b_shape}, c_shape:{c_shape},"
        f" a_stride:{a_stride}, b_stride:{b_stride}, c_stride:{c_stride}, dtype:{dtype}"
PanZezhongQY's avatar
PanZezhongQY committed
90
91
    )

PanZezhong's avatar
PanZezhong committed
92
    # Initialize tensors
PanZezhongQY's avatar
PanZezhongQY committed
93
94
    a = torch.rand(a_shape, dtype=dtype).to(torch_device)
    b = torch.rand(b_shape, dtype=dtype).to(torch_device)
PanZezhong's avatar
PanZezhong committed
95
    c = torch.ones(c_shape, dtype=dtype).to(torch_device)
96
    ans = torch.zeros(c_shape, dtype=dtype).to(torch_device)
PanZezhongQY's avatar
PanZezhongQY committed
97

PanZezhong's avatar
PanZezhong committed
98
    # Compute the PyTorch reference result
99
    gemm(ans, c, beta, a, b, alpha)
PanZezhongQY's avatar
PanZezhongQY committed
100

PanZezhong's avatar
PanZezhong committed
101
102
103
104
105
    a, b, c = [
        rearrange_if_needed(tensor, stride)
        for tensor, stride in zip([a, b, c], [a_stride, b_stride, c_stride])
    ]
    a_tensor, b_tensor, c_tensor = [to_tensor(tensor, lib) for tensor in [a, b, c]]
PanZezhongQY's avatar
PanZezhongQY committed
106

107
108
109
    if sync is not None:
        sync()

PanZezhong's avatar
PanZezhong committed
110
    descriptor = infiniopGemmDescriptor_t()
PanZezhongQY's avatar
PanZezhongQY committed
111
    check_error(
PanZezhong's avatar
PanZezhong committed
112
        lib.infiniopCreateGemmDescriptor(
PanZezhongQY's avatar
PanZezhongQY committed
113
114
            handle,
            ctypes.byref(descriptor),
PanZezhong's avatar
PanZezhong committed
115
            c_tensor.descriptor,
PanZezhongQY's avatar
PanZezhongQY committed
116
117
118
119
120
121
            a_tensor.descriptor,
            b_tensor.descriptor,
        )
    )

    # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
PanZezhong's avatar
PanZezhong committed
122
123
    for tensor in [a_tensor, b_tensor, c_tensor]:
        tensor.destroyDesc(lib)
PanZezhongQY's avatar
PanZezhongQY committed
124

PanZezhong's avatar
PanZezhong committed
125
126
    # Get workspace size and create workspace
    workspace_size = c_uint64(0)
PanZezhongQY's avatar
PanZezhongQY committed
127
    check_error(
PanZezhong's avatar
PanZezhong committed
128
        lib.infiniopGetGemmWorkspaceSize(descriptor, ctypes.byref(workspace_size))
PanZezhongQY's avatar
PanZezhongQY committed
129
    )
PanZezhong's avatar
PanZezhong committed
130
    workspace = create_workspace(workspace_size.value, a.device)
PanZezhongQY's avatar
PanZezhongQY committed
131

PanZezhong's avatar
PanZezhong committed
132
133
    # Execute infiniop gemm operator
    def lib_gemm():
PanZezhongQY's avatar
PanZezhongQY committed
134
        check_error(
PanZezhong's avatar
PanZezhong committed
135
            lib.infiniopGemm(
PanZezhongQY's avatar
PanZezhongQY committed
136
                descriptor,
PanZezhong's avatar
PanZezhong committed
137
138
139
                workspace.data_ptr() if workspace is not None else None,
                workspace_size.value,
                c_tensor.data,
PanZezhongQY's avatar
PanZezhongQY committed
140
141
                a_tensor.data,
                b_tensor.data,
PanZezhong's avatar
PanZezhong committed
142
143
                alpha,
                beta,
PanZezhongQY's avatar
PanZezhongQY committed
144
145
146
147
                None,
            )
        )

PanZezhong's avatar
PanZezhong committed
148
    lib_gemm()
PanZezhongQY's avatar
PanZezhongQY committed
149

PanZezhong's avatar
PanZezhong committed
150
151
152
153
154
    # Validate results
    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
    if DEBUG:
        debug(c, ans, atol=atol, rtol=rtol)
    assert torch.allclose(c, ans, atol=atol, rtol=rtol)
PanZezhongQY's avatar
PanZezhongQY committed
155

PanZezhong's avatar
PanZezhong committed
156
157
    # Profiling workflow
    if PROFILE:
158
        # fmt: off
159
        profile_operation("PyTorch", lambda: gemm(ans, c, beta, a, b, alpha), torch_device, NUM_PRERUN, NUM_ITERATIONS)
PanZezhong's avatar
PanZezhong committed
160
        profile_operation("    lib", lambda: lib_gemm(), torch_device, NUM_PRERUN, NUM_ITERATIONS)
161
        # fmt: on
PanZezhong's avatar
PanZezhong committed
162
    check_error(lib.infiniopDestroyGemmDescriptor(descriptor))
PanZezhongQY's avatar
PanZezhongQY committed
163
164


PanZezhong's avatar
PanZezhong committed
165
166
167
# ==============================================================================
#  Main Execution
# ==============================================================================
PanZezhongQY's avatar
PanZezhongQY committed
168
169
170
171
if __name__ == "__main__":
    args = get_args()
    lib = open_lib()

PanZezhong's avatar
PanZezhong committed
172
173
    lib.infiniopCreateGemmDescriptor.restype = c_int32
    lib.infiniopCreateGemmDescriptor.argtypes = [
PanZezhongQY's avatar
PanZezhongQY committed
174
        infiniopHandle_t,
PanZezhong's avatar
PanZezhong committed
175
        POINTER(infiniopGemmDescriptor_t),
PanZezhongQY's avatar
PanZezhongQY committed
176
177
178
179
180
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
        infiniopTensorDescriptor_t,
    ]

PanZezhong's avatar
PanZezhong committed
181
182
183
184
    lib.infiniopGetGemmWorkspaceSize.restype = c_int32
    lib.infiniopGetGemmWorkspaceSize.argtypes = [
        infiniopGemmDescriptor_t,
        POINTER(c_size_t),
PanZezhongQY's avatar
PanZezhongQY committed
185
186
    ]

PanZezhong's avatar
PanZezhong committed
187
188
189
    lib.infiniopGemm.restype = c_int32
    lib.infiniopGemm.argtypes = [
        infiniopGemmDescriptor_t,
PanZezhongQY's avatar
PanZezhongQY committed
190
191
192
193
194
        c_void_p,
        c_uint64,
        c_void_p,
        c_void_p,
        c_void_p,
PanZezhong's avatar
PanZezhong committed
195
196
        c_float,
        c_float,
PanZezhongQY's avatar
PanZezhongQY committed
197
198
199
        c_void_p,
    ]

PanZezhong's avatar
PanZezhong committed
200
201
202
    lib.infiniopDestroyGemmDescriptor.restype = c_int32
    lib.infiniopDestroyGemmDescriptor.argtypes = [
        infiniopGemmDescriptor_t,
PanZezhongQY's avatar
PanZezhongQY committed
203
204
    ]

PanZezhong's avatar
PanZezhong committed
205
206
207
208
209
210
211
212
213
214
    # Configure testing options
    DEBUG = args.debug
    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
215
    print("\033[92mTest passed!\033[0m")