3 files changed, 55 insertions, 0 deletions

diff --git a/polly/include/polly/ScopInfo.h b/polly/include/polly/ScopInfo.h
index 18ec5c8e247..1f56c94bebc 100644
--- a/polly/include/polly/ScopInfo.h
+++ b/polly/include/polly/ScopInfo.h
@@ -1888,6 +1888,10 @@ private:
   /// in a schedule tree is given in the isl manual.
   isl::schedule Schedule = nullptr;
 
+  /// Whether the schedule has been modified after derived from the CFG by
+  /// ScopBuilder.
+  bool ScheduleModified = false;
+
   /// The set of minimal/maximal accesses for each alias group.
   ///
   /// When building runtime alias checks we look at all memory instructions and
@@ -2997,6 +3001,10 @@ public:
   /// NewSchedule The new schedule (given as schedule tree).
   void setScheduleTree(isl::schedule NewSchedule);
 
+  /// Whether the schedule is the original schedule as derived from the CFG by
+  /// ScopBuilder.
+  bool isOriginalSchedule() const { return !ScheduleModified; }
+
   /// Intersects the domains of all statements in the SCoP.
   ///
   /// @return true if a change was made
diff --git a/polly/lib/Analysis/ScopInfo.cpp b/polly/lib/Analysis/ScopInfo.cpp
index 5b3280c6095..157b4c6fde8 100644
--- a/polly/lib/Analysis/ScopInfo.cpp
+++ b/polly/lib/Analysis/ScopInfo.cpp
@@ -4521,10 +4521,12 @@ void Scop::setSchedule(isl::union_map NewSchedule) {
   auto S = isl::schedule::from_domain(getDomains());
   Schedule = S.insert_partial_schedule(
       isl::multi_union_pw_aff::from_union_map(NewSchedule));
+  ScheduleModified = true;
 }
 
 void Scop::setScheduleTree(isl::schedule NewSchedule) {
   Schedule = NewSchedule;
+  ScheduleModified = true;
 }
 
 bool Scop::restrictDomains(isl::union_set Domain) {
diff --git a/polly/lib/Transform/ForwardOpTree.cpp b/polly/lib/Transform/ForwardOpTree.cpp
index 9c24170bcd5..00d611e5e04 100644
--- a/polly/lib/Transform/ForwardOpTree.cpp
+++ b/polly/lib/Transform/ForwardOpTree.cpp
@@ -396,6 +396,14 @@ public:
     return UseScatter.apply_range(ReachDefTimepoints);
   }
 
+  /// Is @p InnerLoop nested inside @p OuterLoop?
+  static bool isInsideLoop(Loop *OuterLoop, Loop *InnerLoop) {
+    // If OuterLoop is nullptr, we cannot call its contains() method. In this
+    // case OuterLoop represents the 'top level' and therefore contains all
+    // loop.
+    return !OuterLoop || OuterLoop->contains(InnerLoop);
+  }
+
   /// Get a domain translation map from a (scalar) definition to the statement
   /// where the definition is being moved to.
   ///
@@ -422,6 +430,43 @@ public:
           getDomainFor(TargetStmt).get_space().map_from_set());
 
     isl::map &Result = DefToTargetCache[std::make_pair(TargetStmt, DefStmt)];
+
+    // This is a shortcut in case the schedule is still the original and
+    // TargetStmt is in the same or nested inside DefStmt's loop. With the
+    // additional assumption that operand trees do not cross DefStmt's loop
+    // header, then TargetStmt's instance shared coordinated are the same as
+    // DefStmt's coordinated. All TargetStmt instances with this prefix share
+    // the same DefStmt instance.
+    // Model:
+    //
+    //   for (int i < 0; i < N; i+=1) {
+    // DefStmt:
+    //    D = ...;
+    //    for (int j < 0; j < N; j+=1) {
+    // TargetStmt:
+    //      use(D);
+    //    }
+    //  }
+    //
+    // Here, the value used in TargetStmt is defined in the corresponding
+    // DefStmt, i.e.
+    //
+    //   { DefStmt[i] -> TargetStmt[i,j] }
+    //
+    // In practice, this should cover the majority of cases.
+    if (S->isOriginalSchedule() &&
+        isInsideLoop(DefStmt->getSurroundingLoop(),
+                     TargetStmt->getSurroundingLoop())) {
+      isl::set DefDomain = getDomainFor(DefStmt);
+      isl::set TargetDomain = getDomainFor(TargetStmt);
+      assert(DefDomain.dim(isl::dim::set) <= TargetDomain.dim(isl::dim::set));
+
+      Result = isl::map::from_domain_and_range(DefDomain, TargetDomain);
+      for (unsigned i = 0, DefDims = DefDomain.dim(isl::dim::set); i < DefDims;
+           i += 1)
+        Result = Result.equate(isl::dim::in, i, isl::dim::out, i);
+    }
+
     if (!Result) {
       // { DomainDef[] -> DomainTarget[] }
       Result = computeUseToDefFlowDependency(TargetStmt, DefStmt).reverse();