DandinPower
125 lines
5.2 KiB
C++
125 lines
5.2 KiB
C++
#include "sparsity.h"
|
|
|
|
sparse_array_t *allocate_sparse_array(uint16_t num_tokens, uint16_t num_features, float sparse_ratio) {
|
|
if (!num_tokens || !num_features) return NULL;
|
|
if (sparse_ratio < 0.0f || sparse_ratio > 1.0f) return NULL;
|
|
|
|
float raw_sparse = (float)num_features * sparse_ratio;
|
|
uint16_t num_sparse_features = (uint16_t)roundf(raw_sparse);
|
|
|
|
// clamp to valid range
|
|
if (num_sparse_features > num_features) {
|
|
num_sparse_features = num_features;
|
|
} else if (num_sparse_features == 0 && sparse_ratio > 0.0f) {
|
|
num_sparse_features = 1; // Avoid total sparsity if ratio positive;
|
|
}
|
|
|
|
uint32_t sparse_elements = (uint32_t)num_tokens * num_sparse_features;
|
|
uint64_t total = sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t));
|
|
sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, total);
|
|
if (!sparse_array) return NULL;
|
|
|
|
/* initialise the header fields */
|
|
sparse_array->num_tokens = num_tokens;
|
|
sparse_array->num_features = num_features;
|
|
sparse_array->num_sparse_features = num_sparse_features;
|
|
sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1); /* just after the header */
|
|
sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements); /* after the sparse_indices */
|
|
|
|
return sparse_array;
|
|
}
|
|
|
|
void free_sparse_array(sparse_array_t *sparse_array) {
|
|
if (!sparse_array) return;
|
|
free(sparse_array);
|
|
}
|
|
|
|
uint64_t get_sparse_array_size(const sparse_array_t *sparse_array) {
|
|
if (!sparse_array) return 0;
|
|
|
|
uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features;
|
|
|
|
return sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t));
|
|
}
|
|
|
|
sparse_array_t *load_sparse_array_from_buffer(const void *buffer, uint64_t buffer_size) {
|
|
sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, buffer_size);
|
|
if (!sparse_array) return NULL;
|
|
|
|
memcpy(sparse_array, buffer, buffer_size);
|
|
|
|
uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features;
|
|
|
|
sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1);
|
|
sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements);
|
|
|
|
return sparse_array;
|
|
}
|
|
|
|
typedef struct {
|
|
uint16_t index;
|
|
float abs_val;
|
|
} sort_entry_t;
|
|
|
|
static int abs_sort_cmp(const void *a, const void *b) {
|
|
const float abs_a = ((const sort_entry_t *)a)->abs_val;
|
|
const float abs_b = ((const sort_entry_t *)b)->abs_val;
|
|
if (abs_a != abs_b) {
|
|
return (abs_a > abs_b) ? -1 : 1;
|
|
}
|
|
const uint16_t idx_a = ((const sort_entry_t *)a)->index;
|
|
const uint16_t idx_b = ((const sort_entry_t *)b)->index;
|
|
return (int)idx_a - (int)idx_b;
|
|
}
|
|
|
|
int compress(const float *float_array, uint16_t num_tokens, uint16_t num_features, float sparse_ratio, sparse_array_t **sparse_array) {
|
|
if (!float_array || num_tokens == 0 || num_features == 0 || *sparse_array) return 1;
|
|
|
|
/* ---- allocate sparse ------------------------------------------ */
|
|
*sparse_array = allocate_sparse_array(num_tokens, num_features, sparse_ratio);
|
|
if (!*sparse_array) return 1;
|
|
|
|
#pragma omp parallel for
|
|
for (uint16_t cur_token_index = 0; cur_token_index < num_tokens; cur_token_index++) {
|
|
sort_entry_t *entries = (sort_entry_t *)malloc(num_features * sizeof(sort_entry_t));
|
|
|
|
uint32_t dense_base = (uint32_t)cur_token_index * num_features;
|
|
uint32_t sparse_base = (uint32_t)cur_token_index * (*sparse_array)->num_sparse_features;
|
|
|
|
for (uint16_t i = 0; i < num_features; i++) {
|
|
entries[i].index = i;
|
|
entries[i].abs_val = fabsf(float_array[dense_base + i]);
|
|
}
|
|
qsort(entries, num_features, sizeof(sort_entry_t), abs_sort_cmp);
|
|
|
|
for (uint16_t keep_feature_index = 0; keep_feature_index < (*sparse_array)->num_sparse_features; keep_feature_index++) {
|
|
uint16_t orig_index = entries[keep_feature_index].index;
|
|
(*sparse_array)->sparse_indices[sparse_base + keep_feature_index] = orig_index;
|
|
(*sparse_array)->values[sparse_base + keep_feature_index] = float_array[dense_base + orig_index];
|
|
}
|
|
|
|
free(entries);
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int decompress(const sparse_array_t *sparse_array, float *float_array) {
|
|
if (!float_array || !sparse_array) return 1;
|
|
|
|
uint32_t num_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_features;
|
|
memset(float_array, 0, num_elements * sizeof(float));
|
|
|
|
for (uint16_t cur_token_index = 0; cur_token_index < sparse_array->num_tokens; cur_token_index++) {
|
|
uint32_t dense_base = (uint32_t)cur_token_index * sparse_array->num_features;
|
|
uint32_t sparse_base = (uint32_t)cur_token_index * sparse_array->num_sparse_features;
|
|
|
|
for (uint16_t keep_feature_index = 0; keep_feature_index < sparse_array->num_sparse_features; keep_feature_index++) {
|
|
uint16_t original_feature_index = sparse_array->sparse_indices[sparse_base + keep_feature_index];
|
|
float_array[dense_base + original_feature_index] = sparse_array->values[sparse_base + keep_feature_index];
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|