Skip to content

Commit d980858

Browse files
metascroymetascroy
authored
Add 2bit packing routines
Differential Revision: D62133659 Pull Request resolved: #797
1 parent 8c18489 commit d980858

5 files changed

Lines changed: 461 additions & 13 deletions

File tree

‎torchao/experimental/kernels/cpu/aarch64/benchmarks/benchmark_bitpacking.cpp‎

Lines changed: 169 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,144 @@
66

77
#include <arm_neon.h>
88
#include <benchmark/benchmark.h>
9-
#include <iostream>
109

1110
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/bitpack.h>
11+
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint2.h>
1212
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint3.h>
1313
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint4.h>
1414
#include <torchao/experimental/kernels/cpu/aarch64/tests/test_utils.h>
1515
#include <cassert>
1616

1717
namespace {
1818

19+
// Benchmark utility to compare variants of uint2 packing
20+
void pack_uint2_values(
21+
uint8_t* packed,
22+
uint8_t* unpacked,
23+
int packed_size,
24+
int unpacked_size,
25+
int variant) {
26+
constexpr int nbit = 2;
27+
constexpr int bitsPerByte = 8;
28+
assert(unpacked_size * nbit / bitsPerByte == packed_size);
29+
assert(packed_size % variant == 0);
30+
31+
uint8x8_t unpacked0_8x8;
32+
uint8x8_t unpacked1_8x8;
33+
uint8x8_t unpacked2_8x8;
34+
uint8x8_t unpacked3_8x8;
35+
36+
uint8x16_t unpacked0_8x16;
37+
uint8x16_t unpacked1_8x16;
38+
uint8x16_t unpacked2_8x16;
39+
uint8x16_t unpacked3_8x16;
40+
41+
switch (variant) {
42+
case 4:
43+
for (int i = 0; i < unpacked_size; i += 4) {
44+
torchao::bitpacking::internal::pack_4_uint2_values(
45+
packed + ((i * nbit) / bitsPerByte), unpacked + i);
46+
}
47+
break;
48+
case 32:
49+
for (int i = 0; i < unpacked_size; i += 32) {
50+
torchao::bitpacking::internal::vec_load_32_uint8_values(
51+
unpacked0_8x8,
52+
unpacked1_8x8,
53+
unpacked2_8x8,
54+
unpacked3_8x8,
55+
unpacked + i);
56+
torchao::bitpacking::internal::vec_pack_32_uint2_values(
57+
packed + ((i * nbit) / bitsPerByte),
58+
unpacked0_8x8,
59+
unpacked1_8x8,
60+
unpacked2_8x8,
61+
unpacked3_8x8);
62+
}
63+
break;
64+
case 64:
65+
for (int i = 0; i < unpacked_size; i += 64) {
66+
torchao::bitpacking::internal::vec_load_64_uint8_values(
67+
unpacked0_8x16,
68+
unpacked1_8x16,
69+
unpacked2_8x16,
70+
unpacked3_8x16,
71+
unpacked + i);
72+
torchao::bitpacking::internal::vec_pack_64_uint2_values(
73+
packed + ((i * nbit) / bitsPerByte),
74+
unpacked0_8x16,
75+
unpacked1_8x16,
76+
unpacked2_8x16,
77+
unpacked3_8x16);
78+
}
79+
break;
80+
}
81+
}
82+
83+
// Benchmark utility to compare variants of uint2 packing
84+
void unpack_uint2_values(
85+
uint8_t* unpacked,
86+
uint8_t* packed,
87+
int unpacked_size,
88+
int packed_size,
89+
int variant) {
90+
constexpr int nbit = 2;
91+
constexpr int bitsPerByte = 8;
92+
assert(unpacked_size * nbit / bitsPerByte == packed_size);
93+
assert(packed_size % variant == 0);
94+
95+
uint8x8_t unpacked0_8x8;
96+
uint8x8_t unpacked1_8x8;
97+
uint8x8_t unpacked2_8x8;
98+
uint8x8_t unpacked3_8x8;
99+
100+
uint8x16_t unpacked0_8x16;
101+
uint8x16_t unpacked1_8x16;
102+
uint8x16_t unpacked2_8x16;
103+
uint8x16_t unpacked3_8x16;
104+
105+
switch (variant) {
106+
case 4:
107+
for (int i = 0; i < unpacked_size; i += 4) {
108+
torchao::bitpacking::internal::unpack_4_uint2_values(
109+
unpacked + i, packed + ((i * nbit) / bitsPerByte));
110+
}
111+
break;
112+
case 32:
113+
for (int i = 0; i < unpacked_size; i += 32) {
114+
torchao::bitpacking::internal::vec_unpack_32_uint2_values(
115+
unpacked0_8x8,
116+
unpacked1_8x8,
117+
unpacked2_8x8,
118+
unpacked3_8x8,
119+
packed + ((i * nbit) / bitsPerByte));
120+
torchao::bitpacking::internal::vec_store_32_uint8_values(
121+
unpacked + i,
122+
unpacked0_8x8,
123+
unpacked1_8x8,
124+
unpacked2_8x8,
125+
unpacked3_8x8);
126+
}
127+
break;
128+
case 64:
129+
for (int i = 0; i < unpacked_size; i += 64) {
130+
torchao::bitpacking::internal::vec_unpack_64_uint2_values(
131+
unpacked0_8x16,
132+
unpacked1_8x16,
133+
unpacked2_8x16,
134+
unpacked3_8x16,
135+
packed + ((i * nbit) / bitsPerByte));
136+
torchao::bitpacking::internal::vec_store_64_uint8_values(
137+
unpacked + i,
138+
unpacked0_8x16,
139+
unpacked1_8x16,
140+
unpacked2_8x16,
141+
unpacked3_8x16);
142+
}
143+
break;
144+
}
145+
}
146+
19147
// Benchmark utility to compare variants of uint3 packing
20148
void pack_uint3_values(
21149
uint8_t* packed,
@@ -220,6 +348,44 @@ void unpack_uint4_values(
220348

221349
} // namespace
222350

351+
static void benchmark_pack_uint2_values(benchmark::State& state) {
352+
int unpacked_size = state.range(0);
353+
int variant = state.range(1);
354+
int nbit = 2;
355+
356+
assert(unpacked_size % 8 == 0);
357+
int packed_size = (unpacked_size / 8) * nbit;
358+
359+
auto packed = std::vector<uint8_t>(unpacked_size, 0);
360+
auto unpacked = torchao::get_random_lowbit_vector(packed_size, 8);
361+
362+
for (auto _ : state) {
363+
pack_uint2_values(
364+
packed.data(), unpacked.data(), packed_size, unpacked_size, variant);
365+
}
366+
}
367+
368+
static void benchmark_unpack_uint2_values(benchmark::State& state) {
369+
int unpacked_size = state.range(0);
370+
int variant = state.range(1);
371+
int nbit = 2;
372+
373+
assert(unpacked_size % 8 == 0);
374+
int packed_size = (unpacked_size / 8) * nbit;
375+
376+
auto packed = torchao::get_random_lowbit_vector(packed_size, 8);
377+
auto unpacked = std::vector<uint8_t>(unpacked_size, 0);
378+
379+
for (auto _ : state) {
380+
unpack_uint2_values(
381+
unpacked.data(),
382+
packed.data(),
383+
unpacked.size(),
384+
packed.size(),
385+
variant);
386+
}
387+
}
388+
223389
static void benchmark_pack_uint3_values(benchmark::State& state) {
224390
int unpacked_size = state.range(0);
225391
int variant = state.range(1);
@@ -296,6 +462,8 @@ static void benchmark_unpack_uint4_values(benchmark::State& state) {
296462
}
297463
}
298464

465+
BENCHMARK(benchmark_pack_uint2_values)->ArgsProduct({{128}, {4, 32, 64}});
466+
BENCHMARK(benchmark_unpack_uint2_values)->ArgsProduct({{128}, {4, 32, 64}});
299467
BENCHMARK(benchmark_pack_uint3_values)->ArgsProduct({{128}, {8, 64, 128}});
300468
BENCHMARK(benchmark_unpack_uint3_values)->ArgsProduct({{128}, {8, 64, 128}});
301469
BENCHMARK(benchmark_pack_uint4_values)->ArgsProduct({{128}, {2, 16, 32}});

‎torchao/experimental/kernels/cpu/aarch64/benchmarks/benchmark_linear.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,14 +228,20 @@ channelwise_8bit_activation_groupwise_lowbit_weight_1x8x16_f32_neondot(
228228
false>) \
229229
->ArgsProduct(BENCHMARK_PARAMS)
230230

231+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT(
232+
2);
231233
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT(
232234
3);
233235
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT(
234236
4);
237+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
238+
2);
235239
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
236240
3);
237241
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
238242
4);
243+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
244+
2);
239245
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
240246
3);
241247
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(

‎torchao/experimental/kernels/cpu/aarch64/bitpacking/bitpack.h‎

Lines changed: 67 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#pragma once
88
#include <arm_neon.h>
99
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/macro.h>
10+
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint2.h>
1011
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint3.h>
1112
#include <torchao/experimental/kernels/cpu/aarch64/bitpacking/uint4.h>
1213
#include <cassert>
@@ -15,6 +16,30 @@ namespace torchao {
1516
namespace bitpacking {
1617

1718
namespace internal {
19+
TORCHAO_ALWAYS_INLINE inline void vec_store_32_uint8_values(
20+
uint8_t* dest,
21+
const uint8x8_t& vec0,
22+
const uint8x8_t& vec1,
23+
const uint8x8_t& vec2,
24+
const uint8x8_t& vec3) {
25+
vst1_u8(dest, vec0);
26+
vst1_u8(dest + 8, vec1);
27+
vst1_u8(dest + 16, vec2);
28+
vst1_u8(dest + 24, vec3);
29+
}
30+
31+
TORCHAO_ALWAYS_INLINE inline void vec_load_32_uint8_values(
32+
uint8x8_t& vec0,
33+
uint8x8_t& vec1,
34+
uint8x8_t& vec2,
35+
uint8x8_t& vec3,
36+
const uint8_t* src) {
37+
vec0 = vld1_u8(src);
38+
vec1 = vld1_u8(src + 8);
39+
vec2 = vld1_u8(src + 16);
40+
vec3 = vld1_u8(src + 24);
41+
}
42+
1843
TORCHAO_ALWAYS_INLINE inline void vec_store_64_uint8_values(
1944
uint8_t* dest,
2045
const uint8x16_t& vec0,
@@ -49,7 +74,7 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_32_lowbit_values(
4974
static_assert(nbit >= 2);
5075

5176
// Currently supported values
52-
static_assert(nbit >= 3);
77+
static_assert(nbit >= 2);
5378
static_assert(nbit <= 4);
5479

5580
// Shift unpacked values to nonnegative range
@@ -58,6 +83,13 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_32_lowbit_values(
5883
uint8x16_t shifted1 = vreinterpretq_u8_s8(vaddq_s8(unpacked1, shift));
5984

6085
switch (nbit) {
86+
case 2:
87+
torchao::bitpacking::internal::vec_pack_32_uint2_values(
88+
packed,
89+
vget_low_u8(shifted0),
90+
vget_high_u8(shifted0),
91+
vget_low_u8(shifted1),
92+
vget_high_u8(shifted1));
6193
case 3:
6294
uint8_t buffer[32];
6395
vst1q_u8(buffer, shifted0);
@@ -89,13 +121,22 @@ TORCHAO_ALWAYS_INLINE inline void vec_unpack_32_lowbit_values(
89121
static_assert(nbit >= 2);
90122

91123
// Currently supported values
92-
static_assert(nbit >= 3);
124+
static_assert(nbit >= 2);
93125
static_assert(nbit <= 4);
94126

95127
uint8x16_t shifted0;
96128
uint8x16_t shifted1;
97129

98130
switch (nbit) {
131+
case 2:
132+
uint8x8_t shifted0_low;
133+
uint8x8_t shifted0_high;
134+
uint8x8_t shifted1_low;
135+
uint8x8_t shifted1_high;
136+
torchao::bitpacking::internal::vec_unpack_32_uint2_values(
137+
shifted0_low, shifted0_high, shifted1_low, shifted1_high, packed);
138+
shifted0 = vcombine_u8(shifted0_low, shifted0_high);
139+
shifted1 = vcombine_u8(shifted1_low, shifted1_high);
99140
case 3:
100141
uint8_t buffer[32];
101142
torchao::bitpacking::internal::unpack_8_uint3_values(buffer, packed);
@@ -133,7 +174,7 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_64_lowbit_values(
133174
static_assert(nbit >= 2);
134175

135176
// Currently supported values
136-
static_assert(nbit >= 3);
177+
static_assert(nbit >= 2);
137178
static_assert(nbit <= 4);
138179

139180
// Shift unpacked values to nonnegative range
@@ -144,6 +185,10 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_64_lowbit_values(
144185
uint8x16_t shifted3 = vreinterpretq_u8_s8(vaddq_s8(unpacked3, shift));
145186

146187
switch (nbit) {
188+
case 2:
189+
torchao::bitpacking::internal::vec_pack_64_uint2_values(
190+
packed, shifted0, shifted1, shifted2, shifted3);
191+
break;
147192
case 3:
148193
torchao::bitpacking::internal::vec_pack_64_uint3_values(
149194
packed, shifted0, shifted1, shifted2, shifted3);
@@ -170,7 +215,7 @@ TORCHAO_ALWAYS_INLINE inline void vec_unpack_64_lowbit_values(
170215
static_assert(nbit >= 2);
171216

172217
// Currently supported values
173-
static_assert(nbit >= 3);
218+
static_assert(nbit >= 2);
174219
static_assert(nbit <= 4);
175220

176221
uint8x16_t shifted0;
@@ -179,6 +224,10 @@ TORCHAO_ALWAYS_INLINE inline void vec_unpack_64_lowbit_values(
179224
uint8x16_t shifted3;
180225

181226
switch (nbit) {
227+
case 2:
228+
torchao::bitpacking::internal::vec_unpack_64_uint2_values(
229+
shifted0, shifted1, shifted2, shifted3, packed);
230+
break;
182231
case 3:
183232
torchao::bitpacking::internal::vec_unpack_64_uint3_values(
184233
shifted0, shifted1, shifted2, shifted3, packed);
@@ -216,7 +265,7 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_128_lowbit_values(
216265
static_assert(nbit >= 2);
217266

218267
// Currently supported values
219-
static_assert(nbit >= 3);
268+
static_assert(nbit >= 2);
220269
static_assert(nbit <= 4);
221270

222271
// Shift unpacked values to nonnegative range
@@ -231,6 +280,12 @@ TORCHAO_ALWAYS_INLINE inline void vec_pack_128_lowbit_values(
231280
uint8x16_t shifted7 = vreinterpretq_u8_s8(vaddq_s8(unpacked7, shift));
232281

233282
switch (nbit) {
283+
case 2:
284+
torchao::bitpacking::internal::vec_pack_64_uint2_values(
285+
packed, shifted0, shifted1, shifted2, shifted3);
286+
torchao::bitpacking::internal::vec_pack_64_uint2_values(
287+
packed + 16, shifted4, shifted5, shifted6, shifted7);
288+
break;
234289
case 3:
235290
torchao::bitpacking::internal::vec_pack_128_uint3_values(
236291
packed,
@@ -273,7 +328,7 @@ TORCHAO_ALWAYS_INLINE inline void vec_unpack_128_lowbit_values(
273328
static_assert(nbit >= 2);
274329

275330
// Currently supported values
276-
static_assert(nbit >= 3);
331+
static_assert(nbit >= 2);
277332
static_assert(nbit <= 4);
278333

279334
uint8x16_t shifted0;
@@ -286,6 +341,12 @@ TORCHAO_ALWAYS_INLINE inline void vec_unpack_128_lowbit_values(
286341
uint8x16_t shifted7;
287342

288343
switch (nbit) {
344+
case 2:
345+
torchao::bitpacking::internal::vec_unpack_64_uint2_values(
346+
shifted0, shifted1, shifted2, shifted3, packed);
347+
torchao::bitpacking::internal::vec_unpack_64_uint2_values(
348+
shifted4, shifted5, shifted6, shifted7, packed + 16);
349+
break;
289350
case 3:
290351
torchao::bitpacking::internal::vec_unpack_128_uint3_values(
291352
shifted0,

0 commit comments

Comments
 (0)