[SPIR-V] Fix flattened SPIR-V variable codegen. (microsoft#6756)

When the [branch] annotation is used, switches are converted into an
if/else tree.
Issue arise when declaring a variable into a switch case:
- when flattened, the same variable could be traversed twice in case of
a case fall-through.

In addition, there was a small bug in the flattening logic: only breaks
were stopping the handling of further statements, while early-returns
were ignored. This was not an important bug as it only added more
dead-code, but it was wrong.

Fixes microsoft#6718

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1083,6 +1083,9 @@ void DeclResultIdMapper::createCounterVarForDecl(const DeclaratorDecl *decl) {
 SpirvVariable *
 DeclResultIdMapper::createFnVar(const VarDecl *var,
                                 llvm::Optional<SpirvInstruction *> init) {
+  if (astDecls[var].instr != nullptr)
+    return cast<SpirvVariable>(astDecls[var].instr);
+
   const auto type = getTypeOrFnRetType(var);
   const auto loc = var->getLocation();
   const auto name = var->getName();
@@ -1095,10 +1098,7 @@ DeclResultIdMapper::createFnVar(const VarDecl *var,
   bool isAlias = false;
   (void)getTypeAndCreateCounterForPotentialAliasVar(var, &isAlias);
   varInstr->setContainsAliasComponent(isAlias);
-
-  assert(astDecls[var].instr == nullptr);
   astDecls[var].instr = varInstr;
-
   return varInstr;
 }
 

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -13768,7 +13768,9 @@ void SpirvEmitter::processSwitchStmtUsingIfStmts(const SwitchStmt *switchStmt) {
     // current case.
     std::vector<Stmt *> statements;
     unsigned i = curCaseIndex + 1;
-    for (; i < flatSwitch.size() && !isa<BreakStmt>(flatSwitch[i]); ++i) {
+    for (; i < flatSwitch.size() && !isa<BreakStmt>(flatSwitch[i]) &&
+           !isa<ReturnStmt>(flatSwitch[i]);
+         ++i) {
       if (!isa<CaseStmt>(flatSwitch[i]) && !isa<DefaultStmt>(flatSwitch[i]))
         statements.push_back(const_cast<Stmt *>(flatSwitch[i]));
     }

tools/clang/test/CodeGenSPIRV/cf.switch.ifstmt.early-return.hlsl

@@ -0,0 +1,91 @@
+// RUN: %dxc -T cs_6_0 -E main -fcgl  %s -spirv | FileCheck %s
+
+void simple() {
+  uint a = 0;
+
+// CHECK:    [[a:%[0-9]+]] = OpLoad %uint %a
+// CHECK: [[cond:%[0-9]+]] = OpIEqual %bool [[a]] %uint_0
+// CHECK:                    OpBranchConditional [[cond]] %if_true %if_false
+
+// CHECK:         %if_true = OpLabel
+// CHECK:                    OpStore %d1 %uint_0
+// CHECK:                    OpBranch %if_merge_0
+// CHECK:        %if_false = OpLabel
+// CHECK:    [[a:%[0-9]+]] = OpLoad %uint %a
+// CHECK: [[cond:%[0-9]+]] = OpIEqual %bool [[a]] %uint_1
+// CHECK:                    OpBranchConditional [[cond]] %if_true_0 %if_false_0
+
+// CHECK:       %if_true_0 = OpLabel
+// CHECK:                    OpStore %d1_0 %uint_1
+// CHECK:                    OpBranch %if_merge
+// CHECK:      %if_false_0 = OpLabel
+// CHECK:                    OpBranch %if_merge
+
+// CHECK:        %if_merge = OpLabel
+// CHECK:                    OpBranch %if_merge_0
+
+// CHECK:      %if_merge_0 = OpLabel
+// CHECK:                    OpReturn
+  [branch]
+  switch (a) {
+    default:
+      return;
+    case 0: {
+      uint d1 = 0;
+      return;
+    }
+    case 1: {
+      uint d1 = 1;
+    }
+  }
+}
+
+// CHECK:      [[b:%[0-9]+]] = OpLoad %uint %b
+// CHECK:   [[cond:%[0-9]+]] = OpIEqual %bool [[b]] %uint_0
+// CHECK:                      OpBranchConditional [[cond]] %if_true_1 %if_false_1
+
+// CHECK:         %if_true_1 = OpLabel
+// CHECK:                      OpStore %v1 %uint_0
+// CHECK:                      OpStore %v1_0 %uint_2
+// CHECK:                      OpStore %v2 %uint_1
+// CHECK:                      OpBranch %if_merge_2
+// CHECK:        %if_false_1 = OpLabel
+// CHECK:      [[b:%[0-9]+]] = OpLoad %uint %b
+// CHECK:   [[cond:%[0-9]+]] = OpIEqual %bool [[b]] %uint_1
+// CHECK:                      OpBranchConditional [[cond]] %if_true_2 %if_false_2
+
+// CHECK:         %if_true_2 = OpLabel
+// CHECK:                      OpStore %v1_0 %uint_2
+// CHECK:                      OpStore %v2 %uint_1
+// CHECK:                      OpBranch %if_merge_1
+// CHECK:        %if_false_2 = OpLabel
+// CHECK:                      OpBranch %if_merge_1
+
+// CHECK:        %if_merge_1 = OpLabel
+// CHECK:                      OpBranch %if_merge_2
+
+// CHECK:        %if_merge_2 = OpLabel
+// CHECK:                      OpReturn
+void fallthrough() {
+  uint b = 0;
+
+  [branch]
+  switch (b) {
+    default:
+      return;
+    case 0: {
+      uint v1 = 0;
+    }
+    case 1: {
+      uint v1 = 2;
+      uint v2 = 1;
+    }
+  }
+}
+
+[numthreads(1, 1, 1)]
+void main()
+{
+  simple();
+  fallthrough();
+}