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)