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
235 lines
7.5 KiB
Python
235 lines
7.5 KiB
Python
"""Tests for store_global_32, store_global_64, store_global_128, store_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_stg32_codegen():
|
|
"""Test that stg32 generates tl::store_global_32 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg32_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:
|
|
val = T.reinterpret(T.uint32, X[pid])
|
|
T.stg32(Y[pid], val)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg32_kernel(X, Y)
|
|
src = stg32_kernel.get_kernel_source(N=128)
|
|
print("=== stg32 codegen ===")
|
|
print(src)
|
|
# Verify codegen
|
|
assert "store_global_32" in src, "Expected store_global_32 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
torch.testing.assert_close(Y, X, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_stg64_codegen():
|
|
"""Test that stg64 generates tl::store_global_64 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg64_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:
|
|
val = T.ldg64(X[pid * 2 : pid * 2 + 2])
|
|
T.stg64(Y[pid * 2 : pid * 2 + 2], val)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg64_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg64_kernel.get_kernel_source(N=128)
|
|
print("=== stg64 codegen ===")
|
|
print(src)
|
|
assert "store_global_64" in src, "Expected store_global_64 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
torch.testing.assert_close(Y, X, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_stg128_codegen():
|
|
"""Test that stg128 generates tl::store_global_128 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg128_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:
|
|
val = T.ldg128(X[pid * 4 : pid * 4 + 4])
|
|
T.stg128(Y[pid * 4 : pid * 4 + 4], val)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg128_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg128_kernel.get_kernel_source(N=128)
|
|
print("=== stg128 codegen ===")
|
|
print(src)
|
|
assert "store_global_128" in src, "Expected store_global_128 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
torch.testing.assert_close(Y, X, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(10, 0)
|
|
def test_stg256_codegen():
|
|
"""Test that stg256 generates tl::store_global_256 in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg256_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:
|
|
val = T.ldg256(X[pid * 8 : pid * 8 + 8])
|
|
T.stg256(Y[pid * 8 : pid * 8 + 8], val)
|
|
|
|
X = torch.randn(256, dtype=torch.float32, device="cuda")
|
|
Y = torch.empty(256, dtype=torch.float32, device="cuda")
|
|
|
|
stg256_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg256_kernel.get_kernel_source(N=256)
|
|
print("=== stg256 codegen ===")
|
|
print(src)
|
|
assert "store_global_256" in src, "Expected store_global_256 call in generated CUDA source"
|
|
|
|
# Verify correctness
|
|
torch.testing.assert_close(Y, X, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_stg32_predicated_codegen():
|
|
"""Test that stg32 with predicate generates tl::store_global_32_conditional(ptr, val, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg32_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:
|
|
val = T.reinterpret(T.uint32, X[pid])
|
|
# Only store for the first half of elements
|
|
T.stg32(Y[pid], val, pred=pid < N // 2)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg32_pred_kernel(X, Y)
|
|
src = stg32_pred_kernel.get_kernel_source(N=128)
|
|
print("=== stg32 predicated codegen ===")
|
|
print(src)
|
|
# Verify codegen - should have store_global_32 with predicate
|
|
assert "store_global_32" in src, "Expected store_global_32 call in generated CUDA source"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_stg64_predicated_codegen():
|
|
"""Test that stg64 with predicate generates tl::store_global_64_conditional(ptr, val, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg64_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:
|
|
val = T.ldg64(X[pid * 2 : pid * 2 + 2])
|
|
# Only store for the first half of elements
|
|
T.stg64(Y[pid * 2 : pid * 2 + 2], val, pred=pid < N // 4)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg64_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg64_pred_kernel.get_kernel_source(N=128)
|
|
print("=== stg64 predicated codegen ===")
|
|
print(src)
|
|
assert "store_global_64" in src, "Expected store_global_64 call in generated CUDA source"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_stg128_predicated_codegen():
|
|
"""Test that stg128 with predicate generates tl::store_global_128_conditional(ptr, val, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg128_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:
|
|
val = T.ldg128(X[pid * 4 : pid * 4 + 4])
|
|
# Only store for the first half of elements
|
|
T.stg128(Y[pid * 4 : pid * 4 + 4], val, pred=pid < N // 8)
|
|
|
|
X = torch.randn(128, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(128, dtype=torch.float32, device="cuda")
|
|
|
|
stg128_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg128_pred_kernel.get_kernel_source(N=128)
|
|
print("=== stg128 predicated codegen ===")
|
|
print(src)
|
|
assert "store_global_128" in src, "Expected store_global_128 call in generated CUDA source"
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_ge(10, 0)
|
|
def test_stg256_predicated_codegen():
|
|
"""Test that stg256 with predicate generates tl::store_global_256_conditional(ptr, val, pred) in CUDA source."""
|
|
|
|
@tilelang.jit
|
|
def stg256_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:
|
|
val = T.ldg256(X[pid * 8 : pid * 8 + 8])
|
|
# Only store for the first half of elements
|
|
T.stg256(Y[pid * 8 : pid * 8 + 8], val, pred=pid < N // 16)
|
|
|
|
X = torch.randn(256, dtype=torch.float32, device="cuda")
|
|
Y = torch.zeros(256, dtype=torch.float32, device="cuda")
|
|
|
|
stg256_pred_kernel(X, Y)
|
|
|
|
# Verify codegen
|
|
src = stg256_pred_kernel.get_kernel_source(N=256)
|
|
print("=== stg256 predicated codegen ===")
|
|
print(src)
|
|
assert "store_global_256" in src, "Expected store_global_256 call in generated CUDA source"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|