gemm_standalone_xdl_fp16.cpp 7.76 KB
Newer Older
Anthony Chang's avatar
Anthony Chang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#include "gemm_util.hpp"

#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp"

#include "gemm_f16_nn_instance.hpp"
#include "gemm_f16_nt_instance.hpp"
#include "gemm_f16_tn_instance.hpp"
#include "gemm_f16_tt_instance.hpp"

using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;

Anthony Chang's avatar
tidy  
Anthony Chang committed
17
18
19
20
21
22
using PassThrough = ck::tensor_operation::element_wise::PassThrough;
using F16         = ck::half_t;
using ADataType   = F16;
using BDataType   = F16;
using AccDataType = float;
using CDataType   = F16;
Anthony Chang's avatar
Anthony Chang committed
23
24
25
26
27
28
29
30
31

using ALayout = Row;
using BLayout = Col;
using CLayout = Row;

using AElementOp = PassThrough;
using BElementOp = PassThrough;
using CElementOp = PassThrough;

32
using ck::gemm_util::GemmParams;
Anthony Chang's avatar
Anthony Chang committed
33
using ck::tensor_operation::device::BaseOperator;
34
35
using ck::tensor_operation::device::DeviceGemm;
using namespace ck::tensor_operation::device::instance;
Anthony Chang's avatar
Anthony Chang committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

using DeviceGemmNN =
    DeviceGemm<Col, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>;
using DeviceGemmNT =
    DeviceGemm<Col, Row, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>;
using DeviceGemmTN =
    DeviceGemm<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>;
using DeviceGemmTT =
    DeviceGemm<Row, Row, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>;

struct LayoutConfig
{
    bool ARowMajor;
    bool BRowMajor;
    bool CRowMajor;
};

int main(int argc, char* argv[])
{
    // Class DeviceGemm is templated by layout and precision types so it is not an option to contain
    // them in a single vector. Instead we use abstract BaseOperator class and dynamic_cast() it
    // upon invocation.
    // And since DeviceGemm does not expose template arg information, an extra book keeping class
    // LayoutConfig is used for determining which type a BaseOperator instance should be cast to.
    using OpFactoryFn = void (*)(std::vector<std::unique_ptr<BaseOperator>>&);

62
    std::vector<std::tuple<GemmParams, LayoutConfig, OpFactoryFn>> problems = {
Anthony Chang's avatar
Anthony Chang committed
63
64
        // clang-format off
    // 104 tiles
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
    {GemmParams{2048, 3328, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_256x256},
    {GemmParams{2048, 1664, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_256x128},
    {GemmParams{1024, 1664, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_128x128},
    {GemmParams{1024,  832, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_128x64},
    {GemmParams{2048, 3328, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_256x256},
    {GemmParams{2048, 1664, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_256x128},
    {GemmParams{1024, 1664, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_128x128},
    {GemmParams{1024,  832, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_128x64},
    {GemmParams{2048, 3328, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_256x128},
    {GemmParams{2048, 1664, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_256x128},
    {GemmParams{1024, 1664, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_128x128},
    {GemmParams{1024,  832, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_128x64},
    {GemmParams{2048, 3328, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_256x256},
    {GemmParams{2048, 1664, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_256x128},
    {GemmParams{1024, 1664, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_128x128},
    {GemmParams{1024,  832, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_128x64},
Anthony Chang's avatar
Anthony Chang committed
81
    // 110 tiles
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    {GemmParams{2560, 2816, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_256x256},
    {GemmParams{2560, 1408, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_256x128},
    {GemmParams{1280, 1408, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_128x128},
    {GemmParams{1280,  704, 4096, -1, -1, -1}, LayoutConfig{false, false, true}, add_gemm_f16_nn_128x64},
    {GemmParams{2560, 2816, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_256x256},
    {GemmParams{2560, 1408, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_256x128},
    {GemmParams{1280, 1408, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_128x128},
    {GemmParams{1280,  704, 4096, -1, -1, -1}, LayoutConfig{false, true, true}, add_gemm_f16_nt_128x64},
    {GemmParams{2560, 2816, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_256x128},
    {GemmParams{2560, 1408, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_256x128},
    {GemmParams{1280, 1408, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_128x128},
    {GemmParams{1280,  704, 4096, -1, -1, -1}, LayoutConfig{true, false, true}, add_gemm_f16_tn_128x64},
    {GemmParams{2560, 2816, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_256x256},
    {GemmParams{2560, 1408, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_256x128},
    {GemmParams{1280, 1408, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_128x128},
    {GemmParams{1280,  704, 4096, -1, -1, -1}, LayoutConfig{true, true, true}, add_gemm_f16_tt_128x64},
Anthony Chang's avatar
Anthony Chang committed
98
99
100
        // clang-format on
    };

101
102
    bool do_verification = true;
    bool time_kernel     = true;
Anthony Chang's avatar
Anthony Chang committed
103

104
105
106
107
108
109
110
111
112
113
    if(argc == 1)
    {
        // use default
    }
    else if(argc == 3)
    {
        do_verification = std::stoi(argv[1]);
        time_kernel     = std::stoi(argv[2]);
    }
    else
Anthony Chang's avatar
Anthony Chang committed
114
    {
115
116
        std::cerr << "arg1: verification (0=no, 1=yes)" << std::endl
                  << "arg2: time kernel (0=no, 1=yes)" << std::endl;
117
        return 0;
Anthony Chang's avatar
Anthony Chang committed
118
119
    }

120
    bool pass = true;
Anthony Chang's avatar
Anthony Chang committed
121
122
    for(auto& p : problems)
    {
123
        GemmParams& problem_size          = std::get<0>(p);
Anthony Chang's avatar
Anthony Chang committed
124
125
126
127
128
        const LayoutConfig& layout_config = std::get<1>(p);
        const auto& factory               = std::get<2>(p);
        std::vector<std::unique_ptr<BaseOperator>> ops;
        factory(ops);

129
130
131
132
        problem_size.StrideA = layout_config.ARowMajor ? problem_size.K : problem_size.M;
        problem_size.StrideB = layout_config.BRowMajor ? problem_size.N : problem_size.K;
        problem_size.StrideC = layout_config.CRowMajor ? problem_size.N : problem_size.M;

Anthony Chang's avatar
Anthony Chang committed
133
134
135
        if(!layout_config.ARowMajor && !layout_config.BRowMajor)
        {
            auto op_ptr = dynamic_cast<DeviceGemmNN*>(ops[0].get());
136
            pass &= ck::gemm_util::TestGemm<AccDataType>{}(
137
                op_ptr, problem_size, do_verification, time_kernel);
Anthony Chang's avatar
Anthony Chang committed
138
139
140
141
        }
        else if(!layout_config.ARowMajor && layout_config.BRowMajor)
        {
            auto op_ptr = dynamic_cast<DeviceGemmNT*>(ops[0].get());
142
            pass &= ck::gemm_util::TestGemm<AccDataType>{}(
143
                op_ptr, problem_size, do_verification, time_kernel);
Anthony Chang's avatar
Anthony Chang committed
144
145
146
147
        }
        else if(layout_config.ARowMajor && !layout_config.BRowMajor)
        {
            auto op_ptr = dynamic_cast<DeviceGemmTN*>(ops[0].get());
148
            pass &= ck::gemm_util::TestGemm<AccDataType>{}(
149
                op_ptr, problem_size, do_verification, time_kernel);
Anthony Chang's avatar
Anthony Chang committed
150
151
152
153
        }
        else if(layout_config.ARowMajor && layout_config.BRowMajor)
        {
            auto op_ptr = dynamic_cast<DeviceGemmTT*>(ops[0].get());
154
            pass &= ck::gemm_util::TestGemm<AccDataType>{}(
155
                op_ptr, problem_size, do_verification, time_kernel);
Anthony Chang's avatar
Anthony Chang committed
156
157
158
        }
    }

159
160
    std::cout << (pass ? "ALL TESTS PASSED" : "SOME TESTS FAILED") << std::endl;
    return pass ? 0 : 1;
Anthony Chang's avatar
Anthony Chang committed
161
}