Skip to content

cheetah.api

ctx module-attribute 

ctx = Context()

ty module-attribute 

ty = TypeBuilder()

AsmOperand module-attribute 

AsmOperand: TypeAlias = Type | Expr

AsmOption module-attribute 

AsmOption: TypeAlias = AsmOperand | AsmVolatile | AsmMemoryClobber

Type

ctx instance-attribute 

ctx: Context = ctx

id instance-attribute 

id: TypeId = id

__init__ 

__init__(ctx: Context, id: TypeId)

try_to_prim 

try_to_prim() -> Optional[core.Prim]

try_to_ptr 

try_to_ptr() -> Optional[tuple[Optional[Type], core.Mutability]]

try_to_tensor 

try_to_tensor() -> Optional[tuple[Optional[Type], core.Mutability]]

__repr__ 

__repr__()

__eq__ 

__eq__(other)

__hash__ 

__hash__()

TypeBuilder

i8 property 

i8: Type

i16 property 

i16: Type

i32 property 

i32: Type

i64 property 

i64: Type

ix2 property 

ix2: Type

ix4 property 

ix4: Type

u8 property 

u8: Type

u16 property 

u16: Type

u32 property 

u32: Type

u64 property 

u64: Type

bf16 property 

bf16: Type

bf16x2 property 

bf16x2: Type

f16 property 

f16: Type

f32 property 

f32: Type

f32x2 property 

f32x2: Type

f64 property 

f64: Type

bool property 

bool: Type

ptr 

ptr(mut: Mutability, elem: Optional[Type]) -> Type

ptr_const 

ptr_const(elem: Optional[Type]) -> Type

ptr_mut 

ptr_mut(elem: Optional[Type]) -> Type

ptr_volatile 

ptr_volatile(elem: Optional[Type]) -> Type

tensor 

tensor(mut: Mutability, elem: Type) -> Type

tensor_const 

tensor_const(elem: Type) -> Type

tensor_mut 

tensor_mut(elem: Type) -> Type

Expr

ctx instance-attribute 

ctx: Context = ctx

id instance-attribute 

id: VarId = id

val property writable 

val: Expr

__init__ 

__init__(ctx: Context, id: VarId)

type 

type() -> Type

const_like 

const_like(value) -> Expr

__add__ 

__add__(other) -> Expr

__radd__ 

__radd__(other) -> Expr

__sub__ 

__sub__(other) -> Expr

__rsub__ 

__rsub__(other) -> Expr

__mul__ 

__mul__(other) -> Expr

__rmul__ 

__rmul__(other) -> Expr

__truediv__ 

__truediv__(other) -> Expr

__rtruediv__ 

__rtruediv__(other) -> Expr

__floordiv__ 

__floordiv__(other) -> Expr

__rfloordiv__ 

__rfloordiv__(other) -> Expr

__mod__ 

__mod__(other) -> Expr

__rmod__ 

__rmod__(other) -> Expr

__eq__ 

__eq__(other) -> Expr

__ne__ 

__ne__(other) -> Expr

__lt__ 

__lt__(other) -> Expr

__le__ 

__le__(other) -> Expr

__gt__ 

__gt__(other) -> Expr

__ge__ 

__ge__(other) -> Expr

__rshift__ 

__rshift__(other) -> Expr

__lshift__ 

__lshift__(other) -> Expr

id_repr 

id_repr(id)

__repr__ 

__repr__()

offset 

offset(index) -> Expr

__getitem__ 

__getitem__(index) -> Expr

__setitem__ 

__setitem__(index, value) -> None

cast 

cast(cast_ptr_type: Type) -> Expr

IfElseHandle

ctx instance-attribute 

ctx: Context = ctx

__init__ 

__init__(ctx: Context)

then 

then()

else_ 

else_()

AsmVolatile dataclass

volatile class-attribute instance-attribute 

volatile: bool = True

__init__ 

__init__(volatile: bool = True) -> None

AsmMemoryClobber dataclass

memory_clobber class-attribute instance-attribute 

memory_clobber: bool = True

__init__ 

__init__(memory_clobber: bool = True) -> None

Func dataclass

ctx instance-attribute 

ctx: Context

id instance-attribute 

id: FuncId

__init__ 

__init__(ctx: Context, id: FuncId) -> None

void_ptr 

void_ptr() -> Expr

__call__ 

__call__(*args: Expr) -> Expr

Params

params instance-attribute 

params: dict[str, Type] = kwargs

__init__ 

__init__(**kwargs: Type)

const 

const(value, ty: Type) -> Expr

set_name 

set_name(name: str, val, force: bool = False) -> None

if_ 

if_(cond: Expr)

if_else 

if_else(cond) -> Iterator[IfElseHandle]

cond 

cond(cond, then, else_)

for_ 

for_(a: Expr | int, b: Optional[Expr | int] = None, c: Optional[Expr | int] = None, /, *, unroll: Optional[int | Literal[True]] = None) -> Iterator[Expr]

ret 

ret(value: Expr | None = None) -> None

alloc 

alloc(ty: Type, init=None, name=None) -> Expr

alloc_array 

alloc_array(ty: Type, count: int) -> Expr

alloc_shared 

alloc_shared(ty: Type) -> Expr

alloc_shared_array 

alloc_shared_array(ty: Type, count: int) -> Expr

alloc_extern_shared 

alloc_extern_shared(ty: Type) -> Expr

alloc_tensor 

alloc_tensor(ty: Type, count: int) -> Expr

and_ 

and_(*args) -> Expr

or_ 

or_(*args) -> Expr

not_ 

not_(arg)

raw_expr 

raw_expr(src: str, out_ty: Optional[Type] = None, /, *args: Expr, **kwargs: Expr) -> Optional[Expr]

raw_stmt 

raw_stmt(src: str, /, *args: Expr, **kwargs: Expr) -> None

asm 

asm(src: str, *args: AsmOption, **kwargs: AsmOperand) -> tuple[Expr, ...] | dict[str | int, Expr]

asm_volatile 

asm_volatile(src: str, *args: AsmOption, **kwargs: AsmOperand) -> tuple[Expr, ...] | dict[str | int, Expr]

thread_idx_x 

thread_idx_x() -> Expr

thread_idx_y 

thread_idx_y() -> Expr

block_dim_x 

block_dim_x() -> Expr

block_dim_y 

block_dim_y() -> Expr

block_dim_z 

block_dim_z() -> Expr

block_idx_x 

block_idx_x() -> Expr

block_idx_y 

block_idx_y() -> Expr

block_idx_z 

block_idx_z() -> Expr

grid_dim_x 

grid_dim_x() -> Expr

grid_dim_y 

grid_dim_y() -> Expr

grid_dim_z 

grid_dim_z() -> Expr

kernel 

kernel(params, *, raw_attrs: Optional[list[str]] = None)

fn 

fn(params, *, host: bool = False, device: bool = False, force_inline: bool = False, ret: Optional[Type] = None, raw_attrs: Optional[list[str]] = None)

render 

render(entry_points: list[Func | tuple[str, Func]] | tuple[str, Func] | Func, headers: list[str] = [], bind: bool = False, bind_name: Optional[str] = None) -> str

Render the code for the given entry points.

Parameters:

Name Type Description Default
entry_points list[Func | tuple[str, Func]] | tuple[str, Func] | Func

One of: - A Func (no export name) - A tuple of (export_name, Func) - A list of Funcs and/or (export_name, Func) tuples

 required

Returns:

Type Description
str

The generated code as a string.