mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Kuris
7ec5542db9
* [Fix] Remove unnecessary binding in loop variable analysis and add test for issue 1728 * [Test] Add test
26 lines
575 B
Python
26 lines
575 B
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
def test_issue_1728():
|
|
@tilelang.jit()
|
|
def get_qwq(hidden: int):
|
|
num_tokens = T.dynamic("num_tokens")
|
|
num_sms = num_tokens
|
|
|
|
@T.prim_func
|
|
def qwq(A: T.Tensor[(num_tokens,)]):
|
|
with T.Kernel(num_sms) as sm_id:
|
|
stop = sm_id + 1
|
|
for block_idx in T.serial(sm_id, stop):
|
|
_pid_x, _pid_y = (block_idx, hidden)
|
|
|
|
return qwq
|
|
|
|
get_qwq(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|