tilelang.engine.callback¶
Functions¶
|
Register a post-processing function for CUDA code generation. |
|
Register a post-processing function for HIP code generation. |
|
Register a post-processing function for C host code generation. |
|
Register a post-processing function for Metal code generation. |
|
Decorator for registering CUDA post-processing callback function. |
|
Decorator for registering HIP post-processing callback function. |
|
Decorator for registering C host post-processing callback function. |
|
Decorator for registering Metal 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_c_postproc(func, override=True)¶
Register a post-processing function for C host code generation.
This callback intercepts C host code emitted by TileLang just before it is wrapped into a CSourceModule. It should take the generated code string and the Target as inputs, and return the (possibly) modified code.
- 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_metal_postproc(func, override=True)¶
Register a post-processing function for Metal 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): …
- 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): …
- tilelang.engine.callback.register_c_postproc_callback(func=None, override=True)¶
Decorator for registering C host post-processing callback function.
- Can be used with or without parentheses:
@register_c_postproc_callback def func(code, target): …
@register_c_postproc_callback() def func(code, target): …
@register_c_postproc_callback(override=False) def func(code, target): …
- tilelang.engine.callback.register_metal_postproc_callback(func=None, override=True)¶
Decorator for registering Metal post-processing callback function.
- Can be used with or without parentheses:
@register_metal_postproc_callback def func(code, target): …
@register_metal_postproc_callback() def func(code, target): …
@register_metal_postproc_callback(override=False) def func(code, target): …