Merge pull request #445 from tensor-compiler/array_algebra_iter_const

Array algebra iter const

Author: weiya711

src/lower/merge_lattice.cpp

@@ -541,11 +541,53 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
     bool locateLeft = locateFromLeft(left, right);
 
     // Append all combinations of a and b merge points
-    for (auto& leftPoint : left.points()) {
-      for (auto& rightPoint : right.points()) {
-        points.push_back(intersectPoints(leftPoint, rightPoint, locateLeft));
+    struct pointSort {
+      bool operator()(const MergePoint& a, const MergePoint& b) {
+        size_t left_size  = a.iterators().size() + a.locators().size();
+        size_t right_size = b.iterators().size() + b.locators().size();
+        return left_size > right_size;
+      }
+    } pointSorter;
+
+    // Append all combinations of the merge points of a and b
+    auto sorted_apoint = left.points();
+    auto sorted_bpoint = right.points();
+    std::sort(sorted_apoint.begin(), sorted_apoint.end(), pointSorter);
+    std::sort(sorted_bpoint.begin(), sorted_bpoint.end(), pointSorter);
+
+    set<Iterator> apoint_root_set;
+    if (!sorted_apoint.empty())
+      apoint_root_set = sorted_apoint.begin()->tensorRegion();
+
+    set<Iterator>bpoint_root_set;
+    if (!sorted_bpoint.empty())
+      bpoint_root_set = sorted_bpoint.begin()->tensorRegion();
+
+
+    for (auto& apoint : sorted_apoint) {
+      for (auto& bpoint : sorted_bpoint) {
+        bool hasIntersection = true;
+
+        auto apoint_set = apoint.tensorRegion();
+        auto bpoint_set = bpoint.tensorRegion();
+
+        for (auto& it : apoint_set) {
+          if (!std::count(bpoint_set.begin(), bpoint_set.end(), it) &&
+              std::count(bpoint_root_set.begin(), bpoint_root_set.end(), it)) {
+            hasIntersection = false;
+          }
+        }
+        for (auto& it : bpoint_set) {
+          if (!std::count(apoint_set.begin(), apoint_set.end(), it) &&
+              std::count(apoint_root_set.begin(), apoint_root_set.end(), it)) {
+            hasIntersection = false;
+          }
+        }
+        if (hasIntersection)
+          points.push_back(intersectPoints(apoint, bpoint, locateLeft));
       }
     }
+    std::sort(points.begin(), points.end(), pointSorter);
 
     // Correctness: ensures that points produced on BOTH the left and the
     //              right lattices are produced in the final intersection.
@@ -557,7 +599,7 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
     //              points and resolves conflicts arising between omitters and
     //              producers
      points = removeDuplicatedTensorRegions(points, true);
-
+     
     // Optimization: Removed a subLattice of points if the entire subLattice is
     //               made of only omitters
     // points = removeUnnecessaryOmitterPoints(points);
@@ -577,10 +619,49 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
   {
     vector<MergePoint> points;
 
+    struct pointSort {
+      bool operator()(const MergePoint& a, const MergePoint& b) {
+        size_t left_size  = a.iterators().size() + a.locators().size();
+        size_t right_size = b.iterators().size() + b.locators().size();
+        return left_size > right_size;
+      }
+    } pointSorter;
+
     // Append all combinations of the merge points of a and b
-    for (auto& apoint : left.points()) {
-      for (auto& bpoint : right.points()) {
-        points.push_back(unionPoints(apoint, bpoint));
+    auto sorted_apoint = left.points();
+    auto sorted_bpoint = right.points();
+    std::sort(sorted_apoint.begin(), sorted_apoint.end(), pointSorter);
+    std::sort(sorted_bpoint.begin(), sorted_bpoint.end(), pointSorter);
+
+    set<Iterator> apoint_root_set;
+    if (!sorted_apoint.empty())
+      apoint_root_set = sorted_apoint.begin()->tensorRegion();
+
+    set<Iterator>bpoint_root_set;
+    if (!sorted_bpoint.empty())
+      bpoint_root_set = sorted_bpoint.begin()->tensorRegion();
+
+    for (auto& apoint : sorted_apoint) {
+      for (auto& bpoint : sorted_bpoint) {
+        bool hasIntersection = true;
+
+        auto apoint_set = apoint.tensorRegion();
+        auto bpoint_set = bpoint.tensorRegion();
+
+        for (auto& it : apoint_set) {
+          if (!std::count(bpoint_set.begin(), bpoint_set.end(), it) &&
+              std::count(bpoint_root_set.begin(), bpoint_root_set.end(), it)) {
+            hasIntersection = false;
+          }
+        }
+        for (auto& it : bpoint_set) {
+          if (!std::count(apoint_set.begin(), apoint_set.end(), it) &&
+              std::count(apoint_root_set.begin(), apoint_root_set.end(), it)) {
+            hasIntersection = false;
+          }
+        }
+        if (hasIntersection)
+          points.push_back(unionPoints(apoint, bpoint));
       }
     }
 
@@ -590,22 +671,13 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
     // Append the merge points of b
     util::append(points, right.points());
 
-    struct pointSort {
-      bool operator()(const MergePoint& a, const MergePoint& b) {
-        size_t left_size  = a.iterators().size() + a.locators().size();
-        size_t right_size = b.iterators().size() + b.locators().size();
-        return left_size > right_size;
-      }
-    } pointSorter;
-
     std::sort(points.begin(), points.end(), pointSorter);
 
     // Correctness: This ensures that points omitted on BOTH the left and the
     //              right lattices are omitted in the Union. Needed since some
     //              subpoints may produce leading to erroneous producer regions
     points = correctPointTypesAfterUnion(left.points(), right.points(), points);
 
-
     // Correctness: Deduplicate regions that are described by multiple lattice
     //              points and resolves conflicts arising between omitters and
     //              producers
@@ -671,6 +743,7 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
    */
   static MergePoint unionPoints(MergePoint left, MergePoint right)
   {
+
     vector<Iterator> iterators= combine(left.iterators(),right.iterators());
     vector<Iterator> locaters = combine(left.locators(), right.locators());
     vector<Iterator> results  = combine(left.results(),  right.results());
@@ -784,6 +857,21 @@ class MergeLatticeBuilder : public IndexNotationVisitorStrict, public IterationA
     return deduplicates;
   }
 
+  static vector<Iterator>
+  removeDimensionIterators(const vector<Iterator>& iterators)
+  {
+    vector<Iterator> result;
+
+    // Remove all but one of the dense iterators, which are all the same.
+    for (auto& iterator : iterators) {
+      if (!iterator.isDimensionIterator()) {
+        result.push_back(iterator);
+      }
+    }
+    return result;
+  }
+
+
   static vector<MergePoint>
   flipPoints(const vector<MergePoint>& points) {
     vector<MergePoint> flippedPoints;
@@ -1119,8 +1207,18 @@ ostream& operator<<(ostream& os, const MergeLattice& ml) {
 }
 
 bool operator==(const MergeLattice& a, const MergeLattice& b) {
-  auto& apoints = a.points();
-  auto& bpoints = b.points();
+  auto apoints = a.points();
+  auto bpoints = b.points();
+  struct pointSort {
+    bool operator()(const MergePoint& a, const MergePoint& b) {
+      size_t left_size  = a.iterators().size() + a.locators().size();
+      size_t right_size = b.iterators().size() + b.locators().size();
+      return left_size > right_size;
+    }
+  } pointSorter;
+
+  std::sort(apoints.begin(), apoints.end(), pointSorter);
+  std::sort(bpoints.begin(), bpoints.end(), pointSorter);
   if (apoints.size() != bpoints.size()) {
     return false;
   }