quant_gemm.cpp 7.11 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
60
61
62
63
64
65
66
67
68
69
70
71
72
    // if(!inputs.at(1).transposed())
    // {
    //     if (pack_1.empty())
    //     {
    //         pack_1 = allocate_gpu(inputs.at(1));
    //     }
    // }
    // if(inputs.at(0).transposed())
    // {
    //     if (pack_0.empty())
    //     {
    //         pack_0 = allocate_gpu(inputs.at(0));
    //     }
    // }
73

Shucai Xiao's avatar
Shucai Xiao committed
74
75
76
77
78
    check_shapes{input_shapes}.not_broadcasted();
    return op.compute_shape(input_shapes);
}

argument miopen_quant_gemm::compute(context& ctx,
Shucai Xiao's avatar
Shucai Xiao committed
79
80
                                    const shape& output_shape,
                                    const std::vector<argument>& args) const
Shucai Xiao's avatar
Shucai Xiao committed
81
{
82
83
84
85
86
87
88
89
90
91
92
93
94
    // handling the packing of B MUST be before handling that for A
    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];
    rocblas_int ldc = args[2].get_shape().strides()[dim_0];

    if(!transb)
    {
        // use the algorithm to pack A
Shucai Xiao's avatar
Shucai Xiao committed
95
        if(pack_1.empty())
Shucai Xiao's avatar
Shucai Xiao committed
96
97
98
99
        {
            std::cout << "allocate pack_1" << std::endl;
            pack_1 = allocate_gpu(args.at(1).get_shape());
        }
Shucai Xiao's avatar
Shucai Xiao committed
100
        // assert(!pack_1.empty());
Shucai Xiao's avatar
Shucai Xiao committed
101
102
103
        device::pack_a(ctx.get_stream().get(), pack_1, args[1]);
        auto pb = from_gpu(pack_1);
        std::cout << "pb = " << pb << std::endl;
104
105
    }

Shucai Xiao's avatar
Shucai Xiao committed
106
    // need to pack A in this scenario, use the algorithm to pack B in the
107
108
109
    // comment of the API
    if(transa)
    {
Shucai Xiao's avatar
Shucai Xiao committed
110
        if(pack_0.empty())
Shucai Xiao's avatar
Shucai Xiao committed
111
112
113
114
115
        {
            std::cout << "allocate pack_0" << std::endl;
            pack_0 = allocate_gpu(args.at(0).get_shape());
        }
        device::pack_b(ctx.get_stream().get(), pack_0, args[0]);
Shucai Xiao's avatar
Shucai Xiao committed
116
        auto a  = from_gpu(args[0]);
Shucai Xiao's avatar
Shucai Xiao committed
117
118
119
        auto pa = from_gpu(pack_0);
        std::cout << "a = " << a << std::endl;
        std::cout << "pa = " << pa << std::endl;
120
121
    }

Shucai Xiao's avatar
Shucai Xiao committed
122
    bool is_3inputs = (args.size() == 4);
Shucai Xiao's avatar
Shucai Xiao committed
123
    int8_t beta     = 0;
Shucai Xiao's avatar
Shucai Xiao committed
124
125
126
127
128
129
130
131
    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
132
133
134
135
136
137
        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];
138
139
        auto to_pointer = [&](auto&& arg) { return to_rocblas_type(as.from(arg.data())); };
        assert(k % 4 == 0);
Shucai Xiao's avatar
Shucai Xiao committed
140

Shucai Xiao's avatar
Shucai Xiao committed
141
142
143
144
        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
145
            // the rocblas_gemm API handles inputs and output matrices as
Shucai Xiao's avatar
Shucai Xiao committed
146
147
148
            // 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.
149
            generic_rocblas_gemm_ex(ctx.get_stream().get_rocblas(),
150
                                    transb ? rocblas_operation_transpose : rocblas_operation_none,
151
                                    transa ? rocblas_operation_transpose : rocblas_operation_none,
152
                                    n,
153
                                    m,
Shucai Xiao's avatar
Shucai Xiao committed
154
155
                                    k,
                                    &alpha_r,
Shucai Xiao's avatar
Shucai Xiao committed
156
                                    (!transb) ? to_pointer(pack_1) : to_pointer(args[1]),
157
158
                                    rocblas_datatype_i8_r,
                                    ldb,
Shucai Xiao's avatar
Shucai Xiao committed
159
                                    transa ? to_pointer(pack_0) : to_pointer(args[0]),
160
161
                                    rocblas_datatype_i8_r,
                                    lda,
Shucai Xiao's avatar
Shucai Xiao committed
162
163
164
165
166
167
168
169
170
171
172
173
174
                                    &beta_r,
                                    to_pointer(args[2]),
                                    rocblas_datatype_i32_r,
                                    ldc,
                                    (is_3inputs ? to_pointer(args[3]) : to_pointer(args[2])),
                                    rocblas_datatype_i32_r,
                                    ldc,
                                    rocblas_datatype_i32_r,
                                    rocblas_gemm_algo_standard,
                                    0,
                                    0,
                                    nullptr,
                                    nullptr);
Shucai Xiao's avatar
Shucai Xiao committed
175
176
177
        }
        else
        {
178
            generic_rocblas_batched_gemm_ex(
Shucai Xiao's avatar
Shucai Xiao committed
179
180
181
182
183
184
185
                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
186
                (!transb) ? to_pointer(pack_1) : to_pointer(args[1]),
187
                rocblas_datatype_i8_r,
Shucai Xiao's avatar
Shucai Xiao committed
188
189
                ldb,
                k * n,
Shucai Xiao's avatar
Shucai Xiao committed
190
                transa ? to_pointer(pack_0) : to_pointer(args[0]),
191
                rocblas_datatype_i8_r,
Shucai Xiao's avatar
Shucai Xiao committed
192
193
194
                lda,
                m * k,
                &beta_r,
195
196
197
198
                to_pointer(args[2]),
                rocblas_datatype_i32_r,
                ldc,
                m * n,
Shucai Xiao's avatar
Shucai Xiao committed
199
                (is_3inputs ? to_pointer(args[3]) : to_pointer(args[2])),
200
                rocblas_datatype_i32_r,
Shucai Xiao's avatar
Shucai Xiao committed
201
202
                ldc,
                m * n,
203
204
205
                num_matrices,
                rocblas_datatype_i32_r,
                rocblas_gemm_algo_standard,
Shucai Xiao's avatar
Shucai Xiao committed
206
207
208
209
                0,
                0,
                nullptr,
                nullptr);
Shucai Xiao's avatar
Shucai Xiao committed
210
211
212
213
214
215
216
217
218
        }
    });

    return (is_3inputs ? args[3] : args[2]);
}

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