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.PREDEF_HOST_FUNC_PY
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_DESC_INIT_FUNC_PY
tilelang.jit.adapter.wrapper.KERNEL_LAUNCH_FUNC_PY
tilelang.jit.adapter.wrapper.logger
Classes
-------
.. autoapisummary::
tilelang.jit.adapter.wrapper.BaseWrapper
tilelang.jit.adapter.wrapper.TLCUDASourceWrapper
tilelang.jit.adapter.wrapper.TLNVRTCSourceWrapper
tilelang.jit.adapter.wrapper.TLHIPSourceWrapper
tilelang.jit.adapter.wrapper.TLCPUSourceWrapper
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} != CUDA_SUCCESS) {{
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
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
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
Show Value
.. code-block:: python
"""
extern "C" int call({}) {{
{}
return 0;
}}
"""
.. raw:: html
Show Value
.. code-block:: python
"""
import cuda.bindings.driver
import ctypes
_function_names = {}
def call({}):
{}
"""
.. 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
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
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
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
Show Value
.. code-block:: python
"""
{0}_type = cuda.bindings.driver.CUtensorMapDataType({1})
{0}_tensorRank = {2}
{0}_globalAddress = {3}.data_ptr()
{0}_globalDim = [{4}]
{0}_globalStride = [{5}][1:]
{0}_boxDim = [{6}]
{0}_elementStrides = [{7}]
{0}_interleave = cuda.bindings.driver.CUtensorMapInterleave({8})
{0}_swizzle = cuda.bindings.driver.CUtensorMapSwizzle({9})
{0}_l2Promotion = cuda.bindings.driver.CUtensorMapL2promotion({10})
{0}_oobFill = cuda.bindings.driver.CUtensorMapFloatOOBfill({11})
res, {0} = cuda.bindings.driver.cuTensorMapEncodeTiled(
{0}_type,
{0}_tensorRank,
{0}_globalAddress,
{0}_globalDim,
{0}_globalStride,
{0}_boxDim,
{0}_elementStrides,
{0}_interleave,
{0}_swizzle,
{0}_l2Promotion,
{0}_oobFill,
)
if res != cuda.bindings.driver.CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to initialize the TMA descriptor {0}: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
res = cuda.bindings.driver.cuKernelSetAttribute(
cuda.bindings.driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
{7},
kernels["{0}"],
cuda.bindings.driver.CUdevice({10})
)[0]
if res != cuda.bindings.driver.CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to set max dynamic shared memory size to {7} for kernel {0}: {{res}}")
config = cuda.bindings.driver.CUlaunchConfig()
config.gridDimX = {1}
config.gridDimY = {2}
config.gridDimZ = {3}
config.blockDimX = {4}
config.blockDimY = {5}
config.blockDimZ = {6}
config.sharedMemBytes = {7}
config.hStream = stream
arg_values = {8}
arg_types = {9}
res = cuda.bindings.driver.cuLaunchKernelEx(config, kernels["{0}"], (arg_values, arg_types), 0)[0]
if res != cuda.bindings.driver.CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to launch kernel {0}: {{res}}")
"""
.. raw:: html