tilelang.engine.callback¶

Functions¶

register_cuda_postproc(func[, override])

Register a post-processing function for CUDA code generation.

register_hip_postproc(func[, override])

Register a post-processing function for HIP code generation.

register_cuda_postproc_callback([func, override])

Decorator for registering CUDA post-processing callback function.

register_hip_postproc_callback([func, override])

Decorator for registering HIP post-processing callback function.

Module Contents¶

tilelang.engine.callback.register_cuda_postproc(func, override=True)¶

Register a post-processing function for CUDA code generation.

Parameters:

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

  • override (bool) – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_hip_postproc(func, override=True)¶

Register a post-processing function for HIP code generation.

Parameters:

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

  • override (bool) – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_cuda_postproc_callback(func=None, override=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 (Union[Callable, bool]) – The function to be decorated or a boolean override flag

  • override (bool) – Whether to override existing registered function. Defaults to True.

tilelang.engine.callback.register_hip_postproc_callback(func=None, override=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 (Union[Callable, bool]) – The function to be decorated or a boolean override flag

  • override (bool) – Whether to override existing registered function. Defaults to True.