armorpaint/base/tools/iris.c/iris_gguf.c

/*
 * iris_gguf.c - GGUF file format reader implementation
 *
 * Parses the GGUF header/metadata/tensor-directory and fills a
 * safetensors_file_t so downstream code is format-agnostic. The mmap'd payload
 * is shared; for BF16/F16/F32 tensors the bytes match safetensors exactly.
 */

#include "iris_gguf.h"
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

/* GGML tensor types we support. BF16 checkpoints use only F32 + BF16; Q8_0
 * checkpoints additionally use the Q8_0 block-quantized type for the large
 * projection/embedding weights (norms stay F32). */
#define GGML_TYPE_F32  0
#define GGML_TYPE_F16  1
#define GGML_TYPE_Q8_0 8
#define GGML_TYPE_BF16 30

/* GGML Q8_0: 32 quantized elements per block, stored as an fp16 scale (2 bytes)
 * followed by 32 signed int8 quants -> 34 bytes per block. */
#define GGML_Q8_0_BLOCK 32
#define GGML_Q8_0_BYTES 34

/* GGUF metadata value types */
enum {
	GGUF_T_UINT8   = 0,
	GGUF_T_INT8    = 1,
	GGUF_T_UINT16  = 2,
	GGUF_T_INT16   = 3,
	GGUF_T_UINT32  = 4,
	GGUF_T_INT32   = 5,
	GGUF_T_FLOAT32 = 6,
	GGUF_T_BOOL    = 7,
	GGUF_T_STRING  = 8,
	GGUF_T_ARRAY   = 9,
	GGUF_T_UINT64  = 10,
	GGUF_T_INT64   = 11,
	GGUF_T_FLOAT64 = 12
};

/* Cursor over the mmap'd file, with end-of-buffer guarding. */
typedef struct {
	const uint8_t *p;
	const uint8_t *end;
	int            error;
} gguf_cursor_t;

static uint64_t cur_read(gguf_cursor_t *c, void *dst, size_t n) {
	if (c->error || (size_t)(c->end - c->p) < n) {
		c->error = 1;
		return 0;
	}
	if (dst)
		memcpy(dst, c->p, n);
	c->p += n;
	return 1;
}

static uint32_t cur_u32(gguf_cursor_t *c) {
	uint32_t v = 0;
	cur_read(c, &v, 4);
	return v;
}
static uint64_t cur_u64(gguf_cursor_t *c) {
	uint64_t v = 0;
	cur_read(c, &v, 8);
	return v;
}

/* GGUF string: uint64 length + raw bytes (not NUL-terminated). */
static uint64_t cur_str(gguf_cursor_t *c, char *out, size_t max) {
	uint64_t len = cur_u64(c);
	if (c->error || (size_t)(c->end - c->p) < len) {
		c->error = 1;
		return 0;
	}
	if (out) {
		size_t n = len < max - 1 ? (size_t)len : max - 1;
		memcpy(out, c->p, n);
		out[n] = '\0';
	}
	c->p += len;
	return len;
}

static size_t gguf_scalar_size(uint32_t t) {
	switch (t) {
	case GGUF_T_UINT8:
	case GGUF_T_INT8:
	case GGUF_T_BOOL:
		return 1;
	case GGUF_T_UINT16:
	case GGUF_T_INT16:
		return 2;
	case GGUF_T_UINT32:
	case GGUF_T_INT32:
	case GGUF_T_FLOAT32:
		return 4;
	case GGUF_T_UINT64:
	case GGUF_T_INT64:
	case GGUF_T_FLOAT64:
		return 8;
	default:
		return 0;
	}
}

/* Read one metadata value, capturing it only if it is the alignment key. */
static void gguf_read_value(gguf_cursor_t *c, uint32_t type, int want_align, uint32_t *align_out) {
	if (type == GGUF_T_STRING) {
		cur_str(c, NULL, 0);
	}
	else if (type == GGUF_T_ARRAY) {
		uint32_t et = cur_u32(c);
		uint64_t n  = cur_u64(c);
		for (uint64_t i = 0; i < n && !c->error; i++) {
			gguf_read_value(c, et, 0, NULL);
		}
	}
	else {
		size_t sz = gguf_scalar_size(type);
		if (sz == 0) {
			c->error = 1;
			return;
		}
		uint64_t v = 0;
		cur_read(c, &v, sz);
		if (want_align && type == GGUF_T_UINT32 && align_out)
			*align_out = (uint32_t)v;
	}
}

static safetensor_dtype_t gguf_map_dtype(uint32_t ggml_type, size_t *elem_size) {
	switch (ggml_type) {
	case GGML_TYPE_F32:
		*elem_size = 4;
		return DTYPE_F32;
	case GGML_TYPE_F16:
		*elem_size = 2;
		return DTYPE_F16;
	case GGML_TYPE_BF16:
		*elem_size = 2;
		return DTYPE_BF16;
	case GGML_TYPE_Q8_0:
		*elem_size = 0;
		return DTYPE_Q8_0; /* block type, see below */
	default:
		*elem_size = 0;
		return DTYPE_UNKNOWN;
	}
}

safetensors_file_t *gguf_open(const char *path) {
	int fd = open(path, O_RDONLY);
	if (fd < 0) {
		perror("gguf_open: open failed");
		return NULL;
	}

	struct stat st;
	if (fstat(fd, &st) < 0) {
		perror("gguf_open: fstat failed");
		close(fd);
		return NULL;
	}
	size_t file_size = (size_t)st.st_size;
	if (file_size < 24) {
		fprintf(stderr, "gguf_open: file too small\n");
		close(fd);
		return NULL;
	}

	void *data = mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
	close(fd);
	if (data == MAP_FAILED) {
		perror("gguf_open: mmap failed");
		return NULL;
	}

	gguf_cursor_t c = {(const uint8_t *)data, (const uint8_t *)data + file_size, 0};

	char magic[4];
	cur_read(&c, magic, 4);
	if (memcmp(magic, "GGUF", 4) != 0) {
		fprintf(stderr, "gguf_open: not a GGUF file\n");
		munmap(data, file_size);
		return NULL;
	}
	uint32_t version = cur_u32(&c);
	if (version != 2 && version != 3) {
		fprintf(stderr, "gguf_open: unsupported GGUF version %u\n", version);
		munmap(data, file_size);
		return NULL;
	}
	uint64_t n_tensors = cur_u64(&c);
	uint64_t n_kv      = cur_u64(&c);

	if (n_tensors > SAFETENSORS_MAX_TENSORS) {
		fprintf(stderr, "gguf_open: %llu tensors exceeds limit %d\n", (unsigned long long)n_tensors, SAFETENSORS_MAX_TENSORS);
		munmap(data, file_size);
		return NULL;
	}

	/* Metadata: we only care about general.alignment (default 32). */
	uint32_t alignment = 32;
	for (uint64_t i = 0; i < n_kv && !c.error; i++) {
		char key[128];
		cur_str(&c, key, sizeof(key));
		uint32_t type       = cur_u32(&c);
		int      want_align = (strcmp(key, "general.alignment") == 0);
		gguf_read_value(&c, type, want_align, &alignment);
	}
	if (alignment == 0)
		alignment = 32;

	safetensors_file_t *sf = calloc(1, sizeof(safetensors_file_t));
	if (!sf) {
		munmap(data, file_size);
		return NULL;
	}
	sf->path        = strdup(path);
	sf->data        = data;
	sf->file_size   = file_size;
	sf->header_json = NULL; /* not used for GGUF; safetensors_close free(NULL) is safe */

	/* Tensor directory. */
	int nt = 0;
	for (uint64_t i = 0; i < n_tensors && !c.error; i++) {
		safetensor_t *t = &sf->tensors[nt];
		cur_str(&c, t->name, sizeof(t->name));
		uint32_t ndim = cur_u32(&c);
		if (ndim > 8) {
			c.error = 1;
			break;
		}

		/* GGML stores dims fastest-first; reverse to PyTorch (slowest-first). */
		int64_t ne[8];
		for (uint32_t d = 0; d < ndim; d++)
			ne[d] = (int64_t)cur_u64(&c);
		t->ndim = (int)ndim;
		for (uint32_t d = 0; d < ndim; d++)
			t->shape[d] = ne[ndim - 1 - d];

		uint32_t ggml_type = cur_u32(&c);
		uint64_t offset    = cur_u64(&c);

		size_t elem_size = 0;
		t->dtype         = gguf_map_dtype(ggml_type, &elem_size);
		if (t->dtype == DTYPE_UNKNOWN) {
			fprintf(stderr,
			        "gguf_open: %s: unsupported ggml type %u "
			        "(only F32/F16/BF16/Q8_0 supported)\n",
			        t->name, ggml_type);
			c.error = 1;
			break;
		}

		int64_t numel = 1;
		for (uint32_t d = 0; d < ndim; d++)
			numel *= ne[d];
		t->data_offset = (size_t)offset; /* relative to data section */
		if (t->dtype == DTYPE_Q8_0) {
			/* Block-quantized: 34 bytes per 32 elements. ne[0] (the fastest GGML
			 * dim) is a multiple of 32, so blocks tile the tensor exactly. */
			t->data_size = (size_t)(numel / GGML_Q8_0_BLOCK) * GGML_Q8_0_BYTES;
		}
		else {
			t->data_size = (size_t)numel * elem_size;
		}
		nt++;
	}

	if (c.error) {
		fprintf(stderr, "gguf_open: %s: malformed header\n", path);
		safetensors_close(sf);
		return NULL;
	}

	/* Data section begins at the next alignment boundary after the directory. */
	size_t header_end = (size_t)(c.p - (const uint8_t *)data);
	size_t data_start = (header_end + alignment - 1) / alignment * alignment;
	if (data_start < 8)
		data_start = 8;

	/* Reuse safetensors_data()'s arithmetic: it returns
	 * data + 8 + header_size + data_offset. Set header_size so that
	 * 8 + header_size == data_start, leaving data_offset as the GGUF-relative
	 * tensor offset. */
	sf->header_size = data_start - 8;
	sf->num_tensors = nt;

	/* Validate every tensor lies within the file. */
	for (int i = 0; i < nt; i++) {
		safetensor_t *t = &sf->tensors[i];
		if (data_start + t->data_offset + t->data_size > file_size) {
			fprintf(stderr,
			        "gguf_open: %s: tensor '%s' extends past end of file "
			        "(truncated download?)\n",
			        path, t->name);
			safetensors_close(sf);
			return NULL;
		}
	}

	return sf;
}