cgemm_xdl_common.hpp 10.7 KB
Newer Older
1
2
3
4
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#include <cstdlib>
5
6
#include <initializer_list>
#include <numeric>
7
8
9

#include "ck/ck.hpp"
#include "ck/stream_config.hpp"
10
11
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"

12
13
14
15
#include "ck/library/utility/check_err.hpp"
#include "ck/library/utility/device_memory.hpp"
#include "ck/library/utility/host_tensor.hpp"
#include "ck/library/utility/host_tensor_generator.hpp"
16
#include "ck/library/utility/literals.hpp"
17
18
19
20
21
22
23
24
25

template <ck::index_t... Is>
using S = ck::Sequence<Is...>;

using F16   = ck::half_t;
using F32   = float;
using BF16  = ck::bhalf_t;
using INT8  = std::int8_t;
using INT32 = std::int32_t;
26
27
28
#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
using INT4 = ck::int4_t;
#endif
29
30
31
32
33
34
35
36
37
38
39

template <typename ADataType,
          typename BDataType,
          typename CDataType,
          typename ALayout,
          typename BLayout,
          typename CLayout,
          typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation,
          typename DeviceCGemmInstance,
40
41
42
43
44
45
46
47
48
49
50
51
52
          typename ReferenceCGemmInstance,
          typename KernelADataType = ADataType,
          typename KernelBDataType = BDataType,
          typename KernelCDataType = CDataType>
bool run_cgemm_xdl(ck::index_t M,
                   ck::index_t N,
                   ck::index_t K,
                   ck::index_t StrideA,
                   ck::index_t StrideB,
                   ck::index_t StrideC,
                   bool do_verification,
                   int init_method,
                   bool time_kernel)
53
{
54
55
56
57
58
59
60
61
62
63
64
#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
    static_assert(sizeof(ck::int4_t) == sizeof(int8_t),
                  "sizeof ck::int4_t and int8_t is different!");
    static_assert(sizeof(ADataType) == sizeof(KernelADataType),
                  "sizeof ADataType and KernelADataType is different!");
    static_assert(sizeof(BDataType) == sizeof(KernelBDataType),
                  "sizeof BDataType and KernelBDataType is different!");
    static_assert(sizeof(CDataType) == sizeof(KernelCDataType),
                  "sizeof CDataType and KernelCDataType is different!");
#endif

65
66
    using namespace ck::literals;

67
68
    auto f_host_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t stride, auto layout) {
69
            if constexpr(std::is_same_v<decltype(layout), ck::tensor_layout::gemm::RowMajor>)
70
            {
71
                return HostTensorDescriptor({row, col}, {stride, 1_uz});
72
73
74
            }
            else
            {
75
                return HostTensorDescriptor({row, col}, {1_uz, stride});
76
77
78
79
80
81
82
            }
        };

    Tensor<ADataType> a_m_k_real(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
    Tensor<ADataType> a_m_k_imag(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
    Tensor<BDataType> b_k_n_real(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
    Tensor<BDataType> b_k_n_imag(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
83
84
85
86
    Tensor<KernelCDataType> c_m_n_real_device_result(
        f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
    Tensor<KernelCDataType> c_m_n_imag_device_result(
        f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
87

88
89
90
91
92
93
    std::cout << "a_m_k_real: " << a_m_k_real.GetDesc() << std::endl;
    std::cout << "a_m_k_imag: " << a_m_k_imag.GetDesc() << std::endl;
    std::cout << "b_k_n_real: " << b_k_n_real.GetDesc() << std::endl;
    std::cout << "b_k_n_imag: " << b_k_n_imag.GetDesc() << std::endl;
    std::cout << "c_m_n_real: " << c_m_n_real_device_result.GetDesc() << std::endl;
    std::cout << "c_m_n_imag: " << c_m_n_imag_device_result.GetDesc() << std::endl;
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

    switch(init_method)
    {
    case 0: break;
    case 1:
        a_m_k_real.GenerateTensorValue(GeneratorTensor_2<ADataType>{-2, 2});
        a_m_k_imag.GenerateTensorValue(GeneratorTensor_2<ADataType>{-2, 2});
        b_k_n_real.GenerateTensorValue(GeneratorTensor_2<BDataType>{-2, 2});
        b_k_n_imag.GenerateTensorValue(GeneratorTensor_2<BDataType>{-2, 2});
        break;
    default:
        a_m_k_real.GenerateTensorValue(GeneratorTensor_3<ADataType>{-0.5, 0.5});
        a_m_k_imag.GenerateTensorValue(GeneratorTensor_3<ADataType>{-0.5, 0.5});
        b_k_n_real.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
        b_k_n_imag.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
    }

    auto cgemm = DeviceCGemmInstance{};

113
114
115
116
117
118
    DeviceMem a_m_k_real_device_buf(a_m_k_real.GetMemorySize());
    DeviceMem a_m_k_imag_device_buf(a_m_k_imag.GetMemorySize());
    DeviceMem b_k_n_real_device_buf(b_k_n_real.GetMemorySize());
    DeviceMem b_k_n_imag_device_buf(b_k_n_imag.GetMemorySize());
    DeviceMem c_m_n_real_device_buf(c_m_n_real_device_result.GetMemorySize());
    DeviceMem c_m_n_imag_device_buf(c_m_n_imag_device_result.GetMemorySize());
119
120
    DeviceMem workspace_device_buf(cgemm.GetWorkspaceSize(M, N, K, StrideA, StrideB, StrideC));

121
122
123
124
125
126
127
128
#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
    if constexpr(std::is_same_v<ADataType, ck::int4_t>)
    {
        Tensor<KernelADataType> a_m_k_real_converted(a_m_k_real);
        Tensor<KernelADataType> a_m_k_imag_converted(a_m_k_imag);
        Tensor<KernelBDataType> b_k_n_real_converted(b_k_n_real);
        Tensor<KernelBDataType> b_k_n_imag_converted(b_k_n_imag);

129
130
131
132
        a_m_k_real_device_buf.ToDevice(a_m_k_real_converted.data());
        a_m_k_imag_device_buf.ToDevice(a_m_k_imag_converted.data());
        b_k_n_real_device_buf.ToDevice(b_k_n_real_converted.data());
        b_k_n_imag_device_buf.ToDevice(b_k_n_imag_converted.data());
133
134
135
136
    }
    else
#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
    {
137
138
139
140
        a_m_k_real_device_buf.ToDevice(a_m_k_real.data());
        a_m_k_imag_device_buf.ToDevice(a_m_k_imag.data());
        b_k_n_real_device_buf.ToDevice(b_k_n_real.data());
        b_k_n_imag_device_buf.ToDevice(b_k_n_imag.data());
141
    }
142
143
144
145
146
147

    auto a_element_op = AElementwiseOperation{};
    auto b_element_op = BElementwiseOperation{};
    auto c_element_op = CElementwiseOperation{};

    // do GEMM
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    auto invoker  = cgemm.MakeInvoker();
    auto argument = cgemm.MakeArgument(a_m_k_real_device_buf.GetDeviceBuffer(),
                                       a_m_k_imag_device_buf.GetDeviceBuffer(),
                                       b_k_n_real_device_buf.GetDeviceBuffer(),
                                       b_k_n_imag_device_buf.GetDeviceBuffer(),
                                       c_m_n_real_device_buf.GetDeviceBuffer(),
                                       c_m_n_imag_device_buf.GetDeviceBuffer(),
                                       workspace_device_buf.GetDeviceBuffer(),
                                       M,
                                       N,
                                       K,
                                       StrideA,
                                       StrideB,
                                       StrideC,
                                       a_element_op,
                                       b_element_op,
                                       c_element_op);
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

    if(!cgemm.IsSupportedArgument(argument))
    {
        throw std::runtime_error(
            "wrong! device_cgemm with the specified compilation parameters does "
            "not support this CGEMM problem");
    }

    float ave_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});

    std::size_t flop = std::size_t(8) * M * N * K;
    std::size_t num_btype =
        std::size_t(2) *
        (sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N);

180
    float tflops     = static_cast<float>(flop) / 1.E9 / ave_time;
181
182
183
184
185
186
187
188
189
190
191
192
    float gb_per_sec = num_btype / 1.E6 / ave_time;

    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
              << cgemm.GetTypeString() << std::endl;

    if(do_verification)
    {
        Tensor<CDataType> c_m_n_real_host_result(
            f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
        Tensor<CDataType> c_m_n_imag_host_result(
            f_host_tensor_descriptor(M, N, StrideC, CLayout{}));

193
194
        auto ref_cgemm    = ReferenceCGemmInstance{};
        auto ref_invoker  = ref_cgemm.MakeInvoker();
195
196
197
198
199
200
201
202
203
204
205
206
        auto ref_argument = ref_cgemm.MakeArgument(a_m_k_real,
                                                   a_m_k_imag,
                                                   b_k_n_real,
                                                   b_k_n_imag,
                                                   c_m_n_real_host_result,
                                                   c_m_n_imag_host_result,
                                                   a_element_op,
                                                   b_element_op,
                                                   c_element_op);

        ref_invoker.Run(ref_argument);

207
208
        c_m_n_real_device_buf.FromDevice(c_m_n_real_device_result.data());
        c_m_n_imag_device_buf.FromDevice(c_m_n_imag_device_result.data());
209

210
        bool result = true;
211
212
213
214
215
216
#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
        if constexpr(std::is_same_v<ADataType, ck::int4_t>)
        {
            const Tensor<CDataType> c_m_n_real_device_result_converted(c_m_n_real_device_result);
            const Tensor<CDataType> c_m_n_imag_device_result_converted(c_m_n_imag_device_result);

217
218
            result = ck::utils::check_err(c_m_n_real_device_result_converted,
                                          c_m_n_real_host_result,
219
220
221
222
                                          "Verification error: incorrect results in real part!",
                                          1e-2f,
                                          1e-1f);
            result = result && ck::utils::check_err(
223
224
                                   c_m_n_imag_device_result_converted,
                                   c_m_n_imag_host_result,
225
226
227
228
229
230
231
                                   "Verification error: incorrect results in imaginary part!",
                                   1e-2f,
                                   1e-1f);
        }
        else
#endif // CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
        {
232
233
            result = ck::utils::check_err(c_m_n_real_device_result,
                                          c_m_n_real_host_result,
234
235
236
237
                                          "Verification error: incorrect results in real part!",
                                          1e-2f,
                                          1e-1f);
            result = result && ck::utils::check_err(
238
239
                                   c_m_n_imag_device_result,
                                   c_m_n_imag_host_result,
240
241
242
243
244
245
                                   "Verification error: incorrect results in imaginary part!",
                                   1e-2f,
                                   1e-1f);
        }

        return result;
246
    }
247
    return true;
248
}