device_grouped_gemm_splitk.hpp Source File

device_grouped_gemm_splitk.hpp Source File#

Composable Kernel: device_grouped_gemm_splitk.hpp Source File
device_grouped_gemm_splitk.hpp
Go to the documentation of this file.
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
3#pragma once
4
5#include "device_grouped_gemm.hpp"
6
7namespace ck {
8namespace tensor_operation {
9namespace device {
10
11template <typename ALayout,
12 typename BLayout,
13 typename DsLayout,
14 typename ELayout,
15 typename ADataType,
16 typename BDataType,
17 typename DsDataType,
18 typename EDataType,
19 typename AElementwiseOperation,
20 typename BElementwiseOperation,
21 typename CElementwiseOperation>
22struct DeviceGroupedGemmSplitK : public DeviceGroupedGemm<ALayout,
23 BLayout,
24 DsLayout,
25 ELayout,
26 ADataType,
27 BDataType,
28 DsDataType,
29 EDataType,
30 AElementwiseOperation,
31 BElementwiseOperation,
32 CElementwiseOperation>
33{
34 //----------------------------------------------------------------------------------------------
40 virtual void SetKBatchSize(BaseArgument* p_arg, index_t kbatch) const = 0;
41 //----------------------------------------------------------------------------------------------
47 virtual void SetKBatch(BaseArgument* p_arg, index_t kbatch) const
48 {
49 this->SetKBatchSize(p_arg, kbatch);
50 };
51};
52
53} // namespace device
54} // namespace tensor_operation
55} // namespace ck
device_grouped_gemm.hpp
ck::tensor_operation::device
Definition convolution_backward_data_specialization.hpp:8
ck::tensor_operation
Definition convolution_backward_data_specialization.hpp:7
ck
Definition ck.hpp:268
ck::index_t
int32_t index_t
Definition ck.hpp:299
ck::tensor_operation::device::BaseArgument
Definition device_base.hpp:197
ck::tensor_operation::device::DeviceGroupedGemm
Definition device_grouped_gemm.hpp:99
ck::tensor_operation::device::DeviceGroupedGemmSplitK
Definition device_grouped_gemm_splitk.hpp:33
ck::tensor_operation::device::DeviceGroupedGemmSplitK::SetKBatchSize
virtual void SetKBatchSize(BaseArgument *p_arg, index_t kbatch) const =0
Sets the k batch size.
ck::tensor_operation::device::DeviceGroupedGemmSplitK::SetKBatch
virtual void SetKBatch(BaseArgument *p_arg, index_t kbatch) const
Sets the k batch size.
Definition device_grouped_gemm_splitk.hpp:47