Init hidden loop control vars, reset continue var in loop continue block

wpmed92 · Oct 9, 2024 · 4b47402 · 4b47402
1 parent fdd78d3
commit 4b47402
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 53 deletions.
diff --git a/lib/CodeGen/MLIRCodeGen.cpp b/lib/CodeGen/MLIRCodeGen.cpp
@@ -1245,6 +1245,9 @@ void MLIRCodeGen::generateLoop(Statement* initStmt, Expression* conditionExpr, E
       builder.getUnknownLoc(), ptrType, spirv::StorageClass::Function, nullptr)
   );
 
+  setBoolVar(continueStack.back(), false);
+  setBoolVar(breakStack.back(), false);
+
   if (initStmt) {
     initStmt->accept(this);
   }
@@ -1313,13 +1316,6 @@ void MLIRCodeGen::generateLoop(Statement* initStmt, Expression* conditionExpr, E
           builder.setInsertionPointToStart(postGateBlock);
         }
 
-        if (breakDetected) {
-          setBoolVar(breakStack.back(), false);
-        }
-
-        if (continueDetected) {
-          setBoolVar(continueStack.back(), false);
-        }
         breakDetected = false;
         continueDetected = false;
       }
@@ -1330,6 +1326,7 @@ void MLIRCodeGen::generateLoop(Statement* initStmt, Expression* conditionExpr, E
 
   builder.create<spirv::BranchOp>(loc, loopOp.getContinueBlock());
   builder.setInsertionPointToEnd(loopOp.getContinueBlock());
+  setBoolVar(continueStack.back(), false);
 
   if (inductionExpr) {
     inductionExpr->accept(this);

diff --git a/test/CodeGen/cf_loops_for.glsl b/test/CodeGen/cf_loops_for.glsl
@@ -1,28 +1,31 @@
 void main() {
     // CHECK: %2 = spirv.Variable : !spirv.ptr<si32, Function>
     // CHECK: spirv.mlir.loop {
-    // CHECK-NEXT:  spirv.Branch ^bb1
+    // CHECK-NEXT: spirv.Branch ^bb1
     // CHECK-NEXT: ^bb1:  // 2 preds: ^bb0, ^bb3
-    // CHECK-NEXT:  %cst10_si32 = spirv.Constant 10 : si32
-    // CHECK-NEXT:  %3 = spirv.Load "Function" %2 : si32
-    // CHECK-NEXT:  %4 = spirv.SLessThan %3, %cst10_si32 : si32
-    // CHECK-NEXT:  spirv.BranchConditional %4, ^bb2, ^bb4
-    // CHECK-NEXT:^bb2:  // pred: ^bb1
-    // CHECK-NEXT:  %cst1_si32 = spirv.Constant 1 : si32
-    // CHECK-NEXT:  %5 = spirv.Load "Function" %2 : si32
-    // CHECK-NEXT:  %6 = spirv.IAdd %5, %cst1_si32 : si32
-    // CHECK-NEXT:  %7 = spirv.Variable : !spirv.ptr<si32, Function>
-    // CHECK-NEXT:  spirv.Store "Function" %7, %6 : si32
-    // CHECK-NEXT:  spirv.Branch ^bb3
-    // CHECK-NEXT:^bb3:  // pred: ^bb2
-    // CHECK-NEXT:  %8 = spirv.Load "Function" %2 : si32
-    // CHECK-NEXT:  %cst1_si32_0 = spirv.Constant 1 : si32
-    // CHECK-NEXT:  %9 = spirv.IAdd %8, %cst1_si32_0 : si32
-    // CHECK-NEXT:  spirv.Store "Function" %2, %9 : si32
-    // CHECK-NEXT:  spirv.Branch ^bb1
-    // CHECK-NEXT:^bb4:  // pred: ^bb1
-    // CHECK-NEXT:  spirv.mlir.merge
+    // CHECK-NEXT: %cst10_si32 = spirv.Constant 10 : si32
+    // CHECK-NEXT: %3 = spirv.Load "Function" %2 : si32
+    // CHECK-NEXT: %4 = spirv.SLessThan %3, %cst10_si32 : si32
+    // CHECK-NEXT: spirv.BranchConditional %4, ^bb2, ^bb4
+    // CHECK-NEXT: ^bb2:  // pred: ^bb1
+    // CHECK-NEXT: %cst1_si32 = spirv.Constant 1 : si32
+    // CHECK-NEXT: %5 = spirv.Load "Function" %2 : si32
+    // CHECK-NEXT: %6 = spirv.IAdd %5, %cst1_si32 : si32
+    // CHECK-NEXT: %7 = spirv.Variable : !spirv.ptr<si32, Function>
+    // CHECK-NEXT: spirv.Store "Function" %7, %6 : si32
+    // CHECK-NEXT: spirv.Branch ^bb3
+    // CHECK-NEXT: ^bb3:  // pred: ^bb2
+    // CHECK-NEXT: %false_1 = spirv.Constant false
+    // CHECK-NEXT: spirv.Store "Function" %1, %false_1 : i1
+    // CHECK-NEXT: %8 = spirv.Load "Function" %2 : si32
+    // CHECK-NEXT: %cst1_si32_2 = spirv.Constant 1 : si32
+    // CHECK-NEXT: %9 = spirv.IAdd %8, %cst1_si32_2 : si32
+    // CHECK-NEXT: spirv.Store "Function" %2, %9 : si32
+    // CHECK-NEXT: spirv.Branch ^bb1
+    // CHECK-NEXT: ^bb4:  // pred: ^bb1
+    // CHECK-NEXT: spirv.mlir.merge
     // CHECK-NEXT:}
+
     for (int i = 0; i < 10; ++i) {
         int a = i + 1;
     }

diff --git a/test/CodeGen/cf_loops_while.glsl b/test/CodeGen/cf_loops_while.glsl
@@ -18,6 +18,8 @@ void main() {
     }
 
     // CHECK-NEXT: ^bb3:  // pred: ^bb2
+    // CHECK-NEXT:  %false_1 = spirv.Constant false
+    // CHECK-NEXT:  spirv.Store "Function" %3, %false_1 : i1
     // CHECK-NEXT:  spirv.Branch ^bb1
     // CHECK-NEXT: ^bb4:  // pred: ^bb1
     // CHECK-NEXT: spirv.mlir.merge

diff --git a/test/CodeGen/cf_loops_while_break.glsl b/test/CodeGen/cf_loops_while_break.glsl
@@ -10,8 +10,8 @@ void main() {
         int someVarBefore = 1;
 
         // CHECK: ^bb1:  // pred: ^bb0
-        // CHECK-NEXT: %true_2 = spirv.Constant true
-        // CHECK-NEXT: spirv.Store "Function" %0, %true_2 : i1
+        // CHECK-NEXT: %true_4 = spirv.Constant true
+        // CHECK-NEXT: spirv.Store "Function" %0, %true_4 : i1
         if (true) {
             break;
         }
@@ -22,12 +22,7 @@ void main() {
         // CHECK-NEXT: spirv.BranchConditional %3, ^bb5, ^bb3
 
         // CHECK: ^bb3:  // pred: ^bb2
-        // CHECK-NEXT: %false = spirv.Constant false
-
-        // Reset break control var
-
-        // CHECK-NEXT: spirv.Store "Function" %0, %false : i1
-        // CHECK-NEXT: %cst1_si32_1 = spirv.Constant 1 : si32
+        // CHECK-NEXT: %cst1_si32_2 = spirv.Constant 1 : si32
         // CHECK-NEXT: %4 = spirv.Variable : !spirv.ptr<si32, Function>
         int someVarAfter = 1;
     }

diff --git a/test/CodeGen/cf_loops_while_continue.glsl b/test/CodeGen/cf_loops_while_continue.glsl
@@ -10,13 +10,13 @@ void main() {
         int someVarBefore = 1;
 
         // CHECK: ^bb1:  // pred: ^bb0
-        // CHECK-NEXT: %true_3 = spirv.Constant true
-        // CHECK-NEXT: spirv.Store "Function" %1, %true_3 : i1
+        // CHECK-NEXT: %true_4 = spirv.Constant true
+        // CHECK-NEXT: spirv.Store "Function" %1, %true_4 : i1
         if (true) {
             continue;
         // CHECK: ^bb2:  // pred: ^bb0
-        // CHECK-NEXT: %true_4 = spirv.Constant true
-        // CHECK-NEXT: spirv.Store "Function" %0, %true_4 : i1
+        // CHECK-NEXT: %true_5 = spirv.Constant true
+        // CHECK-NEXT: spirv.Store "Function" %0, %true_5 : i1
         } else {
             break;
         }
@@ -29,12 +29,9 @@ void main() {
         // CHECK-NEXT: %4 = spirv.Load "Function" %0 : i1
         // CHECK-NEXT: spirv.BranchConditional %4, ^bb6, ^bb4
 
-        // Reset continue/break control vars
         // CHECK: ^bb4:  // pred: ^bb3
-        // CHECK-NEXT: %false = spirv.Constant false
-        // CHECK-NEXT: spirv.Store "Function" %0, %false : i1
-        // CHECK-NEXT: %false_1 = spirv.Constant false
-        // CHECK-NEXT: spirv.Store "Function" %1, %false_1 : i1
+        // CHECK-NEXT: %cst1_si32_2 = spirv.Constant 1 : si32
+        // CHECK-NEXT: %5 = spirv.Variable : !spirv.ptr<si32, Function>
         int someVarAfter = 1;
     }
 

diff --git a/test/CodeGen/scopes.glsl b/test/CodeGen/scopes.glsl
@@ -14,8 +14,8 @@ void main() {
         // CHECK: %7 = spirv.Variable : !spirv.ptr<si32, Function>
         int a;
 
-        // CHECK: %cst2_si32_1 = spirv.Constant 2 : si32
-        // CHECK-NEXT: spirv.Store "Function" %7, %cst2_si32_1 : si32
+        // CHECK: %cst2_si32_2 = spirv.Constant 2 : si32
+        // CHECK-NEXT: spirv.Store "Function" %7, %cst2_si32_2 : si32
         a = 2;
     } else {
         // CHECK: %8 = spirv.Variable : !spirv.ptr<si32, Function>
@@ -37,7 +37,7 @@ void main() {
      */
 
     // CHECK: %7 = spirv.Load "Function" %0 : si32
-    // CHECK-NEXT: %8 = spirv.IEqual %7, %cst1_si32_1 : si32
+    // CHECK-NEXT: %8 = spirv.IEqual %7, %cst1_si32_2 : si32
     while (a == 1) {
         // CHECK: %9 = spirv.Variable : !spirv.ptr<si32, Function>
         int a;
@@ -59,20 +59,20 @@ void main() {
      */
 
     // CHECK: %5 = spirv.Load "Function" %0 : si32
-    // CHECK-NEXT: %6 = spirv.IEqual %5, %cst1_si32_0 : si32
+    // CHECK-NEXT: %6 = spirv.IEqual %5, %cst1_si32_1 : si32
     if (a == 1) {
         // CHECK: %7 = spirv.Variable : !spirv.ptr<si32, Function>
-        // CHECK-NEXT: %cst1_si32_1 = spirv.Constant 1 : si32
-        // CHECK-NEXT: spirv.Store "Function" %7, %cst1_si32_1 : si32
+        // CHECK-NEXT: %cst1_si32_2 = spirv.Constant 1 : si32
+        // CHECK-NEXT: spirv.Store "Function" %7, %cst1_si32_2 : si32
         int a;
         a = 1;
 
         // CHECK: %8 = spirv.Load "Function" %7 : si32
-        // CHECK-NEXT: %9 = spirv.IEqual %8, %cst2_si32_2 : si32
+        // CHECK-NEXT: %9 = spirv.IEqual %8, %cst2_si32_3 : si32
         if (a == 2) {
             // CHECK: %10 = spirv.Variable : !spirv.ptr<si32, Function>
-            // CHECK-NEXT: %cst2_si32_3 = spirv.Constant 2 : si32
-            // CHECK-NEXT: spirv.Store "Function" %10, %cst2_si32_3 : si32
+            // CHECK-NEXT: %cst2_si32_4 = spirv.Constant 2 : si32
+            // CHECK-NEXT: spirv.Store "Function" %10, %cst2_si32_4 : si32
             int a;
             a = 2;
         }