tilelang.engine.param module

The profiler and convert to torch utils

class tilelang.engine.param.CompiledArtifact(host_mod: IRModule, device_mod: IRModule, params: List[KernelParam], kernel_source: str, rt_mod: Optional[Module] = None)

Bases: object

Represents a compiled kernel artifact containing both host and device code. Stores all necessary components for kernel execution in the TVM runtime.

device_mod: IRModule

host_mod: IRModule

kernel_source: str

params: List[KernelParam]

rt_mod: Optional[Module] = None

class tilelang.engine.param.KernelParam(dtype: torch.dtype, shape: List[Union[int, Var]])

Bases: object

Represents parameters for a kernel operation, storing dtype and shape information. Used to describe tensor or scalar parameters in TVM/PyTorch interop.

dtype: torch.dtype

classmethod from_buffer(buffer: Buffer)

Creates a KernelParam instance from a TVM Buffer object.

Parameters:

buffer – TVM Buffer object containing dtype and shape information

Returns:

KernelParam instance with converted dtype and shape

Raises:

ValueError – If dimension type is not supported (not IntImm or Var)

classmethod from_var(var: Var)

Creates a KernelParam instance from a TVM Variable object. Used for scalar parameters.

Parameters:

var – TVM Variable object containing dtype information

Returns:

KernelParam instance representing a scalar (empty shape)

is_boolean() → bool

Checks if the parameter represents a boolean type.

Returns:

True if parameter is a boolean type, False otherwise

Return type:

bool

is_float8() → bool

Checks if the parameter represents a float8 type.

Returns:

True if parameter is a float8 type, False otherwise

Return type:

bool

is_scalar() → bool

Checks if the parameter represents a scalar value.

Returns:

True if parameter has no dimensions (empty shape), False otherwise

Return type:

bool

is_unsigned() → bool

Checks if the parameter represents an unsigned integer type.

Returns:

True if parameter is an unsigned integer type, False otherwise

Return type:

bool

shape: List[Union[int, Var]]