ckernel_codegen_v2_emit.h File Reference

#include "ckernel_ir_v2.h"
#include "ckernel_mem_plan.h"
#include <stdio.h>

Go to the source code of this file.

Functions
void	ck_codegen_v2_emit_dispatch (FILE *out, const CKIRV2Graph *graph)
int	ck_codegen_v2_emit_preamble (FILE *out)
void	ck_codegen_v2_emit_schedule (FILE *out, const CKIRV2Graph *graph, const char *prefill_runtime, const char *decode_runtime, const char *backward_runtime)
void	ck_codegen_v2_emit_sections (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *prefill_plan, const CKMemPlan *decode_plan, const CKMemPlan *backward_plan)
void	ck_codegen_v2_emit_struct (FILE *out, const CKIRV2Graph *graph, const CKMemPlan *plan, const char *tag)

Function Documentation

ck_codegen_v2_emit_dispatch()

void ck_codegen_v2_emit_dispatch	(	FILE *	out,
		const CKIRV2Graph *	graph
	)

Definition at line 5 of file ckernel_codegen_v2_dispatch.c.

{
    (void)graph;
    fprintf(out,
            "static void ck_v2_dispatch_node(int node_id) {\n"
            "    /* TODO: wire kernel dispatch using node metadata. */\n"
            "    (void)node_id;\n"
            "}\n\n");
}

Referenced by ck_codegen_v2_emit_runtime().

ck_codegen_v2_emit_preamble()

int ck_codegen_v2_emit_preamble

(

FILE *

out

)

Definition at line 5 of file ckernel_codegen_v2_struct.c.

{
    fprintf(out,
            "/* Auto-generated runtime from CKIRV2Graph.\n"
            " * This v2 emitter currently writes the schedule and buffer layout\n"
            " * stubs. Kernel wiring will be layered on as IR v2 grows.\n"
            " */\n\n");
    fprintf(out,
            "#include <stddef.h>\n"
            "#include <stdint.h>\n"
            "#include <stdio.h>\n"
            "#include <string.h>\n\n");
    return 0;
}

Referenced by ck_codegen_v2_emit_runtime().

ck_codegen_v2_emit_schedule()

void ck_codegen_v2_emit_schedule	(	FILE *	out,
		const CKIRV2Graph *	graph,
		const char *	prefill_runtime,
		const char *	decode_runtime,
		const char *	backward_runtime
	)

Definition at line 61 of file ckernel_codegen_v2_schedule.c.

{
    emit_schedule_block(out, graph, "ck_v2_run_prefill", "prefill",
                        prefill_runtime ? prefill_runtime : "ck_v2_prefill_runtime");
    emit_schedule_block(out, graph, "ck_v2_run_decode", "decode",
                        decode_runtime ? decode_runtime : "ck_v2_decode_runtime");
    emit_schedule_block(out, graph, "ck_v2_run_backward", "backward",
                        backward_runtime ? backward_runtime : "ck_v2_backward_runtime");
 
    fprintf(out,
            "void ck_v2_run_forward(void) {\n"
            "    ck_v2_run_prefill();\n"
            "}\n\n");
}

References emit_schedule_block().

Referenced by ck_codegen_v2_emit_runtime().

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().

ck_codegen_v2_emit_struct()

void ck_codegen_v2_emit_struct	(	FILE *	out,
		const CKIRV2Graph *	graph,
		const CKMemPlan *	plan,
		const char *	tag
	)

Definition at line 20 of file ckernel_codegen_v2_struct.c.

{
    const char *suffix = tag ? tag : "prefill";
    fprintf(out, "typedef struct {\n"
                 "    size_t offset_bytes;\n"
                 "    size_t size_bytes;\n"
                 "    int arena;\n"
                 "} CKV2BufferLayout;\n\n");
 
    fprintf(out, "static const CKV2BufferLayout ck_v2_%s_buffers[] = {\n", suffix);
    for (int i = 0; i < graph->num_buffers; ++i) {
        const CKMemSpan *span = &plan->spans[i];
        fprintf(out,
                "    { %zu, %zu, %d }, /* %s */\n",
                span->offset_bytes,
                span->size_bytes,
                (int)span->arena,
                graph->buffers[i].name ? graph->buffers[i].name : "unnamed");
    }
    fprintf(out, "};\n\n");
 
    fprintf(out, "static const size_t ck_v2_%s_total_bytes[] = {\n", suffix);
    for (int i = 0; i < CK_MEM_ARENA_COUNT; ++i) {
        fprintf(out, "    %zu%s\n",
                plan->total_bytes[i],
                (i + 1 == CK_MEM_ARENA_COUNT) ? "" : ",");
    }
    fprintf(out, "};\n\n");
 
    fprintf(out, "typedef struct {\n"
                 "    const CKV2BufferLayout *buffers;\n"
                 "    int num_buffers;\n"
                 "    const size_t *total_bytes;\n"
                 "    size_t alignment_bytes;\n"
                 "} CKV2Runtime;\n\n");
 
    fprintf(out, "static const CKV2Runtime ck_v2_%s_runtime = {\n"
                 "    ck_v2_%s_buffers,\n"
                 "    %d,\n"
                 "    ck_v2_%s_total_bytes,\n"
                 "    %d\n"
                 "};\n\n",
            suffix,
            suffix,
            graph->num_buffers,
            suffix,
            CK_MEM_PLAN_DEFAULT_ALIGN);
}

References CKMemSpan::arena, CKIRV2Graph::buffers, CK_MEM_ARENA_COUNT, CK_MEM_PLAN_DEFAULT_ALIGN, CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKMemSpan::offset_bytes, CKMemSpan::size_bytes, CKMemPlan::spans, and CKMemPlan::total_bytes.

Referenced by ck_codegen_v2_emit_runtime().