ckernel_ir_v2_builder.c File Reference

#include "ckernel_ir_v2.h"
#include "ckernel_kernel_specs.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

Go to the source code of this file.

Functions
int	ck_ir_v2_apply_meta (const char *path, CKIRV2Graph *graph)
static int	ck_ir_v2_apply_weight_dtypes (const char *json, const char *end, 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)
static int	ck_ir_v2_copy_buffer_spec (const CKBufferSpec *spec, CKIRV2Buffer *out)
static void	ck_ir_v2_copy_shape (CKDimToken *dst, const CKDimToken *src)
static int	ck_ir_v2_find_buffer_index (const CKIRV2Graph *graph, const char *name)
static const CKKernelSpec *	ck_ir_v2_find_kernel_spec (const char *name)
static const char *	ck_ir_v2_find_key (const char *json, const char *key, const char *end)
static int	ck_ir_v2_parse_bool (const char *json, const char *key, const char *end, int *out)
static CKDataType	ck_ir_v2_parse_dtype (const char *s)
static int	ck_ir_v2_parse_string (const char *start, const char *end, char **out)
static const char *	ck_ir_v2_select_kernel (const CKKernelSpec *spec, CKDataType dtype, int backward)
static const char *	ck_ir_v2_skip_ws (const char *cur, const char *end)
static char *	ck_ir_v2_strdup (const char *s)

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

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

static

Definition at line 174 of file ckernel_ir_v2_builder.c.

{
    const char *start = ck_ir_v2_find_key(json, "\"weight_dtypes\"", end);
    if (!start) {
        return 0;
    }
    const char *open = strchr(start, '{');
    if (!open || open >= end) {
        return -1;
    }
    const char *cur = open + 1;
    while (cur < end) {
        cur = ck_ir_v2_skip_ws(cur, end);
        if (cur >= end || *cur == '}') {
            break;
        }
        char *key = NULL;
        if (ck_ir_v2_parse_string(cur, end, &key) != 0) {
            break;
        }
        cur = strchr(cur, ':');
        if (!cur) {
            free(key);
            break;
        }
        cur = ck_ir_v2_skip_ws(cur + 1, end);
        char *val = NULL;
        if (ck_ir_v2_parse_string(cur, end, &val) != 0) {
            free(key);
            break;
        }
        int idx = ck_ir_v2_find_buffer_index(graph, key);
        if (idx >= 0) {
            graph->buffers[idx].dtype = ck_ir_v2_parse_dtype(val);
        }
        free(key);
        free(val);
        cur = strchr(cur, ',');
        if (!cur) {
            break;
        }
        cur++;
    }
    return 0;
}

References CKIRV2Graph::buffers, ck_ir_v2_find_buffer_index(), ck_ir_v2_find_key(), ck_ir_v2_parse_dtype(), ck_ir_v2_parse_string(), ck_ir_v2_skip_ws(), CKIRV2Buffer::dtype, end, and start.

Referenced by ck_ir_v2_apply_meta().

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

static int ck_ir_v2_copy_buffer_spec ( const CKBufferSpec *  spec, CKIRV2Buffer *  out  )

static

Definition at line 27 of file ckernel_ir_v2_builder.c.

{
    if (!spec || !out) {
        return -1;
    }
    memset(out, 0, sizeof(*out));
    out->name = ck_ir_v2_strdup(spec->name);
    out->scope = spec->scope;
    out->role = spec->role;
    out->dtype = spec->dtype;
    out->optional = spec->optional;
    out->alias_of = ck_ir_v2_strdup(spec->alias_of);
    out->condition = ck_ir_v2_strdup(spec->condition);
    ck_ir_v2_copy_shape(out->shape, spec->shape);
    return 0;
}

References CKIRV2Buffer::alias_of, CKBufferSpec::alias_of, ck_ir_v2_copy_shape(), ck_ir_v2_strdup(), CKIRV2Buffer::condition, CKBufferSpec::condition, CKIRV2Buffer::dtype, CKBufferSpec::dtype, CKIRV2Buffer::name, CKBufferSpec::name, CKIRV2Buffer::optional, CKBufferSpec::optional, CKIRV2Buffer::role, CKBufferSpec::role, CKIRV2Buffer::scope, CKBufferSpec::scope, CKIRV2Buffer::shape, and CKBufferSpec::shape.

Referenced by ck_ir_v2_build_decoder(), and ck_ir_v2_build_decoder_backward().

ck_ir_v2_copy_shape()

static void ck_ir_v2_copy_shape ( CKDimToken *  dst, const CKDimToken *  src  )

static

Definition at line 22 of file ckernel_ir_v2_builder.c.

{
    memcpy(dst, src, sizeof(CKDimToken) * CK_IR_V2_MAX_DIMS);
}

References CK_IR_V2_MAX_DIMS.

Referenced by ck_ir_v2_copy_buffer_spec().

ck_ir_v2_find_buffer_index()

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

static

Definition at line 44 of file ckernel_ir_v2_builder.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_apply_meta(), ck_ir_v2_apply_weight_dtypes(), ck_ir_v2_build_decoder(), and ck_ir_v2_build_decoder_backward().

ck_ir_v2_find_kernel_spec()

static const CKKernelSpec * ck_ir_v2_find_kernel_spec ( const char *  name )

static

Definition at line 57 of file ckernel_ir_v2_builder.c.

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

References ck_kernel_spec_count, and ck_kernel_specs.

Referenced by ck_ir_v2_build_decoder(), and ck_ir_v2_build_decoder_backward().

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 91 of file ckernel_ir_v2_builder.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_apply_weight_dtypes(), and ck_ir_v2_parse_bool().

ck_ir_v2_parse_bool()

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

static

Definition at line 139 of file ckernel_ir_v2_builder.c.

{
    const char *p = ck_ir_v2_find_key(json, key, end);
    if (!p) {
        return -1;
    }
    p = strchr(p, ':');
    if (!p || p >= end) {
        return -1;
    }
    p = ck_ir_v2_skip_ws(p + 1, end);
    if (p + 4 <= end && memcmp(p, "true", 4) == 0) {
        *out = 1;
        return 0;
    }
    if (p + 5 <= end && memcmp(p, "false", 5) == 0) {
        *out = 0;
        return 0;
    }
    return -1;
}

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

Referenced by ck_ir_v2_apply_meta().

ck_ir_v2_parse_dtype()

static CKDataType ck_ir_v2_parse_dtype ( const char *  s )

static

Definition at line 161 of file ckernel_ir_v2_builder.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_apply_weight_dtypes().

ck_ir_v2_parse_string()

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

static

Definition at line 112 of file ckernel_ir_v2_builder.c.

{
    if (!start || 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 = buf;
    return 0;
}

References end, and start.

Referenced by ck_ir_v2_apply_weight_dtypes().

ck_ir_v2_select_kernel()

static const char * ck_ir_v2_select_kernel ( const CKKernelSpec *  spec, CKDataType  dtype, int  backward  )

static

Definition at line 70 of file ckernel_ir_v2_builder.c.

{
    if (!spec) {
        return NULL;
    }
    if (dtype < 0 || dtype >= CK_DT_COUNT) {
        dtype = spec->default_dtype;
    }
    const char *name = backward ? spec->backward[dtype] : spec->forward[dtype];
    if (name && name[0]) {
        return name;
    }
    for (int i = 0; i < CK_DT_COUNT; ++i) {
        name = backward ? spec->backward[i] : spec->forward[i];
        if (name && name[0]) {
            return name;
        }
    }
    return spec->name;
}

References CKKernelSpec::backward, CK_DT_COUNT, CKKernelSpec::default_dtype, CKKernelSpec::forward, and CKKernelSpec::name.

Referenced by ck_ir_v2_build_decoder(), and ck_ir_v2_build_decoder_backward().

ck_ir_v2_skip_ws()

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

static

Definition at line 104 of file ckernel_ir_v2_builder.c.

{
    while (cur < end && (*cur == ' ' || *cur == '\n' || *cur == '\r' || *cur == '\t')) {
        cur++;
    }
    return cur;
}

References end.

Referenced by ck_ir_v2_apply_weight_dtypes(), and ck_ir_v2_parse_bool().

ck_ir_v2_strdup()

static char * ck_ir_v2_strdup ( const char *  s )

static

Definition at line 8 of file ckernel_ir_v2_builder.c.

{
    if (!s) {
        return NULL;
    }
    size_t len = strlen(s);
    char *out = (char *)malloc(len + 1);
    if (!out) {
        return NULL;
    }
    memcpy(out, s, len + 1);
    return out;
}

Referenced by ck_ir_v2_apply_meta(), ck_ir_v2_build_decoder(), ck_ir_v2_build_decoder_backward(), and ck_ir_v2_copy_buffer_spec().