[TSAR, Transform] Limit scopes of loops for swapping.

Author: yukatan1701

include/tsar/Support/DiagnosticKinds.td

@@ -184,6 +184,8 @@ def warn_loop_swapping_diff_reduction: Warning<"unable to swap loops due to the
 def warn_loop_swapping_true_anti_dependence: Warning<"unable to swap loops due to the true or anti dependence">;
 def warn_loop_swapping_missing_loop: Warning<"not enough loops for swapping">;
 def warn_loop_swapping_redundant_loop: Warning<"too many loops for swapping, ignore redundant">;
+def warn_loop_swapping_no_loop_id: Warning<"cannot find loop ID to perform swapping">;
+
 def error_loop_swapping_lost_loop: Error<"cannot match ForStmt with its IR">;
 def error_loop_swapping_expect_compound: Error<"expected compound statement after pragma">;
-def warn_loop_swapping_no_loop_id: Warning<"cannot find loop ID to perform swapping">;
+def error_loop_swapping_diff_scope: Error<"loops within a pragma must have the same scope">;

lib/Transform/Clang/LoopSwapping.cpp

@@ -54,11 +54,19 @@ using namespace tsar;
 
 namespace {
 
+struct LoopRangeInfo {
+  Loop *LoopPtr;
+  SourceRange Range;
+  CompoundStmt *CompStmtPtr;
+  LoopRangeInfo() : LoopPtr(nullptr), CompStmtPtr(nullptr) {}
+  LoopRangeInfo(llvm::Loop *L, const SourceRange &R, CompoundStmt *S):
+      LoopPtr(L), Range(R), CompStmtPtr(S) {}
+};
+
 /// This provides access to function-level analysis results on server.
 using ClangLoopSwappingProvider =
     FunctionPassAAProvider<DIEstimateMemoryPass, DIDependencyAnalysisPass>;
 using DIAliasTraitVector = std::vector<const DIAliasTrait *>;
-using LoopRangeInfo = std::pair<Loop *, SourceRange>;
 using LoopRangeList = SmallVector<LoopRangeInfo, 2>;
 using PragmaInfoList = SmallVector<std::pair<Stmt *, LoopRangeList>, 2>;
 
@@ -74,7 +82,6 @@ class LoopVisitor : public RecursiveASTVisitor<LoopVisitor> {
     , mLangOpts(Rewr.getLangOpts())
     , mLoopInfo(LM)
     , mState(TraverseState::NONE)
-    , mCurrentLevel(-1)
   {}
 
   bool TraverseStmt(Stmt *S) {
@@ -104,7 +111,6 @@ class LoopVisitor : public RecursiveASTVisitor<LoopVisitor> {
         RemoveEmptyLine.RemoveLineIfEmpty = false;
         /*for (auto SR : ToRemove)
           mRewriter.RemoveText(SR, RemoveEmptyLine);*/
-        mCurrentLevel++;
         mPragmaLoopsInfo.resize(mPragmaLoopsInfo.size() + 1);
         mPragmaLoopsInfo.back().first = S;
         mState = TraverseState::PRAGMA;
@@ -121,22 +127,30 @@ class LoopVisitor : public RecursiveASTVisitor<LoopVisitor> {
   }
 
   bool TraverseCompoundStmt(CompoundStmt *S) {
+    mCompStmtStack.push(S);
     if (mState == TraverseState::PRAGMA) {
       mState = TraverseState::OUTERFOR;
       auto Res = RecursiveASTVisitor::TraverseCompoundStmt(S);
       mState = TraverseState::NONE;
-      mCurrentLevel--;
       return Res;
     }
-    return RecursiveASTVisitor::TraverseCompoundStmt(S);
+    auto Res = RecursiveASTVisitor::TraverseCompoundStmt(S);
+    mCompStmtStack.pop();
+    return Res;
   }
 
   bool TraverseForStmt(ForStmt *S) {
     if (mState == TraverseState::OUTERFOR) {
       auto Match = mLoopInfo.find<AST>(S);
       if (Match != mLoopInfo.end()) {
-        mPragmaLoopsInfo.back().second.push_back(
-            std::make_pair(Match->get<IR>(), S->getSourceRange()));
+        auto &LRL = mPragmaLoopsInfo.back().second;
+        if (!LRL.empty() && LRL.back().CompStmtPtr != mCompStmtStack.top()) {
+          toDiag(mSrcMgr.getDiagnostics(), S->getBeginLoc(),
+            diag::error_loop_swapping_diff_scope);
+            return false;
+        }
+        LRL.push_back(LoopRangeInfo(Match->get<IR>(), S->getSourceRange(),
+            mCompStmtStack.top()));
       } else {
         toDiag(mSrcMgr.getDiagnostics(), S->getBeginLoc(),
             diag::error_loop_swapping_lost_loop);
@@ -164,8 +178,8 @@ class LoopVisitor : public RecursiveASTVisitor<LoopVisitor> {
         ++It, ++N) {
       dbgs() << "\tPragma " << N << " (" << It->first <<"):\n";
       for (const auto &Info : It->second) {
-        const auto LoopPtr = Info.first;
-        const auto &Range = Info.second;
+        const auto LoopPtr = Info.LoopPtr;
+        const auto &Range = Info.Range;
         dbgs() << "\t\t[Range]\n";
         dbgs() << "\t\tBegin:" << Range.getBegin().printToString(mSrcMgr)
             << "\n";
@@ -186,8 +200,7 @@ class LoopVisitor : public RecursiveASTVisitor<LoopVisitor> {
   const LoopMatcherPass::LoopMatcher &mLoopInfo;
   TraverseState mState;
   SmallVector<Stmt *, 1> mClauses;
-
-  int mCurrentLevel;
+  std::stack<CompoundStmt *> mCompStmtStack;
   PragmaInfoList mPragmaLoopsInfo; 
 };
 
@@ -314,15 +327,15 @@ bool ClangLoopSwapping::hasTrueOrAntiDependence(
 
 bool ClangLoopSwapping::isSwappingAvailable(
     const LoopRangeList &LRL, const Stmt *Pragma) const {
-  auto *LoopID0 = mGetLoopID(LRL[0].first->getLoopID());
-  auto *LoopID1 = mGetLoopID(LRL[1].first->getLoopID());
+  auto *LoopID0 = mGetLoopID(LRL[0].LoopPtr->getLoopID());
+  auto *LoopID1 = mGetLoopID(LRL[1].LoopPtr->getLoopID());
   if (!LoopID0) {
-    toDiag(mSrcMgr->getDiagnostics(), LRL[0].second.getBegin(),
+    toDiag(mSrcMgr->getDiagnostics(), LRL[0].Range.getBegin(),
         diag::warn_loop_swapping_no_loop_id);
     return false;
   }
   if (!LoopID1) {
-    toDiag(mSrcMgr->getDiagnostics(), LRL[1].second.getBegin(),
+    toDiag(mSrcMgr->getDiagnostics(), LRL[1].Range.getBegin(),
         diag::warn_loop_swapping_no_loop_id);
     return false;
   }
@@ -363,8 +376,8 @@ void ClangLoopSwapping::swapLoops(const LoopVisitor &Visitor) {
               diag::warn_loop_swapping_redundant_loop);
     }
     if (isSwappingAvailable(Loops, Pragma)) {
-      auto Range0 = Loops[0].second;
-      auto Range1 = Loops[1].second;
+      auto Range0 = Loops[0].Range;
+      auto Range1 = Loops[1].Range;
       Range0.setEnd(GetLoopEnd(Range0));
       Range1.setEnd(GetLoopEnd(Range1));
       auto Range0End = Range0.getEnd();