profile_gemm.cpp 14.4 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
6
#include <iostream>
#include <numeric>
#include <initializer_list>
#include <cstdlib>
#include <stdlib.h>
#include <half.hpp>
Chao Liu's avatar
Chao Liu committed
7

Chao Liu's avatar
Chao Liu committed
8
9
#include "profile_gemm_impl.hpp"

Chao Liu's avatar
Chao Liu committed
10
11
// return true if test pass
bool profile_gemm(int argc, char* argv[])
Chao Liu's avatar
Chao Liu committed
12
{
Chao Liu's avatar
Chao Liu committed
13
14
15
16
17
18
19
    enum struct GemmMatrixLayout
    {
        MK_KN_MN, // 0
        MK_NK_MN, // 1
        KM_KN_MN, // 2
        KM_NK_MN, // 3
    };
Chao Liu's avatar
Chao Liu committed
20

Chao Liu's avatar
Chao Liu committed
21
22
23
24
25
26
27
    enum struct GemmDataType
    {
        F32_F32_F32,    // 0
        F16_F16_F16,    // 1
        BF16_BF16_BF16, // 2
        INT8_INT8_INT8, // 3
    };
Chao Liu's avatar
Chao Liu committed
28

Chao Liu's avatar
Chao Liu committed
29
    if(argc != 14)
Chao Liu's avatar
Chao Liu committed
30
31
    {
        printf("arg1: tensor operation (gemm: GEMM)\n");
32
        printf("arg2: data type (0: fp32; 1: fp16; 2: bf16; 3: int8)\n");
Chao Liu's avatar
Chao Liu committed
33
34
        printf("arg3: matrix layout (0: A[m, k] * B[k, n] = C[m, n];\n");
        printf("                     1: A[m, k] * B[n, k] = C[m, n];\n");
ltqin's avatar
ltqin committed
35
36
        printf("                     2: A[k, m] * B[k, n] = C[m, n];\n");
        printf("                     3: A[k, m] * B[n, k] = C[m, n])\n");
Chao Liu's avatar
Chao Liu committed
37
38
39
40
41
        printf("arg4: verification (0: no; 1: yes)\n");
        printf("arg5: initialization (0: no init; 1: integer value; 2: decimal value)\n");
        printf("arg8: print tensor value (0: no; 1: yes)\n");
        printf("arg7: run kernel # of times (>1)\n");
        printf("arg8 to 13: M, N, K, StrideA, StrideB, StrideC\n");
Chao Liu's avatar
Chao Liu committed
42
        return false;
Chao Liu's avatar
Chao Liu committed
43
44
    }

Chao Liu's avatar
Chao Liu committed
45
46
    const auto data_type       = static_cast<GemmDataType>(std::stoi(argv[2]));
    const auto layout          = static_cast<GemmMatrixLayout>(std::stoi(argv[3]));
Chao Liu's avatar
Chao Liu committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
    const bool do_verification = std::stoi(argv[4]);
    const int init_method      = std::stoi(argv[5]);
    const bool do_log          = std::stoi(argv[6]);
    const int nrepeat          = std::stoi(argv[7]);

    const int M = std::stoi(argv[8]);
    const int N = std::stoi(argv[9]);
    const int K = std::stoi(argv[10]);

    const int StrideA = std::stoi(argv[11]);
    const int StrideB = std::stoi(argv[12]);
    const int StrideC = std::stoi(argv[13]);

    if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
62
63
64
65
66
67
        return ck::profiler::profile_gemm_impl<ck::half_t,
                                               ck::half_t,
                                               ck::half_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
68
69
70
71
72
73
74
75
76
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
77
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
78
79
80
    }
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
81
82
83
84
85
86
        return ck::profiler::profile_gemm_impl<ck::half_t,
                                               ck::half_t,
                                               ck::half_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
87
88
89
90
91
92
93
94
95
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
96
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
97
98
99
    }
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
100
101
102
103
104
105
        return ck::profiler::profile_gemm_impl<ck::half_t,
                                               ck::half_t,
                                               ck::half_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
106
107
108
109
110
111
112
113
114
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
115
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
116
117
118
    }
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
119
120
121
122
123
124
        return ck::profiler::profile_gemm_impl<ck::half_t,
                                               ck::half_t,
                                               ck::half_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
125
126
127
128
129
130
131
132
133
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
134
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
135
136
137
    }
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
138
139
140
141
142
143
        return ck::profiler::profile_gemm_impl<float,
                                               float,
                                               float,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
144
145
146
147
148
149
150
151
152
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
153
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
154
155
156
    }
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
157
158
159
160
161
162
        return ck::profiler::profile_gemm_impl<float,
                                               float,
                                               float,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
163
164
165
166
167
168
169
170
171
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
172
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
173
174
175
    }
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
176
177
178
179
180
181
        return ck::profiler::profile_gemm_impl<float,
                                               float,
                                               float,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
182
183
184
185
186
187
188
189
190
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
191
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
192
193
194
    }
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
195
196
197
198
199
200
        return ck::profiler::profile_gemm_impl<float,
                                               float,
                                               float,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
Chao Liu's avatar
Chao Liu committed
201
202
203
204
205
206
207
208
209
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
210
            (StrideC < 0) ? N : StrideC);
Chao Liu's avatar
Chao Liu committed
211
    }
212
213
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::MK_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
214
215
216
217
218
219
        return ck::profiler::profile_gemm_impl<int8_t,
                                               int8_t,
                                               int8_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
220
221
222
223
224
225
226
227
228
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
229
            (StrideC < 0) ? N : StrideC);
230
    }
231
232
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::MK_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
233
234
235
236
237
238
        return ck::profiler::profile_gemm_impl<int8_t,
                                               int8_t,
                                               int8_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
239
240
241
242
243
244
245
246
247
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
248
            (StrideC < 0) ? N : StrideC);
249
    }
250
251
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::KM_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
252
253
254
255
256
257
        return ck::profiler::profile_gemm_impl<int8_t,
                                               int8_t,
                                               int8_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
258
259
260
261
262
263
264
265
266
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
267
            (StrideC < 0) ? N : StrideC);
268
269
270
    }
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::KM_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
271
272
273
274
275
276
        return ck::profiler::profile_gemm_impl<int8_t,
                                               int8_t,
                                               int8_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
277
278
279
280
281
282
283
284
285
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
286
            (StrideC < 0) ? N : StrideC);
287
288
289
    }
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::MK_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
290
291
292
293
294
295
        return ck::profiler::profile_gemm_impl<ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
296
297
298
299
300
301
302
303
304
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? K : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
305
            (StrideC < 0) ? N : StrideC);
306
    }
307
308
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::MK_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
309
310
311
312
313
314
        return ck::profiler::profile_gemm_impl<ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
315
316
317
318
319
320
321
322
323
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
324
            (StrideC < 0) ? N : StrideC);
325
    }
326
327
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::KM_KN_MN)
    {
Chao Liu's avatar
Chao Liu committed
328
329
330
331
332
333
        return ck::profiler::profile_gemm_impl<ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
334
335
336
337
338
339
340
341
342
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? N : StrideB,
Chao Liu's avatar
Chao Liu committed
343
            (StrideC < 0) ? N : StrideC);
344
345
346
    }
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::KM_NK_MN)
    {
Chao Liu's avatar
Chao Liu committed
347
348
349
350
351
352
        return ck::profiler::profile_gemm_impl<ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::bhalf_t,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::ColumnMajor,
                                               ck::tensor_layout::gemm::RowMajor>(
353
354
355
356
357
358
359
360
361
            do_verification,
            init_method,
            do_log,
            nrepeat,
            M,
            N,
            K,
            (StrideA < 0) ? M : StrideA,
            (StrideB < 0) ? K : StrideB,
Chao Liu's avatar
Chao Liu committed
362
            (StrideC < 0) ? N : StrideC);
363
    }
Chao Liu's avatar
Chao Liu committed
364
365
    else
    {
Chao Liu's avatar
Chao Liu committed
366
        std::cout << "this data_type & layout is not implemented" << std::endl;
Chao Liu's avatar
Chao Liu committed
367

Chao Liu's avatar
Chao Liu committed
368
369
        return true;
    }
Chao Liu's avatar
Chao Liu committed
370
}