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_c_postproc`(func[, override])	Register a post-processing function for C host code generation.
`register_metal_postproc`(func[, override])	Register a post-processing function for Metal 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.
`register_c_postproc_callback`([func, override])	Decorator for registering C host post-processing callback function.
`register_metal_postproc_callback`([func, override])	Decorator for registering Metal post-processing callback function.

Module Contents¶

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

参数:

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)¶

参数:

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)¶

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.

参数:

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)¶

参数:

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): ...

参数:

func (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): ...

参数:

func (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_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): ...

参数:

func (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_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): ...

参数:

func (Callable | bool) -- The function to be decorated or a boolean override flag
override (bool) -- Whether to override existing registered function. Defaults to True.