cpu_gemm_uk.cpp 16 KB
Newer Older
1
2
3
4
5
6
7
8
#include <iostream>
#include <initializer_list>
#include <cstdlib>
#include <stdlib.h>
#include <string>
#include <sstream>
#include <tuple>
#include <memory>
carlushuang's avatar
carlushuang committed
9
#include <omp.h>
carlushuang's avatar
carlushuang committed
10
11
12
13
14
15
16
17
18
#include <string.h>
#include <chrono>
#include "ck/ck.hpp"
#include "ck/library/host_tensor/host_tensor.hpp"
#include "ck/device_utility/kernel_launch.hpp"
#include "ck/library/host_tensor/device_memory.hpp"
#include "ck/utility/print.hpp"
#include "ck/utility/cpuid.hpp"
#include "ck/tensor_operation/cpu/thread/threadwise_gemm_avx2.hpp"
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

#define ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE(FA, FB, FC, TA, TB, NT)   \
    ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 6, 16, TA, TB, NT>,     \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 5, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 4, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 3, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 2, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 1, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 6, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 5, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 4, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 3, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 2, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC, 1, 8, TA, TB, NT>

carlushuang's avatar
carlushuang committed
34
//#define ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE(FA, FB, FC, TA, TB, NT)  \
35
36
//     ck::cpu::ThreadwiseGemmAvx2_MxN_6x16<FA, FB, FC,  6, 16,  TA,  TB,  NT>

carlushuang's avatar
carlushuang committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#define ITERATE_THREAD_GEMM_AVX2_MXN_4X24_INSTANCE(FA, FB, FC, TA, TB, NT)   \
    ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 4, 24, TA, TB, NT>,     \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 3, 24, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 2, 24, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 1, 24, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 4, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 3, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 2, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 1, 16, TA, TB, NT>, \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 4, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 3, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 2, 8, TA, TB, NT>,  \
        ck::cpu::ThreadwiseGemmAvx2_MxN_4x24<FA, FB, FC, 1, 8, TA, TB, NT>

51
52
53
using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;

carlushuang's avatar
carlushuang committed
54
55
56
57
58
// using AType = half_float::half;
// using BType = half_float::half;
using AType = float;
using BType = float;
using CType = float;
carlushuang's avatar
carlushuang committed
59
#define NTStore false
carlushuang's avatar
carlushuang committed
60

61
template <typename ALayout, typename BLayout>
62
63
using thread_gemm_avx2_mxn_6x16_instances = std::tuple<
    // clang-format off
64
    //                                        FloatA FloatB FloatC  ALayout  BLayout NTStore
carlushuang's avatar
carlushuang committed
65
66
67
68
    ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore)
69

carlushuang's avatar
carlushuang committed
70
    // ITERATE_THREAD_GEMM_AVX2_MXN_6X16_INSTANCE(AType, BType, CType,    ALayout,    BLayout, NTStore)
71
72
73
    // clang-format on
    >;

carlushuang's avatar
carlushuang committed
74
75
76
77
template <typename ALayout, typename BLayout>
using thread_gemm_avx2_mxn_4x24_instances = std::tuple<
    // clang-format off
    //                                        FloatA FloatB FloatC  ALayout  BLayout NTStore
carlushuang's avatar
carlushuang committed
78
79
80
81
    ITERATE_THREAD_GEMM_AVX2_MXN_4X24_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_4X24_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_4X24_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore),
    ITERATE_THREAD_GEMM_AVX2_MXN_4X24_INSTANCE( AType, BType, CType, ALayout, BLayout, NTStore)
carlushuang's avatar
carlushuang committed
82
83
84
    // clang-format on
    >;

85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
void dump_cache_hierarchy()
{
    auto dump_cache_type = [&](const ck::cpu::cpuid_cache_type& type) {
        if(type == ck::cpu::cpuid_cache_type_dcache)
            printf("data cache");
        else if(type == ck::cpu::cpuid_cache_type_icache)
            printf("inst cache");
        else if(type == ck::cpu::cpuid_cache_type_unified)
            printf("unif cache");
    };
    auto dump_cache_detail = [&](const ck::cpu::cpuid_cache_detail& detail) {
        dump_cache_type(static_cast<const ck::cpu::cpuid_cache_type>(detail.type));
        printf(" size:%u, cache_line:%u, associativity:%u, sets:%u, partitions:%u, shared by "
               "procs:%u(%u)\n",
               detail.size,
               detail.cache_line_size,
               detail.associativity,
               detail.sets,
               detail.partitions,
               detail.shared_by_procs,
               detail.cores_per_socket);
    };

    ck::cpu::cpuid_cache_hierarchy cache = ck::cpu::cpuid_query_cache();
    if(cache.l1d.size != 0)
    {
        printf("l1 ");
        dump_cache_detail(cache.l1d);
    }
    if(cache.l1i.size != 0)
    {
        printf("l1 ");
        dump_cache_detail(cache.l1i);
    }
    if(cache.l2.size != 0)
    {
        printf("l2 ");
        dump_cache_detail(cache.l2);
    }
    if(cache.l3.size != 0)
    {
        printf("l3 ");
        dump_cache_detail(cache.l3);
    }
    if(cache.l4.size != 0)
    {
        printf("l4 ");
        dump_cache_detail(cache.l4);
    }
}

template <typename T>
void rand_vector(T* v, int elem)
{
    int i;

    static int flag = 0;
    if(!flag)
    {
        srand(time(nullptr));
        flag = 1;
    }

    for(i = 0; i < elem; i++)
    {
        v[i] = (static_cast<T>(rand() % 100)) / 100.0f;
    }
}

bool valid_vector(const float* ref, const float* rhs, uint32_t elem)
{
    float rtol   = 1e-5;
    float atol   = 1e-8;
    uint32_t err = 0;
    for(uint32_t i = 0; i < elem; i++)
    {
        float diff = std::abs(ref[i] - rhs[i]);
        if(diff > atol + rtol * std::abs(ref[i]))
        {
            printf("diff at %u, ref:%f, rhs:%f\n", i, ref[i], rhs[i]);
            err++;
        }
    }

    return err == 0;
}

carlushuang's avatar
carlushuang committed
172
173
174
template <typename FloatA, typename FloatB, typename ALayout, typename BLayout>
void ref_cpu_gemm_uk(
    const FloatA* a, const FloatB* b, float* c, float alpha, uint32_t m, uint32_t n, uint32_t k)
175
{
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
    auto f_host_2d_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
            if(std::is_same<decltype(layout), Row>::value)
            {
                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
                                            std::vector<std::size_t>({stride, 1}));
            }
            else
            {
                return HostTensorDescriptor(std::vector<std::size_t>({row, col}),
                                            std::vector<std::size_t>({1, stride}));
            }
        };

    auto f_host_vectored_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t vec, std::size_t stride) {
            // only valid in row major. stride is for each row, contains vector size
            return HostTensorDescriptor(std::vector<std::size_t>({row, col, vec}),
                                        std::vector<std::size_t>({stride, vec, 1}));
        };

    std::size_t lda = std::is_same<Row, ALayout>::value ? k : m;     // in unit of element
    std::size_t ldb = std::is_same<Row, BLayout>::value ? n : k * 8; // in unit of element
    std::size_t ldc = n;
    HostTensorDescriptor a_m_k = f_host_2d_tensor_descriptor(m, n, lda, ALayout{});
    HostTensorDescriptor b_k_n = std::is_same<Row, BLayout>::value
                                     ? f_host_2d_tensor_descriptor(k, n, ldb, BLayout{})
                                     : f_host_vectored_tensor_descriptor(n / 8, k, 8, ldb);
    HostTensorDescriptor c_m_n = f_host_2d_tensor_descriptor(m, n, ldc, Row{});
205
206
207
208
209
210
211
212

    for(uint32_t im = 0; im < m; im++)
    {
        for(uint32_t in = 0; in < n; in++)
        {
            float acc = .0f;
            for(uint32_t ik = 0; ik < k; ik++)
            {
213
214
215
216
217
                acc += static_cast<float>(a[a_m_k.GetOffsetFromMultiIndex(im, ik)]) *
                       (std::is_same<Row, BLayout>::value
                            ? static_cast<float>(b[b_k_n.GetOffsetFromMultiIndex(ik, in)])
                            : static_cast<float>(
                                  b[b_k_n.GetOffsetFromMultiIndex(in / 8, ik, in % 8)]));
218
219
            }
            acc *= alpha;
220
            c[c_m_n.GetOffsetFromMultiIndex(im, in)] = acc;
221
222
223
224
        }
    }
}

carlushuang's avatar
carlushuang committed
225
template <typename FloatA, typename FloatB, typename ALayout, typename BLayout, typename ukenrel_t>
226
void test_ukernel(ukenrel_t uk,
carlushuang's avatar
carlushuang committed
227
228
                  FloatA* mat_a,
                  FloatB* mat_b,
229
230
231
232
233
234
                  float* mat_c,
                  float alpha,
                  uint32_t m,
                  uint32_t n,
                  uint32_t k)
{
carlushuang's avatar
carlushuang committed
235
236
237
    int max_threads = omp_get_max_threads();

    auto invoke_uk = [&](ck::cpu::ThreadwiseGemmParam& param, float* current_mat_c) {
carlushuang's avatar
carlushuang committed
238
239
240
        assert(m % uk.ThreadMr == 0 && n % uk.ThreadNr == 0);

        for(uint32_t i_m = 0; i_m < m; i_m += uk.ThreadMr)
241
        {
carlushuang's avatar
carlushuang committed
242
            if constexpr(std::is_same<Row, ALayout>::value)
243
            {
carlushuang's avatar
carlushuang committed
244
                param.p_a = mat_a + i_m * k;
245
            }
carlushuang's avatar
carlushuang committed
246
            else
247
            {
carlushuang's avatar
carlushuang committed
248
                param.p_a = mat_a + i_m;
249
            }
250
            for(uint32_t i_n = 0; i_n < n; i_n += uk.ThreadNr)
251
            {
carlushuang's avatar
carlushuang committed
252
253
254
255
256
257
258
259
260
                if constexpr(std::is_same<Row, BLayout>::value)
                {
                    param.p_b = mat_b + i_n;
                }
                else
                {
                    param.p_b = mat_b + i_n * k;
                }
                param.p_c = current_mat_c + i_m * n + i_n;
261
262
263
264
265
                uk.Run(&param);
            }
        }
    };

266
    printf("gemm_uk_%dx%d_%c%c: ", uk.ThreadMr, uk.ThreadNr, ALayout::name[0], BLayout::name[0]);
267
268
269
    fflush(stdout);
    // printf("%s: ", typeid(uk).name());fflush(stdout);

carlushuang's avatar
carlushuang committed
270
    float us = .0f;
271

carlushuang's avatar
carlushuang committed
272
#pragma omp parallel reduction(+ : us)
273
    {
carlushuang's avatar
carlushuang committed
274
275
        int tid = omp_get_thread_num();
        DeviceAlignedMemCPU private_c_mem(m * n * sizeof(float), 32);
276
277
        // float* private_c = reinterpret_cast<float*>(private_c_mem.mpDeviceBuf);
        float* private_c = mat_c + tid * m * n;
278

carlushuang's avatar
carlushuang committed
279
        ck::cpu::ThreadwiseGemmParam param;
280
281
282
283
284
285
286
287
288
        param.p_a         = mat_a;
        param.p_b         = mat_b;
        param.p_c         = private_c;
        param.Kr          = k;
        param.lda         = (std::is_same<Row, ALayout>::value ? k : m) * sizeof(FloatA);
        param.ldb         = (std::is_same<Row, BLayout>::value ? n : k * 8) * sizeof(FloatB);
        param.ldc         = n * sizeof(float);
        param.alpha       = alpha;
        param.accmulate_c = 0;
289

carlushuang's avatar
carlushuang committed
290
291
292
293
294
295
296
297
298
299
        memset(private_c, 0, m * n * sizeof(float));

        int repeat = 7e10 / (2 * m * n * k);

        for(int i = 0; i < (repeat / 5); i++)
        {
            invoke_uk(param, private_c);
        }


carlushuang's avatar
carlushuang committed
300
        auto mStart = std::chrono::high_resolution_clock::now();
carlushuang's avatar
carlushuang committed
301
302
303
304
        for(int i = 0; i < repeat; i++)
        {
            invoke_uk(param, private_c);
        }
carlushuang's avatar
carlushuang committed
305
        auto mStop = std::chrono::high_resolution_clock::now();
carlushuang's avatar
carlushuang committed
306

carlushuang's avatar
carlushuang committed
307
308
        us += static_cast<float>(
                   std::chrono::duration_cast<std::chrono::microseconds>(mStop - mStart).count()) / repeat;
carlushuang's avatar
carlushuang committed
309
310
311
312
313

        memset(private_c, 0, m * n * sizeof(float));
        invoke_uk(param, private_c);

        memcpy(mat_c + tid * m * n, private_c, m * n * sizeof(float));
314
315
    }

carlushuang's avatar
carlushuang committed
316
    us = us / max_threads;
317

carlushuang's avatar
carlushuang committed
318
    float gflops = static_cast<float>(2 * m * n * k * max_threads) * 1e-3 / us;
319
320
321
322
323
324

    printf("m:%u, n:%u, k:%u, alpha:%f, cost:%lfus, GFLOPS:%lf, ", m, n, k, alpha, us, gflops);
    fflush(stdout);
}

// implement small ukernel on L1
carlushuang's avatar
carlushuang committed
325
326
327
328
329
template <typename FloatA,
          typename FloatB,
          typename ALayout,
          typename BLayout,
          typename thread_gemm_instance>
330
331
void test_cpu_ukernel(float alpha, uint32_t m, uint32_t n, uint32_t k)
{
carlushuang's avatar
carlushuang committed
332
    int max_threads = omp_get_max_threads();
333

334
335
    DeviceAlignedMemCPU a_mem(m * k * sizeof(FloatA), 32);
    DeviceAlignedMemCPU b_mem(k * n * sizeof(FloatB), 32);
carlushuang's avatar
carlushuang committed
336
    DeviceAlignedMemCPU c_mem(m * n * sizeof(float) * max_threads, 32);
337
    DeviceAlignedMemCPU c_mem_ref(m * n * sizeof(float), 32);
338

339
340
341
    c_mem_ref.SetZero();
    rand_vector(reinterpret_cast<FloatA*>(a_mem.mpDeviceBuf), m * k);
    rand_vector(reinterpret_cast<FloatB*>(b_mem.mpDeviceBuf), k * n);
342

343
344
345
346
347
348
349
350
    ref_cpu_gemm_uk<FloatA, FloatB, ALayout, BLayout>(
        reinterpret_cast<FloatA*>(a_mem.mpDeviceBuf),
        reinterpret_cast<FloatB*>(b_mem.mpDeviceBuf),
        reinterpret_cast<float*>(c_mem_ref.mpDeviceBuf),
        alpha,
        m,
        n,
        k);
351

352
    // using thread_gemm_instance = thread_gemm_avx2_mxn_6x16_instances<ALayout, BLayout>;
carlushuang's avatar
carlushuang committed
353
354
    // using thread_gemm_instance = thread_gemm_avx2_mxn_4x24_instances<ALayout, BLayout>;
    bool found = false;
355
356
357

    ck::static_for<0, std::tuple_size_v<thread_gemm_instance>, 1>{}([&](auto i) {
        using uk_type = std::tuple_element_t<i, thread_gemm_instance>;
358
        if(m % uk_type::ThreadMr != 0 || n % uk_type::ThreadNr != 0)
359
            return;
carlushuang's avatar
carlushuang committed
360
361
362
363
364
        // if((m != uk_type::ThreadMr && std::is_same<typename uk_type::MatrixALayout, Col>::value)
        // ||
        //    (n != uk_type::ThreadNr && std::is_same<typename uk_type::MatrixBLayout, Row>::value))
        //     // only k is the fast changing dim of A/B can we do muldiplt m, n
        //     return;
365

366
367
368
        if(found)
            return;

369
370
371
372
373
374
375
376
377
378
379
380
        test_ukernel<FloatA, FloatB, ALayout, BLayout>(uk_type{},
                                                       reinterpret_cast<FloatA*>(a_mem.mpDeviceBuf),
                                                       reinterpret_cast<FloatB*>(b_mem.mpDeviceBuf),
                                                       reinterpret_cast<float*>(c_mem.mpDeviceBuf),
                                                       alpha,
                                                       m,
                                                       n,
                                                       k);

        bool is_valid = valid_vector(reinterpret_cast<float*>(c_mem_ref.mpDeviceBuf),
                                     reinterpret_cast<float*>(c_mem.mpDeviceBuf),
                                     m * n);
381
        printf("vald:%s\n", is_valid ? "y" : "n");
382
        found = true;
383
384
385
386
387
    });
}

int main(int argc, char** argv)
{
388
389
    int m       = 4;
    int n       = 24;
390
391
392
393
394
395
396
397
398
399
400
401
402
    int k       = 64;
    float alpha = 1.0f;
    if(argc > 3)
    {
        m = std::atoi(argv[1]);
        n = std::atoi(argv[2]);
        k = std::atoi(argv[3]);
    }
    if(argc > 4)
    {
        alpha = std::atof(argv[4]);
    }
    dump_cache_hierarchy();
carlushuang's avatar
carlushuang committed
403
404
405
406
    if(std::getenv("OMP_NUM_THREADS") == nullptr)
        omp_set_num_threads(1);
    printf("max threads:%d\n", omp_get_max_threads());

carlushuang's avatar
carlushuang committed
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
    test_cpu_ukernel<AType, BType, Row, Row, thread_gemm_avx2_mxn_4x24_instances<Row, Row>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Row, Col, thread_gemm_avx2_mxn_4x24_instances<Row, Col>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Col, Row, thread_gemm_avx2_mxn_4x24_instances<Col, Row>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Col, Col, thread_gemm_avx2_mxn_4x24_instances<Col, Col>>(
        alpha, m, n, k);

    test_cpu_ukernel<AType, BType, Row, Row, thread_gemm_avx2_mxn_6x16_instances<Row, Row>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Row, Col, thread_gemm_avx2_mxn_6x16_instances<Row, Col>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Col, Row, thread_gemm_avx2_mxn_6x16_instances<Col, Row>>(
        alpha, m, n, k);
    test_cpu_ukernel<AType, BType, Col, Col, thread_gemm_avx2_mxn_6x16_instances<Col, Col>>(
        alpha, m, n, k);
424
}