#include "ckernel_ir_v2.h"
#include "ckernel_mem_plan.h"
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

Functions
static size_t	ck_ir_v2_align_up_bytes (size_t n, size_t align)

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

static CKDimKind	ck_ir_v2_dim_kind_from_name (const char *name)

static const char *	ck_ir_v2_dim_name (CKDimKind dim)

static const char *	ck_ir_v2_dtype_name (CKDataType dtype)

static void	ck_ir_v2_emit_dimensions (FILE out, const CKModelConfig cfg, const CKIRV2AlignInfo *align, int tokens_override)

static void	ck_ir_v2_emit_memory_plan (FILE out, const CKIRV2Graph graph, const CKMemPlan *plan)

static void	ck_ir_v2_emit_resolved_shape (FILE out, const CKModelConfig cfg, const CKIRV2AlignInfo align, const CKDimToken shape, int tokens_override)

static int	ck_ir_v2_emit_shape (FILE out, const CKDimToken shape)

static const char *	ck_ir_v2_find_array_end (const char open, const char end)

static int	ck_ir_v2_find_buffer_index (const CKIRV2Graph graph, const char name)

static const CKBufferSpec *	ck_ir_v2_find_buffer_spec (const char *name)

static const char *	ck_ir_v2_find_key (const char json, const char key, const char *end)

void	ck_ir_v2_free (CKIRV2Graph *graph)

static void	ck_ir_v2_free_buffer (CKIRV2Buffer *buf)

static void	ck_ir_v2_free_node (CKIRV2Node *node)

static const char *	ck_ir_v2_mem_arena_name (CKMemArenaKind arena)

static const char *	ck_ir_v2_next_object (const char cur, const char end, const char obj_start, const char obj_end)

static int	ck_ir_v2_parse_bindings (const char obj_start, const char obj_end, CKIRV2Graph graph, CKIRV2Node node)

static int	ck_ir_v2_parse_bool (const char json, const char key, const char end, int out_val)

static int	ck_ir_v2_parse_buffers (const char json, const char end, CKIRV2Graph *graph)

static CKDimKind	ck_ir_v2_parse_dim_kind (const char obj_start, const char obj_end)

static CKDataType	ck_ir_v2_parse_dtype (const char *s)

static int	ck_ir_v2_parse_float (const char json, const char key, const char end, float out_val)

static int	ck_ir_v2_parse_int (const char json, const char key, const char end, int out_val)

int	ck_ir_v2_parse_json (const char path, CKIRV2Graph graph)

static int	ck_ir_v2_parse_nodes (const char json, const char end, CKIRV2Graph *graph)

static CKBufferRole	ck_ir_v2_parse_role (const char *s)

static CKBufferScope	ck_ir_v2_parse_scope (const char *s)

static int	ck_ir_v2_parse_shape (const char obj_start, const char obj_end, CKDimToken *shape_out)

static int	ck_ir_v2_parse_string (const char start, const char end, char **out_str)

static int	ck_ir_v2_parse_string_field (const char json, const char key, const char end, char *out_str)

static void	ck_ir_v2_resolve_align (const CKModelConfig cfg, size_t alignment_bytes, CKIRV2AlignInfo align)

static size_t	ck_ir_v2_resolve_dim_value (const CKModelConfig cfg, const CKIRV2AlignInfo align, CKDimKind dim, int tokens_override)

static const char *	ck_ir_v2_role_name (CKBufferRole role)

static const char *	ck_ir_v2_scope_name (CKBufferScope scope)

int	ck_ir_v2_serialize_json (const CKIRV2Graph graph, const char path)

static int	ck_ir_v2_serialize_json_internal (const CKIRV2Graph graph, const CKMemPlan plan, const char mode, int tokens_override, int base_context_window, const char path)

int	ck_ir_v2_serialize_json_with_plan (const CKIRV2Graph graph, const CKMemPlan plan, const char mode, int tokens_override, int base_context_window, const char path)

static const char *	ck_ir_v2_skip_string (const char cur, const char end)

Function Documentation

◆ ck_ir_v2_align_up_bytes()

static size_t ck_ir_v2_align_up_bytes	(	size_t	n,
		size_t	align
	)

static

Definition at line 164 of file ckernel_ir_v2.c.

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

Referenced by ck_ir_v2_align_up_elems().

◆ ck_ir_v2_align_up_elems()

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

static

Definition at line 172 of file ckernel_ir_v2.c.

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

References ck_ir_v2_align_up_bytes().

Referenced by ck_ir_v2_resolve_align().

◆ ck_ir_v2_dim_kind_from_name()

static CKDimKind ck_ir_v2_dim_kind_from_name ( const char * name )

static

Definition at line 138 of file ckernel_ir_v2.c.

 {
     if (!name) {
         return CK_DIM_END;
     }
     if (strcmp(name, "tokens") == 0) return CK_DIM_TOKENS;
     if (strcmp(name, "embed") == 0) return CK_DIM_EMBED;
     if (strcmp(name, "aligned_embed") == 0) return CK_DIM_ALIGNED_EMBED;
     if (strcmp(name, "head_dim") == 0) return CK_DIM_HEAD_DIM;
     if (strcmp(name, "aligned_head") == 0) return CK_DIM_ALIGNED_HEAD;
     if (strcmp(name, "num_heads") == 0) return CK_DIM_NUM_HEADS;
     if (strcmp(name, "num_kv_heads") == 0) return CK_DIM_NUM_KV_HEADS;
     if (strcmp(name, "aligned_ctx") == 0) return CK_DIM_ALIGNED_CTX;
     if (strcmp(name, "intermediate") == 0) return CK_DIM_INTERMEDIATE;
     if (strcmp(name, "aligned_intermediate") == 0) return CK_DIM_ALIGNED_INTERMEDIATE;
     if (strcmp(name, "vocab") == 0) return CK_DIM_VOCAB;
     return CK_DIM_END;
 }

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, and CK_DIM_VOCAB.

Referenced by ck_ir_v2_parse_dim_kind().

◆ ck_ir_v2_dim_name()

static const char* ck_ir_v2_dim_name ( CKDimKind dim )

static

Definition at line 108 of file ckernel_ir_v2.c.

 {
     switch (dim) {
     case CK_DIM_TOKENS:
         return "tokens";
     case CK_DIM_EMBED:
         return "embed";
     case CK_DIM_ALIGNED_EMBED:
         return "aligned_embed";
     case CK_DIM_HEAD_DIM:
         return "head_dim";
     case CK_DIM_ALIGNED_HEAD:
         return "aligned_head";
     case CK_DIM_NUM_HEADS:
         return "num_heads";
     case CK_DIM_NUM_KV_HEADS:
         return "num_kv_heads";
     case CK_DIM_ALIGNED_CTX:
         return "aligned_ctx";
     case CK_DIM_INTERMEDIATE:
         return "intermediate";
     case CK_DIM_ALIGNED_INTERMEDIATE:
         return "aligned_intermediate";
     case CK_DIM_VOCAB:
         return "vocab";
     default:
         return "unknown";
     }
 }

References CK_DIM_ALIGNED_CTX, CK_DIM_ALIGNED_EMBED, CK_DIM_ALIGNED_HEAD, CK_DIM_ALIGNED_INTERMEDIATE, CK_DIM_EMBED, CK_DIM_HEAD_DIM, CK_DIM_INTERMEDIATE, CK_DIM_NUM_HEADS, CK_DIM_NUM_KV_HEADS, CK_DIM_TOKENS, and CK_DIM_VOCAB.

Referenced by ck_ir_v2_emit_dimensions(), and ck_ir_v2_emit_shape().

◆ ck_ir_v2_dtype_name()

static const char* ck_ir_v2_dtype_name ( CKDataType dtype )

static

Definition at line 86 of file ckernel_ir_v2.c.

 {
     switch (dtype) {
     case CK_DT_FP32:
         return "fp32";
     case CK_DT_BF16:
         return "bf16";
     case CK_DT_FP16:
         return "fp16";
     case CK_DT_Q4_0:
         return "q4_0";
     case CK_DT_Q4_K:
         return "q4_k";
     case CK_DT_Q6_K:
         return "q6_k";
     case CK_DT_Q8_0:
         return "q8_0";
     default:
         return "unknown";
     }
 }

References CK_DT_BF16, CK_DT_FP16, CK_DT_FP32, CK_DT_Q4_0, CK_DT_Q4_K, CK_DT_Q6_K, and CK_DT_Q8_0.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_emit_dimensions()

static void ck_ir_v2_emit_dimensions	(	FILE *	out,
		const CKModelConfig *	cfg,
		const CKIRV2AlignInfo *	align,
		int	tokens_override
	)

static

Definition at line 286 of file ckernel_ir_v2.c.

 {
     fprintf(out, "  \"dimensions\": [\n");
     CKDimKind dims[] = {
         CK_DIM_TOKENS,
         CK_DIM_EMBED,
         CK_DIM_ALIGNED_EMBED,
         CK_DIM_HEAD_DIM,
         CK_DIM_ALIGNED_HEAD,
         CK_DIM_NUM_HEADS,
         CK_DIM_NUM_KV_HEADS,
         CK_DIM_ALIGNED_CTX,
         CK_DIM_INTERMEDIATE,
         CK_DIM_ALIGNED_INTERMEDIATE,
         CK_DIM_VOCAB
     };
     size_t dim_count = sizeof(dims) / sizeof(dims[0]);
     for (size_t i = 0; i < dim_count; ++i) {
         CKDimKind dim = dims[i];
         size_t value = ck_ir_v2_resolve_dim_value(cfg, align, dim, tokens_override);
         fprintf(out,
                 "    {\"id\": %d, \"name\": \"%s\", \"value\": %zu}%s\n",
                 (int)dim,
                 ck_ir_v2_dim_name(dim),
                 value,
                 (i + 1 == dim_count) ? "" : ",");
     }
     fprintf(out, "  ],\n");
 }

References CK_DIM_ALIGNED_CTX, CK_DIM_ALIGNED_EMBED, CK_DIM_ALIGNED_HEAD, CK_DIM_ALIGNED_INTERMEDIATE, CK_DIM_EMBED, CK_DIM_HEAD_DIM, CK_DIM_INTERMEDIATE, CK_DIM_NUM_HEADS, CK_DIM_NUM_KV_HEADS, CK_DIM_TOKENS, CK_DIM_VOCAB, ck_ir_v2_dim_name(), and ck_ir_v2_resolve_dim_value().

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_emit_memory_plan()

static void ck_ir_v2_emit_memory_plan	(	FILE *	out,
		const CKIRV2Graph *	graph,
		const CKMemPlan *	plan
	)

static

Definition at line 333 of file ckernel_ir_v2.c.

 {
     if (!plan) {
         return;
     }
     fprintf(out, "  \"memory_plan\": {\n");
     fprintf(out, "    \"alignment_bytes\": %zu,\n", plan->alignment_bytes);
     fprintf(out, "    \"total_bytes\": {\n");
     fprintf(out, "      \"weights\": %zu,\n", plan->total_bytes[CK_MEM_ARENA_WEIGHTS]);
     fprintf(out, "      \"activations\": %zu,\n", plan->total_bytes[CK_MEM_ARENA_ACTIVATIONS]);
     fprintf(out, "      \"grads\": %zu\n", plan->total_bytes[CK_MEM_ARENA_GRADS]);
     fprintf(out, "    },\n");
     fprintf(out, "    \"buffers\": [\n");
     for (int i = 0; i < graph->num_buffers; ++i) {
         const CKMemSpan *span = &plan->spans[i];
         const char *name = graph->buffers[i].name ? graph->buffers[i].name : "";
         int enabled = span->size_bytes > 0;
         fprintf(out,
                 "      {\"name\": \"%s\", \"arena\": \"%s\", \"offset_bytes\": %zu, \"size_bytes\": %zu, \"enabled\": %s}%s\n",
                 name,
                 ck_ir_v2_mem_arena_name(span->arena),
                 span->offset_bytes,
                 span->size_bytes,
                 enabled ? "true" : "false",
                 (i + 1 == graph->num_buffers) ? "" : ",");
     }
     fprintf(out, "    ]\n");
     fprintf(out, "  },\n");
 }

References CKMemPlan::alignment_bytes, CKMemSpan::arena, CKIRV2Graph::buffers, ck_ir_v2_mem_arena_name(), CK_MEM_ARENA_ACTIVATIONS, CK_MEM_ARENA_GRADS, CK_MEM_ARENA_WEIGHTS, CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKMemSpan::offset_bytes, CKMemSpan::size_bytes, CKMemPlan::spans, and CKMemPlan::total_bytes.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_emit_resolved_shape()

static void ck_ir_v2_emit_resolved_shape	(	FILE *	out,
		const CKModelConfig *	cfg,
		const CKIRV2AlignInfo *	align,
		const CKDimToken *	shape,
		int	tokens_override
	)

static

Definition at line 261 of file ckernel_ir_v2.c.

 {
     fprintf(out, "[");
     int first = 1;
     for (int i = 0; i < CK_IR_V2_MAX_DIMS; ++i) {
         if (shape[i].dim == CK_DIM_END) {
             break;
         }
         size_t dim = ck_ir_v2_resolve_dim_value(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);
         size_t resolved = (div == 0) ? 0 : (dim * mult / div);
         if (!first) {
             fprintf(out, ", ");
         }
         fprintf(out, "%zu", resolved);
         first = 0;
     }
     fprintf(out, "]");
 }

References CK_DIM_END, CK_IR_V2_MAX_DIMS, and ck_ir_v2_resolve_dim_value().

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_emit_shape()

static int ck_ir_v2_emit_shape	(	FILE *	out,
		const CKDimToken *	shape
	)

static

Definition at line 239 of file ckernel_ir_v2.c.

 {
     fprintf(out, "[");
     int first = 1;
     for (int i = 0; i < CK_IR_V2_MAX_DIMS; ++i) {
         if (shape[i].dim == CK_DIM_END) {
             break;
         }
         if (!first) {
             fprintf(out, ", ");
         }
         fprintf(out, "{\"dim\":\"%s\",\"dim_id\":%d,\"mult\":%d,\"div\":%d}",
                 ck_ir_v2_dim_name(shape[i].dim),
                 (int)shape[i].dim,
                 shape[i].mult,
                 shape[i].div);
         first = 0;
     }
     fprintf(out, "]");
     return 0;
 }

References CK_DIM_END, ck_ir_v2_dim_name(), and CK_IR_V2_MAX_DIMS.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_find_array_end()

static const char* ck_ir_v2_find_array_end	(	const char *	open,
		const char *	end
	)

static

Definition at line 610 of file ckernel_ir_v2.c.

 {
     if (!open || open >= end || *open != '[') {
         return NULL;
     }
     int depth = 0;
     for (const char *p = open; p < end; ++p) {
         if (*p == '"') {
             p = ck_ir_v2_skip_string(p, end) - 1;
             continue;
         }
         if (*p == '[') {
             depth++;
             continue;
         }
         if (*p == ']') {
             if (depth > 0) {
                 depth--;
                 if (depth == 0) {
                     return p;
                 }
             }
         }
     }
     return NULL;
 }

References ck_ir_v2_skip_string(), and end.

Referenced by ck_ir_v2_parse_bindings(), ck_ir_v2_parse_buffers(), ck_ir_v2_parse_nodes(), and ck_ir_v2_parse_shape().

◆ ck_ir_v2_find_buffer_index()

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

static

Definition at line 868 of file ckernel_ir_v2.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 ck_ir_v2_parse_bindings().

◆ ck_ir_v2_find_buffer_spec()

static const CKBufferSpec* ck_ir_v2_find_buffer_spec ( const char * name )

static

Definition at line 855 of file ckernel_ir_v2.c.

 {
     if (!name) {
         return NULL;
     }
     for (size_t i = 0; i < ck_decoder_buffer_count; ++i) {
         if (strcmp(ck_decoder_buffers[i].name, name) == 0) {
             return &ck_decoder_buffers[i];
         }
     }
     return NULL;
 }

References ck_decoder_buffer_count, and ck_decoder_buffers.

Referenced by ck_ir_v2_parse_buffers().

◆ ck_ir_v2_find_key()

static const char* ck_ir_v2_find_key	(	const char *	json,
		const char *	key,
		const char *	end
	)

static

Definition at line 637 of file ckernel_ir_v2.c.

 {
     size_t key_len = strlen(key);
     const char *cur = json;
     while (cur + key_len < end) {
         if (memcmp(cur, key, key_len) == 0) {
             return cur;
         }
         cur++;
     }
     return NULL;
 }

References end.

Referenced by ck_ir_v2_parse_bindings(), ck_ir_v2_parse_bool(), ck_ir_v2_parse_buffers(), ck_ir_v2_parse_float(), ck_ir_v2_parse_int(), ck_ir_v2_parse_json(), ck_ir_v2_parse_nodes(), ck_ir_v2_parse_shape(), and ck_ir_v2_parse_string_field().

◆ ck_ir_v2_free()

void ck_ir_v2_free ( CKIRV2Graph * graph )

Definition at line 34 of file ckernel_ir_v2.c.

 {
     if (!graph) {
         return;
     }
     if (graph->buffers) {
         for (int i = 0; i < graph->num_buffers; ++i) {
             ck_ir_v2_free_buffer(&graph->buffers[i]);
         }
         free(graph->buffers);
     }
     if (graph->nodes) {
         for (int i = 0; i < graph->num_nodes; ++i) {
             ck_ir_v2_free_node(&graph->nodes[i]);
         }
         free(graph->nodes);
     }
     memset(graph, 0, sizeof(*graph));
 }

References CKIRV2Graph::buffers, ck_ir_v2_free_buffer(), ck_ir_v2_free_node(), CKIRV2Graph::nodes, CKIRV2Graph::num_buffers, and CKIRV2Graph::num_nodes.

Referenced by ck_ir_v2_build_decoder(), ck_ir_v2_build_decoder_backward(), ck_ir_v2_lower_emit_json(), ck_ir_v2_lower_graph(), ck_ir_v2_parse_json(), and main().

◆ ck_ir_v2_free_buffer()

static void ck_ir_v2_free_buffer ( CKIRV2Buffer * buf )

static

Definition at line 9 of file ckernel_ir_v2.c.

 {
     if (!buf) {
         return;
     }
     free(buf->name);
     free(buf->alias_of);
     free(buf->condition);
     memset(buf, 0, sizeof(*buf));
 }

References CKIRV2Buffer::alias_of, CKIRV2Buffer::condition, and CKIRV2Buffer::name.

Referenced by ck_ir_v2_free().

◆ ck_ir_v2_free_node()

static void ck_ir_v2_free_node ( CKIRV2Node * node )

static

Definition at line 20 of file ckernel_ir_v2.c.

 {
     if (!node) {
         return;
     }
     for (int i = 0; i < node->n_bindings; ++i) {
         free(node->bindings[i].arg);
     }
     free(node->op);
     free(node->kernel);
     free(node->condition);
     memset(node, 0, sizeof(*node));
 }

References CKIRV2Binding::arg, CKIRV2Node::bindings, CKIRV2Node::condition, CKIRV2Node::kernel, CKIRV2Node::n_bindings, and CKIRV2Node::op.

Referenced by ck_ir_v2_free(), and ck_ir_v2_parse_nodes().

◆ ck_ir_v2_mem_arena_name()

static const char* ck_ir_v2_mem_arena_name ( CKMemArenaKind arena )

static

Definition at line 319 of file ckernel_ir_v2.c.

 {
     switch (arena) {
     case CK_MEM_ARENA_WEIGHTS:
         return "weights";
     case CK_MEM_ARENA_ACTIVATIONS:
         return "activations";
     case CK_MEM_ARENA_GRADS:
         return "grads";
     default:
         return "unknown";
     }
 }

References CK_MEM_ARENA_ACTIVATIONS, CK_MEM_ARENA_GRADS, and CK_MEM_ARENA_WEIGHTS.

Referenced by ck_ir_v2_emit_memory_plan().

◆ ck_ir_v2_next_object()

static const char* ck_ir_v2_next_object	(	const char *	cur,
		const char *	end,
		const char **	obj_start,
		const char **	obj_end
	)

static

Definition at line 577 of file ckernel_ir_v2.c.

 {
     int depth = 0;
     const char *start = NULL;
     for (const char *p = cur; p < end; ++p) {
         if (*p == '"') {
             p = ck_ir_v2_skip_string(p, end) - 1;
             continue;
         }
         if (*p == '{') {
             if (depth == 0) {
                 start = p;
             }
             depth++;
             continue;
         }
         if (*p == '}') {
             if (depth > 0) {
                 depth--;
                 if (depth == 0 && start) {
                     *obj_start = start;
                     *obj_end = p;
                     return p + 1;
                 }
             }
         }
     }
     return NULL;
 }

References ck_ir_v2_skip_string(), end, and start.

Referenced by ck_ir_v2_parse_bindings(), ck_ir_v2_parse_buffers(), ck_ir_v2_parse_json(), ck_ir_v2_parse_nodes(), and ck_ir_v2_parse_shape().

◆ ck_ir_v2_parse_bindings()

static int ck_ir_v2_parse_bindings	(	const char *	obj_start,
		const char *	obj_end,
		CKIRV2Graph *	graph,
		CKIRV2Node *	node
	)

static

Definition at line 965 of file ckernel_ir_v2.c.

 {
     const char *start = ck_ir_v2_find_key(obj_start, "\"bindings\"", obj_end);
     if (!start) {
         return 0;
     }
     const char *open = strchr(start, '[');
     const char *close = ck_ir_v2_find_array_end(open, obj_end);
     if (!open || !close || close > obj_end) {
         return -1;
     }
  
     const char *cur = open;
     const char *bstart = NULL;
     const char *bend = NULL;
     while (node->n_bindings < CK_IR_V2_MAX_BINDINGS &&
            (cur = ck_ir_v2_next_object(cur, close, &bstart, &bend)) != NULL) {
  
         char *arg = NULL;
         char *buf = NULL;
         ck_ir_v2_parse_string_field(bstart, "\"arg\"", bend, &arg);
         ck_ir_v2_parse_string_field(bstart, "\"buffer\"", bend, &buf);
  
         if (arg) {
             int buf_idx = ck_ir_v2_find_buffer_index(graph, buf);
             node->bindings[node->n_bindings].arg = arg;
             node->bindings[node->n_bindings].buffer = buf_idx;
             node->n_bindings++;
         } else {
             free(arg);
         }
         free(buf);
     }
  
     return 0;
 }

References CKIRV2Binding::arg, CKIRV2Node::bindings, CKIRV2Binding::buffer, ck_ir_v2_find_array_end(), ck_ir_v2_find_buffer_index(), ck_ir_v2_find_key(), CK_IR_V2_MAX_BINDINGS, ck_ir_v2_next_object(), ck_ir_v2_parse_string_field(), CKIRV2Node::n_bindings, and start.

Referenced by ck_ir_v2_parse_nodes().

◆ ck_ir_v2_parse_bool()

static int ck_ir_v2_parse_bool	(	const char *	json,
		const char *	key,
		const char *	end,
		int *	out_val
	)

static

Definition at line 673 of file ckernel_ir_v2.c.

 {
     const char *p = ck_ir_v2_find_key(json, key, end);
     if (!p) {
         return -1;
     }
     const char *colon = strchr(p, ':');
     if (!colon || colon >= end) {
         return -1;
     }
     const char *cur = colon + 1;
     while (cur < end && (*cur == ' ' || *cur == '\t' || *cur == '\n' || *cur == '\r')) {
         cur++;
     }
     if (cur + 4 <= end && memcmp(cur, "true", 4) == 0) {
         *out_val = 1;
         return 0;
     }
     if (cur + 5 <= end && memcmp(cur, "false", 5) == 0) {
         *out_val = 0;
         return 0;
     }
     return -1;
 }

References ck_ir_v2_find_key(), and end.

Referenced by ck_ir_v2_parse_json().

◆ ck_ir_v2_parse_buffers()

static int ck_ir_v2_parse_buffers	(	const char *	json,
		const char *	end,
		CKIRV2Graph *	graph
	)

static

Definition at line 881 of file ckernel_ir_v2.c.

 {
     const char *start = ck_ir_v2_find_key(json, "\"buffers\"", end);
     if (!start) {
         return -1;
     }
     const char *open = strchr(start, '[');
     const char *close = ck_ir_v2_find_array_end(open, end);
     if (!open || !close) {
         return -1;
     }
  
     int count = 0;
     const char *cur = open;
     const char *obj_start = NULL;
     const char *obj_end = NULL;
     while ((cur = ck_ir_v2_next_object(cur, close, &obj_start, &obj_end)) != NULL) {
         count++;
     }
     if (count <= 0) {
         return -1;
     }
  
     CKIRV2Buffer *buffers = (CKIRV2Buffer *)calloc((size_t)count, sizeof(CKIRV2Buffer));
     if (!buffers) {
         return -1;
     }
     int idx = 0;
     cur = open;
     while (idx < count &&
            (cur = ck_ir_v2_next_object(cur, close, &obj_start, &obj_end)) != NULL) {
  
         CKIRV2Buffer buf = {0};
         char *name = NULL;
         char *scope = NULL;
         char *role = NULL;
         char *dtype = NULL;
         char *alias = NULL;
         char *cond = NULL;
         int optional = 0;
  
         ck_ir_v2_parse_string_field(obj_start, "\"name\"", obj_end, &name);
         ck_ir_v2_parse_string_field(obj_start, "\"scope\"", obj_end, &scope);
         ck_ir_v2_parse_string_field(obj_start, "\"role\"", obj_end, &role);
         ck_ir_v2_parse_string_field(obj_start, "\"dtype\"", obj_end, &dtype);
         ck_ir_v2_parse_string_field(obj_start, "\"alias_of\"", obj_end, &alias);
         ck_ir_v2_parse_string_field(obj_start, "\"condition\"", obj_end, &cond);
         ck_ir_v2_parse_int(obj_start, "\"optional\"", obj_end, &optional);
  
         buf.name = name;
         buf.scope = ck_ir_v2_parse_scope(scope);
         buf.role = ck_ir_v2_parse_role(role);
         buf.dtype = ck_ir_v2_parse_dtype(dtype);
         buf.optional = optional;
         buf.alias_of = alias;
         buf.condition = cond;
  
         if (ck_ir_v2_parse_shape(obj_start, obj_end, buf.shape) != 0) {
             const CKBufferSpec *spec = ck_ir_v2_find_buffer_spec(name);
             if (spec) {
                 memcpy(buf.shape, spec->shape, sizeof(buf.shape));
             } else {
                 for (int i = 0; i < CK_IR_V2_MAX_DIMS; ++i) {
                     buf.shape[i].dim = CK_DIM_END;
                     buf.shape[i].mult = 0;
                     buf.shape[i].div = 0;
                 }
             }
         }
  
         free(scope);
         free(role);
         free(dtype);
  
         buffers[idx++] = buf;
     }
  
     graph->buffers = buffers;
     graph->num_buffers = idx;
     return 0;
 }

References CKIRV2Buffer::alias_of, CKIRV2Graph::buffers, CK_DIM_END, ck_ir_v2_find_array_end(), ck_ir_v2_find_buffer_spec(), ck_ir_v2_find_key(), CK_IR_V2_MAX_DIMS, ck_ir_v2_next_object(), ck_ir_v2_parse_dtype(), ck_ir_v2_parse_int(), ck_ir_v2_parse_role(), ck_ir_v2_parse_scope(), ck_ir_v2_parse_shape(), ck_ir_v2_parse_string_field(), CKIRV2Buffer::condition, CKDimToken::dim, CKDimToken::div, CKIRV2Buffer::dtype, end, CKDimToken::mult, CKIRV2Buffer::name, CKIRV2Graph::num_buffers, CKIRV2Buffer::optional, CKIRV2Buffer::role, CKIRV2Buffer::scope, CKIRV2Buffer::shape, CKBufferSpec::shape, and start.

Referenced by ck_ir_v2_parse_json().

◆ ck_ir_v2_parse_dim_kind()

static CKDimKind ck_ir_v2_parse_dim_kind	(	const char *	obj_start,
		const char *	obj_end
	)

static

Definition at line 785 of file ckernel_ir_v2.c.

 {
     char *dim_str = NULL;
     CKDimKind kind = CK_DIM_END;
     if (ck_ir_v2_parse_string_field(obj_start, "\"dim\"", obj_end, &dim_str) == 0 &&
         dim_str) {
         kind = ck_ir_v2_dim_kind_from_name(dim_str);
     }
     if (dim_str) {
         free(dim_str);
     }
     if (kind != CK_DIM_END) {
         return kind;
     }
     int dim_id = -1;
     if (ck_ir_v2_parse_int(obj_start, "\"dim_id\"", obj_end, &dim_id) == 0) {
         return (CKDimKind)dim_id;
     }
     if (ck_ir_v2_parse_int(obj_start, "\"dim\"", obj_end, &dim_id) == 0) {
         return (CKDimKind)dim_id;
     }
     return CK_DIM_END;
 }

References CK_DIM_END, ck_ir_v2_dim_kind_from_name(), ck_ir_v2_parse_int(), and ck_ir_v2_parse_string_field().

Referenced by ck_ir_v2_parse_shape().

◆ ck_ir_v2_parse_dtype()

static CKDataType ck_ir_v2_parse_dtype ( const char * s )

static

Definition at line 772 of file ckernel_ir_v2.c.

 {
     if (!s) return CK_DT_FP32;
     if (strcmp(s, "fp32") == 0) return CK_DT_FP32;
     if (strcmp(s, "bf16") == 0) return CK_DT_BF16;
     if (strcmp(s, "fp16") == 0) return CK_DT_FP16;
     if (strcmp(s, "q4_0") == 0) return CK_DT_Q4_0;
     if (strcmp(s, "q4_k") == 0) return CK_DT_Q4_K;
     if (strcmp(s, "q6_k") == 0) return CK_DT_Q6_K;
     if (strcmp(s, "q8_0") == 0) return CK_DT_Q8_0;
     return CK_DT_FP32;
 }

References CK_DT_BF16, CK_DT_FP16, CK_DT_FP32, CK_DT_Q4_0, CK_DT_Q4_K, CK_DT_Q6_K, and CK_DT_Q8_0.

Referenced by ck_ir_v2_parse_buffers(), and ck_ir_v2_parse_nodes().

◆ ck_ir_v2_parse_float()

static int ck_ir_v2_parse_float	(	const char *	json,
		const char *	key,
		const char *	end,
		float *	out_val
	)

static

Definition at line 701 of file ckernel_ir_v2.c.

 {
     const char *p = ck_ir_v2_find_key(json, key, end);
     if (!p) {
         return -1;
     }
     const char *colon = strchr(p, ':');
     if (!colon || colon >= end) {
         return -1;
     }
     float value = 0.0f;
     if (sscanf(colon + 1, "%f", &value) != 1) {
         return -1;
     }
     *out_val = value;
     return 0;
 }

References ck_ir_v2_find_key(), and end.

Referenced by ck_ir_v2_parse_json().

◆ ck_ir_v2_parse_int()

static int ck_ir_v2_parse_int	(	const char *	json,
		const char *	key,
		const char *	end,
		int *	out_val
	)

static

Definition at line 652 of file ckernel_ir_v2.c.

 {
     const char *p = ck_ir_v2_find_key(json, key, end);
     if (!p) {
         return -1;
     }
     const char *colon = strchr(p, ':');
     if (!colon || colon >= end) {
         return -1;
     }
     int value = 0;
     if (sscanf(colon + 1, "%d", &value) != 1) {
         return -1;
     }
     *out_val = value;
     return 0;
 }

References ck_ir_v2_find_key(), and end.

Referenced by ck_ir_v2_parse_buffers(), ck_ir_v2_parse_dim_kind(), ck_ir_v2_parse_json(), ck_ir_v2_parse_nodes(), and ck_ir_v2_parse_shape().

◆ ck_ir_v2_parse_json()

int ck_ir_v2_parse_json	(	const char *	path,
		CKIRV2Graph *	graph
	)

Definition at line 1100 of file ckernel_ir_v2.c.

 {
     if (!path || !graph) {
         return -1;
     }
     FILE *f = fopen(path, "rb");
     if (!f) {
         perror("ck_ir_v2_parse_json: fopen");
         return -1;
     }
     if (fseek(f, 0, SEEK_END) != 0) {
         fclose(f);
         return -1;
     }
     long len = ftell(f);
     if (len < 0) {
         fclose(f);
         return -1;
     }
     if (fseek(f, 0, SEEK_SET) != 0) {
         fclose(f);
         return -1;
     }
     char *buf = (char *)malloc((size_t)len + 1);
     if (!buf) {
         fclose(f);
         return -1;
     }
     size_t nread = fread(buf, 1, (size_t)len, f);
     fclose(f);
     buf[nread] = '\0';
  
     CKIRV2Graph tmp = {0};
     const char *end = buf + nread;
  
     tmp.has_pos_emb = 1;
     tmp.tie_word_embeddings = -1;
     tmp.fused_qkv = -1;
     tmp.gated_mlp = -1;
  
     if (ck_ir_v2_parse_int(buf, "\"num_layers\"", end, &tmp.config.num_layers) != 0 ||
         ck_ir_v2_parse_int(buf, "\"hidden_size\"", end, &tmp.config.hidden_size) != 0 ||
         ck_ir_v2_parse_int(buf, "\"intermediate_size\"", end, &tmp.config.intermediate_size) != 0 ||
         ck_ir_v2_parse_int(buf, "\"num_attention_heads\"", end, &tmp.config.num_heads) != 0 ||
         ck_ir_v2_parse_int(buf, "\"num_kv_heads\"", end, &tmp.config.num_kv_heads) != 0) {
         free(buf);
         return -1;
     }
  
     ck_ir_v2_parse_int(buf, "\"vocab_size\"", end, &tmp.config.vocab_size);
     ck_ir_v2_parse_int(buf, "\"context_window\"", end, &tmp.config.context_window);
     ck_ir_v2_parse_float(buf, "\"rms_norm_eps\"", end, &tmp.config.rms_norm_eps);
     ck_ir_v2_parse_float(buf, "\"rope_theta\"", end, &tmp.config.rope_theta);
     const char *meta_key = ck_ir_v2_find_key(buf, "\"meta\"", end);
     if (meta_key) {
         const char *brace = strchr(meta_key, '{');
         const char *obj_start = NULL;
         const char *obj_end = NULL;
         if (brace && ck_ir_v2_next_object(brace, end, &obj_start, &obj_end)) {
             ck_ir_v2_parse_bool(obj_start, "\"has_pos_emb\"", obj_end, &tmp.has_pos_emb);
             ck_ir_v2_parse_bool(obj_start, "\"tie_word_embeddings\"", obj_end, &tmp.tie_word_embeddings);
             ck_ir_v2_parse_bool(obj_start, "\"fused_qkv\"", obj_end, &tmp.fused_qkv);
             ck_ir_v2_parse_bool(obj_start, "\"gated_mlp\"", obj_end, &tmp.gated_mlp);
         }
     }
  
     if (ck_ir_v2_parse_buffers(buf, end, &tmp) != 0 ||
         ck_ir_v2_parse_nodes(buf, end, &tmp) != 0) {
         ck_ir_v2_free(&tmp);
         free(buf);
         return -1;
     }
  
     *graph = tmp;
     free(buf);
     return 0;
 }

References ck_ir_v2_find_key(), ck_ir_v2_free(), ck_ir_v2_next_object(), ck_ir_v2_parse_bool(), ck_ir_v2_parse_buffers(), ck_ir_v2_parse_float(), ck_ir_v2_parse_int(), ck_ir_v2_parse_nodes(), CKIRV2Graph::config, CKModelConfig::context_window, end, CKIRV2Graph::fused_qkv, CKIRV2Graph::gated_mlp, CKIRV2Graph::has_pos_emb, CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, CKModelConfig::num_layers, CKModelConfig::rms_norm_eps, CKModelConfig::rope_theta, CKIRV2Graph::tie_word_embeddings, and CKModelConfig::vocab_size.

Referenced by main().

◆ ck_ir_v2_parse_nodes()

static int ck_ir_v2_parse_nodes	(	const char *	json,
		const char *	end,
		CKIRV2Graph *	graph
	)

static

Definition at line 1005 of file ckernel_ir_v2.c.

 {
     const char *start = ck_ir_v2_find_key(json, "\"nodes\"", end);
     if (!start) {
         return -1;
     }
     const char *open = strchr(start, '[');
     const char *close = ck_ir_v2_find_array_end(open, end);
     if (!open || !close) {
         return -1;
     }
  
     int count = 0;
     const char *cur = open;
     const char *obj_start = NULL;
     const char *obj_end = NULL;
     while ((cur = ck_ir_v2_next_object(cur, close, &obj_start, &obj_end)) != NULL) {
         count++;
     }
     if (count <= 0) {
         return -1;
     }
  
     CKIRV2Node *nodes = (CKIRV2Node *)calloc((size_t)count, sizeof(CKIRV2Node));
     if (!nodes) {
         return -1;
     }
     int idx = 0;
     cur = open;
     while (idx < count &&
            (cur = ck_ir_v2_next_object(cur, close, &obj_start, &obj_end)) != NULL) {
  
         CKIRV2Node node = {0};
         char *op = NULL;
         char *kernel = NULL;
         char *kernel_variant = NULL;
         char *kernel_dtype = NULL;
         char *cond = NULL;
         int layer = 0;
         int flags = 0;
  
         ck_ir_v2_parse_string_field(obj_start, "\"op\"", obj_end, &op);
         ck_ir_v2_parse_string_field(obj_start, "\"kernel\"", obj_end, &kernel);
         ck_ir_v2_parse_string_field(obj_start, "\"kernel_variant\"", obj_end, &kernel_variant);
         ck_ir_v2_parse_string_field(obj_start, "\"kernel_dtype\"", obj_end, &kernel_dtype);
         ck_ir_v2_parse_string_field(obj_start, "\"condition\"", obj_end, &cond);
         ck_ir_v2_parse_int(obj_start, "\"layer\"", obj_end, &layer);
         ck_ir_v2_parse_int(obj_start, "\"flags\"", obj_end, &flags);
  
         if (!kernel && kernel_variant) {
             kernel = kernel_variant;
             kernel_variant = NULL;
         }
  
         node.op = op;
         node.kernel = kernel;
         node.kernel_dtype = ck_ir_v2_parse_dtype(kernel_dtype);
         node.condition = cond;
         node.layer = (uint16_t)layer;
         node.flags = (uint8_t)flags;
         node.n_bindings = 0;
         node.n_inputs = 0;
         node.n_outputs = 0;
  
         if (kernel_variant) {
             free(kernel_variant);
         }
         if (kernel_dtype) {
             free(kernel_dtype);
         }
  
         if (ck_ir_v2_parse_bindings(obj_start, obj_end, graph, &node) != 0) {
             free(node.op);
             free(node.kernel);
             free(node.condition);
             for (int b = 0; b < node.n_bindings; ++b) {
                 free(node.bindings[b].arg);
             }
             for (int j = 0; j < idx; ++j) {
                 ck_ir_v2_free_node(&nodes[j]);
             }
             free(nodes);
             return -1;
         }
  
         nodes[idx++] = node;
     }
  
     graph->nodes = nodes;
     graph->num_nodes = idx;
     return 0;
 }

References CKIRV2Binding::arg, CKIRV2Node::bindings, ck_ir_v2_find_array_end(), ck_ir_v2_find_key(), ck_ir_v2_free_node(), ck_ir_v2_next_object(), ck_ir_v2_parse_bindings(), ck_ir_v2_parse_dtype(), ck_ir_v2_parse_int(), ck_ir_v2_parse_string_field(), CKIRV2Node::condition, end, CKIRV2Node::flags, CKIRV2Node::kernel, CKIRV2Node::kernel_dtype, CKIRV2Node::layer, CKIRV2Node::n_bindings, CKIRV2Node::n_inputs, CKIRV2Node::n_outputs, CKIRV2Graph::nodes, CKIRV2Graph::num_nodes, CKIRV2Node::op, and start.

Referenced by ck_ir_v2_parse_json().

◆ ck_ir_v2_parse_role()

static CKBufferRole ck_ir_v2_parse_role ( const char * s )

static

Definition at line 760 of file ckernel_ir_v2.c.

 {
     if (!s) return CK_ROLE_ACTIVATION;
     if (strcmp(s, "input") == 0) return CK_ROLE_INPUT;
     if (strcmp(s, "output") == 0) return CK_ROLE_OUTPUT;
     if (strcmp(s, "activation") == 0) return CK_ROLE_ACTIVATION;
     if (strcmp(s, "weight") == 0) return CK_ROLE_WEIGHT;
     if (strcmp(s, "scratch") == 0) return CK_ROLE_SCRATCH;
     if (strcmp(s, "grad") == 0) return CK_ROLE_GRAD;
     return CK_ROLE_ACTIVATION;
 }

References CK_ROLE_ACTIVATION, CK_ROLE_GRAD, CK_ROLE_INPUT, CK_ROLE_OUTPUT, CK_ROLE_SCRATCH, and CK_ROLE_WEIGHT.

Referenced by ck_ir_v2_parse_buffers().

◆ ck_ir_v2_parse_scope()

static CKBufferScope ck_ir_v2_parse_scope ( const char * s )

static

Definition at line 752 of file ckernel_ir_v2.c.

 {
     if (!s) return CK_SCOPE_GLOBAL;
     if (strcmp(s, "layer") == 0) return CK_SCOPE_LAYER;
     if (strcmp(s, "global") == 0) return CK_SCOPE_GLOBAL;
     return CK_SCOPE_GLOBAL;
 }

References CK_SCOPE_GLOBAL, and CK_SCOPE_LAYER.

Referenced by ck_ir_v2_parse_buffers().

◆ ck_ir_v2_parse_shape()

static int ck_ir_v2_parse_shape	(	const char *	obj_start,
		const char *	obj_end,
		CKDimToken *	shape_out
	)

static

Definition at line 810 of file ckernel_ir_v2.c.

 {
     if (!shape_out) {
         return -1;
     }
     for (int i = 0; i < CK_IR_V2_MAX_DIMS; ++i) {
         shape_out[i].dim = CK_DIM_END;
         shape_out[i].mult = 0;
         shape_out[i].div = 0;
     }
  
     const char *start = ck_ir_v2_find_key(obj_start, "\"shape\"", obj_end);
     if (!start) {
         return -1;
     }
     const char *open = strchr(start, '[');
     const char *close = ck_ir_v2_find_array_end(open, obj_end);
     if (!open || !close) {
         return -1;
     }
  
     const char *cur = open;
     const char *sstart = NULL;
     const char *send = NULL;
     int idx = 0;
     while (idx < CK_IR_V2_MAX_DIMS &&
            (cur = ck_ir_v2_next_object(cur, close, &sstart, &send)) != NULL) {
         CKDimKind dim = ck_ir_v2_parse_dim_kind(sstart, send);
         if (dim == CK_DIM_END) {
             continue;
         }
         int mult = 1;
         int div = 1;
         ck_ir_v2_parse_int(sstart, "\"mult\"", send, &mult);
         ck_ir_v2_parse_int(sstart, "\"div\"", send, &div);
         shape_out[idx].dim = dim;
         shape_out[idx].mult = mult;
         shape_out[idx].div = div;
         idx++;
     }
     return (idx > 0) ? 0 : -1;
 }

References CK_DIM_END, ck_ir_v2_find_array_end(), ck_ir_v2_find_key(), CK_IR_V2_MAX_DIMS, ck_ir_v2_next_object(), ck_ir_v2_parse_dim_kind(), ck_ir_v2_parse_int(), CKDimToken::dim, CKDimToken::div, CKDimToken::mult, and start.

Referenced by ck_ir_v2_parse_buffers().

◆ ck_ir_v2_parse_string()

static int ck_ir_v2_parse_string	(	const char *	start,
		const char *	end,
		char **	out_str
	)

static

Definition at line 529 of file ckernel_ir_v2.c.

 {
     if (!start || !out_str || start >= end || *start != '"') {
         return -1;
     }
     const char *cur = start + 1;
     while (cur < end && *cur != '"') {
         if (*cur == '\\' && (cur + 1) < end) {
             cur += 2;
             continue;
         }
         cur++;
     }
     if (cur >= end || *cur != '"') {
         return -1;
     }
     size_t len = (size_t)(cur - (start + 1));
     char *buf = (char *)malloc(len + 1);
     if (!buf) {
         return -1;
     }
     memcpy(buf, start + 1, len);
     buf[len] = '\0';
     *out_str = buf;
     return 0;
 }

References end, and start.

Referenced by ck_ir_v2_parse_string_field().

◆ ck_ir_v2_parse_string_field()

static int ck_ir_v2_parse_string_field	(	const char *	json,
		const char *	key,
		const char *	end,
		char **	out_str
	)

static

Definition at line 722 of file ckernel_ir_v2.c.

 {
     const char *p = ck_ir_v2_find_key(json, key, end);
     if (!p) {
         return -1;
     }
     const char *colon = strchr(p, ':');
     if (!colon || colon >= end) {
         return -1;
     }
     const char *cur = colon + 1;
     while (cur < end && (*cur == ' ' || *cur == '\t' || *cur == '\n' || *cur == '\r')) {
         cur++;
     }
     if (cur >= end) {
         return -1;
     }
     if (strncmp(cur, "null", 4) == 0) {
         *out_str = NULL;
         return 0;
     }
     if (*cur != '"') {
         return -1;
     }
     return ck_ir_v2_parse_string(cur, end, out_str);
 }

References ck_ir_v2_find_key(), ck_ir_v2_parse_string(), and end.

Referenced by ck_ir_v2_parse_bindings(), ck_ir_v2_parse_buffers(), ck_ir_v2_parse_dim_kind(), and ck_ir_v2_parse_nodes().

◆ ck_ir_v2_resolve_align()

static void ck_ir_v2_resolve_align	(	const CKModelConfig *	cfg,
		size_t	alignment_bytes,
		CKIRV2AlignInfo *	align
	)

static

Definition at line 179 of file ckernel_ir_v2.c.

 {
     if (!align) {
         return;
     }
     memset(align, 0, sizeof(*align));
     if (!cfg) {
         return;
     }
     if (alignment_bytes == 0) {
         alignment_bytes = CK_MEM_PLAN_DEFAULT_ALIGN;
     }
     size_t elem_bytes = sizeof(float);
     size_t head_dim = (cfg->num_heads > 0) ? (size_t)(cfg->hidden_size / cfg->num_heads) : 0;
     align->aligned_embed = ck_ir_v2_align_up_elems((size_t)cfg->hidden_size, elem_bytes, alignment_bytes);
     align->aligned_head = ck_ir_v2_align_up_elems(head_dim, elem_bytes, alignment_bytes);
     align->aligned_intermediate = ck_ir_v2_align_up_elems((size_t)cfg->intermediate_size,
                                                           elem_bytes, alignment_bytes);
     align->aligned_context = ck_ir_v2_align_up_elems((size_t)cfg->context_window, elem_bytes, alignment_bytes);
 }

References ck_ir_v2_align_up_elems(), CK_MEM_PLAN_DEFAULT_ALIGN, CKModelConfig::context_window, CKModelConfig::hidden_size, CKModelConfig::intermediate_size, and CKModelConfig::num_heads.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_resolve_dim_value()

static size_t ck_ir_v2_resolve_dim_value	(	const CKModelConfig *	cfg,
		const CKIRV2AlignInfo *	align,
		CKDimKind	dim,
		int	tokens_override
	)

static

Definition at line 202 of file ckernel_ir_v2.c.

 {
     switch (dim) {
     case CK_DIM_TOKENS:
         if (tokens_override >= 0) {
             return (size_t)tokens_override;
         }
         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 ck_ir_v2_emit_dimensions(), and ck_ir_v2_emit_resolved_shape().

◆ ck_ir_v2_role_name()

static const char* ck_ir_v2_role_name ( CKBufferRole role )

static

Definition at line 66 of file ckernel_ir_v2.c.

 {
     switch (role) {
     case CK_ROLE_INPUT:
         return "input";
     case CK_ROLE_OUTPUT:
         return "output";
     case CK_ROLE_ACTIVATION:
         return "activation";
     case CK_ROLE_WEIGHT:
         return "weight";
     case CK_ROLE_SCRATCH:
         return "scratch";
     case CK_ROLE_GRAD:
         return "grad";
     default:
         return "unknown";
     }
 }

References CK_ROLE_ACTIVATION, CK_ROLE_GRAD, CK_ROLE_INPUT, CK_ROLE_OUTPUT, CK_ROLE_SCRATCH, and CK_ROLE_WEIGHT.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_scope_name()

static const char* ck_ir_v2_scope_name ( CKBufferScope scope )

static

Definition at line 54 of file ckernel_ir_v2.c.

 {
     switch (scope) {
     case CK_SCOPE_LAYER:
         return "layer";
     case CK_SCOPE_GLOBAL:
         return "global";
     default:
         return "unknown";
     }
 }

References CK_SCOPE_GLOBAL, and CK_SCOPE_LAYER.

Referenced by ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_serialize_json()

int ck_ir_v2_serialize_json	(	const CKIRV2Graph *	graph,
		const char *	path
	)

Definition at line 511 of file ckernel_ir_v2.c.

 {
     return ck_ir_v2_serialize_json_internal(graph, NULL, NULL, -1, -1, path);
 }

References ck_ir_v2_serialize_json_internal().

Referenced by main().

◆ ck_ir_v2_serialize_json_internal()

static int ck_ir_v2_serialize_json_internal	(	const CKIRV2Graph *	graph,
		const CKMemPlan *	plan,
		const char *	mode,
		int	tokens_override,
		int	base_context_window,
		const char *	path
	)

static

Definition at line 365 of file ckernel_ir_v2.c.

 {
     if (!graph || !path) {
         return -1;
     }
     FILE *out = fopen(path, "wb");
     if (!out) {
         fprintf(stderr, "ck_ir_v2_serialize_json: failed to open %s: %s\n",
                 path, strerror(errno));
         return -1;
     }
  
     CKIRV2AlignInfo align = {0};
     ck_ir_v2_resolve_align(&graph->config, CK_MEM_PLAN_DEFAULT_ALIGN, &align);
  
     fprintf(out, "{\n");
     fprintf(out, "  \"version\": 2,\n");
     fprintf(out, "  \"notes\": [\n");
     fprintf(out, "    \"shape.dim uses symbolic names; see dimensions for resolved values\",\n");
     fprintf(out,
             "    \"resolved_shape applies mult/div using alignment_bytes=%d and elem_bytes=4\",\n",
             CK_MEM_PLAN_DEFAULT_ALIGN);
     fprintf(out, "    \"kernel is the selected impl; kernel_dtype records dtype selection\"\n");
     fprintf(out, "  ],\n");
     fprintf(out, "  \"config\": {\n");
     fprintf(out, "    \"num_layers\": %d,\n", graph->config.num_layers);
     fprintf(out, "    \"hidden_size\": %d,\n", graph->config.hidden_size);
     fprintf(out, "    \"intermediate_size\": %d,\n", graph->config.intermediate_size);
     fprintf(out, "    \"num_attention_heads\": %d,\n", graph->config.num_heads);
     fprintf(out, "    \"num_kv_heads\": %d,\n", graph->config.num_kv_heads);
     fprintf(out, "    \"vocab_size\": %d,\n", graph->config.vocab_size);
     fprintf(out, "    \"context_window\": %d,\n", graph->config.context_window);
     fprintf(out, "    \"rms_norm_eps\": %.9g,\n", graph->config.rms_norm_eps);
     fprintf(out, "    \"rope_theta\": %.9g\n", graph->config.rope_theta);
     fprintf(out, "  },\n");
  
     ck_ir_v2_emit_dimensions(out, &graph->config, &align, tokens_override);
  
     fprintf(out, "  \"meta\": {\n");
     fprintf(out, "    \"has_pos_emb\": %s,\n", graph->has_pos_emb ? "true" : "false");
     if (graph->tie_word_embeddings < 0) {
         fprintf(out, "    \"tie_word_embeddings\": null,\n");
     } else {
         fprintf(out, "    \"tie_word_embeddings\": %s,\n", graph->tie_word_embeddings ? "true" : "false");
     }
     if (graph->fused_qkv < 0) {
         fprintf(out, "    \"fused_qkv\": null,\n");
     } else {
         fprintf(out, "    \"fused_qkv\": %s,\n", graph->fused_qkv ? "true" : "false");
     }
     if (graph->gated_mlp < 0) {
         fprintf(out, "    \"gated_mlp\": null\n");
     } else {
         fprintf(out, "    \"gated_mlp\": %s\n", graph->gated_mlp ? "true" : "false");
     }
     fprintf(out, "  },\n");
  
     if (plan) {
         int training = (mode && strcmp(mode, "backward") == 0);
         int tokens = (tokens_override >= 0) ? tokens_override : graph->config.context_window;
         fprintf(out, "  \"lowering\": {\n");
         fprintf(out, "    \"mode\": \"%s\",\n", mode ? mode : "unknown");
         fprintf(out, "    \"training\": %s,\n", training ? "true" : "false");
         fprintf(out, "    \"tokens\": %d", tokens);
         if (base_context_window >= 0) {
             fprintf(out, ",\n    \"base_context_window\": %d\n", base_context_window);
         } else {
             fprintf(out, "\n");
         }
         fprintf(out, "  },\n");
         ck_ir_v2_emit_memory_plan(out, graph, plan);
     }
  
     fprintf(out, "  \"buffers\": [\n");
     for (int i = 0; i < graph->num_buffers; ++i) {
         const CKIRV2Buffer *buf = &graph->buffers[i];
         fprintf(out, "    {\n");
         fprintf(out, "      \"name\": \"%s\",\n", buf->name ? buf->name : "");
         fprintf(out, "      \"scope\": \"%s\",\n", ck_ir_v2_scope_name(buf->scope));
         fprintf(out, "      \"role\": \"%s\",\n", ck_ir_v2_role_name(buf->role));
         fprintf(out, "      \"dtype\": \"%s\",\n", ck_ir_v2_dtype_name(buf->dtype));
         fprintf(out, "      \"optional\": %d,\n", buf->optional ? 1 : 0);
         fprintf(out, "      \"shape\": ");
         ck_ir_v2_emit_shape(out, buf->shape);
         fprintf(out, ",\n");
         fprintf(out, "      \"resolved_shape\": ");
         ck_ir_v2_emit_resolved_shape(out, &graph->config, &align, buf->shape,
                                      tokens_override);
         fprintf(out, ",\n");
         if (buf->alias_of) {
             fprintf(out, "      \"alias_of\": \"%s\",\n", buf->alias_of);
         } else {
             fprintf(out, "      \"alias_of\": null,\n");
         }
         if (buf->condition) {
             fprintf(out, "      \"condition\": \"%s\"\n", buf->condition);
         } else {
             fprintf(out, "      \"condition\": null\n");
         }
         fprintf(out, "    }%s\n", (i + 1 == graph->num_buffers) ? "" : ",");
     }
     fprintf(out, "  ],\n");
  
     fprintf(out, "  \"nodes\": [\n");
     for (int i = 0; i < graph->num_nodes; ++i) {
         const CKIRV2Node *node = &graph->nodes[i];
         fprintf(out, "    {\n");
         fprintf(out, "      \"layer\": %d,\n", (int)node->layer);
         fprintf(out, "      \"op\": \"%s\",\n", node->op ? node->op : "");
         fprintf(out, "      \"kernel\": \"%s\",\n", node->kernel ? node->kernel : "");
         fprintf(out, "      \"kernel_variant\": \"%s\",\n", node->kernel ? node->kernel : "");
         fprintf(out, "      \"kernel_dtype\": \"%s\",\n", ck_ir_v2_dtype_name(node->kernel_dtype));
         fprintf(out, "      \"flags\": %u,\n", (unsigned)node->flags);
         if (node->condition) {
             fprintf(out, "      \"condition\": \"%s\",\n", node->condition);
         } else {
             fprintf(out, "      \"condition\": null,\n");
         }
         fprintf(out, "      \"bindings\": [\n");
         for (int b = 0; b < node->n_bindings; ++b) {
             const CKIRV2Binding *bind = &node->bindings[b];
             const char *buf_name = "";
             if (bind->buffer >= 0 && bind->buffer < graph->num_buffers) {
                 buf_name = graph->buffers[bind->buffer].name ? graph->buffers[bind->buffer].name : "";
             }
             fprintf(out,
                     "        {\"arg\": \"%s\", \"buffer\": \"%s\"}%s\n",
                     bind->arg ? bind->arg : "",
                     buf_name,
                     (b + 1 == node->n_bindings) ? "" : ",");
         }
         fprintf(out, "      ]\n");
         fprintf(out, "    }%s\n", (i + 1 == graph->num_nodes) ? "" : ",");
     }
     fprintf(out, "  ]\n");
     fprintf(out, "}\n");
  
     fclose(out);
     return 0;
 }

References CKIRV2Buffer::alias_of, CKIRV2Binding::arg, CKIRV2Node::bindings, CKIRV2Binding::buffer, CKIRV2Graph::buffers, ck_ir_v2_dtype_name(), ck_ir_v2_emit_dimensions(), ck_ir_v2_emit_memory_plan(), ck_ir_v2_emit_resolved_shape(), ck_ir_v2_emit_shape(), ck_ir_v2_resolve_align(), ck_ir_v2_role_name(), ck_ir_v2_scope_name(), CK_MEM_PLAN_DEFAULT_ALIGN, CKIRV2Buffer::condition, CKIRV2Node::condition, CKIRV2Graph::config, CKModelConfig::context_window, CKIRV2Buffer::dtype, CKIRV2Node::flags, CKIRV2Graph::fused_qkv, CKIRV2Graph::gated_mlp, CKIRV2Graph::has_pos_emb, CKModelConfig::hidden_size, CKModelConfig::intermediate_size, CKIRV2Node::kernel, CKIRV2Node::kernel_dtype, CKIRV2Node::layer, CKIRV2Node::n_bindings, CKIRV2Buffer::name, CKIRV2Graph::nodes, CKIRV2Graph::num_buffers, CKModelConfig::num_heads, CKModelConfig::num_kv_heads, CKModelConfig::num_layers, CKIRV2Graph::num_nodes, CKIRV2Node::op, CKIRV2Buffer::optional, CKModelConfig::rms_norm_eps, CKIRV2Buffer::role, CKModelConfig::rope_theta, CKIRV2Buffer::scope, CKIRV2Buffer::shape, CKIRV2Graph::tie_word_embeddings, and CKModelConfig::vocab_size.

Referenced by ck_ir_v2_serialize_json(), and ck_ir_v2_serialize_json_with_plan().

◆ ck_ir_v2_serialize_json_with_plan()

int ck_ir_v2_serialize_json_with_plan	(	const CKIRV2Graph *	graph,
		const CKMemPlan *	plan,
		const char *	mode,
		int	tokens_override,
		int	base_context_window,
		const char *	path
	)

Definition at line 516 of file ckernel_ir_v2.c.

 {
     return ck_ir_v2_serialize_json_internal(graph, plan, mode,
                                             tokens_override,
                                             base_context_window,
                                             path);
 }

References ck_ir_v2_serialize_json_internal().

◆ ck_ir_v2_skip_string()

static const char* ck_ir_v2_skip_string	(	const char *	cur,
		const char *	end
	)

static

Definition at line 558 of file ckernel_ir_v2.c.

 {
     if (!cur || cur >= end || *cur != '"') {
         return cur;
     }
     cur++;
     while (cur < end) {
         if (*cur == '\\' && (cur + 1) < end) {
             cur += 2;
             continue;
         }
         if (*cur == '"') {
             return cur + 1;
         }
         cur++;
     }
     return end;
 }

References end.

Referenced by ck_ir_v2_find_array_end(), and ck_ir_v2_next_object().

Functions

Function Documentation

◆ ck_ir_v2_align_up_bytes()

◆ ck_ir_v2_align_up_elems()

◆ ck_ir_v2_dim_kind_from_name()

◆ ck_ir_v2_dim_name()

◆ ck_ir_v2_dtype_name()

◆ ck_ir_v2_emit_dimensions()

◆ ck_ir_v2_emit_memory_plan()

◆ ck_ir_v2_emit_resolved_shape()

◆ ck_ir_v2_emit_shape()

◆ ck_ir_v2_find_array_end()

◆ ck_ir_v2_find_buffer_index()

◆ ck_ir_v2_find_buffer_spec()

◆ ck_ir_v2_find_key()

◆ ck_ir_v2_free()

◆ ck_ir_v2_free_buffer()

◆ ck_ir_v2_free_node()

◆ ck_ir_v2_mem_arena_name()

◆ ck_ir_v2_next_object()

◆ ck_ir_v2_parse_bindings()

◆ ck_ir_v2_parse_bool()

◆ ck_ir_v2_parse_buffers()

◆ ck_ir_v2_parse_dim_kind()

◆ ck_ir_v2_parse_dtype()

◆ ck_ir_v2_parse_float()

◆ ck_ir_v2_parse_int()

◆ ck_ir_v2_parse_json()

◆ ck_ir_v2_parse_nodes()

◆ ck_ir_v2_parse_role()

◆ ck_ir_v2_parse_scope()

◆ ck_ir_v2_parse_shape()

◆ ck_ir_v2_parse_string()

◆ ck_ir_v2_parse_string_field()

◆ ck_ir_v2_resolve_align()

◆ ck_ir_v2_resolve_dim_value()

◆ ck_ir_v2_role_name()

◆ ck_ir_v2_scope_name()

◆ ck_ir_v2_serialize_json()

◆ ck_ir_v2_serialize_json_internal()

◆ ck_ir_v2_serialize_json_with_plan()

◆ ck_ir_v2_skip_string()