[BUG] NVRTC wrapper omits host-computed scalar kernel arguments
Required prerequisites
- I have read the documentation https://tilelang.com.
- I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)
What version of TileLang are you using?
0.1.14+cuda.git6ba187e2 (source checkout).
Initially reproduced at eab74a4a; the affected wrapper is unchanged in 6ba187e2.
System information
- Installation: source build with CUDA enabled, CMake/Ninja.
- OS: Windows x64 (
win32). - Python: 3.12.13, 64-bit.
- PyTorch:
2.11.0+cu130. - NVRTC: 13.0.
Problem description
The NVRTC wrapper can omit scalar kernel arguments computed on the host after host/device splitting.
For valid host IR such as:
p = n * 3 + 1
q = p * 2 + 5
T.call_packed("main_kernel", A, B, p, q, 128)the device signature is main_kernel(A, B, p, q), but the generated NVRTC dispatcher contains only:
arg_values = A.data_ptr(), B.data_ptr()
arg_types = ctypes.c_void_p, ctypes.c_void_pSplitHostDevice correctly captures the prepared values. The wrapper subsequently matches CUDA parameter names against the original function's arguments, where p and q do not exist. The existing TVM-FFI path executes the same host preparation correctly.
This may have gone unnoticed because ordinary compilation inlines this example's arithmetic into the kernel before SplitHostDevice: A[i] + p + q becomes A[i] + 9 * n + 8, leaving only the original scalar parameter n to forward. This happens before the NVRTC adapter; writing the bindings outside T.Kernel does not prevent it.
The reproducer preserves the host bindings to exercise the adapter boundary. It does not demonstrate a failure through the default compilation pipeline. Launch-time hoisting such as #3261 could expose this gap if the computed values remain on the host and become additional kernel arguments.
Reproducible example code
Requires a CUDA-enabled TileLang source build. This inspects generated source without launching the incomplete argument list.
import tilelang as tl
import tilelang.language as T
from tilelang import tvm
from tilelang.backend.module import create_backend_context
from tilelang.engine.lower import device_codegen_without_compile, get_device_call, get_host_call
from tilelang.jit.adapter.nvrtc.wrapper import TLNVRTCSourceWrapper
@T.prim_func
def main(A: T.Tensor((128,), "int32"), B: T.Tensor((128,), "int32"), n: T.int32):
p = T.bind(n * 3 + 1)
q = T.bind(p * 2 + 5)
with T.Kernel(1, threads=128):
i = T.get_thread_binding()
B[i] = A[i] + p + q
ctx = create_backend_context({"kind": "cuda", "arch": "sm_80"}, "c", "nvrtc")
original = tvm.IRModule({"main": main})
# Preserve host bindings to exercise the adapter boundary.
mod = tvm.transform.Sequential(
[
tvm.tirx.transform.BindTarget(ctx.target),
tl.transform.MaterializeKernelLaunch(),
tl.transform.LowerOpaqueBlock(),
tl.transform.AnnotateDeviceRegions(),
tl.transform.SplitHostDevice(),
tvm.tirx.transform.AnnotateEntryFunc(),
tl.transform.MakePackedAPI(),
tl.transform.LowerDeviceKernelLaunch(),
]
)(original)
host = tvm.tirx.transform.Filter(get_host_call())(mod)
device = tvm.tirx.transform.Filter(get_device_call())(mod)
source = device_codegen_without_compile(device, ctx).inspect_source()
wrapper = TLNVRTCSourceWrapper(original, source, ctx.target, device_mod=device, host_mod=host)
print(source)
print(wrapper.host_func)Traceback
No traceback is produced by the source-generation reproducer. The malformed dispatcher is inspected, not executed with an incomplete CUDA argument list.
Expected behavior
Resolve launch operands from the actual host call, preserve the required scalar bindings and their execution scope, and marshal every device argument in signature order. Unsupported host computations should produce an explicit diagnostic rather than an incomplete CUDA argument list.
Additional context
The existing TVM-FFI path and a separate IR-driven NVRTC prototype passed changing-input and guarded-launch checks on the same host/device program.
This differs from #2755/#2756, which concerned forwarding original scalar parameters and dynamic strides. The missing values here are introduced by host preparation and are not original function parameters. Related preparation work is discussed in #3261.
Source: tile-ai/tilelang