tilelang.jit.adapter.nvrtc.wrapper
==================================
.. py:module:: tilelang.jit.adapter.nvrtc.wrapper
.. autoapi-nested-parse::
NVRTC Source Wrapper for TileLang.
Generates Python runtime code for launching CUDA kernels compiled via NVRTC.
Why this exists:
- NVRTC compiles kernels at runtime, needs Python launch code (not C++)
- TMA descriptors must be initialized once per unique buffer, not per kernel
- L2 cache policies require explicit CUDA Driver API setup/teardown
Key design:
- Two-pass generation: collect all descriptors first, then generate launches
- Dict-based deduplication ensures TMA descriptors created only once
- Generates pure Python using cuda.bindings.driver for zero C++ dependency
Attributes
----------
.. autoapisummary::
tilelang.jit.adapter.nvrtc.wrapper.PREDEF_HOST_FUNC_PY
tilelang.jit.adapter.nvrtc.wrapper.TMA_DESC_INIT_FUNC_PY
tilelang.jit.adapter.nvrtc.wrapper.TMA_IM2COL_DESC_INIT_FUNC_PY
tilelang.jit.adapter.nvrtc.wrapper.L2_PERSISTENT_MAP_CREATE_HANDLE_PY
tilelang.jit.adapter.nvrtc.wrapper.L2_PERSISTENT_MAP_INIT_FUNC_PY
tilelang.jit.adapter.nvrtc.wrapper.L2_PERSISTENT_MAP_RESET_HANDLE_PY
tilelang.jit.adapter.nvrtc.wrapper.KERNEL_LAUNCH_FUNC_PY
Classes
-------
.. autoapisummary::
tilelang.jit.adapter.nvrtc.wrapper.TLNVRTCSourceWrapper
Module Contents
---------------
.. py:data:: PREDEF_HOST_FUNC_PY
:value: Multiline-String
.. raw:: html
Show Value
.. code-block:: python
"""
from cuda.bindings.driver import (
CUtensorMapDataType,
CUtensorMapInterleave,
CUtensorMapSwizzle,
CUtensorMapL2promotion,
CUtensorMapFloatOOBfill,
cuTensorMapEncodeTiled,
cuTensorMapEncodeIm2col,
CUresult,
cuKernelSetAttribute,
CUfunction_attribute,
CUdevice,
CUlaunchConfig,
cuLaunchKernelEx,
cuuint64_t,
cuuint32_t,
CUkernel,
)
import ctypes
_function_names = {}
def call({}):
{}
"""
.. raw:: html
Show Value
.. code-block:: python
"""
{0}_type = 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 = CUtensorMapInterleave({8})
{0}_swizzle = CUtensorMapSwizzle({9})
{0}_l2Promotion = CUtensorMapL2promotion({10})
{0}_oobFill = CUtensorMapFloatOOBfill({11})
res, {0} = 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 != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to initialize the TMA descriptor {0}: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
{0}_type = CUtensorMapDataType({1})
{0}_tensorRank = {2}
{0}_globalAddress = {3}.data_ptr()
{0}_globalDim = [{4}]
{0}_globalStride = [{5}][1:]
{0}_elementStrides = [{6}]
{0}_lowerCorner = [{7}]
{0}_upperCorner = [{8}]
{0}_channelsPerPixel = {9}
{0}_pixelsPerColumn = {10}
{0}_interleave = CUtensorMapInterleave({11})
{0}_swizzle = CUtensorMapSwizzle({12})
{0}_l2Promotion = CUtensorMapL2promotion({13})
{0}_oobFill = CUtensorMapFloatOOBfill({14})
res, {0} = cuTensorMapEncodeIm2col(
{0}_type,
{0}_tensorRank,
{0}_globalAddress,
{0}_globalDim,
{0}_globalStride,
{0}_lowerCorner,
{0}_upperCorner,
{0}_channelsPerPixel,
{0}_pixelsPerColumn,
{0}_elementStrides,
{0}_interleave,
{0}_swizzle,
{0}_l2Promotion,
{0}_oobFill,
)
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to initialize the TMA descriptor {0}: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
from cuda.bindings.driver import (
CUstreamAttrValue,
CUstreamAttrID,
CUlimit,
CUaccessProperty,
cuCtxGetLimit,
cuCtxSetLimit,
cuStreamSetAttribute,
cuCtxResetPersistingL2Cache,
)
stream_attribute = CUstreamAttrValue()
res, init_persisting_l2_cache_size = cuCtxGetLimit(CUlimit.CU_LIMIT_PERSISTING_L2_CACHE_SIZE)
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to get L2 cache size limit: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
stream_attribute.accessPolicyWindow.hitRatio = {1}
stream_attribute.accessPolicyWindow.hitProp = CUaccessProperty.CU_ACCESS_PROPERTY_PERSISTING
stream_attribute.accessPolicyWindow.missProp = CUaccessProperty.CU_ACCESS_PROPERTY_STREAMING
res = cuCtxSetLimit(CUlimit.CU_LIMIT_PERSISTING_L2_CACHE_SIZE, {2})[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to set L2 cache size limit: {{res}}")
stream_attribute.accessPolicyWindow.base_ptr = {0}.data_ptr()
stream_attribute.accessPolicyWindow.num_bytes = {2}
res = cuStreamSetAttribute(stream, CUstreamAttrID.CU_LAUNCH_ATTRIBUTE_ACCESS_POLICY_WINDOW, stream_attribute)[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to set stream L2 access policy: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
stream_attribute.accessPolicyWindow.num_bytes = 0
res = cuStreamSetAttribute(stream, CUstreamAttrID.CU_LAUNCH_ATTRIBUTE_ACCESS_POLICY_WINDOW, stream_attribute)[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to reset stream L2 access policy: {{res}}")
res = cuCtxResetPersistingL2Cache()[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to reset L2 cache: {{res}}")
res = cuCtxSetLimit(CUlimit.CU_LIMIT_PERSISTING_L2_CACHE_SIZE, init_persisting_l2_cache_size)[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to restore L2 cache size limit: {{res}}")
"""
.. raw:: html
Show Value
.. code-block:: python
"""
res = cuKernelSetAttribute(
CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
{7},
kernels["{0}"],
CUdevice({10})
)[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to set max dynamic shared memory size to {7} for kernel {0}: {{res}}")
config = 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 = cuLaunchKernelEx(config, kernels["{0}"], (arg_values, arg_types), 0)[0]
if res != CUresult.CUDA_SUCCESS:
raise RuntimeError(f"Failed to launch kernel {0}: {{res}}")
"""
.. raw:: html