profile_gemm_impl.hpp 9.52 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.

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

Chao Liu's avatar
Chao Liu committed
6
#include <iomanip>
7
8
#include <iostream>
#include <typeinfo>
9

Chao Liu's avatar
Chao Liu committed
10
11
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/device_gemm.hpp"
12
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
Chao Liu's avatar
Chao Liu committed
13
14
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

15
#include "ck/library/tensor_operation_instance/gpu/gemm.hpp"
Chao Liu's avatar
Chao Liu committed
16

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

namespace ck {
namespace profiler {

27
28
29
30
template <typename ALayout,
          typename BLayout,
          typename CLayout,
          typename ADataType,
31
          typename BDataType,
32
          typename AccDataType,
33
          typename CDataType>
Chao Liu's avatar
Chao Liu committed
34
35
36
37
38
39
40
41
42
43
int profile_gemm_impl(int do_verification,
                      int init_method,
                      bool do_log,
                      bool time_kernel,
                      int M,
                      int N,
                      int K,
                      int StrideA,
                      int StrideB,
                      int StrideC)
44
{
Chao Liu's avatar
Chao Liu committed
45
46
    bool pass = true;

47
48
    using namespace ck::literals;

49
50
    auto f_host_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
51
            if constexpr(is_same_v<decltype(layout), tensor_layout::gemm::RowMajor>)
52
            {
53
                return HostTensorDescriptor({row, col}, {stride, 1_uz});
54
55
56
            }
            else
            {
57
                return HostTensorDescriptor({row, col}, {1_uz, stride});
58
59
60
61
62
            }
        };

    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{}));
Chao Liu's avatar
Chao Liu committed
63
    Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
64
65
    Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));

66
67
68
    std::cout << "a_m_k: " << a_m_k.GetDesc() << std::endl;
    std::cout << "b_k_n: " << b_k_n.GetDesc() << std::endl;
    std::cout << "c_m_n: " << c_m_n_device_result.GetDesc() << std::endl;
69
70
71

    switch(init_method)
    {
Chao Liu's avatar
Chao Liu committed
72
    case 0: break;
73
    case 1:
Chao Liu's avatar
Chao Liu committed
74
75
        a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
        b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
76
77
        break;
    default:
Chao Liu's avatar
Chao Liu committed
78
79
        a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
        b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
80
    }
Chao Liu's avatar
Chao Liu committed
81
82
83
84
85
86
87
88
89

    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{};

90
91
92
    DeviceMem a_device_buf(a_m_k.GetMemorySize());
    DeviceMem b_device_buf(b_k_n.GetMemorySize());
    DeviceMem c_device_buf(c_m_n_device_result.GetMemorySize());
93

94
95
    a_device_buf.ToDevice(a_m_k.data());
    b_device_buf.ToDevice(b_k_n.data());
96

97
98
99
100
101
102
103
104
105
    using DeviceOp = ck::tensor_operation::device::DeviceGemm<ALayout,
                                                              BLayout,
                                                              CLayout,
                                                              ADataType,
                                                              BDataType,
                                                              CDataType,
                                                              AElementOp,
                                                              BElementOp,
                                                              CElementOp>;
106

107
108
109
110
111
    // get device op instances
    const auto op_ptrs = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<
        DeviceOp>::GetInstances();

    std::cout << "found " << op_ptrs.size() << " instances" << std::endl;
Chao Liu's avatar
Chao Liu committed
112

113
    // Run reference op
Chao Liu's avatar
Chao Liu committed
114
    if(do_verification)
115
    {
Chao Liu's avatar
Chao Liu committed
116
117
118
119
120
121
122
        using ReferenceGemmInstance = ck::tensor_operation::host::ReferenceGemm<ADataType,
                                                                                BDataType,
                                                                                CDataType,
                                                                                AccDataType,
                                                                                AElementOp,
                                                                                BElementOp,
                                                                                CElementOp>;
Jianfeng Yan's avatar
Jianfeng Yan committed
123

Chao Liu's avatar
Chao Liu committed
124
125
        auto ref_op      = ReferenceGemmInstance{};
        auto ref_invoker = ref_op.MakeInvoker();
Jianfeng Yan's avatar
Jianfeng Yan committed
126

Chao Liu's avatar
Chao Liu committed
127
128
        auto ref_argument = ref_op.MakeArgument(
            a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);
129

Chao Liu's avatar
Chao Liu committed
130
        ref_invoker.Run(ref_argument);
131
132
    }

Chao Liu's avatar
Chao Liu committed
133
    std::string best_op_name;
134
    float best_avg_time   = 0;
135
136
137
    float best_tflops     = 0;
    float best_gb_per_sec = 0;

138
    // profile device op instances
Chao Liu's avatar
Chao Liu committed
139
    for(auto& op_ptr : op_ptrs)
140
    {
141
142
143
144
145
146
147
148
149
150
151
152
        auto argument_ptr = op_ptr->MakeArgumentPointer(a_device_buf.GetDeviceBuffer(),
                                                        b_device_buf.GetDeviceBuffer(),
                                                        c_device_buf.GetDeviceBuffer(),
                                                        M,
                                                        N,
                                                        K,
                                                        StrideA,
                                                        StrideB,
                                                        StrideC,
                                                        a_element_op,
                                                        b_element_op,
                                                        c_element_op);
Chao Liu's avatar
Chao Liu committed
153
154
155
156

        auto invoker_ptr = op_ptr->MakeInvokerPointer();

        if(op_ptr->IsSupportedArgument(argument_ptr.get()))
157
        {
158
            // re-init C to zero before profiling next kernel
Chao Liu's avatar
Chao Liu committed
159
            c_device_buf.SetZero();
160

Chao Liu's avatar
Chao Liu committed
161
            std::string op_name = op_ptr->GetTypeString();
Chao Liu's avatar
Chao Liu committed
162

163
            float avg_time =
JD's avatar
JD committed
164
                invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
165

166
            std::size_t flop = 2_uz * M * N * K;
Chao Liu's avatar
Chao Liu committed
167

168
            std::size_t num_btype =
Chao Liu's avatar
Chao Liu committed
169
                sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
170

171
            float tflops = static_cast<float>(flop) / 1.E9 / avg_time;
172

173
            float gb_per_sec = num_btype / 1.E6 / avg_time;
174

175
            std::cout << "Perf: " << std::setw(10) << avg_time << " ms, " << tflops << " TFlops, "
Chao Liu's avatar
Chao Liu committed
176
                      << gb_per_sec << " GB/s, " << op_name << std::endl;
177
178
179

            if(tflops > best_tflops)
            {
Chao Liu's avatar
Chao Liu committed
180
                best_op_name    = op_name;
181
                best_tflops     = tflops;
182
                best_avg_time   = avg_time;
183
184
185
186
187
                best_gb_per_sec = gb_per_sec;
            }

            if(do_verification)
            {
188
                c_device_buf.FromDevice(c_m_n_device_result.data());
189

190
                pass = pass & ck::utils::check_err(c_m_n_device_result, c_m_n_host_result);
191
192
193

                if(do_log)
                {
194
195
196
                    LogRangeAsType<float>(std::cout << "a : ", a_m_k, ",") << std::endl;
                    LogRangeAsType<float>(std::cout << "b: ", b_k_n, ",") << std::endl;
                    LogRangeAsType<float>(std::cout << "c_host  : ", c_m_n_host_result, ",")
Chao Liu's avatar
Chao Liu committed
197
                        << std::endl;
198
                    LogRangeAsType<float>(std::cout << "c_device: ", c_m_n_device_result, ",")
199
200
201
202
203
204
                        << std::endl;
                }
            }
        }
        else
        {
Chao Liu's avatar
Chao Liu committed
205
            std::cout << op_ptr->GetTypeString() << " does not support this problem" << std::endl;
206
207
208
        }
    }

209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
    if constexpr(is_same<CDataType, float>::value)
    {
        std::cout << "Best Perf for datatype = f32";
    }
    else if constexpr(is_same<CDataType, half_t>::value)
    {
        std::cout << "Best Perf for datatype = f16";
    }
    else if constexpr(is_same<CDataType, bhalf_t>::value)
    {
        std::cout << "Best Perf for datatype = bf16";
    }
    else if constexpr(is_same<CDataType, int8_t>::value)
    {
        std::cout << "Best Perf for datatype = int8";
    }

    if constexpr(is_same<ALayout, tensor_layout::gemm::RowMajor>::value)
    {
        std::cout << " ALayout =  RowMajor";
    }
    else if constexpr(is_same<ALayout, tensor_layout::gemm::ColumnMajor>::value)
    {
        std::cout << " ALayout =  ColumnMajor";
    }

    if constexpr(is_same<BLayout, tensor_layout::gemm::RowMajor>::value)
    {
        std::cout << " BLayout =  RowMajor";
    }
    else if constexpr(is_same<BLayout, tensor_layout::gemm::ColumnMajor>::value)
    {
        std::cout << " BLayout =  ColumnMajor";
    }

    std::cout << " M = " << M << " N = " << N << " K = " << K << " StrideA = " << StrideA
245
              << " StrideB = " << StrideB << " StrideC = " << StrideC << " : " << best_avg_time
246
              << " ms, " << best_tflops << " TFlops, " << best_gb_per_sec << " GB/s, "
Chao Liu's avatar
Chao Liu committed
247
248
249
              << best_op_name << std::endl;

    return pass ? 0 : 1;
250
251
252
253
}

} // namespace profiler
} // namespace ck