From 5235a47db89379a17927a54d1e49c20e9472419f Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 18 Feb 2025 13:42:00 -0600
Subject: [PATCH] Test vectorization of CL ternaries

---
 test/test_target.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/test/test_target.py b/test/test_target.py
index f29fbeec..b4cb509d 100644
--- a/test/test_target.py
+++ b/test/test_target.py
@@ -842,6 +842,30 @@ def test_to_complex_casts(ctx_factory):
     cl.Program(ctx, code).build()
 
 
+def test_cl_vectorize_ternary(ctx_factory):
+    knl = lp.make_kernel(
+            "{ [i]: 0<=i<n }",
+            """
+            b[i] = a[i]*3 if a[i] < 0 else sin(a[i])
+            """)
+
+    knl = lp.split_array_axis(knl, "a,b", 0, 4)
+    knl = lp.split_iname(knl, "i", 4)
+    knl = lp.tag_inames(knl, {"i_inner": "vec"})
+    knl = lp.tag_array_axes(knl, "a,b", "c,vec")
+    knl = lp.set_options(knl, write_code=True)
+    knl = lp.assume(knl, "n % 4 = 0 and n>0")
+
+    rng = np.random.default_rng(seed=12)
+    a = rng.normal(size=(16, 4))
+    ctx = ctx_factory()
+    queue = cl.CommandQueue(ctx)
+    _evt, (result,) = knl(queue, a=a, n=a.size)
+
+    result_ref = np.where(a < 0, a*3, np.sin(a))
+    assert np.allclose(result, result_ref)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab