mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Lei Wang
8bd0293740
* [Enhancement] Add global memory load/store intrinsics for various vector widths (#1710) * Introduced new intrinsics: ldg32, ldg64, ldg128, ldg256 for loading data from global memory. * Added corresponding store intrinsics: stg32, stg64, stg128, stg256 for storing data to global memory. * Updated built-in functions and templates to support these new operations. * Enhanced documentation for usage examples and type checks in the language module. * Updated CUDA code generation to utilize new intrinsics for improved performance. * lint fix * test fix * Refactor buffer type annotations in global memory load/store intrinsics Updated the type annotations for the `ldg` and `stg` functions to use `tir.Buffer` instead of the generic `Buffer` type. This change enhances type specificity and aligns with the TVM type system for better clarity and correctness in memory operations. * lint fix * [Enhancement] Introduce LDG/STG lowering options and transform pass * Added configuration options to enable/disable lowering of non-predicated and predicated global memory loads/stores to ldg/stg intrinsics. * Implemented the LowerLDGSTG transformation pass to optimize global memory access for better performance. * Updated tests to require a minimum CUDA compute version for relevant functionality. * Enhanced documentation for new configuration options and transformation behavior. * lint fix * Refactor LDG/STG intrinsic selection logic for improved readability * Replaced switch-case statements with if-else conditions for determining the appropriate LDG operation based on byte width. * Enhanced error handling for unsupported byte widths with a clear fatal log message. * lint fix * fix * testfix * Enhance LDGSTG pass to skip processing for CuTeDSL backend * Added a check to skip the transformation if the target has the "cutedsl" key. * Updated tests to require a minimum CUDA compute version of 10.0 for LDG and STG code generation tests. * Refactor global memory load/store functions to include conditional variants * Updated the CUDA code generation to support conditional versions of LDG and STG functions, enhancing flexibility in memory operations. * Modified existing functions to include new conditional variants (e.g., `ldg32_conditional`, `stg32_conditional`) for better control over memory access based on predicates. * Adjusted comments to reflect the changes in function signatures and usage. * Update comments and input handling for LDG intrinsics to include address and optional predicate * Modified comments for `ldg32`, `ldg64`, `ldg128`, and `ldg256` to clarify the function signatures, indicating the use of 'address' and an optional 'predicate'. * Changed the number of inputs for these LDG functions from 2 to -1 to accommodate the new signature format. * Add correctness verification to LDG code generation tests * Introduced assertions to verify the correctness of generated CUDA source for `ldg32`, `ldg64`, `ldg128`, and `ldg256` functions. * Added checks to ensure the output matches the expected reference values within specified tolerances. * Updated test cases to enhance validation of LDG intrinsic functionality. * test fix * Refactor LDG and STG intrinsics to use unified naming convention * Updated all instances of `ldg` and `stg` functions to their new names: `load_global` and `store_global`, respectively, across the codebase. * Modified related test cases to reflect the new function names, ensuring consistency in the testing framework. * Enhanced comments and documentation to clarify the changes in function signatures and usage. * Update reference values in LDG256 predicated code generation test to improve accuracy * lint fix * Rename test function for clarity and add new test for predicated load/store global intrinsics. Update assertions to reflect new naming convention in generated kernel source. * lint fix * subtype cast fix
270 lines
8.7 KiB
Python
270 lines
8.7 KiB
Python
"""Tests for load_global_32, load_global_64, load_global_128, load_global_256 intrinsics codegen using eager jit style."""
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
import torch
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg32_codegen():
|
|
"""Test that ldg32 generates tl::load_global_32 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg32_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N, threads=32) as pid:
|
|
Y[pid] = T.reinterpret(T.float32, T.ldg32(X[pid]))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg32_kernel(X, Y)
|
|
src = ldg32_kernel.get_kernel_source(N=128)
|
|
print("=== ldg32 codegen ===")
|
|
print(src)
|
|
# Verify codegen
|
|
assert "load_global_32" in src, "Expected load_global_32 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = X
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg64_codegen():
|
|
"""Test that ldg64 generates tl::load_global_64 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg64_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 2, threads=32) as pid:
|
|
Y[pid * 2 : pid * 2 + 2] = T.reinterpret(T.float32x2, T.ldg64(X[pid * 2 : pid * 2 + 2]))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg64_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg64_kernel.get_kernel_source(N=128)
|
|
print("=== ldg64 codegen ===")
|
|
print(src)
|
|
assert "load_global_64" in src, "Expected load_global_64 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = X
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg128_codegen():
|
|
"""Test that ldg128 generates tl::load_global_128 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg128_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 4, threads=32) as pid:
|
|
Y[pid * 4 : pid * 4 + 4] = T.reinterpret(T.float32x4, T.ldg128(X[pid * 4 : pid * 4 + 4]))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg128_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg128_kernel.get_kernel_source(N=128)
|
|
print("=== ldg128 codegen ===")
|
|
print(src)
|
|
assert "load_global_128" in src, "Expected load_global_128 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = X
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(10, 0)
|
|
def test_ldg256_codegen():
|
|
"""Test that ldg256 generates tl::load_global_256 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg256_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 8, threads=32) as pid:
|
|
Y[pid * 8 : pid * 8 + 8] = T.reinterpret(T.float32x8, T.ldg256(X[pid * 8 : pid * 8 + 8]))
|
|
|
|
X = torch.randn(256, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(256, dtype=torch.float32, device="cuda")
|
|
|
|
ldg256_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg256_kernel.get_kernel_source(N=256)
|
|
print("=== ldg256 codegen ===")
|
|
print(src)
|
|
assert "load_global_256" in src, "Expected load_global_256 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = X
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg32_predicated_codegen():
|
|
"""Test that ldg32 with predicate generates tl::load_global_32_conditional(ptr, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg32_pred_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N, threads=32) as pid:
|
|
# Only load for the first half of elements
|
|
Y[pid] = T.reinterpret(T.float32, T.ldg32(X[pid], pred=pid < N // 2))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg32_pred_kernel(X, Y)
|
|
src = ldg32_pred_kernel.get_kernel_source(N=128)
|
|
print("=== ldg32 predicated codegen ===")
|
|
print(src)
|
|
# Verify codegen - should have load_global_32_conditional with two arguments and non-trivial predicate
|
|
assert "load_global_32_conditional" in src, "Expected load_global_32_conditional call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
for i in range(128):
|
|
if i < 64:
|
|
Y_ref[i] = X[i]
|
|
else:
|
|
Y_ref[i] = 0
|
|
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg64_predicated_codegen():
|
|
"""Test that ldg64 with predicate generates tl::load_global_64_conditional(ptr, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg64_pred_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 2, threads=32) as pid:
|
|
# Only load for the first half of elements
|
|
Y[pid * 2 : pid * 2 + 2] = T.reinterpret(T.float32x2, T.ldg64(X[pid * 2 : pid * 2 + 2], pred=pid < N // 4))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg64_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg64_pred_kernel.get_kernel_source(N=128)
|
|
print("=== ldg64 predicated codegen ===")
|
|
print(src)
|
|
assert "load_global_64_conditional" in src, "Expected load_global_64_conditional call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
for i in range(128):
|
|
if i < 64:
|
|
Y_ref[i] = X[i]
|
|
else:
|
|
Y_ref[i] = 0
|
|
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_ldg128_predicated_codegen():
|
|
"""Test that ldg128 with predicate generates tl::load_global_128_conditional(ptr, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg128_pred_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 4, threads=32) as pid:
|
|
# Only load for the first half of elements
|
|
Y[pid * 4 : pid * 4 + 4] = T.reinterpret(T.float32x4, T.ldg128(X[pid * 4 : pid * 4 + 4], pred=pid < N // 8))
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
ldg128_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg128_pred_kernel.get_kernel_source(N=128)
|
|
print("=== ldg128 predicated codegen ===")
|
|
print(src)
|
|
assert "load_global_128_conditional" in src, "Expected load_global_128_conditional call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
Y_ref = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
for i in range(128):
|
|
if i < 64:
|
|
Y_ref[i] = X[i]
|
|
else:
|
|
Y_ref[i] = 0
|
|
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(10, 0)
|
|
def test_ldg256_predicated_codegen():
|
|
"""Test that ldg256 with predicate generates tl::load_global_256_conditional(ptr, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def ldg256_pred_kernel(X, Y):
|
|
N = T.const("N")
|
|
X: T.Tensor[[N], T.float32]
|
|
Y: T.Tensor[[N], T.float32]
|
|
|
|
with T.Kernel(N // 8, threads=32) as pid:
|
|
# Only load for the first half of elements
|
|
Y[pid * 8 : pid * 8 + 8] = T.reinterpret(T.float32x8, T.ldg256(X[pid * 8 : pid * 8 + 8], pred=pid < N // 16))
|
|
|
|
X = torch.randn(256, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(256, dtype=torch.float32, device="cuda")
|
|
|
|
ldg256_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = ldg256_pred_kernel.get_kernel_source(N=256)
|
|
print("=== ldg256 predicated codegen ===")
|
|
print(src)
|
|
assert "load_global_256_conditional" in src, "Expected load_global_256_conditional call in generated CUDA source"
|
|
# Verify correctness
|
|
Y_ref = torch.zeros(256, dtype=torch.float32, device="cuda")
|
|
for i in range(256):
|
|
if i < 128:
|
|
Y_ref[i] = X[i]
|
|
else:
|
|
Y_ref[i] = 0
|
|
|
|
torch.testing.assert_close(Y, Y_ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|