diff options
author | Max Ryabinin <mryabinin0@gmail.com> | 2022-07-01 17:16:10 +0300 |
---|---|---|
committer | Max Ryabinin <mryabinin0@gmail.com> | 2022-07-01 17:16:10 +0300 |
commit | 8258b4364a21a4da2572cb644d0926080c3268da (patch) | |
tree | 571e95bc327116fbaba08d14871fb0b224b8a65b /csrc/cpu_ops.cpp | |
parent | 33efe4a09f459832e8beceba70add0695cc485e4 (diff) |
Add a CPU-only build option
Diffstat (limited to 'csrc/cpu_ops.cpp')
-rw-r--r-- | csrc/cpu_ops.cpp | 57 |
1 files changed, 57 insertions, 0 deletions
diff --git a/csrc/cpu_ops.cpp b/csrc/cpu_ops.cpp new file mode 100644 index 0000000..11a2615 --- /dev/null +++ b/csrc/cpu_ops.cpp @@ -0,0 +1,57 @@ +#include <BinSearch.h> +#include <pthread.h> +#include <common.h> + +using namespace BinSearch; + +void dequantize_cpu(float *code, unsigned char *A, float *absmax, float *out, int n) { + for (int block_idx = 0; block_idx < n; block_idx += BLOCK_SIZE) { + int valid_items = n - block_idx >= BLOCK_SIZE ? BLOCK_SIZE : n - block_idx; + int block_end = block_idx + valid_items; + for (int i = block_idx; i < block_end; i++) + out[i] = code[A[i]] * absmax[block_idx / BLOCK_SIZE]; + } +} + +void quantize_cpu(float *code, float *A, float *absmax, unsigned char *out, int n) { + + // the default code is has range [-0.993, 1.0] which can cause an error in the binary search algorithm used below + code[0] = -1.0f; + + int num_blocks = n / BLOCK_SIZE; + num_blocks += n % BLOCK_SIZE == 0 ? 0 : 1; + + pthread_t *threads = (pthread_t *) malloc(sizeof(pthread_t) * num_blocks); + struct quantize_block_args **args = (quantize_block_args **) malloc(num_blocks * sizeof(quantize_block_args *)); + + for (int i = 0; i < num_blocks; i++) + args[i] = (quantize_block_args *) malloc(sizeof(quantize_block_args)); + + const uint32 elements_code = 256; + BinAlgo<Scalar, float, Direct2> bin_searcher(code, elements_code); + + for (int block_idx = 0; block_idx < n; block_idx += BLOCK_SIZE) { + int valid_items = n - block_idx >= BLOCK_SIZE ? BLOCK_SIZE : n - block_idx; + int block_end = block_idx + valid_items; + + struct quantize_block_args *arg = args[block_idx / BLOCK_SIZE]; + arg->bin_searcher = &bin_searcher; + arg->code = code; + arg->A = A; + arg->absmax = absmax; + arg->out = out; + arg->block_end = block_end; + arg->block_idx = block_idx; + arg->threadidx = block_idx / BLOCK_SIZE; + + pthread_create(&threads[block_idx / BLOCK_SIZE], NULL, &quantize_block, (void *) arg); + } + + for (int i = 0; i < num_blocks; i++) + int err = pthread_join(threads[i], NULL); + + free(threads); + for (int i = 0; i < num_blocks; i++) + free(args[i]); + free(args); +}
\ No newline at end of file |