Fixed structurized return when there are phi's from shortcircuiting. (#3981)

tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp

@@ -3494,6 +3494,14 @@ void InitRetValue(BasicBlock *exitBB) {
   }
 }
 
+static void ChangePredBranch(BasicBlock *BB, BasicBlock *NewBB) {
+  for (auto predIt = pred_begin(BB); predIt != pred_end(BB);) {
+    BasicBlock *Pred = *(predIt++);
+    TerminatorInst *TI = Pred->getTerminator();
+    TI->replaceUsesOfWith(BB, NewBB);
+  }
+}
+
 // For functions has multiple returns like
 // float foo(float a, float b, float c) {
 //   float r = c;
@@ -3529,6 +3537,7 @@ void InitRetValue(BasicBlock *exitBB) {
 void StructurizeMultiRetFunction(Function *F, ScopeInfo &ScopeInfo,
                                  bool bWaveEnabledStage,
                                  SmallVector<BranchInst *, 16> &DxBreaks) {
+
   if (ScopeInfo.CanSkipStructurize())
     return;
   // Get bbWithRets.
@@ -3621,8 +3630,10 @@ void StructurizeMultiRetFunction(Function *F, ScopeInfo &ScopeInfo,
           BasicBlock *CmpBB = BasicBlock::Create(BB->getContext(),
                                                  "bReturned.cmp.false", F, BB);
 
-          // Make BB preds go to cmpBB.
-          BB->replaceAllUsesWith(CmpBB);
+          // Make BB preds go to cmpBB. Do this instead of replaceAllUsesWith
+          // because BB could have PHI nodes that reference it.
+          ChangePredBranch(BB, CmpBB);
+
           // Update endscopeBB to CmpBB for scopes which has BB as endscope.
           updateEndScope(ScopeInfo, EndBBToScopeIndexMap, BB, CmpBB);
 
@@ -3641,7 +3652,8 @@ void StructurizeMultiRetFunction(Function *F, ScopeInfo &ScopeInfo,
               BasicBlock::Create(BB->getContext(), "bReturned.cmp.true", F, BB);
           BasicBlock *BreakBB =
               BasicBlock::Create(BB->getContext(), "bReturned.break", F, BB);
-          BB->replaceAllUsesWith(CmpBB);
+          ChangePredBranch(BB, CmpBB);
+
           // Update endscopeBB to CmpBB for scopes which has BB as endscope.
           updateEndScope(ScopeInfo, EndBBToScopeIndexMap, BB, CmpBB);
 

tools/clang/test/HLSLFileCheck/hlsl/control_flow/return/short_circuit_with_phi.hlsl

@@ -0,0 +1,58 @@
+// RUN: %dxc -E main -fcgl -opt-enable structurize-returns -T ps_6_0  -enable-short-circuit %s | FileCheck %s
+// RUN: %dxc -E main -opt-enable structurize-returns -T ps_6_0  -enable-short-circuit %s | FileCheck %s -check-prefix=FULL
+
+// FULL: @main
+
+float main(uint x : X, uint y : Y, uint z : Z) : SV_Target {
+  // CHECK: %[[bReturned:.*]] = alloca i1
+  // CHECK: store i1 false, i1* %[[bReturned]]
+
+  // CHECK-NOT: ret float
+  // CHECK: [[label:.+]] ; preds =
+  if (x) {
+
+    // CHECK: [[label2:.+]] ; preds =
+    if (y) {
+      // CHECK: store float 1.000000e+00
+      // CHECK: store i1 true, i1* %[[bReturned]]
+      // CHECK: br
+      return 1;
+
+      // CHECK: [[label_bRet_cmp_false:.+]]: ; preds =
+      // CHECK: %[[RET:.*]] = load i1, i1* %[[bReturned]]
+      // CHECK: %[[NRET:.*]] = icmp ne i1 %[[RET]], false
+      // CHECK: br i1 %[[NRET]],
+    }
+
+    // CHECK: [[if_end:.+]] ; preds =
+    // CHECK: load i32
+    // CHECK: %[[x_to_bool:.+]] = icmp
+    // CHECK: br i1 %[[x_to_bool]],
+
+    // CHECK: [[land_rhs:.+]] ; preds =
+    // CHECK: load i32
+    // CHECK: icmp
+    // CHECK: br
+    bool cond = x && y;
+
+    // CHECK: [[land_end:.+]] ; preds =
+    // CHECK-NOT: phi.+ %[[label_bRet_cmp_false]]
+    // CHECK: phi
+
+    // CHECK: icmp
+    // CHECK: br i1
+    if (cond) {
+      // CHECK: [[then2:.+]] ; preds =
+      // store float 2.000000e+00
+      // CHECK: store i1 true, i1* %[[bReturned]]
+      return 2;
+
+      // CHECK: [[label_bRet_cmp_false_2:.+]]: ; preds =
+      // CHECK: %[[RET2:.*]] = load i1, i1* %[[bReturned]]
+      // CHECK: %[[NRET2:.*]] = icmp ne i1 %[[RET2]], false
+      // CHECK: br i1 %[[NRET2]],
+    }
+  }
+
+  return 0; 
+}