HAMi-core allocator.c 内部原理

struct allocated_device_memory_struct {
    CUdeviceptr address;                    // 分配的显存地址
    size_t length;                          // 分配的字节数
    CUcontext ctx;                          // 所属的 CUDA 上下文
    CUmemGenericAllocationHandle *allocHandle; // 虚拟内存分配句柄（用于大块内存）
};

struct allocated_list_entry_struct {
    allocated_device_memory *entry;           // 指向分配记录
    struct allocated_list_entry_struct *next;  // 后继节点
    struct allocated_list_entry_struct *prev;  // 前驱节点
};

struct allocated_list_struct {
    allocated_list_entry *head;   // 链表头
    allocated_list_entry *tail;   // 链表尾
    size_t length;                // 节点数量
    size_t limit;                 // 内存池上限（仅 async 使用）
};

struct region_struct {
    size_t start;               // 区域起始地址
    size_t freemark;            // 空闲标记位置
    size_t freed_map;           // 已释放的 chunk 数量
    size_t length;              // 已使用的长度
    CUcontext ctx;              // CUDA 上下文
    allocated_list *region_allocs; // 区域内的子分配列表
    char *bitmap;               // 分配位图
    CUmemGenericAllocationHandle *allocHandle; // 分配句柄
};

region_list *r_list;                      // 区域链表（大内存管理）
allocated_list *device_overallocated;     // 同步分配链表
allocated_list *device_allocasync;        // 异步分配链表

size_t BITSIZE = 512;
size_t IPCSIZE = 2097152;       // 2MB
size_t OVERSIZE = 134217728;    // 128MB
#define ALIGN 2097152            // 2MB
#define CHUNK_SIZE (OVERSIZE/BITSIZE)

int allocate_raw(CUdeviceptr *dptr, size_t size) {
    pthread_mutex_lock(&mutex);
    tmp = add_chunk(dptr, size);
    pthread_mutex_unlock(&mutex);
    return tmp;
}

// 异步分配中的增量计算逻辑
if (poollimit > device_allocasync->limit) {
    allocsize = (poollimit - device_allocasync->limit < size)
              ? poollimit - device_allocasync->limit
              : size;
    add_gpu_device_memory_usage(getpid(), dev, allocsize, 2);
    device_allocasync->limit += allocsize;
    e->entry->length = allocsize;
} else {
    e->entry->length = 0;  // 复用内存池，无增量
}

int free_raw(CUdeviceptr dptr) {
    pthread_mutex_lock(&mutex);
    tmp = remove_chunk(device_overallocated, dptr);
    pthread_mutex_unlock(&mutex);
    return tmp;
}

int oom_check(const int dev, size_t addon) {
    uint64_t limit = get_current_device_memory_limit(d);  // 从共享内存读取限制
    size_t _usage = get_gpu_memory_usage(d);              // 当前使用量

    if (limit == 0) return 0;  // limit=0 表示不限制

    size_t new_allocated = _usage + addon;
    if (new_allocated > limit) {
        // 尝试清理过期的进程槽位
        if (clear_proc_slot_nolock(1) > 0)
            return oom_check(dev, addon);  // 清理成功后重试
        return 1;  // OOM
    }
    return 0;  // 通过
}

pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;

pthread_once_t allocator_allocate_flag = PTHREAD_ONCE_INIT;

#define LIST_INIT(list) {
    list->head = NULL;
    list->tail = NULL;
    list->length = 0;
    list->limit = 0;
}

#define LIST_ADD(list, __entry) {
    if (list->head == NULL) {
        list->head = __entry;
        list->tail = __entry;
        list->length = 1;
    } else {
        __entry->prev = list->tail;
        list->tail->next = __entry;
        list->tail = __entry;
        list->length++;
    }
}

#define LIST_REMOVE(list, val) {
    // 更新前驱后继指针
    if (val->prev != NULL) val->prev->next = val->next;
    if (val->next != NULL) val->next->prev = val->prev;
    if (val == list->tail) list->tail = val->prev;
    if (val == list->head) list->head = val->next;
    // 释放内存
    free(val->entry->allocHandle);
    free(val->entry);
    free(val);
    list->length--;
}

int check_memory_type(CUdeviceptr address) {
    // 遍历 device_overallocated 链表
    // 如果 address 落在某个分配区间内，返回 CU_MEMORYTYPE_DEVICE
    // 否则返回 CU_MEMORYTYPE_HOST
}

CUresult get_listsize(allocated_list *al, size_t *size) {
    size_t count = 0;
    for (val = al->head; val != NULL; val = val->next) {
        count += val->entry->length;
    }
    *size = count;
    return CUDA_SUCCESS;
}

size_t round_up(size_t size, size_t unit) {
    if (size & (unit-1))
        return ((size / unit) + 1) * unit;
    return size;
}