Fix CI failures

src/backend/common/op/atomic_reduce.h

@@ -183,10 +183,11 @@ struct AtomicReduce {
                           level);
     }
     auto loop_layout = par_op->GetLoopLayout();
-    return LowerParallelLoop(
-        fused_loop, loop_layout, lower_args.thread_index, analyzer,
-        lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+    return LowerParallelLoop(fused_loop, loop_layout, lower_args.thread_index,
+                             analyzer, lower_args.layout_map,
+                             par_op->GetPredicate(lower_args.thread_index),
+                             /*parallel_loop=*/true, /*should_vectorize=*/true,
+                             par_op->LoopLayoutRequiresPaddingGuard());
   }
 };
 

src/backend/common/op/transpose.h

@@ -51,7 +51,8 @@ struct Transpose {
     return LowerParallelLoop(
         par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
         lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+        /*parallel_loop=*/true, /*should_vectorize=*/true,
+        par_op->LoopLayoutRequiresPaddingGuard());
   }
 };
 

src/cuda/op/atomic_add.cc

@@ -338,10 +338,11 @@ struct AtomicAdd {
                           level);
     }
     auto loop_layout = par_op->GetLoopLayout();
-    return LowerParallelLoop(
-        fused_loop, loop_layout, lower_args.thread_index, analyzer,
-        lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+    return LowerParallelLoop(fused_loop, loop_layout, lower_args.thread_index,
+                             analyzer, lower_args.layout_map,
+                             par_op->GetPredicate(lower_args.thread_index),
+                             /*parallel_loop=*/true, /*should_vectorize=*/true,
+                             par_op->LoopLayoutRequiresPaddingGuard());
   }
 
   static LayoutMap InferLayout(const AtomicAddNode &op,

src/cuda/op/copy.cc

@@ -865,7 +865,8 @@ Stmt Copy::LowerCPAsync(const CopyNode &op, const LowerArgs &lower_args,
   Stmt lowered_loop = LowerParallelLoop(
       par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
       lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-      /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+      /*parallel_loop=*/true, /*should_vectorize=*/true,
+      par_op->LoopLayoutRequiresPaddingGuard());
 
   auto inject_result =
       InjectPTXAsyncCopy(lowered_loop, /*async_without_async_commit_wait=*/

src/metal/op/transpose.cc

@@ -47,7 +47,8 @@ struct Transpose {
     return LowerParallelLoop(
         par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
         lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+        /*parallel_loop=*/true, /*should_vectorize=*/true,
+        par_op->LoopLayoutRequiresPaddingGuard());
   }
 };
 

src/op/copy.cc

@@ -86,7 +86,8 @@ Stmt LowerNormalCopy(const CopyNode &op, const LowerArgs &lower_args,
   return LowerParallelLoop(
       par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
       lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-      /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+      /*parallel_loop=*/true, /*should_vectorize=*/true,
+      par_op->LoopLayoutRequiresPaddingGuard());
 }
 
 namespace {
@@ -261,7 +262,8 @@ Stmt LowerIm2ColSIMT(const Im2ColOpNode &op, const LowerArgs &lower_args,
   return LowerParallelLoop(
       par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
       lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-      /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+      /*parallel_loop=*/true, /*should_vectorize=*/true,
+      par_op->LoopLayoutRequiresPaddingGuard());
 }
 
 bool RegisterDefaultIm2Col() {

src/rocm/op/atomic_add.cc

@@ -227,10 +227,11 @@ struct AtomicAdd {
                           level);
     }
     auto loop_layout = par_op->GetLoopLayout();
-    return LowerParallelLoop(
-        fused_loop, loop_layout, lower_args.thread_index, analyzer,
-        lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+    return LowerParallelLoop(fused_loop, loop_layout, lower_args.thread_index,
+                             analyzer, lower_args.layout_map,
+                             par_op->GetPredicate(lower_args.thread_index),
+                             /*parallel_loop=*/true, /*should_vectorize=*/true,
+                             par_op->LoopLayoutRequiresPaddingGuard());
   }
 
   static LayoutMap InferLayout(const AtomicAddNode &op,

src/rocm/op/copy.cc

@@ -131,7 +131,8 @@ private:
     Stmt lowered_loop = LowerParallelLoop(
         par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
         lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+        /*parallel_loop=*/true,
+        /*should_vectorize=*/true, par_op->LoopLayoutRequiresPaddingGuard());
 
     auto inject_result =
         InjectROCmAsyncCopy(lowered_loop, /*async_without_async_commit_wait=*/

src/transform/loop_partition.cc

@@ -317,7 +317,7 @@ Stmt LowerParallelLoop(For loop, const Fragment &loop_layout,
                        PrimExpr thread_index, arith::Analyzer *analyzer,
                        const LayoutMap &layout_map,
                        Optional<PrimExpr> predicate, bool parallel_loop,
-                       bool require_padding_guard) {
+                       bool should_vectorize, bool require_padding_guard) {
   // Save analyzer state to prevent conflicted bindings during vectorization
   auto saved_analyzer = analyzer->Clone();
 
@@ -339,9 +339,10 @@ Stmt LowerParallelLoop(For loop, const Fragment &loop_layout,
                                 loop_layout, require_padding_guard);
   }
 
-  // Step 2: Vectorize the loop; the planner picks the size per loop
-  // (1 = scalar) from its access analysis.
-  result_loop = VectorizeLoop(result_loop, saved_analyzer.get(), layout_map);
+  // Step 2: Vectorize the loop (if requested)
+  if (should_vectorize) {
+    result_loop = VectorizeLoop(result_loop, saved_analyzer.get(), layout_map);
+  }
 
   result_loop = PragmaUnrollLoop(result_loop);
 

src/transform/loop_partition.h

@@ -66,13 +66,17 @@ For PragmaUnrollLoop(For stmt);
  * \param parallel_loop Whether this is a true parallel loop requiring thread
  *        partitioning. False for loops that only operate on local/register
  *        buffers. (default true)
+ * \param should_vectorize Whether to vectorize the loop. False when reducer
+ *        combine stores (execution-multiplicity markers) are present or when
+ *        there are no non-local buffer accesses. (default true)
  * \return The lowered statement.
  */
 Stmt LowerParallelLoop(
     For loop, const Fragment &loop_layout, PrimExpr thread_index,
     arith::Analyzer *analyzer, const LayoutMap &layout_map = {},
     ffi::Optional<PrimExpr> predicate = ffi::Optional<PrimExpr>(),
-    bool parallel_loop = true, bool require_padding_guard = false);
+    bool parallel_loop = true, bool should_vectorize = true,
+    bool require_padding_guard = false);
 
 } // namespace tl
 } // namespace tvm

src/transform/lower_tile_op.cc

@@ -1359,9 +1359,47 @@ private:
     // iteration has to be owned by the corresponding thread.
     bool parallel_loop = has_non_local_store || has_fragment_access;
 
+    // Check if there are non-local buffer accesses (for vectorization decision)
+    bool has_non_local = false;
+    PostOrderVisit(for_node->body, [&](const ObjectRef &obj) {
+      if (const auto *load = obj.as<BufferLoadNode>()) {
+        if (!IsLocalBuffer(load->buffer, /*allow_var*/ true) &&
+            !IsFragmentBuffer(load->buffer)) {
+          has_non_local = true;
+        }
+      } else if (const auto *store = obj.as<BufferStoreNode>()) {
+        if (!IsLocalBuffer(store->buffer, /*allow_var*/ true) &&
+            !IsFragmentBuffer(store->buffer)) {
+          has_non_local = true;
+        }
+      }
+    });
+
+    // Check if vectorizable cast operations exist
+    bool has_cast_operations = false;
+    PostOrderVisit(for_node->body, [&](const ObjectRef &obj) {
+      if (const auto *cast = obj.as<CastNode>()) {
+        DataType from_ty = cast->value.dtype();
+        DataType target_ty = cast->dtype;
+        if (IsCudaVectorizableCast(from_ty, target_ty) &&
+            TargetIsCuda(Target::Current())) {
+          has_cast_operations = true;
+        }
+      }
+    });
+
+    // Decide whether to vectorize: only if there are non-local buffers or
+    // vectorizable casts. Reducer combine stores (multiplicity markers were
+    // already lowered by PartitionLoop) need no special exclusion: the
+    // vectorizer's planner keeps stores whose indices do not advance with
+    // the vectorized loop var scalar, which is exactly the reduction-axis
+    // hazard; output-axis contiguous combine stores vectorize like any
+    // other read-modify-write.
+    bool should_vectorize = has_non_local || has_cast_operations;
+    // Lower the parallel loop using the common function
     Stmt lowered = LowerParallelLoop(
         for_node, loop_layout, CurrentThreadIndex(), analyzer_, layout_map_,
-        predicate, parallel_loop, require_padding_guard);
+        predicate, parallel_loop, should_vectorize, require_padding_guard);
 
     // Only parallel-loop lowering needs PTX cp.async injection. Thread-level
     // lowering does not require converting eligible global->shared copies to

src/webgpu/op/transpose.cc

@@ -46,7 +46,8 @@ struct Transpose {
     return LowerParallelLoop(
         par_op->GetRoot(), loop_layout, lower_args.thread_index, analyzer,
         lower_args.layout_map, par_op->GetPredicate(lower_args.thread_index),
-        /*parallel_loop=*/true, par_op->LoopLayoutRequiresPaddingGuard());
+        /*parallel_loop=*/true, /*should_vectorize=*/true,
+        par_op->LoopLayoutRequiresPaddingGuard());
   }
 };
 

testing/python/issue/test_tilelang_issue_1734.py

@@ -10,20 +10,20 @@ def test_issue_1734():
     def kernel():
         @T.prim_func
         def main(
-            A: T.Tensor[(2, 512), T.float32],
-            B: T.Tensor[(2, 512), T.float32],
+            A: T.Tensor[(2, 4096), T.float32],
+            B: T.Tensor[(2, 4096), T.float32],
             C: T.Tensor[(2,), T.float32],
         ):
             with T.Kernel(1, threads=256):
-                A_local = T.alloc_fragment((2, 512), T.float32)
-                B_local = T.alloc_fragment((2, 512), T.float32)
+                A_local = T.alloc_fragment((2, 4096), T.float32)
+                B_local = T.alloc_fragment((2, 4096), T.float32)
                 C_local = T.alloc_fragment((2,), T.float32)
 
                 T.copy(A, A_local)
                 T.copy(C, C_local)
 
-                for i, j in T.Parallel(2, 512):
-                    if C_local[i] >= 0:
+                for i, j in T.Parallel(2, 4096):
+                    if C_local[0] >= 0:
                         B_local[i, j] = A_local[i, j]
 
                 T.copy(B_local, B)
@@ -36,6 +36,7 @@ def test_issue_1734():
     # After hoisting, we should see "if" before "for" pattern
     if_pos = source.find("if (")
     for_pos = source.find("for (")
+    assert if_pos != -1 and for_pos != -1
     assert if_pos < for_pos, "Loop-invariant if should be hoisted outside the loop"
 
 

testing/python/language/test_tilelang_language_parallel.py

@@ -85,8 +85,9 @@ def _parallel_vectorize_local_and_var():
 
 def test_parallel_vectorize_var():
     source = _parallel_vectorize_local_and_var.get_kernel_source()
-    # do not vectorize if the loop only contains local/fragment and var buffer access
-    assert "float2" not in source
+    # register-register loops vectorize too; the var operand becomes a
+    # scalar broadcast inside the packed multiply
+    assert "float2" in source
 
 
 if __name__ == "__main__":