mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-28 07:43:01 +08:00
Yuxi Chi
7fc47066cd
* Port CuTeDSL backend to CUTLASS primitives * Update CuTeDSL dependency to CUTLASS DSL 4.7.0 * Fix CuTeDSL dynamic-shape wrapper generation
33 lines
867 B
Python
33 lines
867 B
Python
"""Regression test for issue #2682.
|
|
|
|
T.vectorized with T.Select incorrectly generates scalar ternary expressions.
|
|
CodeGenTileLangCUDA misses the handling of vectorized SelectNode, causing
|
|
a fallback to CodeGenC in tvm that generates a scalar ternary expression.
|
|
"""
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
NUM_ELEMENTS = 4
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_vectorized_select():
|
|
@T.prim_func
|
|
def kernel(
|
|
source: T.Tensor[(NUM_ELEMENTS,), T.float32],
|
|
destination: T.Tensor[(NUM_ELEMENTS,), T.float32],
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
for index in T.vectorized(NUM_ELEMENTS):
|
|
value = source[index]
|
|
destination[index] = T.Select(value > 0.0, value, value + 1.0)
|
|
|
|
tilelang.compile(kernel)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|