gemm_split_k.cpp 8.76 KB
Newer Older
ltqin's avatar
ltqin committed
1
2
3
#include <iostream>
#include <initializer_list>
#include <cstdlib>
Chao Liu's avatar
Chao Liu committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
#include "ck/tensor_operation/gpu/device/device_gemm_xdl_splitk.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

#include "ck/library/utility/check_err.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/host_tensor/device_memory.hpp"
#include "ck/library/reference_tensor_operation/cpu/reference_gemm.hpp"

#include "ck/library/host_tensor/host_gemm.hpp"
ltqin's avatar
ltqin committed
19

Chao Liu's avatar
Chao Liu committed
20
enum struct GemmMatrixLayout
ltqin's avatar
ltqin committed
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
{
    MK_KN_MN, // 0
    MK_NK_MN, // 1
    KM_KN_MN, // 2
    KM_NK_MN, // 3
};

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

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

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>&);

} // namespace device_gemm_instance
} // namespace device
} // namespace tensor_operation
} // namespace ck

template <typename T>
static bool check_out(const Tensor<T>& ref, const Tensor<T>& result)
{
    float max_diff = 1e-6;

53
    for(std::size_t i = 0; i < ref.mData.size(); ++i)
ltqin's avatar
ltqin committed
54
55
56
57
58
59
60
61
62
63
64
    {
        float diff = std::abs(double(ref.mData[i]) - double(result.mData[i]));
        if(max_diff < diff)
        {
            return false;
        }
    }

    return true;
}

65
struct gemmArgs
ltqin's avatar
ltqin committed
66
{
Chao Liu's avatar
Chao Liu committed
67
    GemmMatrixLayout layout;
68
69
70
71
72
73
74
75
    int M;
    int N;
    int K;
    int StrideA;
    int StrideB;
    int StrideC;
    int KBatch;
};
ltqin's avatar
ltqin committed
76

77
78
int test_gemm(const gemmArgs& args)
{
ltqin's avatar
ltqin committed
79
80
    bool a_row_major, b_row_major, c_row_major;

81
    switch(args.layout)
ltqin's avatar
ltqin committed
82
83
84
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
    {
    case GemmMatrixLayout::MK_KN_MN:
        a_row_major = true;
        b_row_major = true;
        c_row_major = true;
        break;
    case GemmMatrixLayout::MK_NK_MN:
        a_row_major = true;
        b_row_major = false;
        c_row_major = true;
        break;
    case GemmMatrixLayout::KM_KN_MN:
        a_row_major = false;
        b_row_major = true;
        c_row_major = true;
        break;
    case GemmMatrixLayout::KM_NK_MN:
        a_row_major = false;
        b_row_major = false;
        c_row_major = true;
        break;
    default: printf("not supported layout"); return 1;
    }

    auto f_host_tensor_descriptor =
        [](std::size_t row, std::size_t col, std::size_t stride, bool row_major) {
            if(row_major)
            {
                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}));
            }
        };

120
121
    Tensor<float> a_m_k(f_host_tensor_descriptor(args.M, args.K, args.StrideA, a_row_major));
    Tensor<float> b_k_n(f_host_tensor_descriptor(args.K, args.N, args.StrideB, b_row_major));
122
123
124
125
    Tensor<float> c_m_n_host_result(
        f_host_tensor_descriptor(args.M, args.N, args.StrideC, c_row_major));
    Tensor<float> c_m_n_device_result(
        f_host_tensor_descriptor(args.M, args.N, args.StrideC, c_row_major));
ltqin's avatar
ltqin committed
126
127

    // init data
128
    std::size_t num_thread = 1;
ltqin's avatar
ltqin committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    a_m_k.GenerateTensorValue(GeneratorTensor_2<float>{-5, 5}, num_thread);
    b_k_n.GenerateTensorValue(GeneratorTensor_2<float>{-5, 5}, num_thread);
    // set zero to c_device_buf
    c_m_n_device_result.GenerateTensorValue(GeneratorTensor_0<float>{}, num_thread);

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

    DeviceMem a_device_buf(sizeof(float) * a_m_k.mDesc.GetElementSpace());
    DeviceMem b_device_buf(sizeof(float) * b_k_n.mDesc.GetElementSpace());
    DeviceMem c_device_buf(sizeof(float) * 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
    std::vector<DeviceGemmNoOpPtr> gemm_ptrs;

152
    if(args.layout == GemmMatrixLayout::MK_KN_MN)
ltqin's avatar
ltqin committed
153
154
155
156
    {
        ck::tensor_operation::device::device_gemm_instance::
            add_device_gemm_xdl_splitk_f32_f32_f32_mk_kn_mn_instances(gemm_ptrs);
    }
157
    else if(args.layout == GemmMatrixLayout::MK_NK_MN)
ltqin's avatar
ltqin committed
158
159
160
161
    {
        ck::tensor_operation::device::device_gemm_instance::
            add_device_gemm_xdl_splitk_f32_f32_f32_mk_nk_mn_instances(gemm_ptrs);
    }
162
    else if(args.layout == GemmMatrixLayout::KM_KN_MN)
ltqin's avatar
ltqin committed
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    {
        ck::tensor_operation::device::device_gemm_instance::
            add_device_gemm_xdl_splitk_f32_f32_f32_km_kn_mn_instances(gemm_ptrs);
    }
    else
    {
        ck::tensor_operation::device::device_gemm_instance::
            add_device_gemm_xdl_splitk_f32_f32_f32_km_nk_mn_instances(gemm_ptrs);
    }

    bool success = false;
    for(auto& gemm_ptr : gemm_ptrs)
    {
        auto argument_ptr =
            gemm_ptr->MakeArgumentPointer(static_cast<float*>(a_device_buf.GetDeviceBuffer()),
                                          static_cast<float*>(b_device_buf.GetDeviceBuffer()),
                                          static_cast<float*>(c_device_buf.GetDeviceBuffer()),
180
181
182
183
184
185
                                          args.M,
                                          args.N,
                                          args.K,
                                          args.StrideA,
                                          args.StrideB,
                                          args.StrideC,
ltqin's avatar
ltqin committed
186
187
188
                                          ck::tensor_operation::element_wise::PassThrough{},
                                          ck::tensor_operation::element_wise::PassThrough{},
                                          ck::tensor_operation::element_wise::PassThrough{},
189
                                          args.KBatch);
ltqin's avatar
ltqin committed
190
191
192
193
194

        auto invoker_ptr = gemm_ptr->MakeInvokerPointer();

        if(gemm_ptr->IsSupportedArgument(argument_ptr.get()))
        {
JD's avatar
JD committed
195
            invoker_ptr->Run(argument_ptr.get());
ltqin's avatar
ltqin committed
196
197

            c_device_buf.FromDevice(c_m_n_device_result.mData.data());
JD's avatar
JD committed
198

ltqin's avatar
ltqin committed
199
200
201
202
203
204
205
206
            if(!check_out(c_m_n_host_result, c_m_n_device_result))
            {
                success = false;
                break;
            }
            success = true;
        }
    }
207
    auto error_code = 0;
ltqin's avatar
ltqin committed
208
209
210
211
212
213
214
    if(success)
    {
        std::cout << "test split k : Pass" << std::endl;
    }
    else
    {
        std::cout << "test split k: Fail " << std::endl;
215
        error_code = -1; // test needs to report failure
216
217
218
219
220
221
222
223
224
    }
    return error_code;
}

int main(int argc, char* argv[])
{
    std::vector<gemmArgs> test_cases;
    if(argc == 1)
    {
Chao Liu's avatar
Chao Liu committed
225
        test_cases = {{GemmMatrixLayout::MK_KN_MN, 3, 3, 3, 3, 3, 3, 1}};
226
227
228
229
230
        // JD: Populate with more and meaningful
        return 0;
    }
    else if(argc == 9)
    {
Chao Liu's avatar
Chao Liu committed
231
        const auto layout = static_cast<GemmMatrixLayout>(std::stoi(argv[1]));
232

233
234
235
        const int M = std::stoi(argv[2]);
        const int N = std::stoi(argv[3]);
        const int K = std::stoi(argv[4]);
236

237
238
239
240
241
        const int StrideA = std::stoi(argv[5]);
        const int StrideB = std::stoi(argv[6]);
        const int StrideC = std::stoi(argv[7]);
        const int KBatch  = std::stoi(argv[8]);
        test_cases        = {{layout, M, N, K, StrideA, StrideB, StrideC, KBatch}};
242
243
244
245
246
247
248
249
250
251
    }
    else
    {
        printf("arg1: matrix layout (0: A[m, k] * B[k, n] = C[m, n];\n");
        printf("                     1: A[m, k] * B[n, k] = C[m, n];\n");
        printf("                     2: A[k, m] * B[k, n] = C[m, n];\n");
        printf("                     3: A[k, m] * B[n, k] = C[m, n])\n");
        printf("arg2 to 7: M, N, K, StrideA, StrideB, StrideC KBatch\n");
        return -1;
    }
252
    for(const auto& kinder : test_cases)
253
254
255
    {
        const auto res = test_gemm(kinder);
        if(!res)
256
            return -1;
ltqin's avatar
ltqin committed
257
258
259
    }
    return 0;
}