profile_grouped_gemm_impl.hpp 11.7 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

zjing14's avatar
zjing14 committed
4
#pragma once
Chao Liu's avatar
Chao Liu committed
5

zjing14's avatar
zjing14 committed
6
#include <iomanip>
7

Chao Liu's avatar
Chao Liu committed
8
9
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
10
#include "ck/tensor_operation/gpu/device/device_grouped_gemm.hpp"
Chao Liu's avatar
Chao Liu committed
11
12
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

13
14
#include "ck/library/tensor_operation_instance/gpu/grouped_gemm.hpp"

Chao Liu's avatar
Chao Liu committed
15
16
17
18
19
20
#include "ck/library/utility/check_err.hpp"
#include "ck/library/utility/conv_util.hpp"
#include "ck/library/host_tensor/device_memory.hpp"
#include "ck/library/host_tensor/host_tensor.hpp"
#include "ck/library/host_tensor/host_tensor_generator.hpp"
#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"
zjing14's avatar
zjing14 committed
21
22
23
24
25
26

namespace ck {
namespace profiler {

template <typename ADataType,
          typename BDataType,
27
          typename EDataType,
28
          typename AccDataType,
zjing14's avatar
zjing14 committed
29
30
31
          typename ALayout,
          typename BLayout,
          typename CLayout>
32
bool profile_grouped_gemm_impl(int do_verification,
zjing14's avatar
zjing14 committed
33
34
                               int init_method,
                               bool do_log,
JD's avatar
JD committed
35
                               bool time_kernel,
36
37
38
39
40
41
                               const std::vector<int>& Ms,
                               const std::vector<int>& Ns,
                               const std::vector<int>& Ks,
                               const std::vector<int>& StrideAs,
                               const std::vector<int>& StrideBs,
                               const std::vector<int>& StrideCs)
zjing14's avatar
zjing14 committed
42
{
43
44
45

    bool pass = true;

zjing14's avatar
zjing14 committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
    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}));
            }
        };

60
    std::size_t group_count = Ms.size();
zjing14's avatar
zjing14 committed
61
62
63
64
65
66
67
68
69

    if(!(group_count == Ns.size() && group_count == Ks.size() && group_count == StrideAs.size() &&
         group_count == StrideBs.size() && group_count == StrideCs.size()))
    {
        throw std::runtime_error("wrong! inconsistent M/N/Ks, StrideA/B/Cs size\n");
    }

    std::vector<Tensor<ADataType>> a_m_k;
    std::vector<Tensor<BDataType>> b_k_n;
70
    std::vector<Tensor<EDataType>> c_m_n_device_results;
zjing14's avatar
zjing14 committed
71

72
    for(std::size_t i = 0; i < group_count; i++)
zjing14's avatar
zjing14 committed
73
74
75
76
77
78
79
    {
        a_m_k.push_back(
            Tensor<ADataType>(f_host_tensor_descriptor(Ms[i], Ks[i], StrideAs[i], ALayout{})));
        b_k_n.push_back(
            Tensor<BDataType>(f_host_tensor_descriptor(Ks[i], Ns[i], StrideBs[i], BLayout{})));

        c_m_n_device_results.push_back(
80
            Tensor<EDataType>(f_host_tensor_descriptor(Ms[i], Ns[i], StrideCs[i], CLayout{})));
zjing14's avatar
zjing14 committed
81
82
83
84
85

        std::cout << "group: " << i << " a_m_k[" << i << "]:" << a_m_k[i].mDesc << ", b_k_n[" << i
                  << "]:" << b_k_n[i].mDesc << ", c_m_n_device_results[" << i
                  << "]:" << c_m_n_device_results[i].mDesc << std::endl;

86
        std::size_t num_thread = 1;
zjing14's avatar
zjing14 committed
87
88
89
90
91
92
93
94
95
96
97
98
        switch(init_method)
        {
        case 0: break;
        case 1:
            a_m_k[i].GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5}, num_thread);
            b_k_n[i].GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5}, num_thread);
            break;
        default:
            a_m_k[i].GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0}, num_thread);
            b_k_n[i].GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5}, num_thread);
        }

99
        c_m_n_device_results[i].GenerateTensorValue(GeneratorTensor_0<EDataType>{}, num_thread);
zjing14's avatar
zjing14 committed
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
    }

    using AElementOp = ck::tensor_operation::element_wise::PassThrough;
    using BElementOp = ck::tensor_operation::element_wise::PassThrough;
    using CElementOp = ck::tensor_operation::element_wise::PassThrough;

    const auto a_element_op = AElementOp{};
    const auto b_element_op = BElementOp{};
    const auto c_element_op = CElementOp{};

    // if(do_verification)
    // {

    // }

    using DeviceMemPtr = std::unique_ptr<DeviceMem>;
    std::vector<DeviceMemPtr> a_device_buf, b_device_buf, c_device_buf;

    a_device_buf.reserve(group_count);
    b_device_buf.reserve(group_count);
    c_device_buf.reserve(group_count);

    std::vector<const void*> p_a, p_b;
    std::vector<void*> p_c;

    p_a.reserve(group_count);
    p_b.reserve(group_count);
    p_c.reserve(group_count);

129
    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
zjing14's avatar
zjing14 committed
130

131
    gemm_descs.reserve(group_count);
zjing14's avatar
zjing14 committed
132

133
    for(std::size_t i = 0; i < group_count; i++)
zjing14's avatar
zjing14 committed
134
135
    {
        a_device_buf.emplace_back(
zjing14's avatar
zjing14 committed
136
            std::make_unique<DeviceMem>(sizeof(ADataType) * a_m_k[i].mDesc.GetElementSpace()));
zjing14's avatar
zjing14 committed
137
        b_device_buf.emplace_back(
zjing14's avatar
zjing14 committed
138
            std::make_unique<DeviceMem>(sizeof(BDataType) * b_k_n[i].mDesc.GetElementSpace()));
zjing14's avatar
zjing14 committed
139
140

        c_device_buf.emplace_back(std::make_unique<DeviceMem>(
141
            sizeof(EDataType) * c_m_n_device_results[i].mDesc.GetElementSpace()));
zjing14's avatar
zjing14 committed
142
143
144
145
146

        a_device_buf[i]->ToDevice(a_m_k[i].mData.data());
        b_device_buf[i]->ToDevice(b_k_n[i].mData.data());
        c_device_buf[i]->ToDevice(c_m_n_device_results[i].mData.data());

147
        gemm_descs.push_back({Ms[i], Ns[i], Ks[i], StrideAs[i], StrideBs[i], StrideCs[i], {}});
zjing14's avatar
zjing14 committed
148
149
150
151
152
153

        p_a.push_back(a_device_buf[i]->GetDeviceBuffer());
        p_b.push_back(b_device_buf[i]->GetDeviceBuffer());
        p_c.push_back(c_device_buf[i]->GetDeviceBuffer());
    }

154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    using DeviceOp = ck::tensor_operation::device::DeviceGroupedGemm<ALayout,
                                                                     BLayout,
                                                                     CLayout,
                                                                     ADataType,
                                                                     BDataType,
                                                                     ck::Tuple<>,
                                                                     EDataType,
                                                                     AElementOp,
                                                                     BElementOp,
                                                                     CElementOp>;

    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
        DeviceOp>::GetInstances();

    if(op_ptrs.size() <= 0)
zjing14's avatar
zjing14 committed
169
170
171
172
173
174
175
176
177
    {
        throw std::runtime_error("wrong! no device GEMM instance found");
    }

    std::string best_gemm_name;
    float best_ave_time   = 0;
    float best_tflops     = 0;
    float best_gb_per_sec = 0;

178
179
    auto p_ds = std::vector<std::array<const void*, 0>>{};

zjing14's avatar
zjing14 committed
180
    // profile device GEMM instances
181
    for(auto& gemm_ptr : op_ptrs)
zjing14's avatar
zjing14 committed
182
183
184
185
    {
        auto argument_ptr =
            gemm_ptr->MakeArgumentPointer(p_a,
                                          p_b,
186
                                          p_ds,
zjing14's avatar
zjing14 committed
187
                                          p_c,
188
                                          gemm_descs,
zjing14's avatar
zjing14 committed
189
190
191
192
193
194
                                          ck::tensor_operation::element_wise::PassThrough{},
                                          ck::tensor_operation::element_wise::PassThrough{},
                                          ck::tensor_operation::element_wise::PassThrough{});

        auto invoker_ptr = gemm_ptr->MakeInvokerPointer();

195
196
197
198
        DeviceMem gemm_desc_workspace(gemm_ptr->GetWorkSpaceSize(argument_ptr.get()));

        gemm_ptr->SetWorkSpacePointer(argument_ptr.get(), gemm_desc_workspace.GetDeviceBuffer());

zjing14's avatar
zjing14 committed
199
200
201
202
        if(gemm_ptr->IsSupportedArgument(argument_ptr.get()))
        {
            std::string gemm_name = gemm_ptr->GetTypeString();

JD's avatar
JD committed
203
204
            float ave_time =
                invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
zjing14's avatar
zjing14 committed
205
206

            std::size_t flop = 0, num_btype = 0;
207
            for(std::size_t i = 0; i < gemm_descs.size(); i++)
zjing14's avatar
zjing14 committed
208
209
210
211
            {
                flop += std::size_t(2) * Ms[i] * Ns[i] * Ks[i];

                num_btype += sizeof(ADataType) * Ms[i] * Ks[i] + sizeof(BDataType) * Ks[i] * Ns[i] +
212
                             sizeof(EDataType) * Ms[i] * Ns[i];
zjing14's avatar
zjing14 committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
            }

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

            float gb_per_sec = num_btype / 1.E6 / ave_time;
            std::cout << "Perf: " << std::setw(10) << ave_time << " ms, " << tflops << " TFlops, "
                      << gb_per_sec << " GB/s, " << gemm_name << std::endl;

            if(tflops > best_tflops)
            {
                best_gemm_name  = gemm_name;
                best_tflops     = tflops;
                best_ave_time   = ave_time;
                best_gb_per_sec = gb_per_sec;
            }

            if(do_verification)
            {
231
                for(std::size_t i = 0; i < gemm_descs.size(); i++)
zjing14's avatar
zjing14 committed
232
233
234
235
                {

                    c_device_buf[i]->FromDevice(c_m_n_device_results[i].mData.data());

236
                    Tensor<EDataType> c_m_n_host_result(
zjing14's avatar
zjing14 committed
237
238
239
240
241
                        f_host_tensor_descriptor(Ms[i], Ns[i], StrideCs[i], CLayout{}));

                    using ReferenceGemmInstance =
                        ck::tensor_operation::host::ReferenceGemm<ADataType,
                                                                  BDataType,
242
                                                                  EDataType,
243
                                                                  AccDataType,
zjing14's avatar
zjing14 committed
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
                                                                  AElementOp,
                                                                  BElementOp,
                                                                  CElementOp>;

                    auto ref_gemm    = ReferenceGemmInstance{};
                    auto ref_invoker = ref_gemm.MakeInvoker();

                    auto ref_argument = ref_gemm.MakeArgument(a_m_k[i],
                                                              b_k_n[i],
                                                              c_m_n_host_result,
                                                              a_element_op,
                                                              b_element_op,
                                                              c_element_op);

                    ref_invoker.Run(ref_argument);
259
260
                    pass = pass && ck::utils::check_err(c_m_n_device_results[i].mData,
                                                        c_m_n_host_result.mData);
zjing14's avatar
zjing14 committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284

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

    std::cout << "Best Perf: " << best_ave_time << " ms, " << best_tflops << " TFlops, "
              << best_gb_per_sec << " GB/s, " << best_gemm_name << std::endl;
285
286

    return pass;
zjing14's avatar
zjing14 committed
287
288
289
290
} // namespace profiler

} // namespace profiler
} // namespace ck