From 6873e19da2b267d3cb7d0305561ba0bed63feae3 Mon Sep 17 00:00:00 2001
From: Jinpeng Wang <jpwang@smail.nju.edu.cn>
Date: Wed, 5 Aug 2026 13:51:31 +0800
Subject: [PATCH] [BugFix] Handle vectorized SelectNode in codegen_cuda (#2843)

* [BugFix] Handle vectorized SelectNode in codegen_cuda

* address review comments

* [BugFix] Fix cuda codegen vectorized SelectNode when condition is a scalar
---
 src/cuda/codegen/codegen_cuda.cc              | 41 +++++++++++++++++++
 src/cuda/codegen/codegen_cuda.h               |  1 +
 .../python/issue/test_tilelang_issue_2682.py  | 31 ++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 testing/python/issue/test_tilelang_issue_2682.py

diff --git a/src/cuda/codegen/codegen_cuda.cc b/src/cuda/codegen/codegen_cuda.cc
index 99ba54053..9cf40d5c6 100644
--- a/src/cuda/codegen/codegen_cuda.cc
+++ b/src/cuda/codegen/codegen_cuda.cc
@@ -5382,6 +5382,47 @@ void CodeGenTileLangCUDA::VisitStmt_(const BufferStoreNode *op) {
   }
 }
 
+void CodeGenTileLangCUDA::VisitExpr_(const SelectNode *op, std::ostream &os) {
+  // Non-vector cases.
+  if (!op->condition.dtype().is_fixed_length_vector()) {
+    CodeGenC::VisitExpr_(op, os);
+    return;
+  }
+
+  // Codegen vector condition case by serializing the select op.
+  TVM_FFI_ICHECK(op->false_value->dtype == op->dtype &&
+                 op->true_value->dtype == op->dtype &&
+                 op->dtype.lanes() == op->condition.dtype().lanes());
+
+  std::string r_var = name_supply_->FreshName("_");
+  this->PrintIndent();
+  this->PrintType(op->dtype, stream);
+  stream << ' ' << r_var << ";\n";
+  {
+    std::string c_var =
+        SSAGetID(PrintExpr(op->condition), op->condition.dtype());
+    std::string t_var = SSAGetID(PrintExpr(op->true_value), op->dtype);
+    std::string f_var = SSAGetID(PrintExpr(op->false_value), op->dtype);
+
+    // The condition is stored as an ushort vector.
+    int lanes = op->dtype.lanes();
+    DataType memory_ty(DataType::TypeCode::kUInt, 16, lanes);
+
+    for (int i = 0; i < lanes; ++i) {
+      std::ostringstream item;
+      item << "(bool(";
+      PrintVecElemLoad(c_var, memory_ty, i, item);
+      item << ")?";
+      PrintVecElemLoad(t_var, op->dtype, i, item);
+      item << ':';
+      PrintVecElemLoad(f_var, op->dtype, i, item);
+      item << ')';
+      PrintVecElemStore(r_var, op->dtype, i, item.str());
+    }
+  }
+  os << r_var;
+}
+
 void CodeGenTileLangCUDA::VisitExpr_(const ShuffleNode *op,
                                      std::ostream &os) { // NOLINT(*)
   // For bfloat16x2 / float16x2 construction from two scalar lanes, emit a
diff --git a/src/cuda/codegen/codegen_cuda.h b/src/cuda/codegen/codegen_cuda.h
index e188118b5..b2203d716 100644
--- a/src/cuda/codegen/codegen_cuda.h
+++ b/src/cuda/codegen/codegen_cuda.h
@@ -62,6 +62,7 @@ public:
   void VisitStmt_(const AttrStmtNode *op) final;
   void VisitExpr_(const BufferLoadNode *op, std::ostream &os) final;
   void VisitStmt_(const BufferStoreNode *op) final;
+  void VisitExpr_(const SelectNode *op, std::ostream &os) final;
 
   // Override this as a work around for __grid_constant__ parameter
   void AddFunction(const GlobalVar &gvar, const PrimFunc &f);
diff --git a/testing/python/issue/test_tilelang_issue_2682.py b/testing/python/issue/test_tilelang_issue_2682.py
new file mode 100644
index 000000000..6292c1a7b
--- /dev/null
+++ b/testing/python/issue/test_tilelang_issue_2682.py
@@ -0,0 +1,31 @@
+"""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
+
+
+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()