mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
Lei Wang
465c6c5e8c
* [Refactor] Remove unsupported tl_gemm and tl_gemm_sp operations * Removed references to tl_gemm and tl_gemm_sp from various files, including code generation and built-in definitions, as these operations are currently unsupported. * Updated related comments and diagnostics to reflect the removal of these operations, ensuring clarity in the codebase. * This change simplifies the handling of unsupported operations and improves overall code maintainability. * Reduce CUDA template include overhead * Fix CUDA reduce fast min max helpers * lint fix * Fix CUDA swizzle ceil div helper * Restore CUDA wait_wgmma intrinsic helper * Fix CUDA cython TMA host adapter include * Remove unnecessary check for Cutlass fast math header in CUDA code generation * Update LibraryGenerator to correctly handle CUDA gencode flags for shared-library compilation, ensuring compatibility with target architecture and avoiding Hopper-only instruction issues. * Refactor smem_ptr_to_uint function to use __cvta_generic_to_shared for improved clarity and performance. * Refactor example scripts and CUDA math intrinsics for improved performance and clarity - Updated `example_dequant_gemm_fp4_hopper.py` to directly call `run_regression_perf()` and print latency, removing unused argument parsing code. - Modified `example_topk.py` to specify the backend as "cupti" in the benchmarking function. - Enhanced CUDA code generation by introducing `RequiresTileLangMathHeader` to manage math header dependencies for specific functions. - Updated barrier operations in `barrier.h` to use the correct syntax for shared memory barriers. - Improved testing for bfloat16 math intrinsics by replacing pytest skips with decorators for CUDA availability checks. * Refactor CUDA code for improved readability and consistency - Reformatted the `RequiresTileLangMathHeader` function for better readability. - Updated assembly syntax in `barrier.h` for clarity and consistency. - Cleaned up whitespace in the bfloat16 math intrinsics test file to enhance code cleanliness. * Remove unnecessary whitespace in CUDA code for improved readability * Add TileLang intrinsic for prefetching TMA descriptor - Introduced `prefetch_tma_descriptor` intrinsic in C++ and updated relevant files to use `call_intrin` instead of `call_extern`. - Updated CUDA code generation to handle the new intrinsic correctly. - Modified tests to ensure the new intrinsic is called appropriately and that no `call_extern` is used for internal TileLang operations. - Enhanced documentation in SKILL.md regarding the addition of new TileLang intrinsics. * Add bfloat16 hexp overload and update common.h for compatibility - Introduced `hexp` overload for `bfloat16_t` in `common.h` to bridge TileLang and CUDA's handling of exponential functions. - Updated tests to verify the presence of the new `hexp` definition in `common.h`. - Enhanced CUDA code generation to support the new intrinsic for better performance and compatibility. * Update common.h for CUDA compatibility with aligned double4 types - Adjusted preprocessor condition in `common.h` to maintain compatibility with CUDA versions prior to 13 by ensuring proper definition of aligned double4 types. - This change enhances the compatibility of the code with older NVRTC built-in vector headers.
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
from pathlib import Path
|
|
import re
|
|
|
|
import torch
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
def test_cuda_common_h_defines_bfloat16_rsqrt_overload():
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
common_h = repo_root / "src" / "tl_templates" / "cuda" / "common.h"
|
|
source = common_h.read_text()
|
|
|
|
pattern = re.compile(
|
|
r"TL_PATCH\s+TL_DEVICE\s+bfloat16_t\s+hrsqrt\s*"
|
|
r"\(\s*const\s+bfloat16_t\s+x\s*\)\s*\{"
|
|
)
|
|
assert pattern.search(source), "common.h must define hrsqrt(bfloat16_t) for bf16 T.rsqrt codegen"
|
|
|
|
|
|
def test_cuda_common_h_defines_bfloat16_hexp_overload():
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
common_h = repo_root / "src" / "tl_templates" / "cuda" / "common.h"
|
|
source = common_h.read_text()
|
|
|
|
pattern = re.compile(
|
|
r"TL_PATCH\s+TL_DEVICE\s+bfloat16_t\s+hexp\s*"
|
|
r"\(\s*const\s+bfloat16_t\s+x\s*\)\s*\{"
|
|
)
|
|
assert pattern.search(source), "common.h must define hexp(bfloat16_t) for bf16 T.exp codegen"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version(8, 0)
|
|
def test_bfloat16_rsqrt_compiles_and_runs():
|
|
n = 32
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((n,), T.bfloat16), B: T.Tensor((n,), T.bfloat16)):
|
|
with T.Kernel(1, threads=32):
|
|
values = T.alloc_fragment((n,), T.bfloat16)
|
|
for i in T.Parallel(n):
|
|
values[i] = T.rsqrt(A[i])
|
|
for i in T.Parallel(n):
|
|
B[i] = values[i]
|
|
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
inputs = torch.full((n,), 4.0, dtype=torch.bfloat16, device="cuda")
|
|
actual = kernel(inputs)
|
|
|
|
torch.testing.assert_close(actual, torch.full_like(actual, 0.5))
|
|
|
|
|
|
def test_cuda_math_h_bfloat16_fast_exp_does_not_self_recurse():
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
common_h = repo_root / "src" / "tl_templates" / "cuda" / "common.h"
|
|
math_h = repo_root / "src" / "tl_templates" / "cuda" / "math.h"
|
|
common_source = common_h.read_text()
|
|
source = math_h.read_text()
|
|
|
|
# `hexp` is #define'd to cutlass::fast_exp, so the bf16 fast_exp overload
|
|
# must route through float; calling `::hexp(x)` here recurses into itself.
|
|
assert "#include <cutlass/fast_math.h>" not in common_source
|
|
assert "#include <cutlass/fast_math.h>" in source
|
|
assert "#define hexp cutlass::fast_exp" in source
|
|
match = re.search(
|
|
r"bfloat16_t\s+fast_exp\s*\(\s*bfloat16_t\s+x\s*\)\s*\{([^}]*)\}",
|
|
source,
|
|
)
|
|
assert match, "math.h must define fast_exp(bfloat16_t) for bf16 T.exp codegen"
|
|
body = match.group(1)
|
|
assert "float(" in body, "bf16 fast_exp must route through float to avoid self-recursion"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version(8, 0)
|
|
def test_bfloat16_exp_compiles_and_runs():
|
|
n = 32
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((n,), T.bfloat16), B: T.Tensor((n,), T.bfloat16)):
|
|
with T.Kernel(1, threads=32):
|
|
values = T.alloc_fragment((n,), T.bfloat16)
|
|
for i in T.Parallel(n):
|
|
values[i] = T.exp(A[i])
|
|
for i in T.Parallel(n):
|
|
B[i] = values[i]
|
|
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
inputs = torch.full((n,), 1.0, dtype=torch.bfloat16, device="cuda")
|
|
actual = kernel(inputs)
|
|
|
|
expected = torch.exp(inputs)
|
|
torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|