1 files changed, 11 insertions, 0 deletions
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
index cb37137d547..28e89b1be5c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
@@ -27,6 +27,7 @@
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/CodeGen/FastISel.h"
@@ -325,6 +326,8 @@ void SelectionDAGISel::getAnalysisUsage(AnalysisUsage &AU) const {
   if (OptLevel != CodeGenOpt::None)
     AU.addRequired<AAResultsWrapperPass>();
   AU.addRequired<GCModuleInfo>();
+  if (OptLevel != CodeGenOpt::None)
+    AU.addRequired<LoopInfoWrapperPass>();
   AU.addRequired<StackProtector>();
   AU.addPreserved<StackProtector>();
   AU.addPreserved<GCModuleInfo>();
@@ -1415,6 +1418,7 @@ void SelectionDAGISel::SelectAllBasicBlocks(const Function &Fn) {
 
   // Iterate over all basic blocks in the function.
   for (const BasicBlock *LLVMBB : RPOT) {
+    CurDAG->IsDAGPartOfLoop = false;
     if (OptLevel != CodeGenOpt::None) {
       bool AllPredsVisited = true;
       for (const_pred_iterator PI = pred_begin(LLVMBB), PE = pred_end(LLVMBB);
@@ -1592,6 +1596,13 @@ void SelectionDAGISel::SelectAllBasicBlocks(const Function &Fn) {
                                    FunctionBasedInstrumentation);
     }
 
+    if (OptLevel != CodeGenOpt::None) {
+      auto &LIWP = getAnalysis<LoopInfoWrapperPass>();
+      LoopInfo &LI = LIWP.getLoopInfo();
+      if (LI.getLoopFor(LLVMBB))
+        CurDAG->IsDAGPartOfLoop = true;
+    }
+
     if (Begin != BI)
       ++NumDAGBlocks;
     else