Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
device_gemm.hpp 2.9 KB
Newer Older
Chao Liu's avatar
ckProfiler and device-level XDL GEMM operator (#48)  
Chao Liu committed
1
2
3
4
5
6
7
8
9
10 
#ifndef DEVICE_GEMM_HPP
#define DEVICE_GEMM_HPP

#include <iostream>
#include "device_base.hpp"

namespace ck {
namespace tensor_operation {
namespace device {
rocking5566's avatar
Support alpha beta scaling for GEMM (#78)  
rocking5566 committed
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 
template <typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation>
struct DeviceGemmBias : public BaseOperator
{
    virtual std::unique_ptr<BaseArgument>
    MakeArgumentPointer(const void* p_a,
                        const void* p_b,
                        const void* p_bias,
                        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) = 0;

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

template <typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation>
using DeviceGemmBiasPtr = std::unique_ptr<
    DeviceGemmBias<AElementwiseOperation, BElementwiseOperation, CElementwiseOperation>>;
Chao Liu's avatar
GEMM/Conv+BiasAdd+ReLU+Add (#55)  
Chao Liu committed
40
41
42 
template <typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation>
Chao Liu's avatar
ckProfiler and device-level XDL GEMM operator (#48)  
Chao Liu committed
43
44 
struct DeviceGemm : public BaseOperator
{
ltqin's avatar
add split-k GEMM (#59)  
ltqin committed
45
46
47
48
49
50
51
52
53
54
55
56
57 
    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 = 1) = 0;
Chao Liu's avatar
ckProfiler and device-level XDL GEMM operator (#48)  
Chao Liu committed
58
59
60
61 

    virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
};
Chao Liu's avatar
GEMM/Conv+BiasAdd+ReLU+Add (#55)  
Chao Liu committed
62
63
64
65
66 
template <typename AElementwiseOperation,
          typename BElementwiseOperation,
          typename CElementwiseOperation>
using DeviceGemmPtr = std::unique_ptr<
    DeviceGemm<AElementwiseOperation, BElementwiseOperation, CElementwiseOperation>>;
Chao Liu's avatar
ckProfiler and device-level XDL GEMM operator (#48)  
Chao Liu committed
67
68
69
70
71 

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