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

/*******************************************************************************
* Copyright 2026 ZTE 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 "common/nstl.hpp"
#include "common/type_helpers.hpp"
#include "common/utils.hpp"

#include "cpu/rv64/brgemm/brgemm.hpp"
#include "cpu/rv64/brgemm/jit_brgemm_kernel.hpp"

#include "xbyak_riscv/xbyak_riscv_util.hpp"

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

using namespace dnnl::impl::utils;

status_t brgemm_desc_init(brgemm_desc_t *brg, cpu_isa_t isa,
        brgemm_batch_kind_t type, data_type_t dt_a, data_type_t dt_b,
        brgemm_layout_t layout, float alpha, float beta, dim_t LDA, dim_t LDB,
        dim_t LDC, dim_t M, dim_t N, dim_t K, const brgemm_strides_t *strides) {

    if (!brg) return status::invalid_arguments;
    if (M <= 0 || K <= 0) return status::invalid_arguments;

    // Only f32 → f32 supported in the MVP.
    if (!everyone_is(data_type::f32, dt_a, dt_b)) return status::unimplemented;

    *brg = utils::zero<brgemm_desc_t>();

    brg->bcast_dim = M;
    brg->load_dim = N;
    brg->reduce_dim = K;
    brg->LDA = LDA;
    brg->LDB = LDB;
    brg->LDC = LDC;
    brg->alpha = alpha;
    brg->beta = beta;
    brg->type = type;
    brg->layout = layout;
    brg->isa_impl = isa;

    brg->dt_a = dt_a;
    brg->dt_b = dt_b;
    brg->dt_c = data_type::f32;
    brg->typesize_A = static_cast<int>(types::data_type_size(dt_a));
    brg->typesize_B = static_cast<int>(types::data_type_size(dt_b));
    brg->typesize_C = static_cast<int>(types::data_type_size(brg->dt_c));
    brg->is_f32 = true;

    if (strides) {
        brg->stride_a = strides->stride_a;
        brg->stride_b = strides->stride_b;
    }

    // Determine bd_block from VLEN. Using LMUL=m4 so that each logical
    // vector register holds VLEN*4/32 f32 elements.
    const uint32_t vlen_bits
            = Xbyak_riscv::CPU::getInstance().getVlen(); // e.g. 128
    const int vlen_f32 = static_cast<int>(vlen_bits) / 32; // elements per m1
    brg->bd_block = vlen_f32 * 4; // LMUL=m4 → 4× elements

    brg->bdb = static_cast<int>(M) / brg->bd_block;
    brg->bdb_tail = static_cast<int>(M) % brg->bd_block;

    brg->n_step = 4; // process 4 output columns per inner iteration
    brg->rd_block = 4; // K unroll factor
    brg->rdb = static_cast<int>(K) / brg->rd_block;
    brg->rdb_tail = static_cast<int>(K) % brg->rd_block;

    return status::success;
}

status_t brgemm_kernel_create(
        brgemm_kernel_t **brg_kernel, const brgemm_desc_t &brg) {
    if (!brg_kernel) return status::invalid_arguments;
    *brg_kernel = nullptr;

    auto *kernel = new brgemm_kernel_common_t(brg);
    status_t st = kernel->create_kernel();
    if (st != status::success) {
        delete kernel;
        return st;
    }
    *brg_kernel = kernel;
    return status::success;
}

void brgemm_kernel_destroy(brgemm_kernel_t *brg_kernel) {
    delete brg_kernel;
}

void brgemm_kernel_execute(const brgemm_kernel_t *brg_kernel, const void *ptr_A,
        const void *ptr_B, void *ptr_C, dim_t N, float beta) {

    const auto &brg = brg_kernel->get_brg();
    const int ts = brg.typesize_C; // sizeof(float) = 4
    const int bd = brg.bd_block;
    const dim_t K = brg.reduce_dim;
    const dim_t LDA_bytes = brg.LDA * ts;

    const auto *A_base = reinterpret_cast<const char *>(ptr_A);
    const auto *B_base = reinterpret_cast<const char *>(ptr_B);
    auto *C_base = reinterpret_cast<char *>(ptr_C);

    // K-blocking: split the reduction dimension into chunks of BK to keep
    // the A working-set (bd_block × BK × 4 bytes) inside the L1D cache.
    const dim_t BK = BRGEMM_BK;

    for (dim_t kb = 0; kb < K; kb += BK) {
        const dim_t K_inner = nstl::min(BK, K - kb);
        const float beta_kb = (kb == 0) ? beta : 1.0f;

        const char *A_kb = A_base + kb * LDA_bytes;
        const char *B_kb = B_base + kb * ts;

        brgemm_kernel_params_t p;
        p.ptr_B = B_kb;
        p.N = N;
        p.K = K_inner;
        p.beta = beta_kb;

        // Process full bd_block tiles.
        for (int m = 0; m < brg.bdb; m++) {
            p.ptr_A = A_kb + static_cast<dim_t>(m) * bd * ts;
            p.ptr_C = C_base + static_cast<dim_t>(m) * bd * ts;
            p.M = bd;
            (*brg_kernel)(&p);
        }

        // Process M tail (if any).
        if (brg.bdb_tail > 0) {
            p.ptr_A = A_kb + static_cast<dim_t>(brg.bdb) * bd * ts;
            p.ptr_C = C_base + static_cast<dim_t>(brg.bdb) * bd * ts;
            p.M = brg.bdb_tail;
            (*brg_kernel)(&p);
        }
    }
}

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