Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
amd_wmma.hpp 4.18 KB
Newer Older
aska-0096's avatar
wmma_op + unit test  
aska-0096 committed
1
2
3
4
5
6
7 
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#ifndef CK_AMD_WMMA_HPP
#define CK_AMD_WMMA_HPP

#include "data_type.hpp"
aska-0096's avatar
change arch limitation  
aska-0096 committed
8 
// TODO: Add arch limitation
aska-0096's avatar
wmma_op + unit test  
aska-0096 committed
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
56
57
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 
namespace ck {

// wave32 only
// src: fp16, dst: fp32
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_f32_16x16x16_f16_w32;

template <>
struct intrin_wmma_f32_16x16x16_f16_w32<16, 16>
{
    template <class FloatC>
    __device__ static void Run(const half16_t& reg_a, const half16_t& reg_b, FloatC& reg_c)
    {
        reg_c.template AsType<float8_t>()(Number<0>{}) = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(
            reg_a, reg_b, reg_c.template AsType<float8_t>()[Number<0>{}]);
    }
};

// src: bf16, dst: fp32
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_f32_16x16x16_bf16_w32;

template <>
struct intrin_wmma_f32_16x16x16_bf16_w32<16, 16>
{
    template <class FloatC>
    __device__ static void Run(const bhalf16_t& reg_a, const bhalf16_t& reg_b, FloatC& reg_c)
    {
        reg_c.template AsType<float8_t>()(Number<0>{}) =
            __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(
                reg_a, reg_b, reg_c.template AsType<float8_t>()[Number<0>{}]);
    }
};

// src: fp16, dst: fp16
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_f16_16x16x16_f16_w32;

template <>
struct intrin_wmma_f16_16x16x16_f16_w32<16, 16>
{
    template <class FloatC>
    __device__ static void
    Run(const half16_t& reg_a, const half16_t& reg_b, FloatC& reg_c, const bool opsel)
    {
        // opsel usage
        // false: D0.[0:15] = result
        // true : D0.[16:31]= result
        reg_c.template AsType<half16_t>()(Number<0>{}) = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(
            reg_a, reg_b, reg_c.template AsType<half16_t>()[Number<0>{}], opsel);
    }
};

// src: bf16, dst: bf32
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_bf16_16x16x16_bf16_w32;

template <>
struct intrin_wmma_bf16_16x16x16_bf16_w32<16, 16>
{
    template <class FloatC>
    __device__ static void
    Run(const bhalf16_t& reg_a, const bhalf16_t& reg_b, FloatC& reg_c, const bool opsel)
    {
        // opsel usage
        // false: D0.[0:15] = result
        // true : D0.[16:31]= result
        reg_c.template AsType<bhalf16_t>()(Number<0>{}) =
            __builtin_amdgcn_wmma_bf16_16x16x16_bf16_w32(
                reg_a, reg_b, reg_c.template AsType<bhalf16_t>()[Number<0>{}], opsel);
    }
};

// src: iu8, dst: i32
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_i32_16x16x16_iu8_w32;

template <>
struct intrin_wmma_i32_16x16x16_iu8_w32<16, 16>
{
    template <class FloatC>
    __device__ static void Run(const bool neg_a,
                               const int8x16_t& reg_a,
                               const bool neg_b,
                               const int8x16_t& reg_b,
                               FloatC& reg_c,
                               const bool clamp)
    {
        reg_c.template AsType<int32x8_t>()(Number<0>{}) =
            __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(
                neg_a,
                bit_cast<int32x4_t>(reg_a),
                neg_b,
                bit_cast<int32x4_t>(reg_b),
                reg_c.template AsType<int32x8_t>()[Number<0>{}],
                clamp);
    }
};

#ifdef CK_EXPERIMENTAL_BIT_INT_EXTENSION_INT4
// src: iu4, dst: i32
template <index_t MPerWave, index_t NPerWave>
struct intrin_wmma_i32_16x16x16_iu4_w32;

template <>
struct intrin_wmma_i32_16x16x16_iu4_w32<16, 16>
{
    template <class FloatC>
    __device__ static void Run(const bool neg_a,
                               const int4x16_t& reg_a,
                               const bool neg_b,
                               const int4x16_t& reg_b,
                               FloatC& reg_c,
                               const bool clamp)
    {
        reg_c.template AsType<int32x8_t>()(Number<0>{}) =
            __builtin_amdgcn_wmma_i32_16x16x16_iu4_w32(
                neg_a,
                bit_cast<int32x4_t>(reg_a),
                neg_b,
                bit_cast<int32x4_t>(reg_b),
                reg_c.template AsType<int32x8_t>()[Number<0>{}],
                clamp);
    }
};
#endif
} // namespace ck
#endif