ckernel_ir_v2.h File Reference

#include "ckernel_ir.h"
#include "ckernel_kernel_specs.h"
#include <stdint.h>

Go to the source code of this file.

Data Structures
struct	CKIRV2Binding
struct	CKIRV2Buffer
struct	CKIRV2Graph
struct	CKIRV2Node

Macros
#define	CK_IR_V2_MAX_BINDINGS   24
#define	CK_IR_V2_MAX_DIMS   4
#define	CK_IR_V2_MAX_INPUTS   8
#define	CK_IR_V2_MAX_OUTPUTS   4

Enumerations
enum	CKIRV2NodeFlags { CK_IR_V2_NODE_NONE = 0 , CK_IR_V2_NODE_FUSED = 1 << 0 , CK_IR_V2_NODE_INFERENCE_ONLY = 1 << 1 }

Functions
int	ck_ir_v2_apply_meta (const char *path, CKIRV2Graph *graph)
int	ck_ir_v2_build_decoder (const CKModelConfig *cfg, CKIRV2Graph *graph)
int	ck_ir_v2_build_decoder_backward (const CKIRV2Graph *forward, CKIRV2Graph *backward)
void	ck_ir_v2_free (CKIRV2Graph *graph)
int	ck_ir_v2_parse_json (const char *path, CKIRV2Graph *graph)
int	ck_ir_v2_serialize_json (const CKIRV2Graph *graph, const char *path)
int	ck_ir_v2_serialize_json_with_plan (const CKIRV2Graph *graph, const struct CKMemPlan *plan, const char *mode, int tokens_override, int base_context_window, const char *path)

Macro Definition Documentation

CK_IR_V2_MAX_BINDINGS

#define CK_IR_V2_MAX_BINDINGS   24

Definition at line 16 of file ckernel_ir_v2.h.

CK_IR_V2_MAX_DIMS

#define CK_IR_V2_MAX_DIMS   4

Definition at line 13 of file ckernel_ir_v2.h.

CK_IR_V2_MAX_INPUTS

#define CK_IR_V2_MAX_INPUTS   8

Definition at line 14 of file ckernel_ir_v2.h.

CK_IR_V2_MAX_OUTPUTS

#define CK_IR_V2_MAX_OUTPUTS   4

Definition at line 15 of file ckernel_ir_v2.h.

Enumeration Type Documentation

CKIRV2NodeFlags

enum CKIRV2NodeFlags

Enumerator
CK_IR_V2_NODE_NONE
CK_IR_V2_NODE_FUSED
CK_IR_V2_NODE_INFERENCE_ONLY

Definition at line 18 of file ckernel_ir_v2.h.

             {
    CK_IR_V2_NODE_NONE = 0,
    CK_IR_V2_NODE_FUSED = 1 << 0,
    CK_IR_V2_NODE_INFERENCE_ONLY = 1 << 1
} CKIRV2NodeFlags;

Function Documentation

ck_ir_v2_apply_meta()

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

Definition at line 220 of file ckernel_ir_v2_builder.c.

{
    if (!path || !graph) {
        return -1;
    }
    FILE *f = fopen(path, "rb");
    if (!f) {
        perror("ck_ir_v2_apply_meta: 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';
    const char *end = buf + nread;
 
    int val = 0;
    if (ck_ir_v2_parse_bool(buf, "\"has_pos_emb\"", end, &val) == 0) {
        graph->has_pos_emb = val ? 1 : 0;
    }
    if (ck_ir_v2_parse_bool(buf, "\"tie_word_embeddings\"", end, &val) == 0) {
        graph->tie_word_embeddings = val ? 1 : 0;
        if (graph->tie_word_embeddings == 0) {
            int idx = ck_ir_v2_find_buffer_index(graph, "lm_head_weight");
            if (idx >= 0) {
                free(graph->buffers[idx].alias_of);
                graph->buffers[idx].alias_of = NULL;
            }
        } else {
            int idx = ck_ir_v2_find_buffer_index(graph, "lm_head_weight");
            if (idx >= 0) {
                free(graph->buffers[idx].alias_of);
                graph->buffers[idx].alias_of = ck_ir_v2_strdup("token_emb");
            }
        }
    }
    if (ck_ir_v2_parse_bool(buf, "\"fused_qkv\"", end, &val) == 0) {
        graph->fused_qkv = val ? 1 : 0;
    }
    if (ck_ir_v2_parse_bool(buf, "\"gated_mlp\"", end, &val) == 0) {
        graph->gated_mlp = val ? 1 : 0;
    }
    (void)ck_ir_v2_apply_weight_dtypes(buf, end, graph);
 
    free(buf);
    return 0;
}

References CKIRV2Buffer::alias_of, CKIRV2Graph::buffers, ck_ir_v2_apply_weight_dtypes(), ck_ir_v2_find_buffer_index(), ck_ir_v2_parse_bool(), ck_ir_v2_strdup(), end, CKIRV2Graph::fused_qkv, CKIRV2Graph::gated_mlp, CKIRV2Graph::has_pos_emb, and CKIRV2Graph::tie_word_embeddings.

Referenced by main().

ck_ir_v2_build_decoder()

int ck_ir_v2_build_decoder	(	const CKModelConfig *	cfg,
		CKIRV2Graph *	graph
	)

Definition at line 285 of file ckernel_ir_v2_builder.c.

{
    if (!cfg || !graph) {
        return -1;
    }
    memset(graph, 0, sizeof(*graph));
    graph->config = *cfg;
    graph->has_pos_emb = 1;
    graph->tie_word_embeddings = -1;
    graph->fused_qkv = -1;
    graph->gated_mlp = -1;
 
    graph->num_buffers = (int)ck_decoder_buffer_count;
    graph->buffers = (CKIRV2Buffer *)calloc((size_t)graph->num_buffers, sizeof(CKIRV2Buffer));
    if (!graph->buffers) {
        return -1;
    }
    for (int i = 0; i < graph->num_buffers; ++i) {
        if (ck_ir_v2_copy_buffer_spec(&ck_decoder_buffers[i], &graph->buffers[i]) != 0) {
            ck_ir_v2_free(graph);
            return -1;
        }
        if (graph->buffers[i].name && strcmp(graph->buffers[i].name, "pos_emb") == 0) {
            free(graph->buffers[i].condition);
            graph->buffers[i].condition = ck_ir_v2_strdup("has_pos_emb");
        }
    }
 
    int plan_count = (int)ck_decoder_forward_plan_v2_count;
    graph->num_nodes = cfg->num_layers * plan_count;
    graph->nodes = (CKIRV2Node *)calloc((size_t)graph->num_nodes, sizeof(CKIRV2Node));
    if (!graph->nodes) {
        ck_ir_v2_free(graph);
        return -1;
    }
 
    int idx = 0;
    for (int layer = 0; layer < cfg->num_layers; ++layer) {
        for (int p = 0; p < plan_count; ++p) {
            const CKPlanStepV2 *step = &ck_decoder_forward_plan_v2[p];
            const CKKernelSpec *spec = ck_ir_v2_find_kernel_spec(step->kernel);
            CKDataType dtype = spec ? spec->default_dtype : CK_DT_FP32;
            const char *impl = ck_ir_v2_select_kernel(spec, dtype, 0);
            CKIRV2Node *node = &graph->nodes[idx++];
            node->layer = (uint16_t)layer;
            node->op = ck_ir_v2_strdup(step->kernel);
            node->kernel = ck_ir_v2_strdup(impl ? impl : step->kernel);
            node->kernel_dtype = dtype;
            node->condition = ck_ir_v2_strdup(step->condition);
            node->flags = 0;
            node->n_bindings = 0;
            if (step->bindings && step->num_bindings > 0) {
                int limit = (int)step->num_bindings;
                if (limit > CK_IR_V2_MAX_BINDINGS) {
                    limit = CK_IR_V2_MAX_BINDINGS;
                }
                for (int b = 0; b < limit; ++b) {
                    const CKPlanBinding *binding = &step->bindings[b];
                    node->bindings[node->n_bindings].arg = ck_ir_v2_strdup(binding->arg);
                    node->bindings[node->n_bindings].buffer =
                        ck_ir_v2_find_buffer_index(graph, binding->buffer);
                    node->n_bindings++;
                }
            }
            node->n_inputs = 0;
            node->n_outputs = 0;
        }
    }
    return 0;
}

References CKIRV2Binding::arg, CKPlanBinding::arg, CKIRV2Node::bindings, CKPlanStepV2::bindings, CKIRV2Binding::buffer, CKPlanBinding::buffer, CKIRV2Graph::buffers, ck_decoder_buffer_count, ck_decoder_buffers, ck_decoder_forward_plan_v2, ck_decoder_forward_plan_v2_count, CK_DT_FP32, ck_ir_v2_copy_buffer_spec(), ck_ir_v2_find_buffer_index(), ck_ir_v2_find_kernel_spec(), ck_ir_v2_free(), CK_IR_V2_MAX_BINDINGS, ck_ir_v2_select_kernel(), ck_ir_v2_strdup(), CKIRV2Buffer::condition, CKIRV2Node::condition, CKPlanStepV2::condition, CKIRV2Graph::config, CKKernelSpec::default_dtype, CKIRV2Node::flags, CKIRV2Graph::fused_qkv, CKIRV2Graph::gated_mlp, CKIRV2Graph::has_pos_emb, CKIRV2Node::kernel, CKPlanStepV2::kernel, CKIRV2Node::kernel_dtype, CKIRV2Node::layer, CKIRV2Node::n_bindings, CKIRV2Node::n_inputs, CKIRV2Node::n_outputs, CKIRV2Buffer::name, CKIRV2Graph::nodes, CKPlanStepV2::num_bindings, CKIRV2Graph::num_buffers, CKModelConfig::num_layers, CKIRV2Graph::num_nodes, CKIRV2Node::op, and CKIRV2Graph::tie_word_embeddings.

Referenced by main().

ck_ir_v2_build_decoder_backward()

int ck_ir_v2_build_decoder_backward	(	const CKIRV2Graph *	forward,
		CKIRV2Graph *	backward
	)

Definition at line 356 of file ckernel_ir_v2_builder.c.

{
    if (!forward || !backward) {
        return -1;
    }
    memset(backward, 0, sizeof(*backward));
    backward->config = forward->config;
 
    backward->num_buffers = forward->num_buffers;
    backward->buffers = (CKIRV2Buffer *)calloc((size_t)backward->num_buffers, sizeof(CKIRV2Buffer));
    if (!backward->buffers) {
        return -1;
    }
    for (int i = 0; i < backward->num_buffers; ++i) {
        CKBufferSpec spec = {0};
        const CKIRV2Buffer *src = &forward->buffers[i];
        spec.name = src->name;
        spec.scope = src->scope;
        spec.role = src->role;
        spec.dtype = src->dtype;
        spec.optional = src->optional;
        spec.alias_of = src->alias_of;
        spec.condition = src->condition;
        memcpy(spec.shape, src->shape, sizeof(spec.shape));
        if (ck_ir_v2_copy_buffer_spec(&spec, &backward->buffers[i]) != 0) {
            ck_ir_v2_free(backward);
            return -1;
        }
    }
 
    int plan_count = (int)ck_decoder_backward_plan_v2_count;
    backward->num_nodes = forward->config.num_layers * plan_count;
    backward->nodes = (CKIRV2Node *)calloc((size_t)backward->num_nodes, sizeof(CKIRV2Node));
    if (!backward->nodes) {
        ck_ir_v2_free(backward);
        return -1;
    }
 
    int idx = 0;
    for (int layer = 0; layer < forward->config.num_layers; ++layer) {
        for (int p = 0; p < plan_count; ++p) {
            const CKPlanStepV2 *step = &ck_decoder_backward_plan_v2[p];
            const CKKernelSpec *spec = ck_ir_v2_find_kernel_spec(step->kernel);
            CKDataType dtype = spec ? spec->default_dtype : CK_DT_FP32;
            const char *impl = ck_ir_v2_select_kernel(spec, dtype, 1);
            CKIRV2Node *node = &backward->nodes[idx++];
            node->layer = (uint16_t)layer;
            node->op = ck_ir_v2_strdup(step->kernel);
            node->kernel = ck_ir_v2_strdup(impl ? impl : step->kernel);
            node->kernel_dtype = dtype;
            node->condition = ck_ir_v2_strdup(step->condition);
            node->flags = 0;
            node->n_bindings = 0;
            if (step->bindings && step->num_bindings > 0) {
                int limit = (int)step->num_bindings;
                if (limit > CK_IR_V2_MAX_BINDINGS) {
                    limit = CK_IR_V2_MAX_BINDINGS;
                }
                for (int b = 0; b < limit; ++b) {
                    const CKPlanBinding *binding = &step->bindings[b];
                    node->bindings[node->n_bindings].arg = ck_ir_v2_strdup(binding->arg);
                    node->bindings[node->n_bindings].buffer =
                        ck_ir_v2_find_buffer_index(backward, binding->buffer);
                    node->n_bindings++;
                }
            }
            node->n_inputs = 0;
            node->n_outputs = 0;
        }
    }
    return 0;
}

References CKIRV2Buffer::alias_of, CKBufferSpec::alias_of, CKIRV2Binding::arg, CKPlanBinding::arg, CKIRV2Node::bindings, CKPlanStepV2::bindings, CKIRV2Binding::buffer, CKPlanBinding::buffer, CKIRV2Graph::buffers, ck_decoder_backward_plan_v2, ck_decoder_backward_plan_v2_count, CK_DT_FP32, ck_ir_v2_copy_buffer_spec(), ck_ir_v2_find_buffer_index(), ck_ir_v2_find_kernel_spec(), ck_ir_v2_free(), CK_IR_V2_MAX_BINDINGS, ck_ir_v2_select_kernel(), ck_ir_v2_strdup(), CKIRV2Buffer::condition, CKIRV2Node::condition, CKBufferSpec::condition, CKPlanStepV2::condition, CKIRV2Graph::config, CKKernelSpec::default_dtype, CKIRV2Buffer::dtype, CKBufferSpec::dtype, CKIRV2Node::flags, CKIRV2Node::kernel, CKPlanStepV2::kernel, CKIRV2Node::kernel_dtype, CKIRV2Node::layer, CKIRV2Node::n_bindings, CKIRV2Node::n_inputs, CKIRV2Node::n_outputs, CKIRV2Buffer::name, CKBufferSpec::name, CKIRV2Graph::nodes, CKPlanStepV2::num_bindings, CKIRV2Graph::num_buffers, CKModelConfig::num_layers, CKIRV2Graph::num_nodes, CKIRV2Node::op, CKIRV2Buffer::optional, CKBufferSpec::optional, CKIRV2Buffer::role, CKBufferSpec::role, CKIRV2Buffer::scope, CKBufferSpec::scope, CKIRV2Buffer::shape, and CKBufferSpec::shape.

Referenced by main().

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_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_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_with_plan()

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

Referenced by ck_ir_v2_lower_emit_json().