mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-30 08:03:09 +08:00
a84c9b1433
[CuTeDSL] Integrate host codegen call sites
Drive CuTeDSL launcher generation from lowered host tvm_call_packed call sites so multi-kernel wrappers preserve host IR order and repeated launches can be represented.
Expose separate device, host, and generated module sources through the CuTeDSL adapter and persist them separately in the kernel/autotune caches.
Normalize bare cutedsl targets through TileLang CUDA arch detection so Hopper runners use sm_90a for CuTeDSL compilation instead of the generic sm_90 Target("cuda") result.
Add CuTeDSL host-codegen integration tests covering multi-call launcher ordering, source accessors, cache restore behavior, and bare target architecture normalization.
Co-authored-by: Jayce <jayce.su@enflame-tech.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
63 lines
2.3 KiB
Python
63 lines
2.3 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.utils.target import determine_target, normalize_cutedsl_target
|
|
|
|
|
|
def test_cutedsl_dict_target_normalizes_to_cuda_marker(monkeypatch):
|
|
from tilelang.jit.adapter.cutedsl import checks
|
|
|
|
monkeypatch.setattr(checks, "check_cutedsl_available", lambda: None)
|
|
target = determine_target({"kind": "cutedsl", "arch": "sm_80"}, return_object=True)
|
|
|
|
assert target.kind.name == "cuda"
|
|
assert target.attrs["arch"] == "sm_80"
|
|
assert {"cuda", "gpu", "cutedsl"}.issubset(set(target.keys))
|
|
|
|
|
|
def test_cutedsl_string_target_uses_detected_cuda_arch(monkeypatch):
|
|
"""Verify bare CuTeDSL targets use TileLang's CUDA arch normalization."""
|
|
|
|
from tilelang.jit.adapter.cutedsl import checks
|
|
|
|
monkeypatch.setattr(checks, "check_cutedsl_available", lambda: None)
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
|
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (9, 0))
|
|
|
|
target = determine_target("cutedsl", return_object=True)
|
|
|
|
assert target.kind.name == "cuda"
|
|
assert target.attrs["arch"] == "sm_90a"
|
|
assert {"cuda", "gpu", "cutedsl"}.issubset(set(target.keys))
|
|
|
|
|
|
def test_cutedsl_codegen_supports_tl_ptx_cp_async():
|
|
if not tvm.runtime.enabled("cuda"):
|
|
pytest.skip("TileLang CuTeDSL codegen requires TVM built with CUDA support.")
|
|
|
|
build_cutedsl = tvm.ffi.get_global_func("target.build.tilelang_cutedsl_without_compile", allow_missing=True)
|
|
if build_cutedsl is None:
|
|
pytest.skip("TileLang CuTeDSL backend is not enabled in this build.")
|
|
|
|
target = normalize_cutedsl_target({"kind": "cutedsl", "arch": "sm_80"})
|
|
assert target is not None
|
|
|
|
@T.prim_func
|
|
def prog(A: T.Tensor((16,), "uint8"), B: T.Tensor((16,), "uint8")):
|
|
with T.Kernel(1, threads=1):
|
|
A_shared = T.alloc_shared((16,), "uint8", scope="shared")
|
|
T.ptx_cp_async(T.access_ptr(A_shared[0], "w", 16), T.access_ptr(A[0], "r", 16), 16)
|
|
B[0] = A_shared[0]
|
|
|
|
artifact = lower(prog.with_attr("global_symbol", "main"), target=target)
|
|
assert "tl.cp_async_gs(" in artifact.kernel_source
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|