David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/swarm/protocol.c

500 lines
15 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#include "protocol.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <math.h>
#ifdef __linux__
#include <sys/sysinfo.h>
#endif
#ifdef __APPLE__
#include <sys/sysctl.h>
#endif
/* ============================================================= */
/* Message Construction */
/* ============================================================= */
uint8_t* sw_msg_alloc(sw_msg_type type, uint32_t payload_size) {
if (payload_size > SW_MAX_MESSAGE_SIZE - SW_HEADER_SIZE) {
return NULL;
}
uint32_t total = SW_HEADER_SIZE + payload_size;
uint8_t* msg = calloc(1, total);
if (!msg) return NULL;
sw_msg_header* hdr = (sw_msg_header*)msg;
hdr->magic = SW_PROTOCOL_MAGIC;
hdr->type = (uint8_t)type;
hdr->length = payload_size;
return msg;
}
void sw_msg_free(uint8_t* msg) {
free(msg);
}
sw_msg_header* sw_msg_get_header(uint8_t* msg) {
if (!msg) return NULL;
return (sw_msg_header*)msg;
}
void* sw_msg_get_payload(uint8_t* msg) {
if (!msg) return NULL;
return msg + SW_HEADER_SIZE;
}
uint32_t sw_msg_total_size(uint8_t* msg) {
if (!msg) return 0;
sw_msg_header* hdr = sw_msg_get_header(msg);
return SW_HEADER_SIZE + hdr->length;
}
bool sw_msg_validate(uint8_t* msg, size_t received_size) {
if (!msg || received_size < SW_HEADER_SIZE) {
return false;
}
sw_msg_header* hdr = sw_msg_get_header(msg);
/* Check magic */
if (hdr->magic != SW_PROTOCOL_MAGIC) {
return false;
}
/* Check size */
uint32_t expected_size = SW_HEADER_SIZE + hdr->length;
if (received_size < expected_size) {
return false;
}
/* Check payload size limit */
if (hdr->length > SW_MAX_MESSAGE_SIZE - SW_HEADER_SIZE) {
return false;
}
return true;
}
/* ============================================================= */
/* Serialization Helpers */
/* ============================================================= */
/* Simple FP16 conversion (no denormals, limited range) */
static uint16_t float_to_fp16(float f) {
uint32_t x = *(uint32_t*)&f;
uint32_t sign = (x >> 16) & 0x8000;
int32_t exp = ((x >> 23) & 0xFF) - 127 + 15;
uint32_t mant = (x >> 13) & 0x3FF;
if (exp <= 0) {
return (uint16_t)sign; /* Zero or denormal → zero */
} else if (exp >= 31) {
return (uint16_t)(sign | 0x7C00); /* Inf */
}
return (uint16_t)(sign | (exp << 10) | mant);
}
static float fp16_to_float(uint16_t h) {
uint32_t sign = (h & 0x8000) << 16;
int32_t exp = (h >> 10) & 0x1F;
uint32_t mant = h & 0x3FF;
if (exp == 0) {
return 0.0f;
} else if (exp == 31) {
uint32_t inf = sign | 0x7F800000;
return *(float*)&inf;
}
exp = exp - 15 + 127;
uint32_t result = sign | (exp << 23) | (mant << 13);
return *(float*)&result;
}
/* Comparison for qsort - sort by absolute value descending */
typedef struct {
uint32_t index;
float abs_value;
} indexed_value;
static int compare_indexed(const void* a, const void* b) {
float va = ((indexed_value*)a)->abs_value;
float vb = ((indexed_value*)b)->abs_value;
if (va > vb) return -1;
if (va < vb) return 1;
return 0;
}
size_t sw_serialize_tensor(uint8_t* buf, float* data, uint64_t numel,
sw_compress_type compression, float topk_ratio) {
if (!buf || !data || numel == 0) return 0;
size_t written = 0;
switch (compression) {
case SW_COMPRESS_NONE: {
memcpy(buf, data, numel * sizeof(float));
written = numel * sizeof(float);
break;
}
case SW_COMPRESS_FP16: {
uint16_t* out = (uint16_t*)buf;
for (uint64_t i = 0; i < numel; i++) {
out[i] = float_to_fp16(data[i]);
}
written = numel * sizeof(uint16_t);
break;
}
case SW_COMPRESS_TOPK: {
/* Find top K values by absolute magnitude */
uint32_t k = (uint32_t)(numel * topk_ratio);
if (k == 0) k = 1;
if (k > numel) k = (uint32_t)numel;
/* Create indexed array for sorting */
indexed_value* indexed = malloc(numel * sizeof(indexed_value));
if (!indexed) return 0;
for (uint64_t i = 0; i < numel; i++) {
indexed[i].index = (uint32_t)i;
indexed[i].abs_value = fabsf(data[i]);
}
/* Partial sort - only need top K */
qsort(indexed, numel, sizeof(indexed_value), compare_indexed);
/* Write: K (4 bytes), then K pairs of (index, value) */
uint32_t* k_ptr = (uint32_t*)buf;
*k_ptr = k;
written += sizeof(uint32_t);
uint8_t* ptr = buf + sizeof(uint32_t);
for (uint32_t i = 0; i < k; i++) {
uint32_t idx = indexed[i].index;
float val = data[idx];
memcpy(ptr, &idx, sizeof(uint32_t));
ptr += sizeof(uint32_t);
memcpy(ptr, &val, sizeof(float));
ptr += sizeof(float);
written += sizeof(uint32_t) + sizeof(float);
}
free(indexed);
break;
}
case SW_COMPRESS_QSGD: {
/* Quantize to 8-bit with scale factor */
float max_val = 0.0f;
for (uint64_t i = 0; i < numel; i++) {
float absv = fabsf(data[i]);
if (absv > max_val) max_val = absv;
}
/* Write scale factor first */
memcpy(buf, &max_val, sizeof(float));
written += sizeof(float);
/* Quantize to int8 */
int8_t* out = (int8_t*)(buf + sizeof(float));
float scale = (max_val > 0) ? (127.0f / max_val) : 0.0f;
for (uint64_t i = 0; i < numel; i++) {
int32_t q = (int32_t)roundf(data[i] * scale);
if (q > 127) q = 127;
if (q < -127) q = -127;
out[i] = (int8_t)q;
}
written += numel;
break;
}
}
return written;
}
size_t sw_deserialize_tensor(uint8_t* buf, float* data, uint64_t numel,
sw_compress_type compression) {
if (!buf || !data || numel == 0) return 0;
size_t read = 0;
switch (compression) {
case SW_COMPRESS_NONE: {
memcpy(data, buf, numel * sizeof(float));
read = numel * sizeof(float);
break;
}
case SW_COMPRESS_FP16: {
uint16_t* in = (uint16_t*)buf;
for (uint64_t i = 0; i < numel; i++) {
data[i] = fp16_to_float(in[i]);
}
read = numel * sizeof(uint16_t);
break;
}
case SW_COMPRESS_TOPK: {
/* Zero out array first */
memset(data, 0, numel * sizeof(float));
/* Read K */
uint32_t k = *(uint32_t*)buf;
read += sizeof(uint32_t);
/* Read K pairs of (index, value) */
uint8_t* ptr = buf + sizeof(uint32_t);
for (uint32_t i = 0; i < k; i++) {
uint32_t idx;
float val;
memcpy(&idx, ptr, sizeof(uint32_t));
ptr += sizeof(uint32_t);
memcpy(&val, ptr, sizeof(float));
ptr += sizeof(float);
read += sizeof(uint32_t) + sizeof(float);
if (idx < numel) {
data[idx] = val;
}
}
break;
}
case SW_COMPRESS_QSGD: {
/* Read scale factor */
float max_val;
memcpy(&max_val, buf, sizeof(float));
read += sizeof(float);
/* Dequantize */
int8_t* in = (int8_t*)(buf + sizeof(float));
float scale = max_val / 127.0f;
for (uint64_t i = 0; i < numel; i++) {
data[i] = (float)in[i] * scale;
}
read += numel;
break;
}
}
return read;
}
size_t sw_estimate_compressed_size(uint64_t numel, sw_compress_type compression,
float topk_ratio) {
switch (compression) {
case SW_COMPRESS_NONE:
return numel * sizeof(float);
case SW_COMPRESS_FP16:
return numel * sizeof(uint16_t);
case SW_COMPRESS_TOPK: {
uint32_t k = (uint32_t)(numel * topk_ratio);
if (k == 0) k = 1;
/* K header + K*(index + value) */
return sizeof(uint32_t) + k * (sizeof(uint32_t) + sizeof(float));
}
case SW_COMPRESS_QSGD:
/* Scale factor + int8 values */
return sizeof(float) + numel;
}
return numel * sizeof(float);
}
/* ============================================================= */
/* Utility */
/* ============================================================= */
const char* sw_msg_type_name(sw_msg_type type) {
switch (type) {
case SW_MSG_REGISTER: return "REGISTER";
case SW_MSG_REGISTER_ACK: return "REGISTER_ACK";
case SW_MSG_DEREGISTER: return "DEREGISTER";
case SW_MSG_HEARTBEAT: return "HEARTBEAT";
case SW_MSG_HEARTBEAT_ACK: return "HEARTBEAT_ACK";
case SW_MSG_WEIGHTS_REQ: return "WEIGHTS_REQ";
case SW_MSG_WEIGHTS: return "WEIGHTS";
case SW_MSG_WEIGHTS_META: return "WEIGHTS_META";
case SW_MSG_GRADIENTS: return "GRADIENTS";
case SW_MSG_GRADIENTS_ACK: return "GRADIENTS_ACK";
case SW_MSG_SYNC: return "SYNC";
case SW_MSG_CHECKPOINT: return "CHECKPOINT";
case SW_MSG_SHUTDOWN: return "SHUTDOWN";
case SW_MSG_PAUSE: return "PAUSE";
case SW_MSG_RESUME: return "RESUME";
case SW_MSG_STATUS_REQ: return "STATUS_REQ";
case SW_MSG_STATUS: return "STATUS";
case SW_MSG_ERROR: return "ERROR";
default: return "UNKNOWN";
}
}
const char* sw_profile_name(sw_hardware_profile profile) {
switch (profile) {
case SW_PROFILE_RPI3: return "rpi3";
case SW_PROFILE_RPI4_2G: return "rpi4-2g";
case SW_PROFILE_RPI4_4G: return "rpi4-4g";
case SW_PROFILE_RPI4_8G: return "rpi4-8g";
case SW_PROFILE_LAPTOP_LIGHT: return "laptop-light";
case SW_PROFILE_LAPTOP_MID: return "laptop-mid";
case SW_PROFILE_DESKTOP_MID: return "desktop-mid";
case SW_PROFILE_DESKTOP_HEAVY: return "desktop-heavy";
case SW_PROFILE_SERVER: return "server";
default: return "unknown";
}
}
static uint64_t get_total_memory_mb(void) {
#ifdef __linux__
struct sysinfo info;
if (sysinfo(&info) == 0) {
return (info.totalram * info.mem_unit) / (1024 * 1024);
}
#endif
#ifdef __APPLE__
int64_t memsize;
size_t len = sizeof(memsize);
if (sysctlbyname("hw.memsize", &memsize, &len, NULL, 0) == 0) {
return (uint64_t)memsize / (1024 * 1024);
}
#endif
return 0;
}
static int get_num_cores(void) {
#ifdef __linux__
return get_nprocs();
#endif
#ifdef __APPLE__
int ncpu;
size_t len = sizeof(ncpu);
if (sysctlbyname("hw.ncpu", &ncpu, &len, NULL, 0) == 0) {
return ncpu;
}
#endif
return 1;
}
sw_hardware_profile sw_detect_profile(void) {
uint64_t mem_mb = get_total_memory_mb();
int cores = get_num_cores();
/* Check for Raspberry Pi (ARM with specific memory sizes) */
#ifdef __aarch64__
if (mem_mb <= 1024) return SW_PROFILE_RPI3;
if (mem_mb <= 2048) return SW_PROFILE_RPI4_2G;
if (mem_mb <= 4096) return SW_PROFILE_RPI4_4G;
if (mem_mb <= 8192 && cores <= 4) return SW_PROFILE_RPI4_8G;
#endif
/* x86 profiles based on memory and cores */
if (mem_mb >= 64000 && cores >= 16) return SW_PROFILE_SERVER;
if (mem_mb >= 32000 && cores >= 8) return SW_PROFILE_DESKTOP_HEAVY;
if (mem_mb >= 16000 && cores >= 4) return SW_PROFILE_DESKTOP_MID;
if (mem_mb >= 16000) return SW_PROFILE_LAPTOP_MID;
if (mem_mb >= 8000) return SW_PROFILE_LAPTOP_LIGHT;
return SW_PROFILE_UNKNOWN;
}
uint32_t sw_profile_batch_size(sw_hardware_profile profile) {
switch (profile) {
case SW_PROFILE_RPI3: return 2;
case SW_PROFILE_RPI4_2G: return 4;
case SW_PROFILE_RPI4_4G: return 8;
case SW_PROFILE_RPI4_8G: return 16;
case SW_PROFILE_LAPTOP_LIGHT: return 16;
case SW_PROFILE_LAPTOP_MID: return 32;
case SW_PROFILE_DESKTOP_MID: return 64;
case SW_PROFILE_DESKTOP_HEAVY: return 128;
case SW_PROFILE_SERVER: return 256;
default: return 8;
}
}
uint32_t sw_profile_accumulation_steps(sw_hardware_profile profile) {
switch (profile) {
case SW_PROFILE_RPI3: return 100;
case SW_PROFILE_RPI4_2G: return 100;
case SW_PROFILE_RPI4_4G: return 50;
case SW_PROFILE_RPI4_8G: return 50;
case SW_PROFILE_LAPTOP_LIGHT: return 50;
case SW_PROFILE_LAPTOP_MID: return 25;
case SW_PROFILE_DESKTOP_MID: return 25;
case SW_PROFILE_DESKTOP_HEAVY: return 10;
case SW_PROFILE_SERVER: return 5;
default: return 50;
}
}
/* ============================================================= */
/* Compression Helpers */
/* ============================================================= */
void sw_fp32_to_fp16(float* src, uint16_t* dst, size_t n) {
for (size_t i = 0; i < n; i++) {
dst[i] = float_to_fp16(src[i]);
}
}
void sw_fp16_to_fp32(uint16_t* src, float* dst, size_t n) {
for (size_t i = 0; i < n; i++) {
dst[i] = fp16_to_float(src[i]);
}
}
void sw_topk_compress(float* values, size_t n, size_t k,
uint32_t* out_indices, float* out_values) {
if (k > n) k = n;
if (k == 0) return;
/* Create indexed array for sorting */
indexed_value* indexed = malloc(n * sizeof(indexed_value));
if (!indexed) return;
for (size_t i = 0; i < n; i++) {
indexed[i].index = (uint32_t)i;
indexed[i].abs_value = fabsf(values[i]);
}
/* Sort by absolute value (descending) */
qsort(indexed, n, sizeof(indexed_value), compare_indexed);
/* Output top K */
for (size_t i = 0; i < k; i++) {
out_indices[i] = indexed[i].index;
out_values[i] = values[indexed[i].index];
}
free(indexed);
}
void sw_topk_decompress(uint32_t* indices, float* values, size_t k,
float* out, size_t n) {
/* Zero output first */
memset(out, 0, n * sizeof(float));
/* Scatter values */
for (size_t i = 0; i < k; i++) {
if (indices[i] < n) {
out[indices[i]] = values[i];
}
}
}
Reference in New Issue View Git Blame Copy Permalink