← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
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.

6{
7 (void)graph;
8 fprintf(out,
9 "static void ck_v2_dispatch_node(int node_id) {\n"
10 " /* TODO: wire kernel dispatch using node metadata. */\n"
11 " (void)node_id;\n"
12 "}\n\n");
13}

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.

6{
7 fprintf(out,
8 "/* Auto-generated runtime from CKIRV2Graph.\n"
9 " * This v2 emitter currently writes the schedule and buffer layout\n"
10 " * stubs. Kernel wiring will be layered on as IR v2 grows.\n"
11 " */\n\n");
12 fprintf(out,
13 "#include <stddef.h>\n"
14 "#include <stdint.h>\n"
15 "#include <stdio.h>\n"
16 "#include <string.h>\n\n");
17 return 0;
18}

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.

66{
67 emit_schedule_block(out, graph, "ck_v2_run_prefill", "prefill",
68 prefill_runtime ? prefill_runtime : "ck_v2_prefill_runtime");
69 emit_schedule_block(out, graph, "ck_v2_run_decode", "decode",
70 decode_runtime ? decode_runtime : "ck_v2_decode_runtime");
71 emit_schedule_block(out, graph, "ck_v2_run_backward", "backward",
72 backward_runtime ? backward_runtime : "ck_v2_backward_runtime");
73
74 fprintf(out,
75 "void ck_v2_run_forward(void) {\n"
76 " ck_v2_run_prefill();\n"
77 "}\n\n");
78}
static void emit_schedule_block(FILE *out, const CKIRV2Graph *graph, const char *func_name, const char *label, const char *runtime_sym)

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.

314{
315 if (!out || !graph) {
316 return;
317 }
318
319 const int L = graph->config.num_layers;
320
321 fprintf(out,
322 "typedef struct {\n"
323 " size_t offset;\n"
324 " size_t size;\n"
325 "} CKV2Span;\n\n");
326
327 fprintf(out, "typedef struct {\n");
328 emit_header_fields(out, graph, CK_ROLE_WEIGHT, 0);
329 fprintf(out, "} CKV2HeaderWeights;\n\n");
330
331 fprintf(out, "typedef struct {\n");
332 emit_body_fields(out, graph, CK_ROLE_WEIGHT, 0);
333 fprintf(out, "} CKV2LayerWeights;\n\n");
334
335 fprintf(out, "typedef struct {\n");
336 emit_footer_fields(out, graph, CK_ROLE_WEIGHT, 0);
337 fprintf(out, "} CKV2FooterWeights;\n\n");
338
339 fprintf(out, "typedef struct {\n");
341 fprintf(out, "} CKV2HeaderActivations;\n\n");
342
343 fprintf(out, "typedef struct {\n");
344 emit_body_fields(out, graph, CK_ROLE_ACTIVATION, 1);
345 fprintf(out, "} CKV2LayerActivations;\n\n");
346
347 fprintf(out, "typedef struct {\n");
349 fprintf(out, "} CKV2FooterActivations;\n\n");
350
351 fprintf(out, "typedef struct {\n");
352 emit_header_fields(out, graph, CK_ROLE_GRAD, 0);
353 fprintf(out, "} CKV2HeaderGrads;\n\n");
354
355 fprintf(out, "typedef struct {\n");
356 emit_body_fields(out, graph, CK_ROLE_GRAD, 0);
357 fprintf(out, "} CKV2LayerGrads;\n\n");
358
359 fprintf(out, "typedef struct {\n");
360 emit_footer_fields(out, graph, CK_ROLE_GRAD, 0);
361 fprintf(out, "} CKV2FooterGrads;\n\n");
362
363 fprintf(out,
364 "typedef struct {\n"
365 " CKV2HeaderWeights header;\n"
366 " CKV2LayerWeights body;\n"
367 " CKV2FooterWeights footer;\n"
368 " size_t layer_stride_bytes;\n"
369 " size_t total_bytes;\n"
370 "} CKV2WeightLayout;\n\n");
371
372 fprintf(out,
373 "typedef struct {\n"
374 " CKV2HeaderActivations header;\n"
375 " CKV2LayerActivations body;\n"
376 " CKV2FooterActivations footer;\n"
377 " size_t layer_stride_bytes;\n"
378 " size_t total_bytes;\n"
379 "} CKV2ActivationLayout;\n\n");
380
381 fprintf(out,
382 "typedef struct {\n"
383 " CKV2HeaderGrads header;\n"
384 " CKV2LayerGrads body;\n"
385 " CKV2FooterGrads footer;\n"
386 " size_t layer_stride_bytes;\n"
387 " size_t total_bytes;\n"
388 "} CKV2GradLayout;\n\n");
389
390 fprintf(out,
391 "typedef struct {\n"
392 " int num_layers;\n"
393 " int context_window;\n"
394 " int hidden_size;\n"
395 " int intermediate_size;\n"
396 " int num_heads;\n"
397 " int num_kv_heads;\n"
398 " int head_dim;\n"
399 " int vocab_size;\n"
400 "} CKV2ModelConfig;\n\n");
401
402 fprintf(out,
403 "typedef struct {\n"
404 " CKV2WeightLayout weights;\n"
405 " CKV2ActivationLayout activations;\n"
406 " CKV2GradLayout grads;\n"
407 "} CKV2SectionLayout;\n\n");
408
409 fprintf(out,
410 "typedef struct {\n"
411 " CKV2SectionLayout prefill;\n"
412 " CKV2SectionLayout decode;\n"
413 " CKV2SectionLayout backward;\n"
414 "} CKV2SectionLayouts;\n\n");
415
416 fprintf(out,
417 "typedef struct {\n"
418 " CKV2ModelConfig config;\n"
419 " CKV2SectionLayouts decoder;\n"
420 "} CKV2RuntimeLayout;\n\n");
421
422 fprintf(out, "static const CKV2RuntimeLayout ck_v2_layout = {\n");
423 fprintf(out, " .config = {\n");
424 fprintf(out, " .num_layers = %d,\n", graph->config.num_layers);
425 fprintf(out, " .context_window = %d,\n", graph->config.context_window);
426 fprintf(out, " .hidden_size = %d,\n", graph->config.hidden_size);
427 fprintf(out, " .intermediate_size = %d,\n", graph->config.intermediate_size);
428 fprintf(out, " .num_heads = %d,\n", graph->config.num_heads);
429 fprintf(out, " .num_kv_heads = %d,\n", graph->config.num_kv_heads);
430 fprintf(out, " .head_dim = %d,\n",
431 graph->config.num_heads > 0 ? graph->config.hidden_size / graph->config.num_heads : 0);
432 fprintf(out, " .vocab_size = %d\n", graph->config.vocab_size);
433 fprintf(out, " },\n");
434 fprintf(out, " .decoder = {\n");
435
436 const CKMemPlan *plans[3] = { prefill_plan, decode_plan, backward_plan };
437 const char *modes[3] = { "prefill", "decode", "backward" };
438
439 for (int mode = 0; mode < 3; ++mode) {
440 const CKMemPlan *plan = plans[mode];
441 fprintf(out, " .%s = {\n", modes[mode]);
442
443 size_t offset = 0;
444 size_t header_end = 0;
445 size_t layer_stride = 0;
446 size_t footer_base = 0;
447
448 fprintf(out, " .weights = {\n");
449 fprintf(out, " .header = {\n");
450 emit_header_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0, &offset);
451 header_end = offset;
452 fprintf(out, " },\n");
453
454 fprintf(out, " .body = {\n");
455 layer_stride = emit_body_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0);
456 fprintf(out, " },\n");
457
458 footer_base = header_end + layer_stride * (size_t)L;
459 offset = footer_base;
460 fprintf(out, " .footer = {\n");
461 emit_footer_values(out, graph, prefill_plan, CK_ROLE_WEIGHT, 0, &offset);
462 fprintf(out, " },\n");
463
464 fprintf(out, " .layer_stride_bytes = %zu,\n", layer_stride);
465 fprintf(out, " .total_bytes = %zu\n", offset);
466 fprintf(out, " },\n");
467
468 offset = 0;
469 header_end = 0;
470 layer_stride = 0;
471 footer_base = 0;
472
473 fprintf(out, " .activations = {\n");
474 fprintf(out, " .header = {\n");
475 emit_header_values(out, graph, plan, CK_ROLE_ACTIVATION, 1, &offset);
476 header_end = offset;
477 fprintf(out, " },\n");
478
479 fprintf(out, " .body = {\n");
480 layer_stride = emit_body_values(out, graph, plan, CK_ROLE_ACTIVATION, 1);
481 fprintf(out, " },\n");
482
483 footer_base = header_end + layer_stride * (size_t)L;
484 offset = footer_base;
485 fprintf(out, " .footer = {\n");
486 emit_footer_values(out, graph, plan, CK_ROLE_ACTIVATION, 1, &offset);
487 fprintf(out, " },\n");
488
489 fprintf(out, " .layer_stride_bytes = %zu,\n", layer_stride);
490 fprintf(out, " .total_bytes = %zu\n", offset);
491 fprintf(out, " },\n");
492
493 offset = 0;
494 header_end = 0;
495 layer_stride = 0;
496 footer_base = 0;
497
498 fprintf(out, " .grads = {\n");
499 fprintf(out, " .header = {\n");
500 emit_header_values(out, graph, plan, CK_ROLE_GRAD, 0, &offset);
501 header_end = offset;
502 fprintf(out, " },\n");
503
504 fprintf(out, " .body = {\n");
505 layer_stride = emit_body_values(out, graph, plan, CK_ROLE_GRAD, 0);
506 fprintf(out, " },\n");
507
508 footer_base = header_end + layer_stride * (size_t)L;
509 offset = footer_base;
510 fprintf(out, " .footer = {\n");
511 emit_footer_values(out, graph, plan, CK_ROLE_GRAD, 0, &offset);
512 fprintf(out, " },\n");
513
514 fprintf(out, " .layer_stride_bytes = %zu,\n", layer_stride);
515 fprintf(out, " .total_bytes = %zu\n", offset);
516 fprintf(out, " }\n");
517
518 fprintf(out, " }%s\n", mode == 2 ? "" : ",");
519 }
520
521 fprintf(out, " }\n");
522 fprintf(out, "};\n\n");
523}
static void emit_footer_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_header_values(FILE *out, const CKIRV2Graph *graph, const CKMemPlan *plan, CKBufferRole role_filter, int activation_group, size_t *offset)
static void emit_body_fields(FILE *out, const CKIRV2Graph *graph, CKBufferRole role_filter, int activation_group)
static void emit_header_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)
@ CK_ROLE_WEIGHT
@ CK_ROLE_GRAD
@ CK_ROLE_ACTIVATION
CKModelConfig config

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.

24{
25 const char *suffix = tag ? tag : "prefill";
26 fprintf(out, "typedef struct {\n"
27 " size_t offset_bytes;\n"
28 " size_t size_bytes;\n"
29 " int arena;\n"
30 "} CKV2BufferLayout;\n\n");
31
32 fprintf(out, "static const CKV2BufferLayout ck_v2_%s_buffers[] = {\n", suffix);
33 for (int i = 0; i < graph->num_buffers; ++i) {
34 const CKMemSpan *span = &plan->spans[i];
35 fprintf(out,
36 " { %zu, %zu, %d }, /* %s */\n",
37 span->offset_bytes,
38 span->size_bytes,
39 (int)span->arena,
40 graph->buffers[i].name ? graph->buffers[i].name : "unnamed");
41 }
42 fprintf(out, "};\n\n");
43
44 fprintf(out, "static const size_t ck_v2_%s_total_bytes[] = {\n", suffix);
45 for (int i = 0; i < CK_MEM_ARENA_COUNT; ++i) {
46 fprintf(out, " %zu%s\n",
47 plan->total_bytes[i],
48 (i + 1 == CK_MEM_ARENA_COUNT) ? "" : ",");
49 }
50 fprintf(out, "};\n\n");
51
52 fprintf(out, "typedef struct {\n"
53 " const CKV2BufferLayout *buffers;\n"
54 " int num_buffers;\n"
55 " const size_t *total_bytes;\n"
56 " size_t alignment_bytes;\n"
57 "} CKV2Runtime;\n\n");
58
59 fprintf(out, "static const CKV2Runtime ck_v2_%s_runtime = {\n"
60 " ck_v2_%s_buffers,\n"
61 " %d,\n"
62 " ck_v2_%s_total_bytes,\n"
63 " %d\n"
64 "};\n\n",
65 suffix,
66 suffix,
67 graph->num_buffers,
68 suffix,
70}
@ CK_MEM_ARENA_COUNT
#define CK_MEM_PLAN_DEFAULT_ALIGN
CKIRV2Buffer * buffers
CKMemSpan * spans
size_t total_bytes[CK_MEM_ARENA_COUNT]
size_t size_bytes
CKMemArenaKind arena
size_t offset_bytes

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