tilelang.engine.callback module#

tilelang.engine.callback.register_cuda_postproc(func: Callable[[str, Target], str], override: bool = True)#

Register a post-processing function for CUDA code generation.

Parameters:

  • func – A callable that takes generated code (str) and target (Target) as input, and returns the processed code (str).

  • override – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_cuda_postproc_callback(func: Optional[Union[Callable, bool]] = None, override: bool = True)#

Decorator for registering CUDA post-processing callback function.

Can be used with or without parentheses:

@register_cuda_postproc_callback def func(code, target): …

@register_cuda_postproc_callback() def func(code, target): …

@register_cuda_postproc_callback(override=False) def func(code, target): …

Parameters:

  • func – The function to be decorated or a boolean override flag

  • override – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_hip_postproc(func: Callable[[str, Target], str], override: bool = True)#

Register a post-processing function for HIP code generation.

Parameters:

  • func – A callable that takes generated code (str) and target (Target) as input, and returns the processed code (str).

  • override – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_hip_postproc_callback(func: Optional[Union[Callable, bool]] = None, override: bool = True)#

Decorator for registering HIP post-processing callback function.

Can be used with or without parentheses:

@register_hip_postproc_callback def func(code, target): …

@register_hip_postproc_callback() def func(code, target): …

@register_hip_postproc_callback(override=False) def func(code, target): …

Parameters:

  • func – The function to be decorated or a boolean override flag

  • override – Whether to override existing registered function. Defaults to True.