ckernel_codegen_v2_sections.c File Reference

#include "ckernel_codegen_v2_emit.h"
#include "ckernel_dtype.h"
#include <stdio.h>
#include <string.h>

Go to the source code of this file.

Macros
#define	CK_V2_SECTION_ALIGN   64

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 size_t	buffer_bytes (const CKIRV2Buffer *buf, const CKModelConfig *cfg, const CKV2AlignInfo *align)
void	ck_codegen_v2_emit_sections (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *prefill_plan, const CKMemPlan *decode_plan, const CKMemPlan *backward_plan)
static CKV2AlignInfo	compute_align (const CKModelConfig *cfg)
static void	emit_body_fields (FILE *out, const CKIRV2Graph *graph, CKBufferRole role_filter, int activation_group)
static size_t	emit_body_values (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *plan, CKBufferRole role_filter, int activation_group)
static void	emit_footer_fields (FILE *out, const CKIRV2Graph *graph, CKBufferRole role_filter, int activation_group)
static void	emit_footer_values (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *plan, CKBufferRole role_filter, int activation_group, size_t *offset)
static void	emit_header_fields (FILE *out, const CKIRV2Graph *graph, CKBufferRole role_filter, int activation_group)
static void	emit_header_values (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *plan, CKBufferRole role_filter, int activation_group, size_t *offset)
static void	emit_span_field (FILE *out, const char *label)
static void	emit_span_value (FILE *out, const char *label, size_t offset, size_t size, int comma)
static int	is_activation_role (CKBufferRole role)
static int	is_footer_global (const char *name)
static size_t	plan_size (const CKMemPlan *plan, int idx)
static size_t	resolve_dim (const CKModelConfig *cfg, const CKV2AlignInfo *align, CKDimKind kind)
static size_t	resolve_shape_elems (const CKModelConfig *cfg, const CKV2AlignInfo *align, const CKDimToken *shape)

Macro Definition Documentation

CK_V2_SECTION_ALIGN

#define CK_V2_SECTION_ALIGN   64

Definition at line 8 of file ckernel_codegen_v2_sections.c.

Function Documentation

align_up_bytes()

static size_t align_up_bytes ( size_t  n, size_t  align  )

static

Definition at line 17 of file ckernel_codegen_v2_sections.c.

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

Referenced by align_up_elems(), emit_body_values(), emit_footer_values(), and emit_header_values().

align_up_elems()

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

static

Definition at line 25 of file ckernel_codegen_v2_sections.c.

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

References align_up_bytes().

Referenced by compute_align().

buffer_bytes()

static size_t buffer_bytes ( const CKIRV2Buffer *  buf, const CKModelConfig *  cfg, const CKV2AlignInfo *  align  )

static

Definition at line 100 of file ckernel_codegen_v2_sections.c.

{
    size_t elems = resolve_shape_elems(cfg, align, buf->shape);
    return ck_dtype_row_bytes(buf->dtype, elems);
}

References ck_dtype_row_bytes(), CKIRV2Buffer::dtype, resolve_shape_elems(), and CKIRV2Buffer::shape.

ck_codegen_v2_emit_sections()

void ck_codegen_v2_emit_sections	(	FILE *	out,
		const CKIRV2Graph *	graph,
		const CKMemPlan *	prefill_plan,
		const CKMemPlan *	decode_plan,
		const CKMemPlan *	backward_plan
	)

Definition at line 309 of file ckernel_codegen_v2_sections.c.

{
    if (!out || !graph) {
        return;
    }
 
    const int L = graph->config.num_layers;
 
    fprintf(out,
            "typedef struct {\n"
            "    size_t offset;\n"
            "    size_t size;\n"
            "} CKV2Span;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_header_fields(out, graph, CK_ROLE_WEIGHT, 0);
    fprintf(out, "} CKV2HeaderWeights;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_body_fields(out, graph, CK_ROLE_WEIGHT, 0);
    fprintf(out, "} CKV2LayerWeights;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_footer_fields(out, graph, CK_ROLE_WEIGHT, 0);
    fprintf(out, "} CKV2FooterWeights;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_header_fields(out, graph, CK_ROLE_ACTIVATION, 1);
    fprintf(out, "} CKV2HeaderActivations;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_body_fields(out, graph, CK_ROLE_ACTIVATION, 1);
    fprintf(out, "} CKV2LayerActivations;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_footer_fields(out, graph, CK_ROLE_ACTIVATION, 1);
    fprintf(out, "} CKV2FooterActivations;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_header_fields(out, graph, CK_ROLE_GRAD, 0);
    fprintf(out, "} CKV2HeaderGrads;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_body_fields(out, graph, CK_ROLE_GRAD, 0);
    fprintf(out, "} CKV2LayerGrads;\n\n");
 
    fprintf(out, "typedef struct {\n");
    emit_footer_fields(out, graph, CK_ROLE_GRAD, 0);
    fprintf(out, "} CKV2FooterGrads;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2HeaderWeights header;\n"
            "    CKV2LayerWeights body;\n"
            "    CKV2FooterWeights footer;\n"
            "    size_t layer_stride_bytes;\n"
            "    size_t total_bytes;\n"
            "} CKV2WeightLayout;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2HeaderActivations header;\n"
            "    CKV2LayerActivations body;\n"
            "    CKV2FooterActivations footer;\n"
            "    size_t layer_stride_bytes;\n"
            "    size_t total_bytes;\n"
            "} CKV2ActivationLayout;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2HeaderGrads header;\n"
            "    CKV2LayerGrads body;\n"
            "    CKV2FooterGrads footer;\n"
            "    size_t layer_stride_bytes;\n"
            "    size_t total_bytes;\n"
            "} CKV2GradLayout;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    int num_layers;\n"
            "    int context_window;\n"
            "    int hidden_size;\n"
            "    int intermediate_size;\n"
            "    int num_heads;\n"
            "    int num_kv_heads;\n"
            "    int head_dim;\n"
            "    int vocab_size;\n"
            "} CKV2ModelConfig;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2WeightLayout weights;\n"
            "    CKV2ActivationLayout activations;\n"
            "    CKV2GradLayout grads;\n"
            "} CKV2SectionLayout;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2SectionLayout prefill;\n"
            "    CKV2SectionLayout decode;\n"
            "    CKV2SectionLayout backward;\n"
            "} CKV2SectionLayouts;\n\n");
 
    fprintf(out,
            "typedef struct {\n"
            "    CKV2ModelConfig config;\n"
            "    CKV2SectionLayouts decoder;\n"
            "} CKV2RuntimeLayout;\n\n");
 
    fprintf(out, "static const CKV2RuntimeLayout ck_v2_layout = {\n");
    fprintf(out, "    .config = {\n");
    fprintf(out, "        .num_layers = %d,\n", graph->config.num_layers);
    fprintf(out, "        .context_window = %d,\n", graph->config.context_window);
    fprintf(out, "        .hidden_size = %d,\n", graph->config.hidden_size);
    fprintf(out, "        .intermediate_size = %d,\n", graph->config.intermediate_size);
    fprintf(out, "        .num_heads = %d,\n", graph->config.num_heads);
    fprintf(out, "        .num_kv_heads = %d,\n", graph->config.num_kv_heads);
    fprintf(out, "        .head_dim = %d,\n",
            graph->config.num_heads > 0 ? graph->config.hidden_size / graph->config.num_heads : 0);
    fprintf(out, "        .vocab_size = %d\n", graph->config.vocab_size);
    fprintf(out, "    },\n");
    fprintf(out, "    .decoder = {\n");
 
    const CKMemPlan *plans[3] = { prefill_plan, decode_plan, backward_plan };
    const char *modes[3] = { "prefill", "decode", "backward" };
 
    for (int mode = 0; mode < 3; ++mode) {
        const CKMemPlan *plan = plans[mode];
        fprintf(out, "        .%s = {\n", modes[mode]);
 
        size_t offset = 0;
        size_t header_end = 0;
        size_t layer_stride = 0;
        size_t footer_base = 0;
 
        fprintf(out, "            .weights = {\n");
        fprintf(out, "                .header = {\n");
        emit_header_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0, &offset);
        header_end = offset;
        fprintf(out, "                },\n");
 
        fprintf(out, "                .body = {\n");
        layer_stride = emit_body_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0);
        fprintf(out, "                },\n");
 
        footer_base = header_end + layer_stride * (size_t)L;
        offset = footer_base;
        fprintf(out, "                .footer = {\n");
        emit_footer_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0, &offset);
        fprintf(out, "                },\n");
 
        fprintf(out, "                .layer_stride_bytes = %zu,\n", layer_stride);
        fprintf(out, "                .total_bytes = %zu\n", offset);
        fprintf(out, "            },\n");
 
        offset = 0;
        header_end = 0;
        layer_stride = 0;
        footer_base = 0;
 
        fprintf(out, "            .activations = {\n");
        fprintf(out, "                .header = {\n");
        emit_header_values(out, graph, plan, CK_ROLE_ACTIVATION, 1, &offset);
        header_end = offset;
        fprintf(out, "                },\n");
 
        fprintf(out, "                .body = {\n");
        layer_stride = emit_body_values(out, graph, plan, CK_ROLE_ACTIVATION, 1);
        fprintf(out, "                },\n");
 
        footer_base = header_end + layer_stride * (size_t)L;
        offset = footer_base;
        fprintf(out, "                .footer = {\n");
        emit_footer_values(out, graph, plan, CK_ROLE_ACTIVATION, 1, &offset);
        fprintf(out, "                },\n");
 
        fprintf(out, "                .layer_stride_bytes = %zu,\n", layer_stride);
        fprintf(out, "                .total_bytes = %zu\n", offset);
        fprintf(out, "            },\n");
 
        offset = 0;
        header_end = 0;
        layer_stride = 0;
        footer_base = 0;
 
        fprintf(out, "            .grads = {\n");
        fprintf(out, "                .header = {\n");
        emit_header_values(out, graph, plan, CK_ROLE_GRAD, 0, &offset);
        header_end = offset;
        fprintf(out, "                },\n");
 
        fprintf(out, "                .body = {\n");
        layer_stride = emit_body_values(out, graph, plan, CK_ROLE_GRAD, 0);
        fprintf(out, "                },\n");
 
        footer_base = header_end + layer_stride * (size_t)L;
        offset = footer_base;
        fprintf(out, "                .footer = {\n");
        emit_footer_values(out, graph, plan, CK_ROLE_GRAD, 0, &offset);
        fprintf(out, "                },\n");
 
        fprintf(out, "                .layer_stride_bytes = %zu,\n", layer_stride);
        fprintf(out, "                .total_bytes = %zu\n", offset);
        fprintf(out, "            }\n");
 
        fprintf(out, "        }%s\n", mode == 2 ? "" : ",");
    }
 
    fprintf(out, "    }\n");
    fprintf(out, "};\n\n");
}

References CK_ROLE_ACTIVATION, CK_ROLE_GRAD, CK_ROLE_WEIGHT, CKIRV2Graph::config, CKModelConfig::context_window, emit_body_fields(), emit_body_values(), emit_footer_fields(), emit_footer_values(), emit_header_fields(), emit_header_values(), CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, CKModelConfig::num_layers, and CKModelConfig::vocab_size.

Referenced by ck_codegen_v2_emit_runtime().

compute_align()

static CKV2AlignInfo compute_align ( const CKModelConfig *  cfg )

static

Definition at line 32 of file ckernel_codegen_v2_sections.c.

{
    CKV2AlignInfo info = {0};
    if (!cfg) {
        return info;
    }
    size_t elem_bytes = sizeof(float);
    size_t head_dim = (cfg->num_heads > 0) ? (size_t)(cfg->hidden_size / cfg->num_heads) : 0;
    info.aligned_embed = align_up_elems((size_t)cfg->hidden_size, elem_bytes, CK_V2_SECTION_ALIGN);
    info.aligned_head = align_up_elems(head_dim, elem_bytes, CK_V2_SECTION_ALIGN);
    info.aligned_intermediate =
        align_up_elems((size_t)cfg->intermediate_size, elem_bytes, CK_V2_SECTION_ALIGN);
    info.aligned_context =
        align_up_elems((size_t)cfg->context_window, elem_bytes, CK_V2_SECTION_ALIGN);
    return info;
}

References align_up_elems(), CK_V2_SECTION_ALIGN, CKModelConfig::context_window, CKModelConfig::hidden_size, CKModelConfig::intermediate_size, and CKModelConfig::num_heads.

emit_body_fields()

static void emit_body_fields ( FILE *  out, const CKIRV2Graph *  graph, CKBufferRole  role_filter, int  activation_group  )

static

Definition at line 171 of file ckernel_codegen_v2_sections.c.

{
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_LAYER) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        emit_span_field(out, buf->name);
    }
}

References CKIRV2Graph::buffers, CK_SCOPE_LAYER, emit_span_field(), is_activation_role(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_body_values()

static size_t emit_body_values ( FILE *  out, const CKIRV2Graph *  graph, const CKMemPlan *  plan, CKBufferRole  role_filter, int  activation_group  )

static

Definition at line 253 of file ckernel_codegen_v2_sections.c.

{
    size_t layer_offset = 0;
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_LAYER) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        size_t bytes = plan_size(plan, i);
        layer_offset = align_up_bytes(layer_offset, CK_V2_SECTION_ALIGN);
        emit_span_value(out, buf->name, layer_offset, bytes, 1);
        layer_offset += bytes;
    }
    return layer_offset;
}

References align_up_bytes(), CKIRV2Graph::buffers, CK_SCOPE_LAYER, CK_V2_SECTION_ALIGN, emit_span_value(), is_activation_role(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, plan_size(), CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_footer_fields()

static void emit_footer_fields ( FILE *  out, const CKIRV2Graph *  graph, CKBufferRole  role_filter, int  activation_group  )

static

Definition at line 192 of file ckernel_codegen_v2_sections.c.

{
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_GLOBAL) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        if (!is_footer_global(buf->name)) {
            continue;
        }
        emit_span_field(out, buf->name);
    }
}

References CKIRV2Graph::buffers, CK_SCOPE_GLOBAL, emit_span_field(), is_activation_role(), is_footer_global(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_footer_values()

static void emit_footer_values ( FILE *  out, const CKIRV2Graph *  graph, const CKMemPlan *  plan, CKBufferRole  role_filter, int  activation_group, size_t *  offset  )

static

Definition at line 280 of file ckernel_codegen_v2_sections.c.

{
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_GLOBAL) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        if (!is_footer_global(buf->name)) {
            continue;
        }
        size_t bytes = plan_size(plan, i);
        *offset = align_up_bytes(*offset, CK_V2_SECTION_ALIGN);
        emit_span_value(out, buf->name, *offset, bytes, 1);
        *offset += bytes;
    }
}

References align_up_bytes(), CKIRV2Graph::buffers, CK_SCOPE_GLOBAL, CK_V2_SECTION_ALIGN, emit_span_value(), is_activation_role(), is_footer_global(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, plan_size(), CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_header_fields()

static void emit_header_fields ( FILE *  out, const CKIRV2Graph *  graph, CKBufferRole  role_filter, int  activation_group  )

static

Definition at line 147 of file ckernel_codegen_v2_sections.c.

{
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_GLOBAL) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        if (is_footer_global(buf->name)) {
            continue;
        }
        emit_span_field(out, buf->name);
    }
}

References CKIRV2Graph::buffers, CK_SCOPE_GLOBAL, emit_span_field(), is_activation_role(), is_footer_global(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_header_values()

static void emit_header_values ( FILE *  out, const CKIRV2Graph *  graph, const CKMemPlan *  plan, CKBufferRole  role_filter, int  activation_group, size_t *  offset  )

static

Definition at line 224 of file ckernel_codegen_v2_sections.c.

{
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKIRV2Buffer *buf = &graph->buffers[i];
        if (buf->scope != CK_SCOPE_GLOBAL) {
            continue;
        }
        if (activation_group) {
            if (!is_activation_role(buf->role)) {
                continue;
            }
        } else if (buf->role != role_filter) {
            continue;
        }
        if (is_footer_global(buf->name)) {
            continue;
        }
        size_t bytes = plan_size(plan, i);
        *offset = align_up_bytes(*offset, CK_V2_SECTION_ALIGN);
        emit_span_value(out, buf->name, *offset, bytes, 1);
        *offset += bytes;
    }
}

References align_up_bytes(), CKIRV2Graph::buffers, CK_SCOPE_GLOBAL, CK_V2_SECTION_ALIGN, emit_span_value(), is_activation_role(), is_footer_global(), CKIRV2Buffer::name, CKIRV2Graph::num_buffers, plan_size(), CKIRV2Buffer::role, and CKIRV2Buffer::scope.

Referenced by ck_codegen_v2_emit_sections().

emit_span_field()

static void emit_span_field ( FILE *  out, const char *  label  )

static

Definition at line 131 of file ckernel_codegen_v2_sections.c.

{
    fprintf(out, "    CKV2Span %s;\n", label);
}

Referenced by emit_body_fields(), emit_footer_fields(), and emit_header_fields().

emit_span_value()

static void emit_span_value ( FILE *  out, const char *  label, size_t  offset, size_t  size, int  comma  )

static

Definition at line 136 of file ckernel_codegen_v2_sections.c.

{
    fprintf(out, "        .%s = { %zu, %zu }%s\n", label, offset, size, comma ? "," : "");
}

Referenced by emit_body_values(), emit_footer_values(), and emit_header_values().

is_activation_role()

static int is_activation_role ( CKBufferRole  role )

static

Definition at line 141 of file ckernel_codegen_v2_sections.c.

{
    return role == CK_ROLE_INPUT || role == CK_ROLE_OUTPUT ||
           role == CK_ROLE_ACTIVATION || role == CK_ROLE_SCRATCH;
}

References CK_ROLE_ACTIVATION, CK_ROLE_INPUT, CK_ROLE_OUTPUT, and CK_ROLE_SCRATCH.

Referenced by emit_body_fields(), emit_body_values(), emit_footer_fields(), emit_footer_values(), emit_header_fields(), and emit_header_values().

is_footer_global()

static int is_footer_global ( const char *  name )

static

Definition at line 108 of file ckernel_codegen_v2_sections.c.

{
    if (!name) {
        return 0;
    }
    if (strncmp(name, "final_", 6) == 0) {
        return 1;
    }
    if (strcmp(name, "lm_head_weight") == 0) {
        return 1;
    }
    if (strcmp(name, "logits") == 0) {
        return 1;
    }
    if (strncmp(name, "d_final_", 8) == 0) {
        return 1;
    }
    if (strcmp(name, "d_logits") == 0) {
        return 1;
    }
    return 0;
}

Referenced by emit_footer_fields(), emit_footer_values(), emit_header_fields(), and emit_header_values().

plan_size()

static size_t plan_size ( const CKMemPlan *  plan, int  idx  )

static

Definition at line 216 of file ckernel_codegen_v2_sections.c.

{
    if (!plan || !plan->spans || idx < 0 || idx >= plan->num_spans) {
        return 0;
    }
    return plan->spans[idx].size_bytes;
}

References CKMemPlan::num_spans, CKMemSpan::size_bytes, and CKMemPlan::spans.

Referenced by emit_body_values(), emit_footer_values(), and emit_header_values().

resolve_dim()

static size_t resolve_dim ( const CKModelConfig *  cfg, const CKV2AlignInfo *  align, CKDimKind  kind  )

static

Definition at line 49 of file ckernel_codegen_v2_sections.c.

{
    switch (kind) {
    case CK_DIM_TOKENS:
        return cfg ? (size_t)cfg->context_window : 0;
    case CK_DIM_EMBED:
        return cfg ? (size_t)cfg->hidden_size : 0;
    case CK_DIM_ALIGNED_EMBED:
        return align ? align->aligned_embed : 0;
    case CK_DIM_HEAD_DIM:
        return (cfg && cfg->num_heads > 0) ? (size_t)(cfg->hidden_size / cfg->num_heads) : 0;
    case CK_DIM_ALIGNED_HEAD:
        return align ? align->aligned_head : 0;
    case CK_DIM_NUM_HEADS:
        return cfg ? (size_t)cfg->num_heads : 0;
    case CK_DIM_NUM_KV_HEADS:
        return cfg ? (size_t)cfg->num_kv_heads : 0;
    case CK_DIM_ALIGNED_CTX:
        return align ? align->aligned_context : 0;
    case CK_DIM_INTERMEDIATE:
        return cfg ? (size_t)cfg->intermediate_size : 0;
    case CK_DIM_ALIGNED_INTERMEDIATE:
        return align ? align->aligned_intermediate : 0;
    case CK_DIM_VOCAB:
        return cfg ? (size_t)cfg->vocab_size : 0;
    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 CKV2AlignInfo *  align, const CKDimToken *  shape  )

static

Definition at line 82 of file ckernel_codegen_v2_sections.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);
        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 buffer_bytes().