Skip to content

Commit 6080986

Browse files
metascroymetascroy
authored
Add 5bit packing routines
Differential Revision: D62133660 Pull Request resolved: #798
1 parent d980858 commit 6080986

5 files changed

Lines changed: 713 additions & 20 deletions

File tree

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

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,128 @@ void unpack_uint4_values(
346346
}
347347
}
348348

349+
// Benchmark utility to compare variants of uint5 packing
350+
void pack_uint5_values(
351+
uint8_t* packed,
352+
uint8_t* unpacked,
353+
int packed_size,
354+
int unpacked_size,
355+
int variant) {
356+
constexpr int nbit = 5;
357+
constexpr int bitsPerByte = 8;
358+
assert(unpacked_size * nbit / bitsPerByte == packed_size);
359+
assert(packed_size % variant == 0);
360+
361+
uint8x16_t unpacked0;
362+
uint8x16_t unpacked1;
363+
uint8x16_t unpacked2;
364+
uint8x16_t unpacked3;
365+
uint8x16_t unpacked4;
366+
uint8x16_t unpacked5;
367+
uint8x16_t unpacked6;
368+
uint8x16_t unpacked7;
369+
370+
switch (variant) {
371+
case 8:
372+
for (int i = 0; i < unpacked_size; i += 8) {
373+
torchao::bitpacking::internal::pack_8_uint5_values(
374+
packed + ((i * nbit) / bitsPerByte), unpacked + i);
375+
}
376+
break;
377+
case 64:
378+
for (int i = 0; i < unpacked_size; i += 64) {
379+
torchao::bitpacking::internal::vec_load_64_uint8_values(
380+
unpacked0, unpacked1, unpacked2, unpacked3, unpacked + i);
381+
torchao::bitpacking::internal::vec_pack_64_uint5_values(
382+
packed + ((i * nbit) / bitsPerByte),
383+
unpacked0,
384+
unpacked1,
385+
unpacked2,
386+
unpacked3);
387+
}
388+
break;
389+
case 128:
390+
for (int i = 0; i < unpacked_size; i += 128) {
391+
torchao::bitpacking::internal::vec_load_64_uint8_values(
392+
unpacked0, unpacked1, unpacked2, unpacked3, unpacked + i);
393+
torchao::bitpacking::internal::vec_load_64_uint8_values(
394+
unpacked4, unpacked5, unpacked6, unpacked7, unpacked + i + 64);
395+
torchao::bitpacking::internal::vec_pack_128_uint5_values(
396+
packed + ((i * nbit) / bitsPerByte),
397+
unpacked0,
398+
unpacked1,
399+
unpacked2,
400+
unpacked3,
401+
unpacked4,
402+
unpacked5,
403+
unpacked6,
404+
unpacked7);
405+
}
406+
break;
407+
}
408+
}
409+
410+
// Benchmark utility to compare variants of uint5 unpacking
411+
void unpack_uint5_values(
412+
uint8_t* unpacked,
413+
uint8_t* packed,
414+
int unpacked_size,
415+
int packed_size,
416+
int variant) {
417+
constexpr int nbit = 5;
418+
constexpr int bitsPerByte = 8;
419+
assert(unpacked_size * nbit / bitsPerByte == packed_size);
420+
assert(packed_size % variant == 0);
421+
422+
uint8x16_t unpacked0;
423+
uint8x16_t unpacked1;
424+
uint8x16_t unpacked2;
425+
uint8x16_t unpacked3;
426+
uint8x16_t unpacked4;
427+
uint8x16_t unpacked5;
428+
uint8x16_t unpacked6;
429+
uint8x16_t unpacked7;
430+
431+
switch (variant) {
432+
case 8:
433+
for (int i = 0; i < unpacked_size; i += 8) {
434+
torchao::bitpacking::internal::unpack_8_uint5_values(
435+
unpacked + i, packed + ((i * nbit) / bitsPerByte));
436+
}
437+
break;
438+
case 64:
439+
for (int i = 0; i < unpacked_size; i += 64) {
440+
torchao::bitpacking::internal::vec_unpack_64_uint5_values(
441+
unpacked0,
442+
unpacked1,
443+
unpacked2,
444+
unpacked3,
445+
packed + ((i * nbit) / bitsPerByte));
446+
torchao::bitpacking::internal::vec_store_64_uint8_values(
447+
unpacked + i, unpacked0, unpacked1, unpacked2, unpacked3);
448+
}
449+
break;
450+
case 128:
451+
for (int i = 0; i < unpacked_size; i += 128) {
452+
torchao::bitpacking::internal::vec_unpack_128_uint5_values(
453+
unpacked0,
454+
unpacked1,
455+
unpacked2,
456+
unpacked3,
457+
unpacked4,
458+
unpacked5,
459+
unpacked6,
460+
unpacked7,
461+
packed + ((i * nbit) / bitsPerByte));
462+
torchao::bitpacking::internal::vec_store_64_uint8_values(
463+
unpacked + i, unpacked0, unpacked1, unpacked2, unpacked3);
464+
torchao::bitpacking::internal::vec_store_64_uint8_values(
465+
unpacked + i + 64, unpacked4, unpacked5, unpacked6, unpacked7);
466+
}
467+
break;
468+
}
469+
}
470+
349471
} // namespace
350472

351473
static void benchmark_pack_uint2_values(benchmark::State& state) {
@@ -462,12 +584,52 @@ static void benchmark_unpack_uint4_values(benchmark::State& state) {
462584
}
463585
}
464586

587+
static void benchmark_pack_uint5_values(benchmark::State& state) {
588+
int unpacked_size = state.range(0);
589+
int variant = state.range(1);
590+
int nbit = 5;
591+
592+
assert(unpacked_size % 8 == 0);
593+
int packed_size = (unpacked_size / 8) * nbit;
594+
595+
auto packed = std::vector<uint8_t>(unpacked_size, 0);
596+
auto unpacked = torchao::get_random_lowbit_vector(packed_size, 8);
597+
598+
for (auto _ : state) {
599+
pack_uint5_values(
600+
packed.data(), unpacked.data(), packed_size, unpacked_size, variant);
601+
}
602+
}
603+
604+
static void benchmark_unpack_uint5_values(benchmark::State& state) {
605+
int unpacked_size = state.range(0);
606+
int variant = state.range(1);
607+
int nbit = 5;
608+
609+
assert(unpacked_size % 8 == 0);
610+
int packed_size = (unpacked_size / 8) * nbit;
611+
612+
auto packed = torchao::get_random_lowbit_vector(packed_size, 8);
613+
auto unpacked = std::vector<uint8_t>(unpacked_size, 0);
614+
615+
for (auto _ : state) {
616+
unpack_uint5_values(
617+
unpacked.data(),
618+
packed.data(),
619+
unpacked.size(),
620+
packed.size(),
621+
variant);
622+
}
623+
}
624+
465625
BENCHMARK(benchmark_pack_uint2_values)->ArgsProduct({{128}, {4, 32, 64}});
466626
BENCHMARK(benchmark_unpack_uint2_values)->ArgsProduct({{128}, {4, 32, 64}});
467627
BENCHMARK(benchmark_pack_uint3_values)->ArgsProduct({{128}, {8, 64, 128}});
468628
BENCHMARK(benchmark_unpack_uint3_values)->ArgsProduct({{128}, {8, 64, 128}});
469629
BENCHMARK(benchmark_pack_uint4_values)->ArgsProduct({{128}, {2, 16, 32}});
470630
BENCHMARK(benchmark_unpack_uint4_values)->ArgsProduct({{128}, {2, 16, 32}});
631+
BENCHMARK(benchmark_pack_uint5_values)->ArgsProduct({{128}, {8, 64, 128}});
632+
BENCHMARK(benchmark_unpack_uint5_values)->ArgsProduct({{128}, {8, 64, 128}});
471633

472634
// Run the benchmark
473635
BENCHMARK_MAIN();

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,18 +234,24 @@ BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT
234234
3);
235235
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT(
236236
4);
237+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x1x32_F32_NEONDOT(
238+
5);
237239
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
238240
2);
239241
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
240242
3);
241243
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
242244
4);
245+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x4x16_F32_NEONDOT(
246+
5);
243247
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
244248
2);
245249
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
246250
3);
247251
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
248252
4);
253+
BENCHMARK_CHANNELWISE_8BIT_ACTIVATION_GROUPWISE_LOWBIT_WEIGHT_1x8x16_F32_NEONDOT(
254+
5);
249255

250256
// Run the benchmark
251257
BENCHMARK_MAIN();

0 commit comments

Comments
 (0)