gemm_xdl_bf16.cpp 4.42 KB
Newer Older
1
2
3
4
5
6
7
#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <numeric>
#include <tuple>
#include <vector>

Chao Liu's avatar
Chao Liu committed
8
9
10
11
12
13
14
15
16
17
18
19
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
#include "ck/tensor_operation/gpu/device/device_gemm.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/reference_tensor_operation/cpu/reference_gemm.hpp"

#include "test/gemm/gemm_util.hpp"
20
21
22

using PassThrough = ck::tensor_operation::element_wise::PassThrough;

23
using DeviceGemmNoOpPtr =
24
25
26
27
28
29
30
31
    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 {
Chao Liu's avatar
Chao Liu committed
32
33
34
35
36
37
38
39
void add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_kn_mn_instances(
    std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_nk_mn_instances(
    std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_mk_nk_mn_instances(
    std::vector<DeviceGemmNoOpPtr>&);
void add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_mk_kn_mn_instances(
    std::vector<DeviceGemmNoOpPtr>&);
40
} // namespace device_gemm_instance
41
42
43
44
} // namespace device
} // namespace tensor_operation
} // namespace ck

45
int main()
46
{
47
48
    using RowMajor    = ck::tensor_layout::gemm::RowMajor;
    using ColumnMajor = ck::tensor_layout::gemm::ColumnMajor;
49

50
51
    bool res = true;
    std::vector<DeviceGemmNoOpPtr> gemmPtrs;
52

53
54
    ck::tensor_operation::device::device_gemm_instance::
        add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_kn_mn_instances(gemmPtrs);
55

56
57
58
59
60
61
62
63
64
65
    for(auto& gemmPtr : gemmPtrs)
    {
        res &= ck::gemm_util::TestGemmBF16<DeviceGemmNoOpPtr,
                                           ColumnMajor,
                                           RowMajor,
                                           RowMajor,
                                           PassThrough,
                                           PassThrough,
                                           PassThrough>{}(gemmPtr);
    }
66

67
68
69
    gemmPtrs.clear();
    ck::tensor_operation::device::device_gemm_instance::
        add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_km_nk_mn_instances(gemmPtrs);
70

71
72
73
74
75
76
77
78
79
80
    for(auto& gemmPtr : gemmPtrs)
    {
        res &= ck::gemm_util::TestGemmBF16<DeviceGemmNoOpPtr,
                                           ColumnMajor,
                                           ColumnMajor,
                                           RowMajor,
                                           PassThrough,
                                           PassThrough,
                                           PassThrough>{}(gemmPtr);
    }
81

82
83
84
    gemmPtrs.clear();
    ck::tensor_operation::device::device_gemm_instance::
        add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_mk_kn_mn_instances(gemmPtrs);
85

86
87
88
89
90
91
92
93
94
95
    for(auto& gemmPtr : gemmPtrs)
    {
        res &= ck::gemm_util::TestGemmBF16<DeviceGemmNoOpPtr,
                                           RowMajor,
                                           RowMajor,
                                           RowMajor,
                                           PassThrough,
                                           PassThrough,
                                           PassThrough>{}(gemmPtr);
    }
96

97
    gemmPtrs.clear();
98
99
100
101
102
    ck::tensor_operation::device::device_gemm_instance::
        add_device_gemm_xdl_c_shuffle_bf16_bf16_bf16_mk_nk_mn_instances(gemmPtrs);

    for(auto& gemmPtr : gemmPtrs)
    {
103
104
105
106
107
108
109
        res &= ck::gemm_util::TestGemmBF16<DeviceGemmNoOpPtr,
                                           RowMajor,
                                           ColumnMajor,
                                           RowMajor,
                                           PassThrough,
                                           PassThrough,
                                           PassThrough>{}(gemmPtr);
110
111
112
    }

    std::cout << "TestGemm ..... " << (res ? "SUCCESS" : "FAILURE") << std::endl;
rocking5566's avatar
rocking5566 committed
113
    return res ? 0 : 1;
114
}