/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/device/device_gemm_splitk.hpp Source File

• develop
• include
• ck
• tensor_operation
• gpu
• device

device_gemm_splitk.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

2 // Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.

3 

4 #pragma once

5 

6 #include <iostream>

7 #include <vector>

8 

9 #include "device_base.hpp"

10 

11 namespace ck {

12 namespace tensor_operation {

13 namespace device {

14 

15 template <typename ALayout,

16  typename BLayout,

17  typename CLayout,

18  typename ADataType,

19  typename BDataType,

20  typename CDataType,

21  typename AElementwiseOperation,

22  typename BElementwiseOperation,

23  typename CElementwiseOperation,

24  typename ComputeType = CDataType>

25 struct DeviceGemmSplitK : public BaseOperator

26 {

27  virtual std::unique_ptr<BaseArgument> MakeArgumentPointer(const void* p_a,

28  const void* p_b,

29  void* p_c,

30  ck::index_t M,

31  ck::index_t N,

32  ck::index_t K,

33  ck::index_t StrideA,

34  ck::index_t StrideB,

35  ck::index_t StrideC,

36  AElementwiseOperation a_element_op,

37  BElementwiseOperation b_element_op,

38  CElementwiseOperation c_element_op,

39  ck::index_t KBatch) = 0;

40 

41  virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;

42 };

43 

44 template <typename ALayout,

45  typename BLayout,

46  typename CLayout,

47  typename ADataType,

48  typename BDataType,

49  typename CDataType,

50  typename AElementwiseOperation,

51  typename BElementwiseOperation,

52  typename CElementwiseOperation,

53  typename ComputeType = CDataType>

54 using DeviceGemmSplitKPtr = std::unique_ptr<DeviceGemmSplitK<ALayout,

55  BLayout,

56  CLayout,

57  ADataType,

58  BDataType,

59  CDataType,

60  AElementwiseOperation,

61  BElementwiseOperation,

62  CElementwiseOperation,

63  ComputeType>>;

64 

65 } // namespace device

66 } // namespace tensor_operation

67 } // namespace ck

device_base.hpp

ck::tensor_operation::device::DeviceGemmSplitKPtr

std::unique_ptr< DeviceGemmSplitK< ALayout, BLayout, CLayout, ADataType, BDataType, CDataType, AElementwiseOperation, BElementwiseOperation, CElementwiseOperation, ComputeType > > DeviceGemmSplitKPtr

Definition: device_gemm_splitk.hpp:63

ck

Definition: ck.hpp:268

ck::index_t

int32_t index_t

Definition: ck.hpp:299

ck::tensor_operation::device::BaseOperator

Definition: device_base.hpp:223

ck::tensor_operation::device::DeviceGemmSplitK

Definition: device_gemm_splitk.hpp:26

ck::tensor_operation::device::DeviceGemmSplitK::MakeInvokerPointer

virtual std::unique_ptr< BaseInvoker > MakeInvokerPointer()=0

ck::tensor_operation::device::DeviceGemmSplitK::MakeArgumentPointer

virtual std::unique_ptr< BaseArgument > MakeArgumentPointer(const void *p_a, const void *p_b, void *p_c, ck::index_t M, ck::index_t N, ck::index_t K, ck::index_t StrideA, ck::index_t StrideB, ck::index_t StrideC, AElementwiseOperation a_element_op, BElementwiseOperation b_element_op, CElementwiseOperation c_element_op, ck::index_t KBatch)=0