profile_gemm_impl.hpp 15.7 KB
Newer Older
1
#pragma once
Chao Liu's avatar
Chao Liu committed
2
3
//#include "device_gemm_instance.hpp"
//#include "device_gemm_splitk_xdl_instance.hpp"
ltqin's avatar
ltqin committed
4
5
6
7
8
9

namespace ck {
namespace tensor_operation {
namespace device {
namespace device_gemm_instance {

Jing Zhang's avatar
Jing Zhang committed
10
11
12
13
14
using DeviceGemmNoOpPtr =
    ck::tensor_operation::device::DeviceGemmPtr<ck::tensor_operation::element_wise::PassThrough,
                                                ck::tensor_operation::element_wise::PassThrough,
                                                ck::tensor_operation::element_wise::PassThrough>;

Chao Liu's avatar
Chao Liu committed
15
#if 0
ltqin's avatar
ltqin committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
template <>
void add_device_gemm_instance<float,
                              float,
                              float,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<float,
                              float,
                              float,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<float,
                              float,
                              float,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<float,
                              float,
                              float,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<ck::half_t,
                              ck::half_t,
                              ck::half_t,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<ck::half_t,
                              ck::half_t,
                              ck::half_t,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<ck::half_t,
                              ck::half_t,
                              ck::half_t,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);

template <>
void add_device_gemm_instance<ck::half_t,
                              ck::half_t,
                              ck::half_t,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::ColumnMajor,
                              ck::tensor_layout::gemm::RowMajor>(std::vector<DeviceGemmNoOpPtr>&);
Chao Liu's avatar
Chao Liu committed
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#else
void add_device_gemm_xdl_f16_f16_f16_mk_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f16_f16_f16_mk_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f16_f16_f16_km_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f16_f16_f16_km_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);

void add_device_gemm_xdl_f32_f32_f32_mk_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f32_f32_f32_mk_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f32_f32_f32_km_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_f32_f32_f32_km_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);

void add_device_gemm_xdl_splitk_f32_f32_f32_mk_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_splitk_f32_f32_f32_mk_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_splitk_f32_f32_f32_km_kn_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_splitk_f32_f32_f32_km_nk_mn_instances(std::vector<DeviceGemmNoOpPtr>&);
#endif
ltqin's avatar
ltqin committed
95
96
97
98
99

} // namespace device_gemm_instance
} // namespace device
} // namespace tensor_operation
} // namespace ck
100
101
102
103
104
105
106
107
108
109

namespace ck {
namespace profiler {

template <typename ADataType,
          typename BDataType,
          typename CDataType,
          typename ALayout,
          typename BLayout,
          typename CLayout>
Chao Liu's avatar
Chao Liu committed
110
111
112
113
114
115
116
117
118
void profile_gemm_impl(int do_verification,
                       int init_method,
                       bool do_log,
                       int nrepeat,
                       int M,
                       int N,
                       int K,
                       int StrideA,
                       int StrideB,
119
                       int StrideC,
ltqin's avatar
ltqin committed
120
                       int KBatch = 1)
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
{
    auto f_host_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
            if(is_same<decltype(layout), tensor_layout::gemm::RowMajor>::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}));
            }
        };

    Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
    Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
    Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
    Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));

    std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
    std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
    std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;

ltqin's avatar
ltqin committed
145
    std::size_t num_thread = std::thread::hardware_concurrency();
146
147
148
149
    switch(init_method)
    {
    case 0: break;
    case 1:
ltqin's avatar
ltqin committed
150
151
        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5}, num_thread);
        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5}, num_thread);
152
153
        break;
    default:
ltqin's avatar
ltqin committed
154
155
        a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0}, num_thread);
        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5}, num_thread);
156
    }
ltqin's avatar
ltqin committed
157
158
    // set zero to c_device_buf
    c_m_n_device_result.GenerateTensorValue(GeneratorTensor_0<CDataType>{}, num_thread);
159
160
161

    if(do_verification)
    {
Chao Liu's avatar
Chao Liu committed
162
163
164
165
166
167
        host_gemm_mk_kn_mn(a_m_k,
                           b_k_n,
                           c_m_n_host_result,
                           ck::tensor_operation::element_wise::PassThrough{},
                           ck::tensor_operation::element_wise::PassThrough{},
                           ck::tensor_operation::element_wise::PassThrough{});
168
169
170
171
172
173
174
175
176
177
178
    }

    DeviceMem a_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpace());
    DeviceMem b_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpace());
    DeviceMem c_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpace());

    a_device_buf.ToDevice(a_m_k.mData.data());
    b_device_buf.ToDevice(b_k_n.mData.data());
    c_device_buf.ToDevice(c_m_n_device_result.mData.data());

    // add device GEMM instances
Chao Liu's avatar
Chao Liu committed
179
    std::vector<ck::tensor_operation::device::device_gemm_instance::DeviceGemmNoOpPtr> gemm_ptrs;
180

Jing Zhang's avatar
Jing Zhang committed
181
182
    if constexpr(is_same<ADataType, float>::value && is_same<BDataType, float>::value &&
                 is_same<CDataType, float>::value)
183
    {
Jing Zhang's avatar
Jing Zhang committed
184
185
186
        if constexpr(is_same<ALayout, tensor_layout::gemm::RowMajor>::value &&
                     is_same<BLayout, tensor_layout::gemm::RowMajor>::value &&
                     is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
Jing Zhang's avatar
Jing Zhang committed
187
        {
Jing Zhang's avatar
Jing Zhang committed
188
            if(KBatch > 1)
Jing Zhang's avatar
Jing Zhang committed
189
190
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
191
                    add_device_gemm_xdl_splitk_f32_f32_f32_mk_kn_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
192
            }
Jing Zhang's avatar
Jing Zhang committed
193
            else
Jing Zhang's avatar
Jing Zhang committed
194
195
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
196
                    add_device_gemm_xdl_f32_f32_f32_mk_kn_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
197
            }
Jing Zhang's avatar
Jing Zhang committed
198
199
200
201
202
203
        }
        else if constexpr(is_same<ALayout, tensor_layout::gemm::RowMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            if(KBatch > 1)
Jing Zhang's avatar
Jing Zhang committed
204
205
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
206
                    add_device_gemm_xdl_splitk_f32_f32_f32_mk_nk_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
207
            }
Jing Zhang's avatar
Jing Zhang committed
208
            else
Jing Zhang's avatar
Jing Zhang committed
209
210
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
211
                    add_device_gemm_xdl_f32_f32_f32_mk_nk_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
212
213
            }
        }
Jing Zhang's avatar
Jing Zhang committed
214
215
216
        else if constexpr(is_same<ALayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::RowMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
Jing Zhang's avatar
Jing Zhang committed
217
        {
Jing Zhang's avatar
Jing Zhang committed
218
            if(KBatch > 1)
Jing Zhang's avatar
Jing Zhang committed
219
            {
Jing Zhang's avatar
Jing Zhang committed
220

Jing Zhang's avatar
Jing Zhang committed
221
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
222
                    add_device_gemm_xdl_splitk_f32_f32_f32_km_kn_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
223
            }
Jing Zhang's avatar
Jing Zhang committed
224
            else
Jing Zhang's avatar
Jing Zhang committed
225
226
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
227
                    add_device_gemm_xdl_f32_f32_f32_km_kn_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
228
            }
Jing Zhang's avatar
Jing Zhang committed
229
230
231
232
233
234
        }
        else if constexpr(is_same<ALayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            if(KBatch > 1)
Jing Zhang's avatar
Jing Zhang committed
235
236
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
237
                    add_device_gemm_xdl_splitk_f32_f32_f32_km_nk_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
238
            }
Jing Zhang's avatar
Jing Zhang committed
239
            else
Jing Zhang's avatar
Jing Zhang committed
240
241
            {
                ck::tensor_operation::device::device_gemm_instance::
Jing Zhang's avatar
Jing Zhang committed
242
                    add_device_gemm_xdl_f32_f32_f32_km_nk_mn_instances(gemm_ptrs);
Jing Zhang's avatar
Jing Zhang committed
243
244
            }
        }
245
    }
Jing Zhang's avatar
Jing Zhang committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    else if constexpr(is_same<ADataType, half_t>::value && is_same<BDataType, half_t>::value &&
                      is_same<CDataType, half_t>::value)
    {
        if constexpr(is_same<ALayout, tensor_layout::gemm::RowMajor>::value &&
                     is_same<BLayout, tensor_layout::gemm::RowMajor>::value &&
                     is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            ck::tensor_operation::device::device_gemm_instance::
                add_device_gemm_xdl_f16_f16_f16_mk_kn_mn_instances(gemm_ptrs);
        }
        else if constexpr(is_same<ALayout, tensor_layout::gemm::RowMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            ck::tensor_operation::device::device_gemm_instance::
                add_device_gemm_xdl_f16_f16_f16_mk_nk_mn_instances(gemm_ptrs);
        }
        else if constexpr(is_same<ALayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::RowMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            ck::tensor_operation::device::device_gemm_instance::
                add_device_gemm_xdl_f16_f16_f16_km_kn_mn_instances(gemm_ptrs);
        }
        else if constexpr(is_same<ALayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<BLayout, tensor_layout::gemm::ColumnMajor>::value &&
                          is_same<CLayout, tensor_layout::gemm::RowMajor>::value)
        {
            ck::tensor_operation::device::device_gemm_instance::
                add_device_gemm_xdl_f16_f16_f16_km_nk_mn_instances(gemm_ptrs);
        }
    }
278
279
280
281
282
283

    if(gemm_ptrs.size() <= 0)
    {
        throw std::runtime_error("wrong! no device GEMM instance found");
    }

Chao Liu's avatar
Chao Liu committed
284
    std::string best_gemm_name;
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
    float best_ave_time   = 0;
    float best_tflops     = 0;
    float best_gb_per_sec = 0;

    // profile device GEMM instances
    for(auto& gemm_ptr : gemm_ptrs)
    {
        auto argument_ptr =
            gemm_ptr->MakeArgumentPointer(static_cast<ADataType*>(a_device_buf.GetDeviceBuffer()),
                                          static_cast<BDataType*>(b_device_buf.GetDeviceBuffer()),
                                          static_cast<CDataType*>(c_device_buf.GetDeviceBuffer()),
                                          M,
                                          N,
                                          K,
                                          StrideA,
                                          StrideB,
Chao Liu's avatar
Chao Liu committed
301
302
303
                                          StrideC,
                                          ck::tensor_operation::element_wise::PassThrough{},
                                          ck::tensor_operation::element_wise::PassThrough{},
304
                                          ck::tensor_operation::element_wise::PassThrough{},
ltqin's avatar
ltqin committed
305
                                          KBatch);
306
307
308
309
310

        auto invoker_ptr = gemm_ptr->MakeInvokerPointer();

        if(gemm_ptr->IsSupportedArgument(argument_ptr.get()))
        {
Chao Liu's avatar
Chao Liu committed
311
312
            std::string gemm_name = gemm_ptr->GetTypeString();

313
314
315
            float ave_time = invoker_ptr->Run(argument_ptr.get(), nrepeat);

            std::size_t flop = std::size_t(2) * M * N * K;
Chao Liu's avatar
Chao Liu committed
316

317
318
319
320
321
322
323
324
            std::size_t num_btype =
                sizeof(ADataType) * M * K + sizeof(BDataType) * K * M + sizeof(CDataType) * M * N;

            float tflops = static_cast<float>(flop) / 1.E9 / ave_time;

            float gb_per_sec = num_btype / 1.E6 / ave_time;

            std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec
Chao Liu's avatar
Chao Liu committed
325
                      << " GB/s, " << gemm_name << std::endl;
326
327
328

            if(tflops > best_tflops)
            {
Chao Liu's avatar
Chao Liu committed
329
                best_gemm_name  = gemm_name;
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
                best_tflops     = tflops;
                best_ave_time   = ave_time;
                best_gb_per_sec = gb_per_sec;
            }

            if(do_verification)
            {
                c_device_buf.FromDevice(c_m_n_device_result.mData.data());

                check_error(c_m_n_host_result, c_m_n_device_result);

                if(do_log)
                {
                    LogRangeAsType<float>(std::cout << "a : ", a_m_k.mData, ",") << std::endl;
                    LogRangeAsType<float>(std::cout << "b: ", b_k_n.mData, ",") << std::endl;
                    LogRangeAsType<float>(std::cout << "c_host  : ", c_m_n_host_result.mData, ",")
                        << std::endl;
                    LogRangeAsType<float>(std::cout << "c_device: ", c_m_n_device_result.mData, ",")
                        << std::endl;
                }
            }
        }
        else
        {
            std::cout << "this device GEMM instance does not support this GEMM problem"
                      << std::endl;
        }
    }

    std::cout << "Best Perf: " << best_ave_time << " ms, " << best_tflops << " TFlops, "
Chao Liu's avatar
Chao Liu committed
360
              << best_gb_per_sec << " GB/s, " << best_gemm_name << std::endl;
361
362
363
364
}

} // namespace profiler
} // namespace ck