mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Chaofan Lin
77cbe6da87
* [WIP][Feature] Introduce CUDARunner * refactor * refactor to T.SourceCodeKernel api * remove streamk change * remove runner * remove prim func params * enhance API * add entry_name * add IsCodeBlockKey * remove misc file * cleanup code * move parse logic to python through callback * fix * comments
88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
import os
|
|
import re
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
import torch
|
|
|
|
|
|
CUDA_SOURCE = """
|
|
extern "C" __global__ void external_copy(float* A, float* B, int n) {
|
|
int i = (int)(blockIdx.x * blockDim.x + threadIdx.x);
|
|
if (i < n) {
|
|
B[i] = A[i];
|
|
}
|
|
}
|
|
"""
|
|
|
|
|
|
def make_source_kernel(source_code_or_path: str | os.PathLike[str], entry_name: str):
|
|
N = T.dynamic("N")
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((N,), T.float32),
|
|
B: T.Tensor((N,), T.float32),
|
|
):
|
|
T.CUDASourceCodeKernel(T.ceildiv(N, 128), threads=128, source_code_or_path=source_code_or_path, entry_name=entry_name)
|
|
|
|
return main
|
|
|
|
|
|
def get_single_device_function_name(device_mod) -> str:
|
|
function_names = [g_var.name_hint for g_var in device_mod.functions]
|
|
assert len(function_names) == 1
|
|
return function_names[0]
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_source_kernel_inline_codegen():
|
|
artifact = tilelang.lower(make_source_kernel(CUDA_SOURCE, entry_name="external_copy"), target="cuda")
|
|
function_name = get_single_device_function_name(artifact.device_mod)
|
|
|
|
assert re.search(
|
|
rf"__global__\s+void\s+(?:__launch_bounds__\([^\)]*\)\s+)?{re.escape(function_name)}\s*\(",
|
|
artifact.kernel_source,
|
|
)
|
|
assert "B[i] = A[i];" in artifact.kernel_source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_source_kernel_run():
|
|
kernel = tilelang.compile(make_source_kernel(CUDA_SOURCE, entry_name="external_copy"), target="cuda")
|
|
print(kernel.get_kernel_source())
|
|
print(kernel.get_host_source())
|
|
a = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
b = torch.empty_like(a)
|
|
kernel(a, b)
|
|
torch.testing.assert_close(b, a)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_source_kernel_loads_from_file():
|
|
with tempfile.NamedTemporaryFile("w", suffix=".cu", delete=False, encoding="utf-8") as f:
|
|
f.write(CUDA_SOURCE)
|
|
source_path = f.name
|
|
|
|
try:
|
|
artifact = tilelang.lower(make_source_kernel(Path(source_path), entry_name="external_copy"), target="cuda")
|
|
finally:
|
|
os.unlink(source_path)
|
|
|
|
assert "B[i] = A[i];" in artifact.kernel_source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_source_kernel_invalid_entry_name_fails_in_lower():
|
|
with pytest.raises(Exception, match=r"Available entries: external_copy"):
|
|
tilelang.lower(make_source_kernel(CUDA_SOURCE, entry_name="main_kernel"), target="cuda")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|