(tvm) branch main updated: [SPIRV] Fix forloop codegen in vulkan (#18715)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 739027de42 [SPIRV] Fix forloop codegen in vulkan (#18715)
739027de42 is described below

commit 739027de423d5087b3a5ff59d98bd55ebd7ea9bb
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Thu Feb 5 16:16:08 2026 -0500

    [SPIRV] Fix forloop codegen in vulkan (#18715)
    
    This PR fixes a regression in vulkan codegen where we need to create phi
    node after the start label.
---
 src/target/spirv/codegen_spirv.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/target/spirv/codegen_spirv.cc 
b/src/target/spirv/codegen_spirv.cc
index 136f969896..b1d3ae9cc2 100644
--- a/src/target/spirv/codegen_spirv.cc
+++ b/src/target/spirv/codegen_spirv.cc
@@ -676,13 +676,12 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
   spirv::Value init_value = MakeValue(op->min);
   PrimExpr end = is_zero(op->min) ? op->extent : analyzer_->Simplify(op->min + 
op->extent);
   spirv::Value end_value = MakeValue(end);
-  spirv::PhiValue loop_var = builder_->MakePhi(init_value.stype, 2);
 
   // loop step
   spirv::Value step;
   if (op->HasTrivialStep()) {
-    step = op->loop_var.dtype().is_int() ? builder_->IntImm(loop_var.stype, 1)
-                                         : builder_->UIntImm(loop_var.stype, 
1);
+    step = op->loop_var.dtype().is_int() ? builder_->IntImm(init_value.stype, 
1)
+                                         : builder_->UIntImm(init_value.stype, 
1);
   } else {
     step = MakeValue(tvm::cast(end->dtype, *op->step));
   }
@@ -699,8 +698,9 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
   builder_->SetName(merge_label, "for_loop_merge");
   builder_->MakeInst(spv::OpBranch, head_label);
 
-  // Loop head
+  // Loop head - Phi must be created AFTER StartLabel so it's in the head block
   builder_->StartLabel(head_label);
+  spirv::PhiValue loop_var = builder_->MakePhi(init_value.stype, 2);
   loop_var.SetIncoming(0, init_value, init_label);
   spirv::Value loop_cond = builder_->LT(loop_var, end_value);
   uint32_t control =