#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";
     }
 }