mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
* add hoist_broadcast_values pass * only hoist intImm or floatImm * refactor * add test for hoist_broadcast_values * add a test for issue 1601 * add transform test for hoist_broadcast_value --------- Co-authored-by: kurisu6912 <227995639+kurisu6912@users.noreply.github.com>
25 lines
540 B
Python
25 lines
540 B
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
def test_issue_1601():
|
|
@tilelang.jit
|
|
def qwq():
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((8,), T.float8_e4m3fn),
|
|
):
|
|
with T.Kernel(1, threads=32):
|
|
for i in T.vectorized(8):
|
|
A[i] = 0
|
|
|
|
return main
|
|
|
|
kernel = qwq()
|
|
assert "fp8_e4_t broadcast_var = fp8_e4_t(0x0p+0f/*0.000000e+00*/);" in kernel.get_kernel_source()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|