ckernel_mem_plan.h File Reference

#include "ckernel_ir_v2.h"
#include <stddef.h>

Go to the source code of this file.

Data Structures
struct	CKMemPlan
struct	CKMemSpan

Macros
#define	CK_MEM_PLAN_DEFAULT_ALIGN   64

Enumerations
enum	CKMemArenaKind { CK_MEM_ARENA_WEIGHTS = 0 , CK_MEM_ARENA_ACTIVATIONS = 1 , CK_MEM_ARENA_GRADS = 2 , CK_MEM_ARENA_COUNT }

Functions
int	ck_mem_plan_build_inference (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes)
int	ck_mem_plan_build_inference_with_tokens (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes, int tokens_override)
int	ck_mem_plan_build_training (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes)
int	ck_mem_plan_build_training_with_tokens (const CKIRV2Graph *graph, CKMemPlan *plan, size_t alignment_bytes, int tokens_override)
void	ck_mem_plan_free (CKMemPlan *plan)

Macro Definition Documentation

CK_MEM_PLAN_DEFAULT_ALIGN

#define CK_MEM_PLAN_DEFAULT_ALIGN   64

Definition at line 12 of file ckernel_mem_plan.h.

Enumeration Type Documentation

CKMemArenaKind

enum CKMemArenaKind

Enumerator
CK_MEM_ARENA_WEIGHTS
CK_MEM_ARENA_ACTIVATIONS
CK_MEM_ARENA_GRADS
CK_MEM_ARENA_COUNT

Definition at line 14 of file ckernel_mem_plan.h.

             {
    CK_MEM_ARENA_WEIGHTS = 0,
    CK_MEM_ARENA_ACTIVATIONS = 1,
    CK_MEM_ARENA_GRADS = 2,
    CK_MEM_ARENA_COUNT
} CKMemArenaKind;

Function Documentation

ck_mem_plan_build_inference()

int ck_mem_plan_build_inference	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes
	)

Definition at line 216 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 0, -1);
}

References build_plan().

ck_mem_plan_build_inference_with_tokens()

int ck_mem_plan_build_inference_with_tokens	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes,
		int	tokens_override
	)

Definition at line 230 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 0, tokens_override);
}

References build_plan().

Referenced by ck_codegen_v2_emit_runtime(), and ck_ir_v2_lower_graph().

ck_mem_plan_build_training()

int ck_mem_plan_build_training	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes
	)

Definition at line 223 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 1, -1);
}

References build_plan().

ck_mem_plan_build_training_with_tokens()

int ck_mem_plan_build_training_with_tokens	(	const CKIRV2Graph *	graph,
		CKMemPlan *	plan,
		size_t	alignment_bytes,
		int	tokens_override
	)

Definition at line 238 of file ckernel_mem_plan.c.

{
    return build_plan(graph, plan, alignment_bytes, 1, tokens_override);
}

References build_plan().

Referenced by ck_codegen_v2_emit_runtime(), and ck_ir_v2_lower_graph().

ck_mem_plan_free()

void ck_mem_plan_free

(

CKMemPlan *

plan

)

Definition at line 246 of file ckernel_mem_plan.c.

{
    if (!plan) {
        return;
    }
    free(plan->spans);
    memset(plan, 0, sizeof(*plan));
}

References CKMemPlan::spans.

Referenced by ck_codegen_v2_emit_runtime(), ck_ir_v2_lower_emit_json(), and ck_ir_v2_lower_graph().