mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Lei Wang
c770a58f63
* [Enhancement] Add matrix multiplication functions for integer and float variables in Cython JIT - Introduced `matmul_int_variable` and `matmul_float_variable` functions to support matrix multiplication with dynamic shapes and additional parameters. - Implemented corresponding `run_matmul_int_variable` and `run_matmul_float_variable` functions for testing. - Updated test cases to validate the new matrix multiplication implementations. - Enhanced error handling in library initialization and compilation processes across various modules. - Improved dynamic memory handling in CUDA kernel initialization to provide better error reporting. * lint fix * optimize * Support var defiine * lint fix * Update TVM submodule and add alloc_variable function to allocate local variables in TileLang - Updated the TVM submodule to the latest commit. - Introduced `alloc_variable` function in `allocate.py` to support local variable allocation with specified data types and scopes. * lint fix * Refactor variable allocation functions for consistency - Renamed `alloc_variable` to `alloc_var` across multiple files for improved consistency. - Updated corresponding test functions to reflect the new naming convention. - Adjusted imports in `__init__.py` to align with the changes.
86 lines
1.7 KiB
Python
86 lines
1.7 KiB
Python
import tilelang.testing
|
|
|
|
|
|
def alloc_var(
|
|
N,
|
|
block_N,
|
|
dtype,
|
|
):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Buffer((N,), dtype),
|
|
B: T.Buffer((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=block_N) as bx:
|
|
A_shared = T.alloc_shared([block_N], dtype)
|
|
tmp = T.alloc_var(dtype)
|
|
tmp = 1 # noqa: F841
|
|
T.copy(A[bx * block_N], A_shared)
|
|
T.copy(A_shared, B[bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_alloc_var(
|
|
N,
|
|
block_N,
|
|
dtype,
|
|
min=None,
|
|
max=None,
|
|
):
|
|
program = alloc_var(N, block_N, dtype)
|
|
|
|
kernel = tilelang.compile(program, out_idx=[1])
|
|
code = kernel.get_kernel_source()
|
|
assert "tmp =" in code
|
|
|
|
|
|
def test_alloc_var():
|
|
run_alloc_var(1024, 128, "float16")
|
|
|
|
|
|
def alloc_var_add(
|
|
N,
|
|
block_N,
|
|
dtype,
|
|
):
|
|
import tilelang.language as T
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Buffer((N,), dtype),
|
|
B: T.Buffer((N,), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), threads=block_N) as bx:
|
|
A_shared = T.alloc_shared([block_N], dtype)
|
|
tmp = T.alloc_var(dtype)
|
|
tmp = 1 # noqa: F841
|
|
T.copy(A[bx * block_N], A_shared)
|
|
for i in T.Parallel(block_N):
|
|
A_shared[i] = A_shared[i] + tmp
|
|
T.copy(A_shared, B[bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_alloc_var_add(
|
|
N,
|
|
block_N,
|
|
dtype,
|
|
):
|
|
program = alloc_var_add(N, block_N, dtype)
|
|
|
|
kernel = tilelang.compile(program, out_idx=[1])
|
|
code = kernel.get_kernel_source()
|
|
assert "tmp =" in code
|
|
|
|
|
|
def test_alloc_var_add():
|
|
run_alloc_var_add(1024, 128, "float16")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|