mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-31 08:13:15 +08:00
Chennes
39b3b3e5db
T.__ldg failed to compile for float16 and bfloat16 buffers because CUDA provides __ldg overloads for __half and __nv_bfloat16, while TileLang uses CUTLASS half_t and bfloat16_t for these scalar types. Add forwarding overloads in common.h that reinterpret the bit-compatible pointers as the native CUDA types and wrap the loaded values back in the CUTLASS types. The native 16-bit __ldg is available on SM60+, so no architecture-specific fallback is needed. Update the stale 32-bit-only API wording and add exact round-trip coverage for both dtypes.
57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_language_ldg_codegen():
|
|
N = 128
|
|
|
|
@T.prim_func
|
|
def main(
|
|
x: T.Tensor((N,), T.float32),
|
|
y: T.Tensor((N,), T.float32),
|
|
):
|
|
with T.Kernel(N, threads=32) as pid:
|
|
# Explicitly request read-only cache load for x[pid]
|
|
y[pid] = T.__ldg(x[pid]) + 1.0
|
|
|
|
# Compile for CUDA and retrieve generated CUDA source
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
src = kernel.get_kernel_source()
|
|
print(src)
|
|
# Assert that codegen uses __ldg on CUDA backend
|
|
# We look for the intrinsic call with address-of argument
|
|
assert "__ldg(" in src, "Expected __ldg call in generated CUDA source"
|
|
assert "__ldg(&" in src or "__ldg(&(" in src, "Expected address-of form in __ldg call"
|
|
|
|
|
|
def run_ldg_roundtrip(dtype, N=128):
|
|
@T.prim_func
|
|
def main(
|
|
x: T.Tensor((N,), dtype),
|
|
y: T.Tensor((N,), dtype),
|
|
):
|
|
with T.Kernel(N, threads=32) as pid:
|
|
y[pid] = T.__ldg(x[pid])
|
|
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
# Match the address-of form emitted by CUDA codegen.
|
|
assert "__ldg(&(" in kernel.get_kernel_source()
|
|
|
|
x = torch.randn(N, device="cuda").to(dtype.as_torch())
|
|
assert torch.equal(kernel(x), x)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", [T.float16, T.bfloat16])
|
|
def test_language_ldg_roundtrip(dtype):
|
|
run_ldg_roundtrip(dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|