ckernel_mem_plan.c File Reference

#include "ckernel_mem_plan.h"
#include "ckernel_dtype.h"
#include <stdlib.h>
#include <string.h>

Go to the source code of this file.

Functions
static size_t	align_up_bytes (size_t n, size_t align)
static size_t	align_up_elems (size_t elems, size_t elem_bytes, size_t align_bytes)
static CKMemArenaKind	arena_for_role (CKBufferRole role)
static int	buffer_enabled (const CKIRV2Graph *graph, const CKIRV2Buffer *buf, int training_enabled)
static int	build_plan (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes, int training_enabled, int tokens_override)
int	ck_mem_plan_build_inference (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes)
int	ck_mem_plan_build_inference_with_tokens (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes, int tokens_override)
int	ck_mem_plan_build_training (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes)
int	ck_mem_plan_build_training_with_tokens (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes, int tokens_override)
void	ck_mem_plan_free (CKMemPlan *plan)
static int	find_buffer_by_name (const CKIRV2Graph *graph, const char *name)
static size_t	resolve_dim (const CKModelConfig *cfg, const CKIRV2AlignInfo *align, CKDimKind kind, int tokens_override)
static size_t	resolve_shape_elems (const CKModelConfig *cfg, const CKIRV2AlignInfo *align, const CKDimToken *shape, int tokens_override)

Function Documentation

align_up_bytes()

static size_t align_up_bytes ( size_t  n, size_t  align  )

static

Definition at line 15 of file ckernel_mem_plan.c.

{
    if (align == 0) return n;
    return (n + align - 1) & ~(align - 1);
}

Referenced by align_up_elems(), and build_plan().

align_up_elems()

static size_t align_up_elems ( size_t  elems, size_t  elem_bytes, size_t  align_bytes  )

static

Definition at line 21 of file ckernel_mem_plan.c.

{
    size_t bytes = elems * elem_bytes;
    bytes = align_up_bytes(bytes, align_bytes);
    return bytes / elem_bytes;
}

References align_up_bytes().

Referenced by build_plan().

arena_for_role()

static CKMemArenaKind arena_for_role ( CKBufferRole  role )

static

Definition at line 84 of file ckernel_mem_plan.c.

{
    switch (role) {
    case CK_ROLE_WEIGHT:
        return CK_MEM_ARENA_WEIGHTS;
    case CK_ROLE_GRAD:
        return CK_MEM_ARENA_GRADS;
    case CK_ROLE_SCRATCH:
    case CK_ROLE_ACTIVATION:
    case CK_ROLE_INPUT:
    case CK_ROLE_OUTPUT:
    default:
        return CK_MEM_ARENA_ACTIVATIONS;
    }
}

References CK_MEM_ARENA_ACTIVATIONS, CK_MEM_ARENA_GRADS, CK_MEM_ARENA_WEIGHTS, CK_ROLE_ACTIVATION, CK_ROLE_GRAD, CK_ROLE_INPUT, CK_ROLE_OUTPUT, CK_ROLE_SCRATCH, and CK_ROLE_WEIGHT.

Referenced by build_plan().

buffer_enabled()

static int buffer_enabled ( const CKIRV2Graph *  graph, const CKIRV2Buffer *  buf, int  training_enabled  )

static

Definition at line 100 of file ckernel_mem_plan.c.

{
    if (!graph || !buf) {
        return 0;
    }
    if (buf->condition && strcmp(buf->condition, "rope_theta") == 0) {
        if (graph->config.rope_theta <= 0.0f) {
            return 0;
        }
    }
    if (buf->condition && strcmp(buf->condition, "rope_disabled") == 0) {
        if (graph->config.rope_theta > 0.0f) {
            return 0;
        }
    }
    if (buf->condition && strcmp(buf->condition, "has_pos_emb") == 0) {
        if (graph->has_pos_emb == 0) {
            return 0;
        }
    }
    if (buf->condition && strcmp(buf->condition, "training_enabled") == 0) {
        if (!training_enabled) {
            return 0;
        }
    }
    if (buf->role == CK_ROLE_GRAD && !training_enabled) {
        return 0;
    }
    return 1;
}

References CK_ROLE_GRAD, CKIRV2Buffer::condition, CKIRV2Graph::config, CKIRV2Graph::has_pos_emb, CKIRV2Buffer::role, and CKModelConfig::rope_theta.

Referenced by build_plan().

build_plan()

static int build_plan ( const CKIRV2Graph *  graph, CKMemPlan *  plan, size_t  alignment_bytes, int  training_enabled, int  tokens_override  )

static

Definition at line 146 of file ckernel_mem_plan.c.

{
    if (!graph || !plan) {
        return -1;
    }
    memset(plan, 0, sizeof(*plan));
    if (alignment_bytes == 0) {
        alignment_bytes = CK_MEM_PLAN_DEFAULT_ALIGN;
    }
    plan->alignment_bytes = alignment_bytes;
 
    plan->num_spans = graph->num_buffers;
    plan->spans = (CKMemSpan *)calloc((size_t)graph->num_buffers, sizeof(CKMemSpan));
    if (!plan->spans) {
        return -1;
    }
 
    size_t elem_bytes = sizeof(float);
    CKIRV2AlignInfo align = {0};
    align.aligned_embed = align_up_elems((size_t)graph->config.hidden_size, elem_bytes, alignment_bytes);
    align.aligned_head = align_up_elems((size_t)(graph->config.hidden_size / graph->config.num_heads),
                                        elem_bytes, alignment_bytes);
    align.aligned_intermediate = align_up_elems((size_t)graph->config.intermediate_size,
                                               elem_bytes, alignment_bytes);
    align.aligned_context = align_up_elems((size_t)graph->config.context_window, elem_bytes, alignment_bytes);
 
    size_t arena_offsets[CK_MEM_ARENA_COUNT] = {0};
 
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        CKMemSpan *span = &plan->spans[i];
        span->buffer_id = i;
        span->arena = arena_for_role(buf->role);
        span->offset_bytes = 0;
        span->size_bytes = 0;
 
        if (!buffer_enabled(graph, buf, training_enabled)) {
            continue;
        }
 
        if (buf->alias_of) {
            int alias_id = find_buffer_by_name(graph, buf->alias_of);
            if (alias_id >= 0) {
                span->offset_bytes = plan->spans[alias_id].offset_bytes;
                span->size_bytes = plan->spans[alias_id].size_bytes;
                span->arena = plan->spans[alias_id].arena;
                continue;
            }
        }
 
        size_t n_elems = resolve_shape_elems(&graph->config, &align, buf->shape, tokens_override);
        size_t bytes = ck_dtype_row_bytes(buf->dtype, n_elems);
        size_t aligned = align_up_bytes(bytes, alignment_bytes);
 
        span->offset_bytes = arena_offsets[span->arena];
        span->size_bytes = bytes;
        arena_offsets[span->arena] += aligned;
    }
 
    for (int i = 0; i < CK_MEM_ARENA_COUNT; ++i) {
        plan->total_bytes[i] = arena_offsets[i];
    }
 
    return 0;
}

References CKIRV2Buffer::alias_of, align_up_bytes(), align_up_elems(), CKMemPlan::alignment_bytes, CKMemSpan::arena, arena_for_role(), buffer_enabled(), CKMemSpan::buffer_id, CKIRV2Graph::buffers, ck_dtype_row_bytes(), CK_MEM_ARENA_COUNT, CK_MEM_PLAN_DEFAULT_ALIGN, CKIRV2Graph::config, CKModelConfig::context_window, CKIRV2Buffer::dtype, find_buffer_by_name(), CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKIRV2Graph::num_buffers, CKModelConfig::num_heads, CKMemPlan::num_spans, CKMemSpan::offset_bytes, resolve_shape_elems(), CKIRV2Buffer::role, CKIRV2Buffer::shape, CKMemSpan::size_bytes, CKMemPlan::spans, and CKMemPlan::total_bytes.

Referenced by ck_mem_plan_build_inference(), ck_mem_plan_build_inference_with_tokens(), ck_mem_plan_build_training(), and ck_mem_plan_build_training_with_tokens().

ck_mem_plan_build_inference()

int ck_mem_plan_build_inference	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes
	)

Definition at line 216 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 0, -1);
}

References build_plan().

ck_mem_plan_build_inference_with_tokens()

int ck_mem_plan_build_inference_with_tokens	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes,
		int	tokens_override
	)

Definition at line 230 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 0, tokens_override);
}

References build_plan().

Referenced by ck_codegen_v2_emit_runtime(), and ck_ir_v2_lower_graph().

ck_mem_plan_build_training()

int ck_mem_plan_build_training	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes
	)

Definition at line 223 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 1, -1);
}

References build_plan().

ck_mem_plan_build_training_with_tokens()

int ck_mem_plan_build_training_with_tokens	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes,
		int	tokens_override
	)

Definition at line 238 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 1, tokens_override);
}

References build_plan().

Referenced by ck_codegen_v2_emit_runtime(), and ck_ir_v2_lower_graph().

ck_mem_plan_free()

void ck_mem_plan_free

(

CKMemPlan *

plan

)

Definition at line 246 of file ckernel_mem_plan.c.

{
    if (!plan) {
        return;
    }
    free(plan->spans);
    memset(plan, 0, sizeof(*plan));
}

References CKMemPlan::spans.

Referenced by ck_codegen_v2_emit_runtime(), ck_ir_v2_lower_emit_json(), and ck_ir_v2_lower_graph().

find_buffer_by_name()

static int find_buffer_by_name ( const CKIRV2Graph *  graph, const char *  name  )

static

Definition at line 133 of file ckernel_mem_plan.c.

{
    if (!graph || !name) {
        return -1;
    }
    for (int i = 0; i < graph->num_buffers; ++i) {
        if (graph->buffers[i].name && strcmp(graph->buffers[i].name, name) == 0) {
            return i;
        }
    }
    return -1;
}

References CKIRV2Graph::buffers, CKIRV2Buffer::name, and CKIRV2Graph::num_buffers.

Referenced by build_plan().

resolve_dim()

static size_t resolve_dim ( const CKModelConfig *  cfg, const CKIRV2AlignInfo *  align, CKDimKind  kind, int  tokens_override  )

static

Definition at line 28 of file ckernel_mem_plan.c.

{
    switch (kind) {
    case CK_DIM_TOKENS:
        if (tokens_override >= 0) {
            return (size_t)tokens_override;
        }
        return (size_t)cfg->context_window;
    case CK_DIM_EMBED:
        return (size_t)cfg->hidden_size;
    case CK_DIM_ALIGNED_EMBED:
        return align->aligned_embed;
    case CK_DIM_HEAD_DIM:
        return (size_t)(cfg->hidden_size / cfg->num_heads);
    case CK_DIM_ALIGNED_HEAD:
        return align->aligned_head;
    case CK_DIM_NUM_HEADS:
        return (size_t)cfg->num_heads;
    case CK_DIM_NUM_KV_HEADS:
        return (size_t)cfg->num_kv_heads;
    case CK_DIM_ALIGNED_CTX:
        return align->aligned_context;
    case CK_DIM_INTERMEDIATE:
        return (size_t)cfg->intermediate_size;
    case CK_DIM_ALIGNED_INTERMEDIATE:
        return align->aligned_intermediate;
    case CK_DIM_VOCAB:
        return (size_t)cfg->vocab_size;
    case CK_DIM_END:
    default:
        return 0;
    }
}

References CK_DIM_ALIGNED_CTX, CK_DIM_ALIGNED_EMBED, CK_DIM_ALIGNED_HEAD, CK_DIM_ALIGNED_INTERMEDIATE, CK_DIM_EMBED, CK_DIM_END, CK_DIM_HEAD_DIM, CK_DIM_INTERMEDIATE, CK_DIM_NUM_HEADS, CK_DIM_NUM_KV_HEADS, CK_DIM_TOKENS, CK_DIM_VOCAB, CKModelConfig::context_window, CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, and CKModelConfig::vocab_size.

Referenced by resolve_shape_elems().

resolve_shape_elems()

static size_t resolve_shape_elems ( const CKModelConfig *  cfg, const CKIRV2AlignInfo *  align, const CKDimToken *  shape, int  tokens_override  )

static

Definition at line 65 of file ckernel_mem_plan.c.

{
    size_t total = 1;
    for (int i = 0; i < CK_IR_V2_MAX_DIMS; ++i) {
        if (shape[i].dim == CK_DIM_END) {
            break;
        }
        size_t dim = resolve_dim(cfg, align, shape[i].dim, tokens_override);
        size_t mult = (size_t)(shape[i].mult > 0 ? shape[i].mult : 1);
        size_t div = (size_t)(shape[i].div > 0 ? shape[i].div : 1);
        if (div == 0) div = 1;
        total = total * dim * mult / div;
    }
    return total;
}

References CK_DIM_END, CK_IR_V2_MAX_DIMS, and resolve_dim().

Referenced by build_plan().