DxilValueCache::MayBranchTo: handle phi from switch properly (#4306)

Without this change, during DxilValueCache::ProcessAndSimplify_PHI,
MayBranchTo will return false for switch. If all but one case or default
block are collapsed into the switch's successor's phi (by simplifycfg), this
causes DxilValueCache to assume the remaining branch is the only predecessor
to the phi. If the corresponding incoming value is constant, DxilValueCache
will assume that's the only possible value, replacing the phi with the
constant.

With the change, MayBranchTo will return false only when a switch with a
constant condition would branch to a different successor.

lib/Analysis/DxilValueCache.cpp

@@ -11,6 +11,7 @@
 //
 
 
+#include "dxc/Support/Global.h"
 #include "llvm/Pass.h"
 #include "dxc/DXIL/DxilConstants.h"
 #include "llvm/Analysis/DxilSimplify.h"
@@ -68,9 +69,8 @@ bool DxilValueCache::IsAlwaysReachable_(BasicBlock *BB) {
 }
 
 bool DxilValueCache::MayBranchTo(BasicBlock *A, BasicBlock *B) {
-  BranchInst *Br = dyn_cast<BranchInst>(A->getTerminator());
-  if (!Br) return false;
-
+  TerminatorInst *Term = A->getTerminator();
+  if (BranchInst *Br = dyn_cast<BranchInst>(Term)) {
     if (Br->isUnconditional() && Br->getSuccessor(0) == B)
       return true;
 
@@ -78,6 +78,25 @@ bool DxilValueCache::MayBranchTo(BasicBlock *A, BasicBlock *B) {
       unsigned SuccIndex = C->getLimitedValue() != 0 ? 0 : 1;
       return Br->getSuccessor(SuccIndex) == B;
     }
+
+  } else if (SwitchInst *Sw = dyn_cast<SwitchInst>(Term)) {
+    if (ConstantInt *C =
+            dyn_cast<ConstantInt>(TryGetCachedValue(Sw->getCondition()))) {
+      for (auto Case : Sw->cases()) {
+        if (Case.getCaseValue() == C)
+          return Case.getCaseSuccessor() == B;
+      }
+      return Sw->getDefaultDest() == B;
+    }
+
+  } else if (isa<ReturnInst>(Term) || isa<UnreachableInst>(Term)) {
+    return false;
+
+  } else {
+    // Should not see: IndirectBrInst, InvokeInst, ResumeInst
+    DXASSERT(false, "otherwise, unexpected terminator instruction.");
+  }
+
   return true;
 }
 

tools/clang/test/HLSLFileCheck/hlsl/control_flow/switch/switch4.hlsl

@@ -0,0 +1,30 @@
+// RUN: %dxc -E main -T cs_6_0 %s | FileCheck %s
+
+// Ensure values through non-default cases are not eliminated
+// CHECK: call %dx.types.CBufRet.i32 @dx.op.cbufferLoadLegacy.i32(i32 59,
+
+cbuffer CBStruct : register(b0)
+{
+  uint value;
+};
+
+RWStructuredBuffer<uint> Output : register(u0);
+
+[numthreads(64, 1, 1)]
+void main(uint DTid : SV_DispatchThreadID)
+{
+  uint copyOfValue = value;
+
+  uint data = 0;
+  switch (DTid) {
+    case 0:
+    case 2:
+      data = copyOfValue;
+      break;
+
+    default:
+      data = 0;
+      break;
+  }
+  Output[DTid] = data;
+}