mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Lei Wang
79730b1123
* [Refactor] Enhance TLVectorizer with loop vectorization convenience method and improve let variable handling * lint fix * let test fix * lint fix
24 lines
665 B
Python
24 lines
665 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: T.float32x4 = 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.imported_modules[0].get_source()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|