tilelang.jit.adapter.wrapper
============================
.. py:module:: tilelang.jit.adapter.wrapper
Attributes
----------
.. autoapisummary::
tilelang.jit.adapter.wrapper.PREDEF_ATTRIBUTE_SET_DYNAMIC_MEMORY
tilelang.jit.adapter.wrapper.PREDEF_ATTRIBUTE_SET_DYNAMIC_MEMORY_HIP
tilelang.jit.adapter.wrapper.PREDEF_INIT_FUNC
tilelang.jit.adapter.wrapper.PREDEF_HOST_FUNC
tilelang.jit.adapter.wrapper.L2_PERSISTENT_MAP_CREATE_HANDLE
tilelang.jit.adapter.wrapper.L2_PERSISTENT_MAP_INIT_FUNC
tilelang.jit.adapter.wrapper.L2_PERSISTENT_MAP_RESET_HANDLE
tilelang.jit.adapter.wrapper.TMA_DESC_INIT_FUNC
tilelang.jit.adapter.wrapper.TMA_IM2COL_DESC_INIT_FUNC
tilelang.jit.adapter.wrapper.logger
Classes
-------
.. autoapisummary::
tilelang.jit.adapter.wrapper.BaseWrapper
tilelang.jit.adapter.wrapper.TLCUDASourceWrapper
tilelang.jit.adapter.wrapper.TLHIPSourceWrapper
tilelang.jit.adapter.wrapper.TLCPUSourceWrapper
tilelang.jit.adapter.wrapper.TLMetalSourceWrapper
tilelang.jit.adapter.wrapper.TLWrapper
tilelang.jit.adapter.wrapper.TLPyWrapper
Module Contents
---------------
.. py:data:: PREDEF_ATTRIBUTE_SET_DYNAMIC_MEMORY
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
cudaError_t result_{0} = cudaFuncSetAttribute({0}, cudaFuncAttributeMaxDynamicSharedMemorySize, {1});
if (result_{0} != cudaSuccess) {{
snprintf(error_buf, ERROR_BUF_SIZE, "Failed to set the allowed dynamic shared memory size to %d with error: %s", {1}, cudaGetErrorString(result_{0}));
return -1;
}}
"""
.. raw:: html
.. py:data:: PREDEF_ATTRIBUTE_SET_DYNAMIC_MEMORY_HIP
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
if ({1} > 65536) {{
snprintf(error_buf, ERROR_BUF_SIZE, "Failed to set the allowed dynamic shared memory size for {0} to %d", {1});
return -1;
}}
return 0;
"""
.. raw:: html
.. py:data:: PREDEF_INIT_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
#define ERROR_BUF_SIZE 1024
static char error_buf[ERROR_BUF_SIZE];
extern "C" const char* get_last_error() {{
return error_buf;
}}
extern "C" int init() {{
error_buf[0] = '\0';
{0}
return 0;
}}
"""
.. raw:: html
.. py:data:: PREDEF_HOST_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
extern "C" int call({}) {{
{}
return 0;
}}
"""
.. raw:: html
.. py:data:: L2_PERSISTENT_MAP_CREATE_HANDLE
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
cudaStreamAttrValue stream_attribute;
size_t init_persisting_l2_cache_size;
cudaDeviceGetLimit(&init_persisting_l2_cache_size, cudaLimitPersistingL2CacheSize);
"""
.. raw:: html
.. py:data:: L2_PERSISTENT_MAP_INIT_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
stream_attribute.accessPolicyWindow.hitRatio = {1};
stream_attribute.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting;
stream_attribute.accessPolicyWindow.missProp = cudaAccessPropertyStreaming;
cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, {2});
stream_attribute.accessPolicyWindow.base_ptr = (void*)({0});
stream_attribute.accessPolicyWindow.num_bytes = {2};
cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);
"""
.. raw:: html
.. py:data:: L2_PERSISTENT_MAP_RESET_HANDLE
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
stream_attribute.accessPolicyWindow.num_bytes = 0;
cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &stream_attribute);
cudaCtxResetPersistingL2Cache();
cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, init_persisting_l2_cache_size);
"""
.. raw:: html
.. py:data:: TMA_DESC_INIT_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
CUtensorMap {0};
CUtensorMapDataType {0}_type= (CUtensorMapDataType){1};
cuuint32_t {0}_tensorRank= {2};
void *{0}_globalAddress= {3};
cuuint64_t {0}_globalDim[{2}]= {{{4}}};
cuuint64_t {0}_globalStride[{2}]= {{{5}}};
cuuint32_t {0}_boxDim[{2}]= {{{6}}};
cuuint32_t {0}_elementStrides[{2}]= {{{7}}};
CUtensorMapInterleave {0}_interleave= (CUtensorMapInterleave){8};
CUtensorMapSwizzle {0}_swizzle= (CUtensorMapSwizzle){9};
CUtensorMapL2promotion {0}_l2Promotion= (CUtensorMapL2promotion){10};
CUtensorMapFloatOOBfill {0}_oobFill= (CUtensorMapFloatOOBfill){11};
CUresult {0}_result = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)(
&{0}, {0}_type, {0}_tensorRank, {0}_globalAddress, {0}_globalDim, {0}_globalStride + 1, {0}_boxDim, {0}_elementStrides, {0}_interleave, {0}_swizzle, {0}_l2Promotion, {0}_oobFill);
if ({0}_result != CUDA_SUCCESS) {{
std::stringstream ss;
ss << "Error: Failed to initialize the TMA descriptor {0}";
snprintf(error_buf, ERROR_BUF_SIZE, "%s", ss.str().c_str());
return -1;
}}
"""
.. raw:: html
.. py:data:: TMA_IM2COL_DESC_INIT_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
CUtensorMap {0};
CUtensorMapDataType {0}_type= (CUtensorMapDataType){1};
cuuint32_t {0}_tensorRank= {2};
void *{0}_globalAddress= {3};
cuuint64_t {0}_globalDim[{2}]= {{{4}}};
cuuint64_t {0}_globalStride[{2}]= {{{5}}};
cuuint32_t {0}_elementStrides[{2}]= {{{6}}};
int {0}_lowerCorner[{2} - 2]= {{{7}}};
int {0}_upperCorner[{2} - 2]= {{{8}}};
cuuint32_t {0}_channelsPerPixel= {9};
cuuint32_t {0}_pixelsPerColumn= {10};
CUtensorMapInterleave {0}_interleave= (CUtensorMapInterleave){11};
CUtensorMapSwizzle {0}_swizzle= (CUtensorMapSwizzle){12};
CUtensorMapL2promotion {0}_l2Promotion= (CUtensorMapL2promotion){13};
CUtensorMapFloatOOBfill {0}_oobFill= (CUtensorMapFloatOOBfill){14};
CUresult {0}_result = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeIm2col)(
&{0}, {0}_type, {0}_tensorRank, {0}_globalAddress, {0}_globalDim, {0}_globalStride + 1,
{0}_lowerCorner, {0}_upperCorner, {0}_channelsPerPixel, {0}_pixelsPerColumn, {0}_elementStrides, {0}_interleave, {0}_swizzle, {0}_l2Promotion, {0}_oobFill);
if ({0}_result != CUDA_SUCCESS) {{
std::stringstream ss;
ss << "Error: Failed to initialize the TMA descriptor {0}";
snprintf(error_buf, ERROR_BUF_SIZE, "%s", ss.str().c_str());
return -1;
}}
"""
.. raw:: html
.. py:class:: BaseWrapper
Bases: :py:obj:`abc.ABC`
Helper class that provides a standard way to create an ABC using
inheritance.
.. py:method:: wrap(*args, **kwargs)
:abstractmethod:
.. py:data:: logger
.. py:class:: TLCUDASourceWrapper(scheduled_ir_module, source, target, device_mod = None, host_mod = None, pass_configs = None)
.. py:attribute:: backend
:value: 'tl'
.. py:attribute:: device_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: host_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: pass_configs
:type: dict[str, Any] | None
:value: None
.. py:attribute:: mod
.. py:attribute:: target
.. py:attribute:: source
.. py:attribute:: function_names
:type: str | None
:value: None
.. py:attribute:: dynamic_smem_buf
:type: int | None
:value: None
.. py:attribute:: block_info
:type: list[int] | dict
:value: [1, 1, 1]
.. py:attribute:: grid_info
:type: list[int] | dict
:value: [1, 1, 1]
.. py:attribute:: tma_descriptor_args
:type: dict | None
:value: None
.. py:attribute:: l2_persistent_map
:type: dict[str, dict] | None
.. py:attribute:: srcpath
:type: str | None
:value: None
.. py:attribute:: libpath
:type: str | None
:value: None
.. py:attribute:: lib_code
:type: str | None
.. py:method:: is_tma_descriptor_arg(arg_name)
.. py:method:: create_dispatch_func(code, function_informations)
.. py:method:: generate_l2_persistent_map(function_name)
.. py:method:: generate_tma_descriptor_args(desc_name_map, desc_name_var_map)
.. py:method:: parse_source_information()
.. py:method:: get_dynamic_symbolic_set(prim_func)
.. py:method:: get_init_func()
.. py:method:: update_lib_code(code)
.. py:method:: get_stream_type()
.. py:property:: prim_func
.. py:property:: device_func
.. py:property:: host_func
.. py:class:: TLHIPSourceWrapper(scheduled_ir_module, source, target, device_mod = None, host_mod = None, pass_configs = None)
Bases: :py:obj:`TLCUDASourceWrapper`
A wrapper class for the TileLang HIP backend.
.. py:method:: get_init_func()
.. py:method:: get_stream_type()
.. py:class:: TLCPUSourceWrapper(scheduled_ir_module, source, target, device_mod = None, host_mod = None, pass_configs = None)
.. py:attribute:: INIT_FUNC
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
#define ERROR_BUF_SIZE 1024
static char error_buf[ERROR_BUF_SIZE];
extern "C" const char* get_last_error() {
return error_buf;
}
extern "C" int init() {
error_buf[0] = '\0';
return 0;
}
"""
.. raw:: html
.. py:attribute:: CALL_PREFIX
.. py:attribute:: backend
:value: 'tl'
.. py:attribute:: device_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: host_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: pass_configs
:type: dict[str, Any] | None
:value: None
.. py:attribute:: mod
.. py:attribute:: target
.. py:attribute:: source
.. py:attribute:: function_names
:type: str | None
:value: None
.. py:attribute:: dynamic_smem_buf
:type: int | None
:value: None
.. py:attribute:: srcpath
:type: str | None
:value: None
.. py:attribute:: libpath
:type: str | None
:value: None
.. py:attribute:: lib_code
:type: str | None
.. py:method:: create_call_func(code, function_informations)
.. py:method:: parse_source_information()
.. py:method:: get_dynamic_symbolic_set(prim_func)
.. py:method:: get_cpu_init_func()
.. py:method:: update_lib_code(code)
.. py:property:: prim_func
.. py:class:: TLMetalSourceWrapper(scheduled_ir_module, source, target, device_mod = None, host_mod = None, pass_configs = None)
.. py:attribute:: mod
.. py:attribute:: target
.. py:attribute:: source
.. py:attribute:: pass_configs
:value: None
.. py:attribute:: device_mod
:value: None
.. py:attribute:: host_mod
:value: None
.. py:attribute:: lib_code
.. py:method:: update_lib_code(code)
.. py:class:: TLWrapper(target)
Bases: :py:obj:`BaseWrapper`
A wrapper class for the TileLang backend.
.. py:attribute:: device_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: host_mod
:type: tvm.IRModule | None
:value: None
.. py:attribute:: pass_configs
:type: dict[str, Any] | None
:value: None
.. py:attribute:: target
:type: tvm.target.Target | None
:value: None
.. py:attribute:: lib
:type: object | None
:value: None
.. py:attribute:: scheduled_ir_module
:value: None
.. py:method:: assign_optimized_module(scheduled_ir_module)
.. py:method:: assign_pass_configs(pass_configs)
.. py:method:: assign_host_module(host_mod)
.. py:method:: assign_device_module(device_mod)
.. py:method:: wrap(c_source)
.. py:class:: TLPyWrapper(target)
Bases: :py:obj:`TLWrapper`
A wrapper class for the TileLang backend.
.. py:method:: wrap(c_source)