Check flatten attribute in SimplifyCFG. (#2311)

* Check flatten attribute in SimplifyCFG.

include/dxc/DXIL/DxilMetadataHelper.h

@@ -370,6 +370,8 @@ public:
   void LoadDxilViewIdState(std::vector<unsigned> &SerializedState);
   // Control flow hints.
   static llvm::MDNode *EmitControlFlowHints(llvm::LLVMContext &Ctx, std::vector<DXIL::ControlFlowHint> &hints);
+  static unsigned GetControlFlowHintMask(const llvm::Instruction *I);
+  static bool HasControlFlowHintToPreventFlatten(const llvm::Instruction *I);
 
   // Subobjects
   void EmitSubobjects(const DxilSubobjects &Subobjects);

lib/DXIL/DxilMetadataHelper.cpp

@@ -1369,6 +1369,35 @@ MDNode *DxilMDHelper::EmitControlFlowHints(llvm::LLVMContext &Ctx, std::vector<D
   return hintsNode;
 }
 
+unsigned DxilMDHelper::GetControlFlowHintMask(const Instruction *I) {
+  // Check that there are control hint to use
+  // branch.
+  MDNode *MD = I->getMetadata(hlsl::DxilMDHelper::kDxilControlFlowHintMDName);
+  if (!MD)
+    return 0;
+
+  if (MD->getNumOperands() < 3)
+    return 0;
+  unsigned mask = 0;
+  for (unsigned i = 2; i < MD->getNumOperands(); i++) {
+    Metadata *Op = MD->getOperand(2).get();
+    auto ConstOp = cast<ConstantAsMetadata>(Op);
+    unsigned hint = ConstOp->getValue()->getUniqueInteger().getLimitedValue();
+    mask |= 1 << hint;
+  }
+  return mask;
+}
+
+bool DxilMDHelper::HasControlFlowHintToPreventFlatten(
+    const llvm::Instruction *I) {
+  unsigned mask = GetControlFlowHintMask(I);
+  const unsigned BranchMask =
+      1 << (unsigned)(DXIL::ControlFlowHint::Branch) |
+      1 << (unsigned)(DXIL::ControlFlowHint::Call) |
+      1 << (unsigned)(DXIL::ControlFlowHint::ForceCase);
+  return mask & BranchMask;
+}
+
 void DxilMDHelper::EmitSubobjects(const DxilSubobjects &Subobjects) {
   NamedMDNode *pSubobjectsNamedMD = m_pModule->getNamedMetadata(kDxilSubobjectsMDName);
   IFTBOOL(pSubobjectsNamedMD == nullptr, DXC_E_INCORRECT_DXIL_METADATA);

lib/Transforms/Utils/SimplifyCFG.cpp

@@ -48,6 +48,9 @@
 #include <algorithm>
 #include <map>
 #include <set>
+
+#include "dxc/DXIL/DxilMetadataHelper.h" // HLSL Change - control flow hint.
+
 using namespace llvm;
 using namespace PatternMatch;
 
@@ -1488,7 +1491,7 @@ static bool SpeculativelyExecuteBB(BranchInst *BI, BasicBlock *ThenBB,
                                    const TargetTransformInfo &TTI) {
   // HLSL Change Begins.
   // Skip block with control flow hint.
-  if (BI->hasMetadataOtherThanDebugLoc()) {
+  if (hlsl::DxilMDHelper::HasControlFlowHintToPreventFlatten(BI)) {
     return false;
   }
   // HLSL Change Ends.
@@ -1911,7 +1914,7 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
   Instruction *InsertPt = DomBlock->getTerminator();
   // HLSL Change Begins.
   // Skip block with control flow hint.
-  if (InsertPt->hasMetadataOtherThanDebugLoc()) {
+  if (hlsl::DxilMDHelper::HasControlFlowHintToPreventFlatten(InsertPt)) {
     return false;
   }
   // HLSL Change Ends.

tools/clang/test/CodeGenHLSL/batch/control_flow/branch.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// Make sure if is not flattened.
+// CHECK:br
+// CHECK:dx.controlflow.hints
+
+float c;
+float main(float2 a:A) : SV_Target {
+    float x = a.x;
+    [branch]
+    if (c > 2)
+      x = x + 1;
+    else
+      x = x *3;
+    return x;
+}

tools/clang/test/CodeGenHLSL/batch/control_flow/flatten.hlsl

@@ -0,0 +1,15 @@
+// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// Make sure if is flattened.
+// CHECK-NOT:br
+// CHECK-NOT:dx.controlflow.hints
+
+float c;
+float main(float2 a:A) : SV_Target {
+    float x = a.x;
+    [flatten]
+    if (c > 2)
+      x = x + 1;
+    else
+      x = x *3;
+    return x;
+}

tools/clang/test/CodeGenHLSL/batch/misc/if2.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// RUN: %dxc -E main -T ps_6_0 -Od %s | FileCheck %s
 
 // CHECK: !"dx.controlflow.hints", i32 2