profile_batched_gemm.cpp 7.46 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
#include <cstdint>
zjing14's avatar
zjing14 committed
5
6
7
8
#include <iostream>
#include <numeric>
#include <initializer_list>
#include <cstdlib>
Chao Liu's avatar
Chao Liu committed
9

10
11
#include "profiler/profile_batched_gemm_impl.hpp"
#include "profiler_operation_registry.hpp"
zjing14's avatar
zjing14 committed
12

Chao Liu's avatar
Chao Liu committed
13
enum struct GemmMatrixLayout
zjing14's avatar
zjing14 committed
14
15
16
17
18
19
20
{
    MK_KN_MN, // 0
    MK_NK_MN, // 1
    KM_KN_MN, // 2
    KM_NK_MN, // 3
};

Chao Liu's avatar
Chao Liu committed
21
enum struct GemmDataType
zjing14's avatar
zjing14 committed
22
{
23
24
    F32_F32_F32,    // 0
    F16_F16_F16,    // 1
Jianfeng Yan's avatar
Jianfeng Yan committed
25
26
    BF16_BF16_BF16, // 2
    INT8_INT8_INT8, // 3
zjing14's avatar
zjing14 committed
27
28
};

29
30
31
#define OP_NAME "batched_gemm"
#define OP_DESC "Batched GEMM"

zjing14's avatar
zjing14 committed
32
33
int profile_batched_gemm(int argc, char* argv[])
{
Chao Liu's avatar
Chao Liu committed
34
    if(argc != 18)
zjing14's avatar
zjing14 committed
35
    {
Chao Liu's avatar
Chao Liu committed
36
        // clang-format off
37
        printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n");
Jianfeng Yan's avatar
Jianfeng Yan committed
38
        printf("arg2: data type (0: fp32; 1: fp16, 2: bf16, 3: int8)\n");
zjing14's avatar
zjing14 committed
39
40
41
42
43
44
        printf("arg3: matrix layout (0: A[g, m, k] * B[g, k, n] = C[g, m, n];\n");
        printf("                     1: A[g, m, k] * B[g, n, k] = C[g, m, n];\n");
        printf("                     2: A[g, k, m] * B[g, k, n] = C[g, m, n];\n");
        printf("                     3: A[g, k, m] * B[g, n, k] = C[g, m, n])\n");
        printf("arg4: verification (0: no; 1: yes)\n");
        printf("arg5: initialization (0: no init; 1: integer value; 2: decimal value)\n");
JD's avatar
JD committed
45
46
        printf("arg6: print tensor value (0: no; 1: yes)\n");
        printf("arg7: time kernel (0=n0, 1=yes)\n");
Chao Liu's avatar
Chao Liu committed
47
48
        printf("arg8 to 17: M, N, K, StrideA, StrideB, StrideC, BatchStrideA, BatchStrideB, BatchStrideC, BatchCount\n");
        // clang-format on
zjing14's avatar
zjing14 committed
49
50
51
        exit(1);
    }

Chao Liu's avatar
Chao Liu committed
52
53
    const auto data_type       = static_cast<GemmDataType>(std::stoi(argv[2]));
    const auto layout          = static_cast<GemmMatrixLayout>(std::stoi(argv[3]));
zjing14's avatar
zjing14 committed
54
55
56
    const bool do_verification = std::stoi(argv[4]);
    const int init_method      = std::stoi(argv[5]);
    const bool do_log          = std::stoi(argv[6]);
JD's avatar
JD committed
57
    const bool time_kernel     = std::stoi(argv[7]);
zjing14's avatar
zjing14 committed
58
59
60
61
62
63
64
65
66

    const int M = std::stoi(argv[8]);
    const int N = std::stoi(argv[9]);
    const int K = std::stoi(argv[10]);

    const int StrideA = std::stoi(argv[11]);
    const int StrideB = std::stoi(argv[12]);
    const int StrideC = std::stoi(argv[13]);

Chao Liu's avatar
Chao Liu committed
67
68
69
70
71
    const int BatchStrideA = std::stoi(argv[14]);
    const int BatchStrideB = std::stoi(argv[15]);
    const int BatchStrideC = std::stoi(argv[16]);

    const int BatchCount = std::stoi(argv[17]);
zjing14's avatar
zjing14 committed
72

Chao Liu's avatar
Chao Liu committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
    using F32  = float;
    using F16  = ck::half_t;
    using BF16 = ck::bhalf_t;
    using INT8 = int8_t;

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

    auto profile = [&](auto a_type,
                       auto b_type,
                       auto c_type,
                       auto a_layout,
                       auto b_layout,
                       auto c_layout) {
        using ADataType = decltype(a_type);
        using BDataType = decltype(b_type);
        using CDataType = decltype(c_type);

        using ALayout = decltype(a_layout);
        using BLayout = decltype(b_layout);
        using CLayout = decltype(c_layout);

        const int DefaultStrideA = ck::is_same_v<ALayout, Row> ? K : M;
        const int DefaultStrideB = ck::is_same_v<BLayout, Row> ? N : K;
        const int DefaultStrideC = ck::is_same_v<CLayout, Row> ? N : M;

99
100
101
102
        const int StrideA_ = (StrideA < 0) ? DefaultStrideA : StrideA;
        const int StrideB_ = (StrideB < 0) ? DefaultStrideB : StrideB;
        const int StrideC_ = (StrideC < 0) ? DefaultStrideC : StrideC;

Chao Liu's avatar
Chao Liu committed
103
104
105
106
107
108
109
        const int DefaultBatchStrideA = (ck::is_same_v<ALayout, Row> ? M : K) * StrideA_;
        const int DefaultBatchStrideB = (ck::is_same_v<BLayout, Row> ? K : N) * StrideB_;
        const int DefaultBatchStrideC = (ck::is_same_v<CLayout, Row> ? M : N) * StrideC_;

        const int BatchStrideA_ = (BatchStrideA < 0) ? DefaultBatchStrideA : BatchStrideA;
        const int BatchStrideB_ = (BatchStrideB < 0) ? DefaultBatchStrideB : BatchStrideB;
        const int BatchStrideC_ = (BatchStrideC < 0) ? DefaultBatchStrideC : BatchStrideC;
110

Chao Liu's avatar
Chao Liu committed
111
112
113
114
115
116
117
118
119
        bool pass = ck::profiler::
            profile_batched_gemm_impl<ADataType, BDataType, CDataType, ALayout, BLayout, CLayout>(
                do_verification,
                init_method,
                do_log,
                time_kernel,
                M,
                N,
                K,
Chao Liu's avatar
Chao Liu committed
120
121
122
                BatchStrideA_,
                BatchStrideB_,
                BatchStrideC_,
123
124
125
                StrideA_,
                StrideB_,
                StrideC_,
Chao Liu's avatar
Chao Liu committed
126
127
128
129
130
131
                BatchCount);

        return pass ? 0 : 1;
    };

    if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_KN_MN)
zjing14's avatar
zjing14 committed
132
    {
Chao Liu's avatar
Chao Liu committed
133
        return profile(F32{}, F32{}, F32{}, Row{}, Row{}, Row{});
zjing14's avatar
zjing14 committed
134
    }
Chao Liu's avatar
Chao Liu committed
135
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::MK_NK_MN)
zjing14's avatar
zjing14 committed
136
    {
Chao Liu's avatar
Chao Liu committed
137
        return profile(F32{}, F32{}, F32{}, Row{}, Col{}, Row{});
zjing14's avatar
zjing14 committed
138
    }
Chao Liu's avatar
Chao Liu committed
139
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_KN_MN)
zjing14's avatar
zjing14 committed
140
    {
Chao Liu's avatar
Chao Liu committed
141
        return profile(F32{}, F32{}, F32{}, Col{}, Row{}, Row{});
zjing14's avatar
zjing14 committed
142
    }
Chao Liu's avatar
Chao Liu committed
143
    else if(data_type == GemmDataType::F32_F32_F32 && layout == GemmMatrixLayout::KM_NK_MN)
zjing14's avatar
zjing14 committed
144
    {
Chao Liu's avatar
Chao Liu committed
145
        return profile(F32{}, F32{}, F32{}, Col{}, Col{}, Row{});
zjing14's avatar
zjing14 committed
146
    }
Chao Liu's avatar
Chao Liu committed
147
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_KN_MN)
Jianfeng Yan's avatar
Jianfeng Yan committed
148
    {
Chao Liu's avatar
Chao Liu committed
149
        return profile(F16{}, F16{}, F16{}, Row{}, Row{}, Row{});
Jianfeng Yan's avatar
Jianfeng Yan committed
150
    }
Chao Liu's avatar
Chao Liu committed
151
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::MK_NK_MN)
Jianfeng Yan's avatar
Jianfeng Yan committed
152
    {
Chao Liu's avatar
Chao Liu committed
153
        return profile(F16{}, F16{}, F16{}, Row{}, Col{}, Row{});
Jianfeng Yan's avatar
Jianfeng Yan committed
154
    }
Chao Liu's avatar
Chao Liu committed
155
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_KN_MN)
Jianfeng Yan's avatar
Jianfeng Yan committed
156
    {
Chao Liu's avatar
Chao Liu committed
157
        return profile(F16{}, F16{}, F16{}, Col{}, Row{}, Row{});
Jianfeng Yan's avatar
Jianfeng Yan committed
158
    }
Chao Liu's avatar
Chao Liu committed
159
    else if(data_type == GemmDataType::F16_F16_F16 && layout == GemmMatrixLayout::KM_NK_MN)
Jianfeng Yan's avatar
Jianfeng Yan committed
160
    {
Chao Liu's avatar
Chao Liu committed
161
        return profile(F16{}, F16{}, F16{}, Col{}, Col{}, Row{});
Jianfeng Yan's avatar
Jianfeng Yan committed
162
    }
Chao Liu's avatar
Chao Liu committed
163
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::MK_KN_MN)
164
    {
Chao Liu's avatar
Chao Liu committed
165
        return profile(BF16{}, BF16{}, BF16{}, Row{}, Row{}, Row{});
166
    }
Chao Liu's avatar
Chao Liu committed
167
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::MK_NK_MN)
168
    {
Chao Liu's avatar
Chao Liu committed
169
        return profile(BF16{}, BF16{}, BF16{}, Row{}, Col{}, Row{});
170
    }
Chao Liu's avatar
Chao Liu committed
171
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::KM_KN_MN)
172
    {
Chao Liu's avatar
Chao Liu committed
173
        return profile(BF16{}, BF16{}, BF16{}, Col{}, Row{}, Row{});
174
    }
Chao Liu's avatar
Chao Liu committed
175
    else if(data_type == GemmDataType::BF16_BF16_BF16 && layout == GemmMatrixLayout::KM_NK_MN)
176
    {
Chao Liu's avatar
Chao Liu committed
177
        return profile(BF16{}, BF16{}, BF16{}, Col{}, Col{}, Row{});
178
    }
Jianfeng Yan's avatar
Jianfeng Yan committed
179
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::MK_KN_MN)
180
    {
Chao Liu's avatar
Chao Liu committed
181
        return profile(INT8{}, INT8{}, INT8{}, Row{}, Row{}, Row{});
182
    }
Jianfeng Yan's avatar
Jianfeng Yan committed
183
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::MK_NK_MN)
184
    {
Chao Liu's avatar
Chao Liu committed
185
        return profile(INT8{}, INT8{}, INT8{}, Row{}, Col{}, Row{});
186
    }
Jianfeng Yan's avatar
Jianfeng Yan committed
187
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::KM_KN_MN)
188
    {
Chao Liu's avatar
Chao Liu committed
189
        return profile(INT8{}, INT8{}, INT8{}, Col{}, Row{}, Row{});
190
    }
Jianfeng Yan's avatar
Jianfeng Yan committed
191
    else if(data_type == GemmDataType::INT8_INT8_INT8 && layout == GemmMatrixLayout::KM_NK_MN)
192
    {
Chao Liu's avatar
Chao Liu committed
193
        return profile(INT8{}, INT8{}, INT8{}, Col{}, Col{}, Row{});
194
    }
zjing14's avatar
zjing14 committed
195
196
    else
    {
Chao Liu's avatar
Chao Liu committed
197
        std::cout << "this data_type & layout is not implemented" << std::endl;
zjing14's avatar
zjing14 committed
198

Chao Liu's avatar
Chao Liu committed
199
200
        return 1;
    }
zjing14's avatar
zjing14 committed
201
}
202
203

REGISTER_PROFILER_OPERATION(OP_NAME, OP_DESC, profile_batched_gemm);