"vscode:/vscode.git/clone" did not exist on "b4bc91c1f884f4c859ac7871a5f631a3731c734a"
device_grouped_gemm_xdl.hpp 21.4 KB
Newer Older
Jing Zhang's avatar
Jing Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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
53
54
55
#ifndef DEVICE_GROUPED_GEMM_XDL_HPP
#define DEVICE_GROUPED_GEMM_XDL_HPP

#include <iostream>
#include <sstream>
#include "device.hpp"
#include "device_base.hpp"
#include "device_gemm.hpp"
#include "common_header.hpp"
#include "tensor_layout.hpp"
#include "tensor_descriptor.hpp"
#include "tensor_descriptor_helper.hpp"
#include "gridwise_grouped_gemm_xdlops_v2r3.hpp"
#include "gemm_specialization.hpp"

namespace ck {
namespace tensor_operation {
namespace device {

template <typename ADataType,
          typename BDataType,
          typename CDataType,
          typename AccDataType,
          typename ALayout,
          typename BLayout,
          typename CLayout,
          typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation,
          GemmSpecialization_t GemmSpecialization,
          ck::index_t BlockSize,
          ck::index_t MPerBlock,
          ck::index_t NPerBlock,
          ck::index_t K0PerBlock,
          ck::index_t K1,
          ck::index_t MPerXDL,
          ck::index_t NPerXDL,
          ck::index_t MXdlPerWave,
          ck::index_t NXdlPerWave,
          typename ABlockTransferThreadClusterLengths_K0_M_K1,
          typename ABlockTransferThreadClusterArrangeOrder,
          typename ABlockTransferSrcAccessOrder,
          ck::index_t ABlockTransferSrcVectorDim,
          ck::index_t ABlockTransferSrcScalarPerVector,
          ck::index_t ABlockTransferDstScalarPerVector_K1,
          bool ABlockLdsAddExtraM,
          typename BBlockTransferThreadClusterLengths_K0_N_K1,
          typename BBlockTransferThreadClusterArrangeOrder,
          typename BBlockTransferSrcAccessOrder,
          ck::index_t BBlockTransferSrcVectorDim,
          ck::index_t BBlockTransferSrcScalarPerVector,
          ck::index_t BBlockTransferDstScalarPerVector_K1,
          bool BBlockLdsAddExtraN,
          ck::index_t CThreadTransferSrcDstVectorDim,
          ck::index_t CThreadTransferDstScalarPerVector,
Jing Zhang's avatar
Jing Zhang committed
56
57
          ck::index_t NumPrefetch   = 1,
          ck::index_t MaxGroupCount = 5>
Jing Zhang's avatar
Jing Zhang committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
struct DeviceGroupedGemmXdl
    : public DeviceGroupedGemm<AElementwiseOperation, BElementwiseOperation, CElementwiseOperation>
{
    static constexpr auto I0 = Number<0>{};
    static constexpr auto I1 = Number<1>{};
    static constexpr auto I2 = Number<2>{};

    static constexpr auto K1Number = Number<K1>{};

    static auto MakeAGridDescriptor_K0_M_K1(index_t M, index_t K, index_t StrideA)
    {
        assert(K % K1 == 0);

        const index_t K0 = K / K1;

        const auto a_grid_desc_m_k = [&]() {
            if constexpr(is_same<tensor_layout::gemm::RowMajor, ALayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(StrideA, I1));
            }
            else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, ALayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(M, K), make_tuple(I1, StrideA));
            }
        }();

        if constexpr(GemmSpecialization == GemmSpecialization_t::MNPadding)
        {
            const auto PadM = (MPerBlock - M % MPerBlock) % MPerBlock;

            return transform_tensor_descriptor(
                a_grid_desc_m_k,
                make_tuple(make_unmerge_transform(make_tuple(K0, K1Number)),
                           make_right_pad_transform(M, PadM)),
                make_tuple(Sequence<1>{}, Sequence<0>{}),
                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
        }
        else
        {
            return transform_tensor_descriptor(
                a_grid_desc_m_k,
                make_tuple(make_unmerge_transform(make_tuple(K0, K1Number)),
                           make_pass_through_transform(M)),
                make_tuple(Sequence<1>{}, Sequence<0>{}),
                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
        }
    }

    static auto MakeBGridDescriptor_K0_N_K1(index_t K, index_t N, index_t StrideB)
    {
        assert(K % K1 == 0);

        const index_t K0 = K / K1;

        const auto b_grid_desc_k_n = [&]() {
            if constexpr(is_same<tensor_layout::gemm::RowMajor, BLayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(K, N), make_tuple(StrideB, I1));
            }
            else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, BLayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(K, N), make_tuple(I1, StrideB));
            }
        }();

        if constexpr(GemmSpecialization == GemmSpecialization_t::MNPadding)
        {
            const auto PadN = (NPerBlock - N % NPerBlock) % NPerBlock;

            return transform_tensor_descriptor(
                b_grid_desc_k_n,
                make_tuple(make_unmerge_transform(make_tuple(K0, K1Number)),
                           make_right_pad_transform(N, PadN)),
                make_tuple(Sequence<0>{}, Sequence<1>{}),
                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
        }
        else
        {
            return transform_tensor_descriptor(
                b_grid_desc_k_n,
                make_tuple(make_unmerge_transform(make_tuple(K0, K1Number)),
                           make_pass_through_transform(N)),
                make_tuple(Sequence<0>{}, Sequence<1>{}),
                make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
        }
    }

    static auto MakeCGridDescriptor_M_N(index_t M, index_t N, index_t StrideC)
    {
        const auto c_grid_desc_m_n = [&]() {
            if constexpr(is_same<tensor_layout::gemm::RowMajor, CLayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(StrideC, I1));
            }
            else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, CLayout>::value)
            {
                return make_naive_tensor_descriptor(make_tuple(M, N), make_tuple(I1, StrideC));
            }
        }();

        if constexpr(GemmSpecialization == GemmSpecialization_t::MNPadding)
        {
            const auto PadM = (MPerBlock - M % MPerBlock) % MPerBlock;
            const auto PadN = (NPerBlock - N % NPerBlock) % NPerBlock;

            return transform_tensor_descriptor(
                c_grid_desc_m_n,
                make_tuple(make_right_pad_transform(M, PadM), make_right_pad_transform(N, PadN)),
                make_tuple(Sequence<0>{}, Sequence<1>{}),
                make_tuple(Sequence<0>{}, Sequence<1>{}));
        }
        else
        {

            return transform_tensor_descriptor(
                c_grid_desc_m_n,
                make_tuple(make_pass_through_transform(M), make_pass_through_transform(N)),
                make_tuple(Sequence<0>{}, Sequence<1>{}),
                make_tuple(Sequence<0>{}, Sequence<1>{}));
        }
    }

    using AGridDesc_K0_M_K1 = decltype(MakeAGridDescriptor_K0_M_K1(1, 1, 1));
    using BGridDesc_K0_N_K1 = decltype(MakeBGridDescriptor_K0_N_K1(1, 1, 1));
    using CGridDesc_M_N     = decltype(MakeCGridDescriptor_M_N(1, 1, 1));

    // GridwiseGemm
    using GridwiseGemm = GridwiseGroupedGemm_k0mk1_k0nk1_mn_xdlops_v2r3<
        BlockSize,
        ADataType, // TODO: distinguish A/B datatype
        AccDataType,
        CDataType,
        InMemoryDataOperationEnum_t::Set,
        AGridDesc_K0_M_K1,
        BGridDesc_K0_N_K1,
        CGridDesc_M_N,
        AElementwiseOperation,
        BElementwiseOperation,
        CElementwiseOperation,
        MPerBlock,
        NPerBlock,
        K0PerBlock,
        MPerXDL,
        NPerXDL,
        K1,
        MXdlPerWave,
        NXdlPerWave,
        ABlockTransferThreadClusterLengths_K0_M_K1,
        ABlockTransferThreadClusterArrangeOrder,
        ABlockTransferSrcAccessOrder,
        ABlockTransferSrcVectorDim,
        ABlockTransferSrcScalarPerVector,
        ABlockTransferDstScalarPerVector_K1,
        false, // AThreadTransferSrcResetCoordinateAfterRun,
        ABlockLdsAddExtraM,
        BBlockTransferThreadClusterLengths_K0_N_K1,
        BBlockTransferThreadClusterArrangeOrder,
        BBlockTransferSrcAccessOrder,
        BBlockTransferSrcVectorDim,
        BBlockTransferSrcScalarPerVector,
        BBlockTransferDstScalarPerVector_K1,
        false, // BThreadTransferSrcResetCoordinateAfterRun,
        BBlockLdsAddExtraN,
        Sequence<0, 2, 4, 5, 6, 1, 3, 7>, // CThreadTransferSrcDstAccessOrder,
        CThreadTransferSrcDstVectorDim,
        CThreadTransferDstScalarPerVector,
        NumPrefetch>;

Jing Zhang's avatar
Jing Zhang committed
226
227
228
229
230
231
232
233
234
235
236
237
238
239
    struct GemmDesc
    {
        AGridDesc_K0_M_K1 a_grid_desc_k0_m_k1_;
        BGridDesc_K0_N_K1 b_grid_desc_k0_n_k1_;
        CGridDesc_M_N c_grid_desc_m_n_;

        typename GridwiseGemm::CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2
            c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_;
        typename GridwiseGemm::DefaultBlock2CTileMap block_2_ctile_map_;

        ck::index_t OffsetA, OffsetB, OffsetC;
        ck::index_t BlockStart, BlockEnd;
    };

Jing Zhang's avatar
Jing Zhang committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    // Argument
    struct Argument : public BaseArgument
    {
        Argument(const ADataType* p_a_grid,
                 const BDataType* p_b_grid,
                 CDataType* p_c_grid,
                 std::vector<gemm_desc> gemm_shapes,
                 index_t M01,
                 index_t N01,
                 AElementwiseOperation a_element_op,
                 BElementwiseOperation b_element_op,
                 CElementwiseOperation c_element_op)
            : p_a_grid_{p_a_grid},
              p_b_grid_{p_b_grid},
              p_c_grid_{p_c_grid},
Jing Zhang's avatar
Jing Zhang committed
255
              gemm_shapes_{gemm_shapes},
Jing Zhang's avatar
Jing Zhang committed
256
257
258
259
260
261
262
              M01_{M01},
              N01_{N01},
              a_element_op_{a_element_op},
              b_element_op_{b_element_op},
              c_element_op_{c_element_op}
        {

Jing Zhang's avatar
Jing Zhang committed
263
264
            grid_size = 0;

Jing Zhang's avatar
Jing Zhang committed
265
266
267
268
269
270
271
272
273
274
275
            static_for<0, MaxGroupCount, 1>{}([&](auto i) {
                if(i < gemm_shapes_.size())
                {
                    const index_t M = gemm_shapes_[i].M;
                    const index_t N = gemm_shapes_[i].N;
                    const index_t K = gemm_shapes_[i].K;

                    const index_t StrideA = gemm_shapes_[i].StrideA;
                    const index_t StrideB = gemm_shapes_[i].StrideB;
                    const index_t StrideC = gemm_shapes_[i].StrideC;

Jing Zhang's avatar
Jing Zhang committed
276
                    gemm_desc_(i).a_grid_desc_k0_m_k1_ =
Jing Zhang's avatar
Jing Zhang committed
277
                        DeviceGroupedGemmXdl::MakeAGridDescriptor_K0_M_K1(M, K, StrideA);
Jing Zhang's avatar
Jing Zhang committed
278
                    gemm_desc_(i).b_grid_desc_k0_n_k1_ =
Jing Zhang's avatar
Jing Zhang committed
279
                        DeviceGroupedGemmXdl::MakeBGridDescriptor_K0_N_K1(K, N, StrideB);
Jing Zhang's avatar
Jing Zhang committed
280
                    gemm_desc_(i).c_grid_desc_m_n_ =
Jing Zhang's avatar
Jing Zhang committed
281
282
                        DeviceGroupedGemmXdl::MakeCGridDescriptor_M_N(M, N, StrideC);

Jing Zhang's avatar
Jing Zhang committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
                    const index_t grid_size_grp =
                        GridwiseGemm::CalculateGridSize(gemm_desc_[i].c_grid_desc_m_n_);

                    gemm_desc_(i).BlockStart = grid_size;
                    gemm_desc_(i).BlockEnd   = grid_size + grid_size_grp;

                    grid_size += grid_size_grp;

                    gemm_desc_(i).OffsetA = gemm_shapes_[i].OffsetA;
                    gemm_desc_(i).OffsetB = gemm_shapes_[i].OffsetB;
                    gemm_desc_(i).OffsetC = gemm_shapes_[i].OffsetC;

                    if(GridwiseGemm::CheckValidity(gemm_desc_[i].a_grid_desc_k0_m_k1_,
                                                   gemm_desc_[i].b_grid_desc_k0_n_k1_,
                                                   gemm_desc_[i].c_grid_desc_m_n_,
Jing Zhang's avatar
Jing Zhang committed
298
299
300
                                                   M01_,
                                                   N01_))
                    {
Jing Zhang's avatar
Jing Zhang committed
301
                        gemm_desc_(i).c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_ =
Jing Zhang's avatar
Jing Zhang committed
302
                            GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(
Jing Zhang's avatar
Jing Zhang committed
303
                                gemm_desc_[i].c_grid_desc_m_n_);
Jing Zhang's avatar
Jing Zhang committed
304

Jing Zhang's avatar
Jing Zhang committed
305
306
                        gemm_desc_(i).block_2_ctile_map_ = GridwiseGemm::MakeDefaultBlock2CTileMap(
                            gemm_desc_[i].c_grid_desc_m_n_, M01, N01);
Jing Zhang's avatar
Jing Zhang committed
307
308
                    }
                }
Jing Zhang's avatar
Jing Zhang committed
309
310
311
312
313
                else
                {
                    gemm_desc_(i).BlockStart = -1;
                    gemm_desc_(i).BlockEnd   = -1;
                }
Jing Zhang's avatar
Jing Zhang committed
314
            });
Jing Zhang's avatar
Jing Zhang committed
315
316
317
318
319
320
        }

        //  private:
        const ADataType* p_a_grid_;
        const BDataType* p_b_grid_;
        CDataType* p_c_grid_;
Jing Zhang's avatar
Jing Zhang committed
321
        std::vector<gemm_desc> gemm_shapes_;
Jing Zhang's avatar
Jing Zhang committed
322
323
324
325
326
        index_t M01_;
        index_t N01_;
        AElementwiseOperation a_element_op_;
        BElementwiseOperation b_element_op_;
        CElementwiseOperation c_element_op_;
Jing Zhang's avatar
Jing Zhang committed
327
328
329
330

        StaticallyIndexedArray<GemmDesc, MaxGroupCount> gemm_desc_;

        index_t grid_size;
Jing Zhang's avatar
Jing Zhang committed
331
332
333
334
335
336
337
338
339
    };

    // Invoker
    struct Invoker : public BaseInvoker
    {
        using Argument = DeviceGroupedGemmXdl::Argument;

        float Run(const Argument& arg, int nrepeat = 1)
        {
Jing Zhang's avatar
Jing Zhang committed
340
341
342
343
            static_for<0, MaxGroupCount, 1>{}([&](auto i) {
                if(i < arg.gemm_shapes_.size())
                {
                    std::cout << "arg.a_grid_desc_k0_m_k1_{"
Jing Zhang's avatar
Jing Zhang committed
344
345
346
347
                              << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I0) << ", "
                              << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I1) << ", "
                              << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I2) << "}"
                              << std::endl;
Jing Zhang's avatar
Jing Zhang committed
348
349

                    std::cout << "arg.b_grid_desc_k0_n_k1_{"
Jing Zhang's avatar
Jing Zhang committed
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
                              << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I0) << ", "
                              << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I1) << ", "
                              << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I2) << "}"
                              << std::endl;

                    std::cout << "arg.c_grid_desc_m_n_{ "
                              << arg.gemm_desc_[i].c_grid_desc_m_n_.GetLength(I0) << ", "
                              << arg.gemm_desc_[i].c_grid_desc_m_n_.GetLength(I1) << "}"
                              << std::endl;

                    std::cout << "Block: " << arg.gemm_desc_[i].BlockStart << ", "
                              << arg.gemm_desc_[i].BlockEnd << std::endl;

                    if(!GridwiseGemm::CheckValidity(arg.gemm_desc_[i].a_grid_desc_k0_m_k1_,
                                                    arg.gemm_desc_[i].b_grid_desc_k0_n_k1_,
                                                    arg.gemm_desc_[i].c_grid_desc_m_n_,
Jing Zhang's avatar
Jing Zhang committed
366
367
368
369
370
371
                                                    arg.M01_,
                                                    arg.N01_))
                    {
                        throw std::runtime_error(
                            "wrong! GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3 has invalid setting");
                    }
Jing Zhang's avatar
Jing Zhang committed
372
                }
Jing Zhang's avatar
Jing Zhang committed
373
            });
Jing Zhang's avatar
Jing Zhang committed
374

Jing Zhang's avatar
Jing Zhang committed
375
            const auto K0 = arg.gemm_desc_[I0].a_grid_desc_k0_m_k1_.GetLength(I0);
Jing Zhang's avatar
Jing Zhang committed
376
377
378
379
380
381
382

            const bool has_main_k0_block_loop = GridwiseGemm::CalculateHasMainK0BlockLoop(K0);

            float ave_time = 0;

            if(has_main_k0_block_loop)
            {
Jing Zhang's avatar
Jing Zhang committed
383
384
385
386
387
388
389
390
391
392
                const auto kernel =
                    kernel_grouped_gemm_xdlops_v2r3<GridwiseGemm,
                                                    ADataType, // TODO: distiguish A/B datatype
                                                    CDataType,
                                                    remove_reference_t<GemmDesc>,
                                                    AElementwiseOperation,
                                                    BElementwiseOperation,
                                                    CElementwiseOperation,
                                                    true,
                                                    MaxGroupCount>;
Jing Zhang's avatar
Jing Zhang committed
393
394
395

                ave_time = launch_and_time_kernel(kernel,
                                                  nrepeat,
Jing Zhang's avatar
Jing Zhang committed
396
                                                  dim3(arg.grid_size),
Jing Zhang's avatar
Jing Zhang committed
397
398
399
400
401
                                                  dim3(BlockSize),
                                                  0,
                                                  arg.p_a_grid_,
                                                  arg.p_b_grid_,
                                                  arg.p_c_grid_,
Jing Zhang's avatar
Jing Zhang committed
402
                                                  arg.gemm_desc_,
Jing Zhang's avatar
Jing Zhang committed
403
                                                  arg.gemm_shapes_.size(),
Jing Zhang's avatar
Jing Zhang committed
404
405
                                                  arg.a_element_op_,
                                                  arg.b_element_op_,
Jing Zhang's avatar
Jing Zhang committed
406
                                                  arg.c_element_op_);
Jing Zhang's avatar
Jing Zhang committed
407
408
409
            }
            else
            {
Jing Zhang's avatar
Jing Zhang committed
410
411
412
413
414
415
416
417
418
419
                const auto kernel =
                    kernel_grouped_gemm_xdlops_v2r3<GridwiseGemm,
                                                    ADataType, // TODO: distiguish A/B datatype
                                                    CDataType,
                                                    remove_reference_t<GemmDesc>,
                                                    AElementwiseOperation,
                                                    BElementwiseOperation,
                                                    CElementwiseOperation,
                                                    false,
                                                    MaxGroupCount>;
Jing Zhang's avatar
Jing Zhang committed
420
421
422

                ave_time = launch_and_time_kernel(kernel,
                                                  nrepeat,
Jing Zhang's avatar
Jing Zhang committed
423
                                                  dim3(arg.grid_size),
Jing Zhang's avatar
Jing Zhang committed
424
425
426
427
428
                                                  dim3(BlockSize),
                                                  0,
                                                  arg.p_a_grid_,
                                                  arg.p_b_grid_,
                                                  arg.p_c_grid_,
Jing Zhang's avatar
Jing Zhang committed
429
                                                  arg.gemm_desc_,
Jing Zhang's avatar
Jing Zhang committed
430
                                                  arg.gemm_shapes_.size(),
Jing Zhang's avatar
Jing Zhang committed
431
432
                                                  arg.a_element_op_,
                                                  arg.b_element_op_,
Jing Zhang's avatar
Jing Zhang committed
433
                                                  arg.c_element_op_);
Jing Zhang's avatar
Jing Zhang committed
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
            }

            return ave_time;
        }

        // polymorphic
        float Run(const BaseArgument* p_arg, int nrepeat = 1) override
        {
            return Run(*dynamic_cast<const Argument*>(p_arg), nrepeat);
        }
    };

    static constexpr bool IsValidCompilationParameter()
    {
        // TODO: properly implement this check
        return true;
    }

    static bool IsSupportedArgument(const Argument& arg)
    {
Jing Zhang's avatar
Jing Zhang committed
454
455
456
        return GridwiseGemm::CheckValidity(arg.gemm_desc_[Number<0>{}].a_grid_desc_k0_m_k1_,
                                           arg.gemm_desc_[Number<0>{}].b_grid_desc_k0_n_k1_,
                                           arg.gemm_desc_[Number<0>{}].c_grid_desc_m_n_,
Jing Zhang's avatar
Jing Zhang committed
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
                                           arg.M01_,
                                           arg.N01_);
    }

    // polymorphic
    bool IsSupportedArgument(const BaseArgument* p_arg) override
    {
        return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
    }

    static auto MakeArgument(const ADataType* p_a,
                             const BDataType* p_b,
                             CDataType* p_c,
                             std::vector<gemm_desc> gemm_shapes,
                             AElementwiseOperation a_element_op,
                             BElementwiseOperation b_element_op,
                             CElementwiseOperation c_element_op)
    {
        return Argument{p_a, p_b, p_c, gemm_shapes, 1, 1, a_element_op, b_element_op, c_element_op};
    }

    static auto MakeInvoker() { return Invoker{}; }

    // polymorphic
    std::unique_ptr<BaseArgument> MakeArgumentPointer(const void* p_a,
                                                      const void* p_b,
                                                      void* p_c,
                                                      std::vector<gemm_desc> gemm_shapes,
                                                      AElementwiseOperation a_element_op,
                                                      BElementwiseOperation b_element_op,
                                                      CElementwiseOperation c_element_op,
                                                      index_t /* KBatch */ = 1) override
    {
        return std::make_unique<Argument>(static_cast<const ADataType*>(p_a),
                                          static_cast<const BDataType*>(p_b),
                                          static_cast<CDataType*>(p_c),
                                          gemm_shapes,
                                          1,
                                          1,
                                          a_element_op,
                                          b_element_op,
                                          c_element_op);
    }

    // polymorphic
    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
    {
        return std::make_unique<Invoker>(Invoker{});
    }

    // polymorphic
    std::string GetTypeString() const override
    {
        auto str = std::stringstream();

        // clang-format off
        str << "DeviceGroupedGemmXdl"
            << "<"
            << BlockSize << ", "
            << MPerBlock << ", "
            << NPerBlock << ", "
            << K0PerBlock << ", "
            << K1 << ", "
            << MPerXDL << ", "
            << NPerXDL << ", "
            << MXdlPerWave << ", "
            << NXdlPerWave
            << ">";
        // clang-format on

        return str.str();
    }
};

} // namespace device
} // namespace tensor_operation
} // namespace ck
#endif