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
/*******************************************************************************
* Copyright 2023 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/x64/jit_avx512_sparse_decompress_kernel.hpp"
#define GET_OFF(field) offsetof(call_params_t, field)
namespace dnnl {
namespace impl {
namespace cpu {
namespace x64 {
using namespace Xbyak;
Xbyak::Opmask jit_avx512_sparse_decompress_kernel_t::get_opmask(int idx) {
switch (idx) {
case 0: return k1; break;
case 1: return k2; break;
case 2: return k3; break;
case 3: return k4; break;
default: assert(!"incorrect index"); return k0;
}
}
Xbyak::Opmask jit_avx512_sparse_decompress_kernel_t::get_load_mask(int idx) {
// This function prepares a load mask for loading packed values. The minimum
// and maximum number of bits that can be set is 0 and 64 respectively.
// Since `shl` instruction doesn't work when the `count` operand is 64
// (the instruction actually uses `count % 64` instead of just `count`)
// we have to split 1 shift into 2 shifts (3 if there is a tail).
//
//
// The following pseudo-code is implemented in JIT below:
// shift = reg_popcnt / 2;
// res = 1;
// res = res << shift;
// res = res << shift;
// shift_tail = reg_popcnt % 2;
// res = res << shift_tail;
// Save original number of bits set in 1.
mov(reg_popcnt_tmp, reg_popcnt);
mov(reg_tmp, 1);
// shift = reg_popcnt / 2
shr(reg_popcnt, 1);
// Apply shift two times.
shl(reg_tmp, reg_popcnt.cvt8());
shl(reg_tmp, reg_popcnt.cvt8());
// Calculate shift_tail = reg_popcnt % 2.
mov(reg_popcnt, reg_popcnt_tmp);
and_(reg_popcnt, 1);
// Apply shift_tail.
shl(reg_tmp, reg_popcnt.cvt8());
sub(reg_tmp, 1);
// Restore the value (used to advance the pointer to packed values).
mov(reg_popcnt, reg_popcnt_tmp);
auto opmask = get_opmask(idx);
kmovq(opmask, reg_tmp);
return opmask;
}
Xbyak::Opmask jit_avx512_sparse_decompress_kernel_t::get_expand_mask(int idx) {
return get_opmask(idx);
}
Xbyak::Reg64 jit_avx512_sparse_decompress_kernel_t::get_reg_mask_tmp(int idx) {
switch (idx) {
case 0: return r13;
case 1: return r14;
case 2: return r15;
case 3: return rax;
default: assert(!"incorrect index"); return Xbyak::Reg64(0);
}
}
Xbyak::Zmm jit_avx512_sparse_decompress_kernel_t::get_zmm(int idx) {
switch (idx) {
case 0: return Xbyak::Zmm(25);
case 1: return Xbyak::Zmm(26);
case 2: return Xbyak::Zmm(27);
case 3: return Xbyak::Zmm(28);
default: assert(!"incorrect index"); return Xbyak::Zmm(0);
}
}
void jit_avx512_sparse_decompress_kernel_t::generate() {
preamble();
mov(reg_bitmask_ptr, ptr[param1 + GET_OFF(bitmask_ptr)]);
mov(reg_dst_ptr, ptr[param1 + GET_OFF(dst_ptr)]);
mov(reg_src_ptr, ptr[param1 + GET_OFF(src_ptr)]);
assert(unroll_factor() == 4);
for (int b = 0; b < nblks_to_decompress_; b++) {
const int blk_offset = b * blk_sz_;
const int bitmask_off = blk_offset / CHAR_BIT;
const int nbytes_per_load = 64;
for (int i = 0; i < b_blk_sz_; i += unroll_factor()) {
for (int uf = 0; uf < unroll_factor(); uf++) {
auto reg_mask_tmp = get_reg_mask_tmp(uf);
mov(reg_mask_tmp,
ptr[reg_bitmask_ptr + (i + uf) * sizeof(uint64_t)
+ bitmask_off]);
popcnt(reg_popcnt, reg_mask_tmp);
auto load_mask = get_load_mask(uf);
auto zmm_reg = get_zmm(uf);
vmovdqu8(zmm_reg | load_mask | T_z, ptr[reg_src_ptr]);
add(reg_src_ptr, reg_popcnt);
auto expand_mask = get_expand_mask(uf);
kmovq(expand_mask, reg_mask_tmp);
vpexpandb(zmm_reg | expand_mask | T_z, zmm_reg);
vmovdqu8(ptr[reg_dst_ptr + blk_offset
+ (i + uf) * nbytes_per_load],
zmm_reg);
}
}
}
postamble();
}
#undef GET_OFF
} // namespace x64
} // namespace cpu
} // namespace impl
} // namespace dnnl