#include "ckernel_ir_v2.h"
#include "ckernel_mem_plan.h"

Enumerations
enum	CKIRV2LowerMode { CK_IR_V2_LOWER_PREFILL = 0 , CK_IR_V2_LOWER_DECODE = 1 , CK_IR_V2_LOWER_BACKWARD = 2 }

Functions
int	ck_ir_v2_lower_emit_json (const CKIRV2Graph input, CKIRV2LowerMode mode, const char path)

int	ck_ir_v2_lower_graph (const CKIRV2Graph input, CKIRV2LowerMode mode, CKIRV2Graph output, CKMemPlan *plan)

int	ck_ir_v2_lower_mode_from_string (const char name, CKIRV2LowerMode out_mode)

const char *	ck_ir_v2_lower_mode_name (CKIRV2LowerMode mode)

Enumeration Type Documentation

◆ CKIRV2LowerMode

enum CKIRV2LowerMode

Enumerator
CK_IR_V2_LOWER_PREFILL
CK_IR_V2_LOWER_DECODE
CK_IR_V2_LOWER_BACKWARD

Definition at line 11 of file ckernel_ir_v2_lower.h.

             {
    CK_IR_V2_LOWER_PREFILL = 0,
    CK_IR_V2_LOWER_DECODE = 1,
    CK_IR_V2_LOWER_BACKWARD = 2
} CKIRV2LowerMode;

Function Documentation

◆ ck_ir_v2_lower_emit_json()

int ck_ir_v2_lower_emit_json	(	const CKIRV2Graph *	input,
		CKIRV2LowerMode	mode,
		const char *	path
	)

Definition at line 219 of file ckernel_ir_v2_lower.c.

{
    if (!input || !path) {
        return -1;
    }
    CKIRV2Graph lowered = {0};
    CKMemPlan plan = {0};
    if (ck_ir_v2_lower_graph(input, mode, &lowered, &plan) != 0) {
        return -1;
    }
    int tokens_override = (mode == CK_IR_V2_LOWER_DECODE) ? 1 : -1;
    int base_context = (tokens_override >= 0) ? input->config.context_window : -1;
    int rc = ck_ir_v2_serialize_json_with_plan(&lowered, &plan,
                                               ck_ir_v2_lower_mode_name(mode),
                                               tokens_override,
                                               base_context,
                                               path);
    ck_ir_v2_free(&lowered);
    ck_mem_plan_free(&plan);
    return rc;
}

References ck_ir_v2_free(), CK_IR_V2_LOWER_DECODE, ck_ir_v2_lower_graph(), ck_ir_v2_lower_mode_name(), ck_ir_v2_serialize_json_with_plan(), ck_mem_plan_free(), CKIRV2Graph::config, and CKModelConfig::context_window.

Referenced by main().

◆ ck_ir_v2_lower_graph()

int ck_ir_v2_lower_graph	(	const CKIRV2Graph *	input,
		CKIRV2LowerMode	mode,
		CKIRV2Graph *	output,
		CKMemPlan *	plan
	)

Definition at line 177 of file ckernel_ir_v2_lower.c.

{
    if (!input || !output || !plan) {
        return -1;
    }
    memset(output, 0, sizeof(*output));
    memset(plan, 0, sizeof(*plan));
 
    output->config = input->config;
    output->has_pos_emb = input->has_pos_emb;
    output->tie_word_embeddings = input->tie_word_embeddings;
    output->fused_qkv = input->fused_qkv;
    output->gated_mlp = input->gated_mlp;
 
    if (ck_ir_v2_lower_copy_buffers(input, output) != 0 ||
        ck_ir_v2_lower_copy_nodes(input, mode, output) != 0) {
        ck_ir_v2_free(output);
        return -1;
    }
 
    int tokens_override = (mode == CK_IR_V2_LOWER_DECODE) ? 1 : -1;
    int rc = 0;
    if (mode == CK_IR_V2_LOWER_BACKWARD) {
        rc = ck_mem_plan_build_training_with_tokens(output, plan,
                                                    CK_MEM_PLAN_DEFAULT_ALIGN,
                                                    tokens_override);
    } else {
        rc = ck_mem_plan_build_inference_with_tokens(output, plan,
                                                     CK_MEM_PLAN_DEFAULT_ALIGN,
                                                     tokens_override);
    }
    if (rc != 0) {
        ck_ir_v2_free(output);
        ck_mem_plan_free(plan);
        return -1;
    }
    return 0;
}

References ck_ir_v2_free(), CK_IR_V2_LOWER_BACKWARD, ck_ir_v2_lower_copy_buffers(), ck_ir_v2_lower_copy_nodes(), CK_IR_V2_LOWER_DECODE, ck_mem_plan_build_inference_with_tokens(), ck_mem_plan_build_training_with_tokens(), CK_MEM_PLAN_DEFAULT_ALIGN, ck_mem_plan_free(), CKIRV2Graph::config, CKIRV2Graph::fused_qkv, CKIRV2Graph::gated_mlp, CKIRV2Graph::has_pos_emb, and CKIRV2Graph::tie_word_embeddings.

Referenced by ck_ir_v2_lower_emit_json().

◆ ck_ir_v2_lower_mode_from_string()

int ck_ir_v2_lower_mode_from_string	(	const char *	name,
		CKIRV2LowerMode *	out_mode
	)

Definition at line 36 of file ckernel_ir_v2_lower.c.

{
    if (!name || !out_mode) {
        return -1;
    }
    if (strcmp(name, "prefill") == 0) {
        *out_mode = CK_IR_V2_LOWER_PREFILL;
        return 0;
    }
    if (strcmp(name, "decode") == 0) {
        *out_mode = CK_IR_V2_LOWER_DECODE;
        return 0;
    }
    if (strcmp(name, "backward") == 0) {
        *out_mode = CK_IR_V2_LOWER_BACKWARD;
        return 0;
    }
    return -1;
}

References CK_IR_V2_LOWER_BACKWARD, CK_IR_V2_LOWER_DECODE, and CK_IR_V2_LOWER_PREFILL.

Referenced by ck_ir_v2_lower_node_enabled(), and main().

◆ ck_ir_v2_lower_mode_name()

const char * ck_ir_v2_lower_mode_name ( CKIRV2LowerMode mode )

Definition at line 22 of file ckernel_ir_v2_lower.c.

{
    switch (mode) {
    case CK_IR_V2_LOWER_PREFILL:
        return "prefill";
    case CK_IR_V2_LOWER_DECODE:
        return "decode";
    case CK_IR_V2_LOWER_BACKWARD:
        return "backward";
    default:
        return "unknown";
    }
}