mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
5f202fe5c1
* tilelang frontend v2 * syntax sugar: defining a local var by annotation * [Refactor] fix type linting warning like `T.float32` * Add tl.local_var_init for new tl.float32 * allow passing default argument as function annotation * allow default arguments as annotation * fix lint error * minor fix * [Refactor] refactor tilelang.jit and tilelang.autotune * minor fix * minor fix * minor fix * fix metal get function name * add par_compile impl and tests * Type consistency on tvm datatype 1. isinstance(tl.float32, tvm.DataType) == True 2. Allow `tl.float32` as function annotations 3. Allow `tl.float32` as argument to be passed to `tl.alloc` or other functions * fix lint error * add more warning in frontend * update tvm version * Minor fix on tvm_ffi annotations * add document and examples * fix lint error * Simplify index calculations in example_chunk_o_bwd.py Refactor index calculations for dg_last_fragment assignment. * minor fix * lint fix --------- Co-authored-by: Lei Wang <leiwang1999@outlook.com> Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
24 lines
647 B
Python
24 lines
647 B
Python
import tilelang.testing
|
|
from tilelang import tvm as tvm
|
|
from tilelang import language as T
|
|
|
|
|
|
def test_let_vectorize_load():
|
|
|
|
@T.prim_func
|
|
def main(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)
|
|
|
|
for _blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for _threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
b = A[0, 0:4]
|
|
A[0, 4:8] = b
|
|
|
|
mod = tvm.IRModule({"main": main})
|
|
mod = tvm.compile(mod, target="cuda")
|
|
assert "float4 b" in mod.mod.imports[0].inspect_source()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|