[mlir][Vector][Affine] Fix heap-use-after-free in vectorizer

[dcaballe]

Summary

This patch fixes a heap-use-after-free introduced by the recent changes
in the vectorizer: https://reviews.llvm.org/rG95db7b4aeaad590f37720898e339a6d54313422f
The problem is due to the way candidate loops are visited. All candidate loops
are pattern-matched beforehand using the 'NestedMatch' utility. These matches may
intersect with each other so it may happen that we try to vectorize a loop that
was previously vectorized. The new vectorization algorithm replaces the original
loops that are vectorized with new loops and, therefore, any reference to the
original loops in the pre-computed matches becomes invalid.

This patch fixes the problem by classifying the candidate matches into buckets
before vectorization. Each bucket contains all the matches that intersect. The
vectorizer uses these buckets to make sure that we only vectorize *one* match from
each bucket, at most.

[ftynse]

Thanks for the fix! Don't forget to reland the patches I rolled back.

mlir/include/mlir/Analysis/NestedMatcher.h

 public:
   static NestedMatch build(Operation *operation,
                            ArrayRef<NestedMatch> nestedMatches);
   NestedMatch(const NestedMatch &) = default;
   NestedMatch &operator=(const NestedMatch &) = default;
 
   explicit operator bool() { return matchedOperation != nullptr; }
 
-  Operation *getMatchedOperation() { return matchedOperation; }
+  Operation *getMatchedOperation() const { return matchedOperation; }
   ArrayRef<NestedMatch> getMatchedChildren() { return matchedChildren; }
 
 private:
   friend class NestedPattern;
   friend class NestedPatternContext;
 
   /// Underlying global bump allocator managed by a NestedPatternContext.
   static llvm::BumpPtrAllocator *&allocator();

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

 static FilterFunctionType
 isVectorizableLoopPtrFactory(const DenseSet<Operation *> &parallelLoops,
                              int fastestVaryingMemRefDimension);
 
 /// Creates a vectorization pattern from the command line arguments.
 /// Up to 3-D patterns are supported.
 /// If the command line argument requests a pattern of higher order, returns an
 /// empty pattern list which will conservatively result in no vectorization.
-static std::vector<NestedPattern>
-makePatterns(const DenseSet<Operation *> &parallelLoops, int vectorRank,
-             ArrayRef<int64_t> fastestVaryingPattern) {
+static Optional<NestedPattern>
+makePattern(const DenseSet<Operation *> &parallelLoops, int vectorRank,
   using matcher::For;
   int64_t d0 = fastestVaryingPattern.empty() ? -1 : fastestVaryingPattern[0];
   int64_t d1 = fastestVaryingPattern.size() < 2 ? -1 : fastestVaryingPattern[1];
   int64_t d2 = fastestVaryingPattern.size() < 3 ? -1 : fastestVaryingPattern[2];
   switch (vectorRank) {
   case 1:
-    return {For(isVectorizableLoopPtrFactory(parallelLoops, d0))};
+    return For(isVectorizableLoopPtrFactory(parallelLoops, d0));
   case 2:
-    return {For(isVectorizableLoopPtrFactory(parallelLoops, d0),
-                For(isVectorizableLoopPtrFactory(parallelLoops, d1)))};
+    return For(isVectorizableLoopPtrFactory(parallelLoops, d0),
+               For(isVectorizableLoopPtrFactory(parallelLoops, d1)));
   case 3:
-    return {For(isVectorizableLoopPtrFactory(parallelLoops, d0),
-                For(isVectorizableLoopPtrFactory(parallelLoops, d1),
-                    For(isVectorizableLoopPtrFactory(parallelLoops, d2))))};
+    return For(isVectorizableLoopPtrFactory(parallelLoops, d0),
+                   For(isVectorizableLoopPtrFactory(parallelLoops, d2))));
   default: {
-    return std::vector<NestedPattern>();
+    return llvm::None;
   }
   }
 }
 
 static NestedPattern &vectorTransferPattern() {
   static auto pattern = matcher::Op([](Operation &op) {
     return isa<vector::TransferReadOp, vector::TransferWriteOp>(op);
   });

 /// below it fails.
 static LogicalResult vectorizeRootMatch(NestedMatch m,
                                         const VectorizationStrategy &strategy) {
   std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize;
   getMatchedAffineLoops(m, loopsToVectorize);
   return vectorizeLoopNest(loopsToVectorize, strategy);
 }
 
+/// Traverses all the loop matches and classifies them into intersection
+/// buckets. Two matches intersect if any of them encloses the other one. A
+/// match intersects with a bucket if the match intersects with the root
+/// (outermost) loop in that bucket.
+static void computeIntersectionBuckets(
+    ArrayRef<NestedMatch> matches,
+    std::vector<SmallVector<NestedMatch, 8>> &intersectionBuckets) {
+  assert(intersectionBuckets.empty() && "Expected empty output");
+  // Keeps track of the root (outermost) loop of each bucket.
+  SmallVector<AffineForOp, 8> bucketRoots;
+
+  for (const NestedMatch &match : matches) {
+    AffineForOp matchRoot = cast<AffineForOp>(match.getMatchedOperation());
+    bool intersects = false;
+    for (int i = 0, end = intersectionBuckets.size(); i < end; ++i) {
+      AffineForOp bucketRoot = bucketRoots[i];
+      // Add match to the bucket if the bucket root encloses the match root.
+      if (bucketRoot->isAncestor(matchRoot)) {
+        intersectionBuckets[i].push_back(match);
+        intersects = true;
+        break;
+      }
+      // Add match to the bucket if the match root encloses the bucket root. The
+      // match root becomes the new bucket root.
+      if (matchRoot->isAncestor(bucketRoot)) {
+        bucketRoots[i] = matchRoot;
+        intersectionBuckets[i].push_back(match);
+        intersects = true;
+        break;
+      }
+    }
+
+    // Match doesn't intersect with any existing bucket. Create a new bucket for
+    // it.
+    if (!intersects) {
+      bucketRoots.push_back(matchRoot);
+      intersectionBuckets.push_back(SmallVector<NestedMatch, 8>());
+      intersectionBuckets.back().push_back(match);
+    }
+  }
+}
+
 /// Internal implementation to vectorize affine loops in 'loops' using the n-D
 /// vectorization factors in 'vectorSizes'. By default, each vectorization
 /// factor is applied inner-to-outer to the loops of each loop nest.
 /// 'fastestVaryingPattern' can be optionally used to provide a different loop
 /// vectorization order.
 static void vectorizeLoops(Operation *parentOp, DenseSet<Operation *> &loops,
                            ArrayRef<int64_t> vectorSizes,
                            ArrayRef<int64_t> fastestVaryingPattern) {
-  for (auto &pat :
-       makePatterns(loops, vectorSizes.size(), fastestVaryingPattern)) {
-    LLVM_DEBUG(dbgs() << "\n******************************************");
-    LLVM_DEBUG(dbgs() << "\n******************************************");
-    LLVM_DEBUG(dbgs() << "\n[early-vect] new pattern on parent op\n");
-    LLVM_DEBUG(parentOp->print(dbgs()));
+  // Compute 1-D, 2-D or 3-D loop pattern to be matched on the target loops.
+  Optional<NestedPattern> pattern =
+      makePattern(loops, vectorSizes.size(), fastestVaryingPattern);
+  if (!pattern.hasValue()) {
+    LLVM_DEBUG(dbgs() << "\n[early-vect] pattern couldn't be computed\n");
+    return;
+  }
+
+  LLVM_DEBUG(dbgs() << *parentOp << "\n");
 
-    unsigned patternDepth = pat.getDepth();
+  unsigned patternDepth = pattern->getDepth();
 
-    SmallVector<NestedMatch, 8> matches;
-    pat.match(parentOp, &matches);
-    // Iterate over all the top-level matches and vectorize eagerly.
-    // This automatically prunes intersecting matches.
-    for (auto m : matches) {
+  // Compute all the pattern matches and classify them into buckets of
+  // intersecting matches.
+  SmallVector<NestedMatch, 32> allMatches;
+  pattern->match(parentOp, &allMatches);
+  std::vector<SmallVector<NestedMatch, 8>> intersectionBuckets;
+  computeIntersectionBuckets(allMatches, intersectionBuckets);
+
+  // Iterate over all buckets and vectorize the matches eagerly. We can only
+  // vectorize one match from each bucket since all the matches within a bucket
+  // intersect.
+  for (auto &intersectingMatches : intersectionBuckets) {
+    for (NestedMatch &match : intersectingMatches) {
       VectorizationStrategy strategy;
       // TODO: depending on profitability, elect to reduce the vector size.
       strategy.vectorSizes.assign(vectorSizes.begin(), vectorSizes.end());
-      if (failed(analyzeProfitability(m.getMatchedChildren(), 1, patternDepth,
-                                      &strategy))) {
+      if (failed(analyzeProfitability(match.getMatchedChildren(), 1,
+                                      patternDepth, &strategy))) {
         continue;
       }
-      vectorizeLoopIfProfitable(m.getMatchedOperation(), 0, patternDepth,
+      vectorizeLoopIfProfitable(match.getMatchedOperation(), 0, patternDepth,
                                 &strategy);
-      // TODO: if pattern does not apply, report it; alter the
-      // cost/benefit.
-      (void)vectorizeRootMatch(m, strategy);
+      // Vectorize match. Skip the rest of intersecting matches in the bucket if
+      // vectorization succeeded.
+      // TODO: if pattern does not apply, report it; alter the cost/benefit.
       // TODO: some diagnostics if failure to vectorize occurs.
+      if (succeeded(vectorizeRootMatch(match, strategy)))
+        break;
     }
   }
+
   LLVM_DEBUG(dbgs() << "\n");
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
   return std::make_unique<Vectorize>(virtualVectorSize);
 }
 std::unique_ptr<OperationPass<FuncOp>> createSuperVectorizePass() {