onednn-src 0.1.13

Source of oneAPI Deep Neural Network Library (oneDNN)
Documentation 
1
2
3
4
5
6
7
8
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211

/*******************************************************************************
* Copyright 2019 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/
#include "cpu/rv64/rvv_matmul.hpp"
#include "common/c_types_map.hpp"
#include "common/dnnl_thread.hpp"
#include "common/utils.hpp"
#include "cpu/rv64/gemm/rvv_gemm_f32.hpp"
#include "cpu/rv64/rvv_postops.hpp"
#include <riscv_vector.h>

namespace dnnl {
namespace impl {
namespace cpu {
namespace rv64 {
namespace matmul {

status_t rvv_matmul_t::execute(const exec_ctx_t &ctx) const {
    auto src = CTX_IN_MEM(const float *, DNNL_ARG_SRC);
    auto weights = CTX_IN_MEM(const float *, DNNL_ARG_WEIGHTS);
    auto dst = CTX_OUT_MEM(float *, DNNL_ARG_DST);

    const memory_desc_wrapper src_d(pd()->src_md());
    const memory_desc_wrapper weights_d(pd()->weights_md());
    const memory_desc_wrapper dst_d(pd()->dst_md());
    const memory_desc_wrapper bias_d(pd()->desc()->bias_desc);

    const post_ops_t &post_ops = pd()->attr()->post_ops_;
    rvv_postops_t postops_handler(post_ops);

    const float *bias = CTX_IN_MEM(const float *, DNNL_ARG_BIAS);
    const int ndims = src_d.ndims();
    const int wei_ndims = weights_d.ndims();
    const dim_t *src_dims = src_d.dims();
    const dim_t *wei_dims = weights_d.dims();

    const dim_t batch = pd()->batch_;
    const dim_t M = pd()->M_;
    const dim_t K = pd()->K_;
    const dim_t N = pd()->N_;
    const bool weights_col_major = pd()->weights_col_major_;

    // row-major [M, K] / [M, N] are reinterpreted as column-major matrices and
    // mapped to GEMM as:
    //   C(N x M) = A(N x K) * B(K x M)
    // where:
    //   - A = W^T, with W logical layout [K, N] row-major. Since row-major
    //     [K, N] is equivalent in memory to column-major [N, K] (leading dim N),
    //     we pass transa = 'N' and lda = N.
    //   - B = src, where src logical layout [M, K] row-major is equivalent to
    //     column-major [K, M] with leading dim K, so transb = 'N', ldb = K.
    //   - C is viewed as column-major [N, M] with leading dim N, which matches
    //     row-major [M, N] in memory.
    char transa = weights_col_major ? 'T' : 'N';
    char transb = 'N';
    dim_t M_gemm = N;
    dim_t N_gemm = M;
    dim_t K_gemm = K;
    // weights: col-major [K, N] uses leading K; row-major [K, N] is equivalent
    // to col-major [N, K] so leading dim is N.
    dim_t lda = weights_col_major ? K : N;
    dim_t ldb = K; // src: row-major [M, K] -> col-major [K, M]
    dim_t ldc = N; // dst: row-major [M, N] -> col-major [N, M]
    float alpha = 1.0f;
    float beta = 0.0f;

    // batch strides for row-major src/dst
    const dim_t src_batch_stride = M * K;
    const dim_t dst_batch_stride = M * N;

    const int src_batch_ndims = ndims > 2 ? ndims - 2 : 0;
    const int wei_batch_ndims = wei_ndims > 2 ? wei_ndims - 2 : 0;
    const int batch_dim_shift = src_batch_ndims - wei_batch_ndims;
    const dim_t K_dim = wei_dims[wei_ndims - 2];
    const dim_t N_dim = wei_dims[wei_ndims - 1];
    const dim_t wei_matrix_stride = K_dim * N_dim;

    // GEMM compute
    if (pd()->weights_are_broadcast_) {
        //   C(N x (batch * M)) = A(N x K) * B(K x (batch * M))
        dim_t M_gemm_all = M_gemm; // N
        dim_t N_gemm_all = batch * N_gemm; // batch * M

        status_t st = rvv_gemm_f32(&transa, &transb, &M_gemm_all, &N_gemm_all,
                &K_gemm, &alpha, weights, &lda, src, &ldb, &beta, dst, &ldc,
                /*bias=*/nullptr);
        assert(st == status::success || st == status::unimplemented);
        MAYBE_UNUSED(st);

    } else {
        // one GEMM per logical batch
        parallel_nd(batch, [&](dim_t b) {
            const float *src_base = src + b * src_batch_stride;
            float *dst_base = dst + b * dst_batch_stride;

            dim_t batch_indices[DNNL_MAX_NDIMS] = {};
            if (src_batch_ndims > 0) {
                utils::l_dims_by_l_offset(
                        batch_indices, b, src_dims, src_batch_ndims);
            }

            dim_t weight_batch_index = 0;
            if (wei_batch_ndims > 0) {
                for (int d = 0; d < wei_batch_ndims; ++d) {
                    const int src_dim_idx = d + batch_dim_shift;
                    dim_t idx = (src_dim_idx >= 0) ? batch_indices[src_dim_idx]
                                                   : dim_t(0);
                    const dim_t wei_dim = wei_dims[d];
                    idx = (wei_dim == 1) ? dim_t(0) : idx;
                    weight_batch_index = weight_batch_index * wei_dim + idx;
                }
            }

            const float *wei_base
                    = weights + weight_batch_index * wei_matrix_stride;

            status_t st = rvv_gemm_f32(&transa, &transb, &M_gemm, &N_gemm,
                    &K_gemm, &alpha, wei_base, &lda, src_base, &ldb, &beta,
                    dst_base, &ldc,
                    /*bias=*/nullptr);
            assert(st == status::success || st == status::unimplemented);
            MAYBE_UNUSED(st);
        });
    }

    if (!bias && post_ops.len() == 0) return status::success;

    const int dst_ndims = dst_d.ndims();
    const int bias_ndims = bias_d.ndims();
    const dim_t *bias_dims = bias_d.dims();

    parallel_nd(batch, [&](dim_t b) {
        float *dst_base = dst + b * dst_batch_stride;

        dim_t dst_idx_prefix[DNNL_MAX_NDIMS] = {};
        size_t bias_strides[DNNL_MAX_NDIMS] = {};

        if (bias && bias_ndims > 1) {
            bias_strides[bias_ndims - 1] = 1;
            for (int d = bias_ndims - 2; d >= 0; --d)
                bias_strides[d]
                        = bias_strides[d + 1] * (size_t)bias_dims[d + 1];
        }

        for (dim_t m = 0; m < M; ++m) {
            if (ndims > 2) {
                utils::l_dims_by_l_offset(
                        dst_idx_prefix, b, src_dims, ndims - 2);
            }
            dst_idx_prefix[ndims - 2] = m;

            float *row_dst = dst_base + m * N;

            for (dim_t n0 = 0; n0 < N;) {
                size_t vl = __riscv_vsetvl_e32m1(N - n0);
                vfloat32m1_t acc = __riscv_vle32_v_f32m1(row_dst + n0, vl);

                if (bias) {
                    if (bias_d.nelems() == 1) {
                        acc = __riscv_vfadd_vf_f32m1(acc, bias[0], vl);
                    } else {
                        size_t base_bias_off = 0;
                        if (bias_ndims > 1) {
                            for (int d = 0; d < bias_ndims - 1; ++d) {
                                int dst_dim_idx = d + (dst_ndims - bias_ndims);
                                dim_t idx = (bias_dims[d] == 1)
                                        ? 0
                                        : dst_idx_prefix[dst_dim_idx];
                                base_bias_off += idx * bias_strides[d];
                            }
                        }

                        if (bias_dims[bias_ndims - 1] == 1) {
                            acc = __riscv_vfadd_vf_f32m1(
                                    acc, bias[base_bias_off], vl);
                        } else {
                            const float *bias_ptr = bias + base_bias_off + n0;
                            vfloat32m1_t bias_vec
                                    = __riscv_vle32_v_f32m1(bias_ptr, vl);
                            acc = __riscv_vfadd_vv_f32m1(acc, bias_vec, vl);
                        }
                    }
                }

                acc = postops_handler.apply(acc, vl);
                __riscv_vse32_v_f32m1(row_dst + n0, acc, vl);
                n0 += vl;
            }
        }
    });

    return status::success;
}

} // namespace matmul
} // namespace rv64
} // namespace cpu
} // namespace impl
} // namespace dnnl