tilelang.jit.adapter.wrapper module

class tilelang.jit.adapter.wrapper.BaseWrapper

Bases: ABC

abstract wrap(*args, **kwargs)

class tilelang.jit.adapter.wrapper.TLCPUSourceWrapper(scheduled_ir_module: IRModule, source: str, target: Target, device_mod: Optional[IRModule] = None, host_mod: Optional[IRModule] = None, pass_configs: Optional[Dict[str, Any]] = None)

Bases: object

CALL_PREFIX = '\n#ifdef __cplusplus\nextern "C"\n#endif\nint32_t call({}) {{\n  return {};\n}}\n'

INIT_FUNC = '\n#ifdef __cplusplus\nextern "C"\n#endif\nint32_t init() {\n    return 0;\n}\n'

backend = 'tl'

create_call_func(code, function_informations)

device_mod: Optional[IRModule] = None

get_cpu_init_func()

get_dynamic_symbolic_set(prim_func)

host_mod: Optional[IRModule] = None

parse_source_information()

pass_configs: Optional[Dict[str, Any]] = None

property prim_func

update_lib_code(code: str)

class tilelang.jit.adapter.wrapper.TLCUDASourceWrapper(scheduled_ir_module: IRModule, source: str, target: Target, device_mod: Optional[IRModule] = None, host_mod: Optional[IRModule] = None, pass_configs: Optional[Dict[str, Any]] = None)

Bases: object

backend = 'tl'

create_dispatch_func(code, function_informations)

device_mod: Optional[IRModule] = None

generate_tma_descriptor_args(desc_name_map: Dict[str, str]) → str

get_cuda_init_func()

get_dynamic_symbolic_set(prim_func)

host_mod: Optional[IRModule] = None

is_tma_descriptor_arg(arg_name: str) → bool

parse_source_information()

pass_configs: Optional[Dict[str, Any]] = None

property prim_func

update_lib_code(code: str)

class tilelang.jit.adapter.wrapper.TLHIPSourceWrapper(scheduled_ir_module: IRModule, source: str, target: Target, device_mod: Optional[IRModule] = None, host_mod: Optional[IRModule] = None, pass_configs: Optional[Dict[str, Any]] = None)

Bases: TLCUDASourceWrapper

A wrapper class for the TileLang HIP backend.

get_hip_init_func()

get_stream_type(function_args)

class tilelang.jit.adapter.wrapper.TLWrapper(target: Target)

Bases: BaseWrapper

A wrapper class for the TileLang backend.

assign_device_module(device_mod: IRModule)

assign_host_module(host_mod: IRModule)

assign_optimized_module(scheduled_ir_module: IRModule)

assign_pass_configs(pass_configs: Dict[str, Any])

device_mod: Optional[IRModule] = None

host_mod: Optional[IRModule] = None

lib: Optional[object] = None

pass_configs: Optional[Dict[str, Any]] = None

target: Optional[Target] = None

wrap(c_source: str)