quant_gemm.cpp 6.55 KB
Newer Older
Shucai Xiao's avatar
Shucai Xiao committed
1
#include <migraphx/gpu/quant_gemm.hpp>
2
#include <migraphx/gpu/device/pack.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
3
#include <migraphx/gpu/context.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
4
#include <migraphx/generate.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
5
6
7
8
9
10

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {

template <class... Ts>
11
rocblas_status generic_rocblas_gemm_ex(Ts&&... xs)
Shucai Xiao's avatar
Shucai Xiao committed
12
{
13
    return rocblas_gemm_ex(std::forward<Ts>(xs)...);
Shucai Xiao's avatar
Shucai Xiao committed
14
15
16
}

template <class... Ts>
17
rocblas_status generic_rocblas_batched_gemm_ex(Ts&&... xs)
Shucai Xiao's avatar
Shucai Xiao committed
18
{
19
    return rocblas_gemm_strided_batched_ex(std::forward<Ts>(xs)...);
Shucai Xiao's avatar
Shucai Xiao committed
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
}

template <class T>
struct compute_rocblas_type
{
    using type = T;
};

template <class T>
struct compute_rocblas_type<const T>
{
    using type = const typename compute_rocblas_type<T>::type;
};

template <>
struct compute_rocblas_type<half>
{
    using type = rocblas_half;
};

template <class T>
using rb_type = typename compute_rocblas_type<T>::type;

template <class T>
rb_type<T> to_rocblas_type(T x)
{
    return reinterpret_cast<const rb_type<T>&>(x);
}

template <class T>
rb_type<T>* to_rocblas_type(T* x)
{
    return reinterpret_cast<rb_type<T>*>(x);
}

shape miopen_quant_gemm::compute_shape(const std::vector<shape>& inputs) const
{
57
58
    std::vector<shape> input_shapes(inputs);
    input_shapes.pop_back();
Shucai Xiao's avatar
Shucai Xiao committed
59
    check_shapes{input_shapes}.not_broadcasted();
Shucai Xiao's avatar
Shucai Xiao committed
60
61
    bool transa = inputs[0].transposed();
    bool transb = inputs[1].transposed();
Shucai Xiao's avatar
Shucai Xiao committed
62
63
64
65
66

    if(!transb)
    {
        if(arg_b.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
67
            auto* p_this  = const_cast<miopen_quant_gemm*>(this);
Shucai Xiao's avatar
Shucai Xiao committed
68
            p_this->arg_b = allocate_gpu(inputs[1]);
Shucai Xiao's avatar
Shucai Xiao committed
69
70
71
72
73
74
75
        }
    }

    if(transa)
    {
        if(arg_a.empty())
        {
Shucai Xiao's avatar
Shucai Xiao committed
76
            auto* p_this  = const_cast<miopen_quant_gemm*>(this);
Shucai Xiao's avatar
Shucai Xiao committed
77
            p_this->arg_a = allocate_gpu(inputs[0]);
Shucai Xiao's avatar
Shucai Xiao committed
78
79
80
        }
    }

Shucai Xiao's avatar
Shucai Xiao committed
81
82
83
84
    return op.compute_shape(input_shapes);
}

argument miopen_quant_gemm::compute(context& ctx,
Shucai Xiao's avatar
Shucai Xiao committed
85
86
                                    const shape& output_shape,
                                    const std::vector<argument>& args) const
Shucai Xiao's avatar
Shucai Xiao committed
87
{
88
89
90
91
92
93
94
    bool transa     = args[0].get_shape().transposed();
    bool transb     = args[1].get_shape().transposed();
    auto n_dim      = output_shape.lens().size();
    auto dim_1      = n_dim - 1;
    auto dim_0      = n_dim - 2;
    rocblas_int lda = args[0].get_shape().strides()[transa ? dim_1 : dim_0];
    rocblas_int ldb = args[1].get_shape().strides()[transb ? dim_1 : dim_0];
Shucai Xiao's avatar
Shucai Xiao committed
95
    rocblas_int ldc = args[2].get_shape().strides()[dim_0];
96
97
98

    if(!transb)
    {
Shucai Xiao's avatar
Shucai Xiao committed
99
        device::pack_a(ctx.get_stream().get(), arg_b, args[1]);
100
101
    }

Shucai Xiao's avatar
Shucai Xiao committed
102
    // need to pack A in this scenario, use the algorithm to pack B in the
103
104
105
    // comment of the API
    if(transa)
    {
Shucai Xiao's avatar
Shucai Xiao committed
106
        device::pack_b(ctx.get_stream().get(), arg_a, args[0]);
107
108
    }

Shucai Xiao's avatar
Shucai Xiao committed
109
    bool is_3inputs = (args.size() == 4);
Shucai Xiao's avatar
Shucai Xiao committed
110
    int32_t beta    = 0;
Shucai Xiao's avatar
Shucai Xiao committed
111
112
113
114
115
116
117
118
    if(is_3inputs)
    {
        beta = op.beta;
    }

    auto a_lens = args[0].get_shape().lens();
    auto b_lens = args[1].get_shape().lens();
    output_shape.visit_type([&](auto as) {
Shucai Xiao's avatar
Shucai Xiao committed
119
120
121
122
123
124
        auto alpha_r    = to_rocblas_type(as(op.alpha));
        auto beta_r     = to_rocblas_type(as(beta));
        auto out_lens   = output_shape.lens();
        rocblas_int m   = out_lens[dim_0];
        rocblas_int n   = out_lens[dim_1];
        rocblas_int k   = args[0].get_shape().lens()[dim_1];
125
126
        auto to_pointer = [&](auto&& arg) { return to_rocblas_type(as.from(arg.data())); };
        assert(k % 4 == 0);
Shucai Xiao's avatar
Shucai Xiao committed
127

Shucai Xiao's avatar
Shucai Xiao committed
128
129
130
131
        auto num_matrices = std::accumulate(
            out_lens.rbegin() + 2, out_lens.rend(), std::size_t{1}, std::multiplies<std::size_t>());
        if(num_matrices == 1)
        {
Shucai Xiao's avatar
Shucai Xiao committed
132
            // the rocblas_gemm API handles inputs and output matrices as
Shucai Xiao's avatar
Shucai Xiao committed
133
134
135
            // column-major format. When doing a C = A * B, we actually do
            // C^T = (B^T) * (A^T). That is the reason we input args[1] as
            // A and args[0] as B in calling the rocblas_gemm.
136
            generic_rocblas_gemm_ex(ctx.get_stream().get_rocblas(),
137
                                    transb ? rocblas_operation_transpose : rocblas_operation_none,
138
                                    transa ? rocblas_operation_transpose : rocblas_operation_none,
139
                                    n,
140
                                    m,
Shucai Xiao's avatar
Shucai Xiao committed
141
142
                                    k,
                                    &alpha_r,
Shucai Xiao's avatar
Shucai Xiao committed
143
                                    (!transb) ? to_pointer(arg_b) : to_pointer(args.at(1)),
144
145
                                    rocblas_datatype_i8_r,
                                    ldb,
Shucai Xiao's avatar
Shucai Xiao committed
146
                                    transa ? to_pointer(arg_a) : to_pointer(args.at(0)),
147
148
                                    rocblas_datatype_i8_r,
                                    lda,
Shucai Xiao's avatar
Shucai Xiao committed
149
150
151
152
                                    &beta_r,
                                    to_pointer(args[2]),
                                    rocblas_datatype_i32_r,
                                    ldc,
Shucai Xiao's avatar
Shucai Xiao committed
153
                                    is_3inputs ? to_pointer(args.at(3)) : to_pointer(args[2]),
Shucai Xiao's avatar
Shucai Xiao committed
154
155
156
157
158
159
160
161
                                    rocblas_datatype_i32_r,
                                    ldc,
                                    rocblas_datatype_i32_r,
                                    rocblas_gemm_algo_standard,
                                    0,
                                    0,
                                    nullptr,
                                    nullptr);
Shucai Xiao's avatar
Shucai Xiao committed
162
163
164
        }
        else
        {
165
            generic_rocblas_batched_gemm_ex(
Shucai Xiao's avatar
Shucai Xiao committed
166
167
168
169
170
171
172
                ctx.get_stream().get_rocblas(),
                transb ? rocblas_operation_transpose : rocblas_operation_none,
                transa ? rocblas_operation_transpose : rocblas_operation_none,
                n,
                m,
                k,
                &alpha_r,
Shucai Xiao's avatar
Shucai Xiao committed
173
                (!transb) ? to_pointer(arg_b) : to_pointer(args.at(1)),
174
                rocblas_datatype_i8_r,
Shucai Xiao's avatar
Shucai Xiao committed
175
176
                ldb,
                k * n,
Shucai Xiao's avatar
Shucai Xiao committed
177
                transa ? to_pointer(arg_a) : to_pointer(args.at(0)),
178
                rocblas_datatype_i8_r,
Shucai Xiao's avatar
Shucai Xiao committed
179
180
181
                lda,
                m * k,
                &beta_r,
182
183
184
185
                to_pointer(args[2]),
                rocblas_datatype_i32_r,
                ldc,
                m * n,
Shucai Xiao's avatar
Shucai Xiao committed
186
                is_3inputs ? to_pointer(args.at(3)) : to_pointer(args[2]),
187
                rocblas_datatype_i32_r,
Shucai Xiao's avatar
Shucai Xiao committed
188
189
                ldc,
                m * n,
190
191
192
                num_matrices,
                rocblas_datatype_i32_r,
                rocblas_gemm_algo_standard,
Shucai Xiao's avatar
Shucai Xiao committed
193
194
195
196
                0,
                0,
                nullptr,
                nullptr);
Shucai Xiao's avatar
Shucai Xiao committed
197
198
199
        }
    });

Shucai Xiao's avatar
Shucai Xiao committed
200
    return is_3inputs ? args.at(3) : args[2];
Shucai Xiao's avatar
Shucai Xiao committed
201
202
203
204
205
}

} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx