[MLIR][Affine][VectorOps] Utility to vectorize loop nest using strategy

[dcaballe]

Summary

This patch adds a utility based on SuperVectorizer to vectorize an
affine loop nest using a given vectorization strategy. This strategy allows
targeting specific loops for vectorization instead of relying of the
SuperVectorizer analysis to choose the right loops to vectorize.

As part of the refactoring needed in SuperVectorizer to make the necessary
pieces public, we replaced NestedMatcher with SmallVector in some functions as
the mechanism to pass loop candidates over the algorithm. This allows us not to
expose NestedMatcher in the public API so that the client can provide loop
candidates using any other approach.

Please, let me know if you have any other ideas to keep the public API generic.
Another option could be using 'std::vector<SmallVector<AffineForOp, 2>>', to
preserve the parent-children relationship (we use something similar in
utility 'gatherLoops'[1].

[1] https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Transforms/LoopUtils.h#L273

[aartbik]

Thanks for your CL, Diego. I am not sure if Nicolas will respond soon since he is traveling.

I will try to start a review in the meantime, but I need a little bit more time than Nicolas would to read this carefully :-)

[aartbik]

Thanks for your patience. So far I only looked at lowering structured ops into vectorized code, not really into this "super" vectorizer that extracts vector code from scalar loops (even though I am of course a big fan of that :-). It looks like an interesting beginning, and refactoring loop nest vectorization into a public API seems useful.

mlir/include/mlir/Dialect/Affine/Utils.h

39

on a single loop nest

40

since this is now a "public" API, an example in the top level comment would help
(e.g. nest with two loops, and fill in the vector/map below with sample values)

40

Nicolas had the following TODO on the original anonymous structure, you may want to keep this here too, or check with him if this is still the plan

TODO: Hoist to a VectorizationStrategy.cpp when appropriate.

45

typo: vectorize -> vectorized

90

Nest's loops reads a bit awkward. How about

The loops in the loop nest must be ...

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

857

there is a TAB char here, replace with spaces

862

here and below, again TAB chars

884

having a

} // end for

will improve readability a bit (that, and fixing the TABs)

1253

break after "after" to even this out better

1310

Function -> function

1346

// end for

mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir

127

can we do better on the map capture?

[dcaballe]

Thanks, Aart! Sorry for the delay. I've been OOO. I'll take a look at your comments and I'll also try to use 'std::vector<SmallVector<AffineForOp, 2>>' for the API so that we can preserve some kind of loop nest structure.

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

862

Sorry about this! Not sure what happened... my editor shouldn't generate tabs at all...

[aartbik]

In D85869#2243647, @dcaballe wrote:

Thanks, Aart! Sorry for the delay. I've been OOO. I'll take a look at your comments

No worries. If you wait a bit longer, Nicolas will be back too :-)

[dcaballe]

Addressed all the comments and replaced the DFS loop container in the API with
an 'std::vector<SmallVector<AffineForOp, 2>>' (we use something similar in other
loop utilities). This container is a bit more friendly for a public API than
providing a plain list of loops in DFS order and also explicitly represents the
nesting level relationship of the loops that we want to vectorized, which is
useful for the vectorization algorithm itself. The vectorizer is therefore updated
to work on this container.

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

857

Not sure what's going on. I don't see any tab in my editor. I'll reformat the line, just in case. Maybe my patch is replacing an existing tab with spaces or some other weird character? Clang-format seems to be happy with the new version.

862

Same here. Hopefully it's fixed now.

1253

reformatted

1346

This loop is really small, right? Hopefully we don't have to add this comment to all the loops. Is this in the coding guidelines? If that's the case, maybe it's better to 'enforce' that with clang-format (like for namespaces)?

[nicolasvasilache]

Thanks for this @dcaballe, I'm. a big fan of more composable APIs and less blackbox heuristics.

mlir/include/mlir/Dialect/Affine/Utils.h

111

Could you please beef up this example a bit to give some intuition about how things are specified in the second dimension?
The examples so far only show 1 loop in the second dimension.
I think I see how this can represent imperfectly nested loops of different depths but more doc always helps.

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

863

While you are at it maybe cleanup trivial braces ?
The style was not established at the time this code was written so there are likely a few more in this file.

1142

I agree it is nicer to extract this recursive iteration into its own logic and populate a more easy to use auxiliary data structure.
Thanks for this cleanup!

1351

I think it should be possible to simplify the logic by using isAncestor / isProperAncestor ?

1357

Shouldn't you also assert that for each pair of loops "m, n" in loops[i], they are not nested ?
I think I can see ways to construct 2-D vectors that may break things.

[dcaballe]

Addressed feedback. Thanks!

[dcaballe]

Thanks! Any other comments? @aartbik?

mlir/include/mlir/Dialect/Affine/Utils.h

111

Something like this?

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

863

Sure!

1351

Ok! I think it would be a bit less efficient since to use isAncestor we have to traverse the "ancestry chain" multiple times (the current approach only traverses it once), but I don't expect the number of nesting levels and loops to be huge so it should be fine. The second check suggested below fits better following isAncestor approach here.

[nicolasvasilache]

mlir/include/mlir/Dialect/Affine/Utils.h

111

I was thinking of something like:

///   affine.for %i0 = 0 to 64 {
///     affine.for %i1 = 0 to 128 {
///       affine.for %i2 = 0 to 512 {
///         %ld = affine.load %in[%i0, %i1, %i2] : memref<64x128x512xf32>
///         affine.store %ld, %out[%i0, %i1, %i2] : memref<64x128x512xf32>
///       }
///     }
///     affine.for %i3 = 0 to 128 {
///       affine.for %i4 = 0 to 1024 {
///         %ld = affine.load %in[%i0, %i3, %i4] : memref<64x128x512xf32>
///         affine.store %ld, %out[%i0, %i3, %i4] : memref<64x128x512xf32>
///       }
///     }
///   }

and then:

/// loops = {{%i0}, {%i2, %i4}}, to vectorize the first, third and fifth loops,

Assuming I got it right .. :) ?

[dcaballe]

Sorry for the delay and sorry, Nicolas, I had written the example that
you were asking for in the .cpp doc but I forgot to copy it to the header
file. This patch is doing that. Please, let me know if there are any other
comments. Thanks!

[nicolasvasilache]

Thanks @dcaballe !

mlir/include/mlir/Dialect/Affine/Utils.h

 //===- Utils.h - Affine dialect utilities -----------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file declares a set of utilities for the affine dialect ops.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_UTILS_H
 #define MLIR_DIALECT_AFFINE_UTILS_H
 
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 class AffineForOp;
 class AffineIfOp;
 class AffineParallelOp;
 struct LogicalResult;
 class Operation;
 
 /// Replaces parallel affine.for op with 1-d affine.parallel op.
 /// mlir::isLoopParallel detect the parallel affine.for ops.
 /// There is no cost model currently used to drive this parallelization.
 void affineParallelize(AffineForOp forOp);
 
 /// Hoists out affine.if/else to as high as possible, i.e., past all invariant
 /// affine.fors/parallel's. Returns success if any hoisting happened; folded` is
 /// set to true if the op was folded or erased. This hoisting could lead to
 /// significant code expansion in some cases.
 LogicalResult hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr);
 
+/// Holds parameters to perform n-D vectorization on a single loop nest.
+/// For example, for the following loop nest:
+///
+/// func @vec2d(%in: memref<64x128x512xf32>, %out: memref<64x128x512xf32>) {
+///   affine.for %i0 = 0 to 64 {
+///     affine.for %i1 = 0 to 128 {
+///       affine.for %i2 = 0 to 512 {
+///         %ld = affine.load %in[%i0, %i1, %i2] : memref<64x128x512xf32>
+///         affine.store %ld, %out[%i0, %i1, %i2] : memref<64x128x512xf32>
+///       }
+///     }
+///   }
+///   return
+/// }
+///
+/// and VectorizationStrategy = 'vectorSizes = {8, 4}', 'loopToVectorDim =
+/// {{i1->0}, {i2->1}}', SuperVectorizer will generate:
+///
+///  func @vec2d(%arg0: memref<64x128x512xf32>, %arg1: memref<64x128x512xf32>) {
+///    affine.for %arg2 = 0 to 64 {
+///      affine.for %arg3 = 0 to 128 step 8 {
+///        affine.for %arg4 = 0 to 512 step 4 {
+///          %cst = constant 0.000000e+00 : f32
+///          %0 = vector.transfer_read %arg0[%arg2, %arg3, %arg4], %cst : ...
+///          vector.transfer_write %0, %arg1[%arg2, %arg3, %arg4] : ...
+///        }
+///      }
+///    }
+///    return
+///  }
+// TODO: Hoist to a VectorizationStrategy.cpp when appropriate.
+struct VectorizationStrategy {
+  // Vectorization factors to apply to each target vector dimension.
+  // Each factor will be applied to a different loop.
+  SmallVector<int64_t, 8> vectorSizes;
+  // Maps each AffineForOp vectorization candidate with its vector dimension.
+  // The candidate will be vectorized using the vectorization factor in
+  // 'vectorSizes' for that dimension.
+  DenseMap<Operation *, unsigned> loopToVectorDim;
+};
+
 /// Vectorizes 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.
 void vectorizeAffineLoops(
     Operation *parentOp,
     llvm::DenseSet<Operation *, DenseMapInfo<Operation *>> &loops,
     ArrayRef<int64_t> vectorSizes, ArrayRef<int64_t> fastestVaryingPattern);
 
+/// External utility to vectorize affine loops from a single loop nest using an
+/// n-D vectorization strategy (see doc in VectorizationStrategy definition).
+/// Loops are provided in a 2D vector container. The first dimension represents
+/// the nesting level relative to the loops to be vectorized. The second
+/// dimension contains the loops. This means that:
+///   a) every loop in 'loops[i]' must have a parent loop in 'loops[i-1]',
+///   b) a loop in 'loops[i]' may or may not have a child loop in 'loops[i+1]'.
+///
+/// For example, for the following loop nest:
+///
+///   func @vec2d(%in0: memref<64x128x512xf32>, %in1: memref<64x128x128xf32>,
+///               %out0: memref<64x128x512xf32>,
+///               %out1: memref<64x128x128xf32>) {
+///     affine.for %i0 = 0 to 64 {
+///       affine.for %i1 = 0 to 128 {
+///         affine.for %i2 = 0 to 512 {
+///           %ld = affine.load %in0[%i0, %i1, %i2] : memref<64x128x512xf32>
+///           affine.store %ld, %out0[%i0, %i1, %i2] : memref<64x128x512xf32>
+///         }
+///         affine.for %i3 = 0 to 128 {
+///           %ld = affine.load %in1[%i0, %i1, %i3] : memref<64x128x128xf32>
+///           affine.store %ld, %out1[%i0, %i1, %i3] : memref<64x128x128xf32>
+///         }
+///       }
+///     }
+///     return
+///   }
+///
+/// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two
+/// innermost loops;
+/// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost
+/// loops;
+/// loops = {{%i2}}, to vectorize only the first innermost loop;
+/// loops = {{%i3}}, to vectorize only the second innermost loop;
+/// loops = {{%i1}}, to vectorize only the middle loop.
+LogicalResult
+vectorizeAffineLoopNest(const std::vector<SmallVector<AffineForOp, 2>> &loops,
+                        const VectorizationStrategy &strategy);
+
 /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1.
 /// As currently implemented, this transformation cannot fail and will return
 /// early if the op is already in a normalized form.
 void normalizeAffineParallel(AffineParallelOp op);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_UTILS_H

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

 ///     loops that have a. conformable loop annotations attached (e.g. parallel,
 ///     reduction, vectorizable, ...) as well as b. all contiguous load/store
 ///     operations along a specified minor dimension (not necessarily the
 ///     fastest varying) ;
 ///  2. analyzing those patterns for profitability (TODO: and
 ///     interference);
 ///  3. Then, for each pattern in order:
 ///    a. applying iterative rewriting of the loop and the load operations in
-///       DFS postorder. Rewriting is implemented by coarsening the loops and
-///       turning load operations into opaque vector.transfer_read ops;
+///       inner-to-outer order. Rewriting is implemented by coarsening the loops
+///       and turning load operations into opaque vector.transfer_read ops;
 ///    b. keeping track of the load operations encountered as "roots" and the
 ///       store operations as "terminals";
 ///    c. traversing the use-def chains starting from the roots and iteratively
 ///       propagating vectorized values. Scalar values that are encountered
 ///       during this process must come from outside the scope of the current
 ///       pattern (TODO: enforce this and generalize). Such a scalar value
 ///       is vectorized only if it is a constant (into a vector splat). The
 ///       non-constant case is not supported for now and results in the pattern

 };
 
 } // end anonymous namespace
 
 Vectorize::Vectorize(ArrayRef<int64_t> virtualVectorSize) {
   vectorSizes = virtualVectorSize;
 }
 
-/////// TODO: Hoist to a VectorizationStrategy.cpp when appropriate.
-/////////
-namespace {
-
-struct VectorizationStrategy {
-  SmallVector<int64_t, 8> vectorSizes;
-  DenseMap<Operation *, unsigned> loopToVectorDim;
-};
-
-} // end anonymous namespace
-
 static void vectorizeLoopIfProfitable(Operation *loop, unsigned depthInPattern,
                                       unsigned patternDepth,
                                       VectorizationStrategy *strategy) {
   assert(patternDepth > depthInPattern &&
          "patternDepth is greater than depthInPattern");
   if (patternDepth - depthInPattern > strategy->vectorSizes.size()) {
     // Don't vectorize this loop
     return;
auto vectorizableBody =
         isVectorizableLoopBody(loop, &memRefDim, vectorTransferPattern());
     if (!vectorizableBody)
       return false;
     return memRefDim == -1 || fastestVaryingMemRefDimension == -1 ||
            memRefDim == fastestVaryingMemRefDimension;
   };
 }
 
-/// Apply vectorization of `loop` according to `state`. This is only triggered
-/// if all vectorizations in `childrenMatches` have already succeeded
-/// recursively in DFS post-order.
+/// Apply vectorization of `loop` according to `state`. `loops` are processed in
+/// inner-to-outer order to ensure that all the children loops have already been
+/// vectorized before vectorizing the parent loop.
 static LogicalResult
-vectorizeLoopsAndLoadsRecursively(NestedMatch oneMatch,
-                                  VectorizationState *state) {
-  auto *loopInst = oneMatch.getMatchedOperation();
-  auto loop = cast<AffineForOp>(loopInst);
-  auto childrenMatches = oneMatch.getMatchedChildren();
-
-  // 1. DFS postorder recursion, if any of my children fails, I fail too.
-  for (auto m : childrenMatches) {
-    if (failed(vectorizeLoopsAndLoadsRecursively(m, state))) {
-      return failure();
-    }
-  }
-
-  // 2. This loop may have been omitted from vectorization for various reasons
-  // (e.g. due to the performance model or pattern depth > vector size).
-  auto it = state->strategy->loopToVectorDim.find(loopInst);
-  if (it == state->strategy->loopToVectorDim.end()) {
-    return success();
-  }
+vectorizeLoopsAndLoads(std::vector<SmallVector<AffineForOp, 2>> &loops,
+  // Vectorize loops in inner-to-outer order. If any children fails, the parent
+  // will fail too.
+  for (auto &loopsInLevel : llvm::reverse(loops)) {
+    for (AffineForOp loop : loopsInLevel) {
+      // 1. This loop may have been omitted from vectorization for various
+      // reasons (e.g. due to the performance model or pattern depth > vector
+      // size).
+      auto it = state->strategy->loopToVectorDim.find(loop.getOperation());
+      if (it == state->strategy->loopToVectorDim.end())
+        continue;
 
-  // 3. Actual post-order transformation.
+      // 2. Actual inner-to-outer transformation.
       auto vectorDim = it->second;
       assert(vectorDim < state->strategy->vectorSizes.size() &&
              "vector dim overflow");
       //   a. get actual vector size
       auto vectorSize = state->strategy->vectorSizes[vectorDim];
       //   b. loop transformation for early vectorization is still subject to
-  //     exploratory tradeoffs (see top of the file). Apply coarsening, i.e.:
+      //     exploratory tradeoffs (see top of the file). Apply coarsening,
+      //     i.e.:
       //        | ub -> ub
       //        | step -> step * vectorSize
       LLVM_DEBUG(dbgs() << "\n[early-vect] vectorizeForOp by " << vectorSize
-                    << " : ");
-  LLVM_DEBUG(loopInst->print(dbgs()));
-  return vectorizeAffineForOp(loop, loop.getStep() * vectorSize, state);
+                        << " : \n"
+                        << loop);
+      if (failed(
+              vectorizeAffineForOp(loop, loop.getStep() * vectorSize, state)))
+        return failure();
+    } // end for.
+  }
+
+  return success();
 }
 
 /// Tries to transform a scalar constant into a vector splat of that constant.
 /// Returns the vectorized splat operation if the constant is a valid vector
 /// element type.
 /// If `type` is not a valid vector type or if the scalar constant is not a
 /// valid vector element type, returns nullptr.
 static Value vectorizeConstant(Operation *op, ConstantOp constant, Type type) {
for (unsigned i = 0; i < worklist.size(); ++i) {
     //    Note that we cannot just call replaceAllUsesWith because it may
     //    result in ops with mixed types, for ops whose operands have not all
     //    yet been vectorized. This would be invalid IR.
     state->registerReplacement(op, vectorizedInst);
   }
   return success();
 }
 
-/// Vectorization is a recursive procedure where anything below can fail.
-/// The root match thus needs to maintain a clone for handling failure.
-/// Each root may succeed independently but will otherwise clean after itself if
-/// anything below it fails.
-static LogicalResult vectorizeRootMatch(NestedMatch m,
-                                        VectorizationStrategy *strategy) {
-  auto loop = cast<AffineForOp>(m.getMatchedOperation());
-  OperationFolder folder(loop.getContext());
+/// Recursive implementation to convert all the nested loops in 'match' to a 2D
+/// vector container that preserves the relative nesting level of each loop with
+/// respect to the others in 'match'. 'currentLevel' is the nesting level that
+/// will be assigned to the loop in the current 'match'.
+static void
+getMatchedAffineLoopsRec(NestedMatch match, unsigned currentLevel,
+                         std::vector<SmallVector<AffineForOp, 2>> &loops) {
+  // Add a new empty level to the output if it doesn't exist already.
+  assert(currentLevel <= loops.size() && "Unexpected currentLevel");
+  if (currentLevel == loops.size())
+    loops.push_back(SmallVector<AffineForOp, 2>());
+
+  // Add current match and recursively visit its children.
+  loops[currentLevel].push_back(cast<AffineForOp>(match.getMatchedOperation()));
+  for (auto childMatch : match.getMatchedChildren()) {
+    getMatchedAffineLoopsRec(childMatch, currentLevel + 1, loops);
+  }
+}
+
+/// Converts all the nested loops in 'match' to a 2D vector container that
+/// preserves the relative nesting level of each loop with respect to the others
+/// in 'match'. This means that every loop in 'loops[i]' will have a parent loop
+/// in 'loops[i-1]'. A loop in 'loops[i]' may or may not have a child loop in
+/// 'loops[i+1]'.
+static void
+getMatchedAffineLoops(NestedMatch match,
+                      std::vector<SmallVector<AffineForOp, 2>> &loops) {
+  getMatchedAffineLoopsRec(match, /*currLoopDepth=*/0, loops);
+}
+
+/// Internal implementation to vectorize affine loops from a single loop nest
+/// using an n-D vectorization strategy.
+static LogicalResult
+vectorizeLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
+                  const VectorizationStrategy &strategy) {
+  assert(loops[0].size() == 1 && "Expected single root loop");
+  AffineForOp rootLoop = loops[0][0];
+  OperationFolder folder(rootLoop.getContext());
   VectorizationState state;
-  state.strategy = strategy;
+  state.strategy = &strategy;
   state.folder = &folder;
 
   // Since patterns are recursive, they can very well intersect.
   // Since we do not want a fully greedy strategy in general, we decouple
   // pattern matching, from profitability analysis, from application.
   // As a consequence we must check that each root pattern is still
   // vectorizable. If a pattern is not vectorizable anymore, we just skip it.
   // TODO: implement a non-greedy profitability analysis that keeps only
   // non-intersecting patterns.
-  if (!isVectorizableLoopBody(loop, vectorTransferPattern())) {
+  if (!isVectorizableLoopBody(rootLoop, vectorTransferPattern())) {
     LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ loop is not vectorizable");
     return failure();
   }
 
   /// Sets up error handling for this root loop. This is how the root match
   /// maintains a clone for handling failure and restores the proper state via
   /// RAII.
-  auto *loopInst = loop.getOperation();
+  auto *loopInst = rootLoop.getOperation();
   OpBuilder builder(loopInst);
   auto clonedLoop = cast<AffineForOp>(builder.clone(*loopInst));
   struct Guard {
     LogicalResult failure() {
       loop.getInductionVar().replaceAllUsesWith(clonedLoop.getInductionVar());
       loop.erase();
       return mlir::failure();
     }
     LogicalResult success() {
       clonedLoop.erase();
       return mlir::success();
     }
     AffineForOp loop;
     AffineForOp clonedLoop;
-  } guard{loop, clonedLoop};
+  } guard{rootLoop, clonedLoop};
 
   //////////////////////////////////////////////////////////////////////////////
   // Start vectorizing.
   // From now on, any error triggers the scope guard above.
   //////////////////////////////////////////////////////////////////////////////
-  // 1. Vectorize all the loops matched by the pattern, recursively.
+  // 1. Vectorize all the loop candidates, in inner-to-outer order.
   // This also vectorizes the roots (AffineLoadOp) as well as registers the
   // terminals (AffineStoreOp) for post-processing vectorization (we need to
   // wait for all use-def chains into them to be vectorized first).
-  if (failed(vectorizeLoopsAndLoadsRecursively(m, &state))) {
+  if (failed(vectorizeLoopsAndLoads(loops, &state))) {
     LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ failed root vectorizeLoop");
     return guard.failure();
   }
 
   // 2. Vectorize operations reached by use-def chains from root except the
   // terminals (store operations) that need to be post-processed separately.
   // TODO: add more as we expand.
   if (failed(vectorizeNonTerminals(&state))) {
vectorizeLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
   }
 
   // 4. Finish this vectorization pattern.
   LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ success vectorizing pattern");
   state.finishVectorizationPattern();
   return guard.success();
 }
 
-/// Applies vectorization to the current Function by searching over a bunch of
-/// predetermined patterns.
-void Vectorize::runOnFunction() {
-  FuncOp f = getFunction();
-  if (!fastestVaryingPattern.empty() &&
-      fastestVaryingPattern.size() != vectorSizes.size()) {
-    f.emitRemark("Fastest varying pattern specified with different size than "
-                 "the vector size.");
-    return signalPassFailure();
-  }
-
-  DenseSet<Operation *> parallelLoops;
-  f.walk([&parallelLoops](AffineForOp loop) {
-    if (isLoopParallel(loop))
-      parallelLoops.insert(loop);
-  });
-
-  vectorizeAffineLoops(f, parallelLoops, vectorSizes, fastestVaryingPattern);
+/// Vectorization is a recursive procedure where anything below can fail. The
+/// root match thus needs to maintain a clone for handling failure. Each root
+/// may succeed independently but will otherwise clean after itself if anything
+/// below it fails.
+static LogicalResult vectorizeRootMatch(NestedMatch m,
+                                        const VectorizationStrategy &strategy) {
+  std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize;
+  getMatchedAffineLoops(m, loopsToVectorize);
+  return vectorizeLoopNest(loopsToVectorize, strategy);
 }
 
-namespace mlir {
-
-/// Vectorizes 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.
-void vectorizeAffineLoops(Operation *parentOp, DenseSet<Operation *> &loops,
+/// 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) {
-  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
-  NestedPatternContext mlContext;
-
   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()));
 
     unsigned patternDepth = pat.getDepth();
for (auto m : matches) {
       if (failed(analyzeProfitability(m.getMatchedChildren(), 1, patternDepth,
                                       &strategy))) {
         continue;
       }
       vectorizeLoopIfProfitable(m.getMatchedOperation(), 0, patternDepth,
                                 &strategy);
       // TODO: if pattern does not apply, report it; alter the
       // cost/benefit.
-      vectorizeRootMatch(m, &strategy);
+      vectorizeRootMatch(m, strategy);
       // TODO: some diagnostics if failure to vectorize occurs.
     }
   }
   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() {
   return std::make_unique<Vectorize>();
 }
 
+/// Applies vectorization to the current function by searching over a bunch of
+/// predetermined patterns.
+void Vectorize::runOnFunction() {
+  FuncOp f = getFunction();
+  if (!fastestVaryingPattern.empty() &&
+      fastestVaryingPattern.size() != vectorSizes.size()) {
+    f.emitRemark("Fastest varying pattern specified with different size than "
+                 "the vector size.");
+    return signalPassFailure();
+  }
+
+  DenseSet<Operation *> parallelLoops;
+  f.walk([&parallelLoops](AffineForOp loop) {
+    if (isLoopParallel(loop))
+      parallelLoops.insert(loop);
+  });
+
+  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
+  NestedPatternContext mlContext;
+  vectorizeLoops(f, parallelLoops, vectorSizes, fastestVaryingPattern);
+}
+
+/// Verify that affine loops in 'loops' meet the nesting criteria expected by
+/// SuperVectorizer:
+///   * There must be at least one loop.
+///   * There must be a single root loop (nesting level 0).
+///   * Each loop at a given nesting level must be nested in a loop from a
+///     previous nesting level.
+static void
+verifyLoopNesting(const std::vector<SmallVector<AffineForOp, 2>> &loops) {
+  assert(!loops.empty() && "Expected at least one loop");
+  assert(!loops[0].size() && "Expected only one root loop");
+
+  // Traverse loops outer-to-inner to check some invariants.
+  for (int i = 1, end = loops.size(); i < end; ++i) {
+    for (AffineForOp loop : loops[i]) {
+      //  Check that each loop at this level is nested in one of the loops from
+      //  the previous level.
+      bool parentFound = false;
+      for (AffineForOp maybeParent : loops[i - 1]) {
+        if (maybeParent.getOperation()->isProperAncestor(loop)) {
+          parentFound = true;
+          break;
+        }
+      }
+      assert(parentFound && "Child loop not nested in any parent loop");
+
+      //  Check that each loop at this level is not nested in another loop from
+      //  this level.
+      for (AffineForOp sibling : loops[i])
+        assert(!sibling.getOperation()->isProperAncestor(loop) &&
+               "Loops at the same level are nested");
+    }
+  }
+}
+
+namespace mlir {
+
+/// External utility 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.
+void vectorizeAffineLoops(Operation *parentOp, DenseSet<Operation *> &loops,
+                          ArrayRef<int64_t> vectorSizes,
+                          ArrayRef<int64_t> fastestVaryingPattern) {
+  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
+  NestedPatternContext mlContext;
+  vectorizeLoops(parentOp, loops, vectorSizes, fastestVaryingPattern);
+}
+
+/// External utility to vectorize affine loops from a single loop nest using an
+/// n-D vectorization strategy (see doc in VectorizationStrategy definition).
+/// Loops are provided in a 2D vector container. The first dimension represents
+/// the nesting level relative to the loops to be vectorized. The second
+/// dimension contains the loops. This means that:
+///   a) every loop in 'loops[i]' must have a parent loop in 'loops[i-1]',
+///   b) a loop in 'loops[i]' may or may not have a child loop in 'loops[i+1]'.
+///
+/// For example, for the following loop nest:
+///
+///   func @vec2d(%in0: memref<64x128x512xf32>, %in1: memref<64x128x128xf32>,
+///               %out0: memref<64x128x512xf32>,
+///               %out1: memref<64x128x128xf32>) {
+///     affine.for %i0 = 0 to 64 {
+///       affine.for %i1 = 0 to 128 {
+///         affine.for %i2 = 0 to 512 {
+///           %ld = affine.load %in0[%i0, %i1, %i2] : memref<64x128x512xf32>
+///           affine.store %ld, %out0[%i0, %i1, %i2] : memref<64x128x512xf32>
+///         }
+///         affine.for %i3 = 0 to 128 {
+///           %ld = affine.load %in1[%i0, %i1, %i3] : memref<64x128x128xf32>
+///           affine.store %ld, %out1[%i0, %i1, %i3] : memref<64x128x128xf32>
+///         }
+///       }
+///     }
+///     return
+///   }
+///
+/// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two
+/// innermost loops;
+/// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost
+/// loops;
+/// loops = {{%i2}}, to vectorize only the first innermost loop;
+/// loops = {{%i3}}, to vectorize only the second innermost loop;
+/// loops = {{%i1}}, to vectorize only the middle loop.
+LogicalResult
+vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
+                        const VectorizationStrategy &strategy) {
+  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
+  NestedPatternContext mlContext;
+  verifyLoopNesting(loops);
+  return vectorizeLoopNest(loops, strategy);
+}
+
+std::unique_ptr<OperationPass<FuncOp>>
+createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
+  return std::make_unique<Vectorize>(virtualVectorSize);
+}
+std::unique_ptr<OperationPass<FuncOp>> createSuperVectorizePass() {
+  return std::make_unique<Vectorize>();
+}
+
 } // namespace mlir

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

 // RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128 test-fastest-varying=0" | FileCheck %s
 
 // Permutation maps used in vectorization.
-// CHECK: #[[$map_proj_d0d1_0:map[0-9]+]] = affine_map<(d0, d1) -> (0)>
+// CHECK-DAG: #[[$map_proj_d0d1_0:map[0-9]+]] = affine_map<(d0, d1) -> (0)>
+// CHECK-DAG: #[[$map_id1:map[0-9]+]] = affine_map<(d0) -> (d0)>
 
 #map0 = affine_map<(d0) -> (d0)>
 #mapadd1 = affine_map<(d0) -> (d0 + 1)>
 #mapadd2 = affine_map<(d0) -> (d0 + 2)>
 #mapadd3 = affine_map<(d0) -> (d0 + 3)>
 #set0 = affine_set<(i) : (i >= 0)>
 
 // Maps introduced to vectorize fastest varying memory index.
 // CHECK-LABEL: func @vec1d_1
 func @vec1d_1(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
 // CHECK-DAG: %[[C0:.*]] = constant 0 : index
 // CHECK-DAG: %[[C1:.*]] = constant 1 : index
 // CHECK-DAG: %[[C2:.*]] = constant 2 : index
 // CHECK-DAG: [[ARG_M:%[0-9]+]] = dim %{{.*}}, %[[C0]] : memref<?x?xf32>
 // CHECK-DAG: [[ARG_N:%[0-9]+]] = dim %{{.*}}, %[[C1]] : memref<?x?xf32>
 // CHECK-DAG: [[ARG_P:%[0-9]+]] = dim %{{.*}}, %[[C2]] : memref<?x?x?xf32>
    %c0 = constant 0 : index
    %c1 = constant 1 : index
    %c2 = constant 2 : index
    %M = dim %A, %c0 : memref<?x?xf32>
    %N = dim %A, %c1 : memref<?x?xf32>
    %P = dim %B, %c2 : memref<?x?x?xf32>
 
 // CHECK: for {{.*}} step 128
-// CHECK-NEXT: %{{.*}} = affine.apply #map0(%[[C0]])
-// CHECK-NEXT: %{{.*}} = affine.apply #map0(%[[C0]])
+// CHECK-NEXT: %{{.*}} = affine.apply #[[$map_id1]](%[[C0]])
+// CHECK-NEXT: %{{.*}} = affine.apply #[[$map_id1]](%[[C0]])
 // CHECK-NEXT: %{{.*}} = constant 0.0{{.*}}: f32
 // CHECK-NEXT: {{.*}} = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {permutation_map = #[[$map_proj_d0d1_0]]} : memref<?x?xf32>, vector<128xf32>
    affine.for %i0 = 0 to %M { // vectorized due to scalar -> vector
      %a0 = affine.load %A[%c0, %c0] : memref<?x?xf32>
    }
    return
 }
 
// CHECK-DAG: [[ARG_P:%[0-9]+]] = dim %{{.*}}, %[[C2]] : memref<?x?x?xf32>
    %c1 = constant 1 : index
    %c2 = constant 2 : index
    %M = dim %A, %c0 : memref<?x?xf32>
    %N = dim %A, %c1 : memref<?x?xf32>
    %P = dim %B, %c2 : memref<?x?x?xf32>
 
 // CHECK: affine.for %{{.*}}{{[0-9]*}} = 0 to %{{[0-9]*}} {
 // CHECK:   for [[IV18:%[a-zA-Z0-9]+]] = 0 to [[ARG_M]] step 128
-// CHECK:     %{{.*}} = affine.apply #map0(%{{.*}})
-// CHECK:     %{{.*}} = affine.apply #map0(%{{.*}})
+// CHECK:     %{{.*}} = affine.apply #[[$map_id1]](%{{.*}})
+// CHECK:     %{{.*}} = affine.apply #[[$map_id1]](%{{.*}})
 // CHECK:     %{{.*}} = constant 0.0{{.*}}: f32
 // CHECK:     {{.*}} = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {permutation_map = #[[$map_proj_d0d1_0]]} : memref<?x?xf32>, vector<128xf32>
    affine.for %i17 = 0 to %M { // not vectorized, the 1-D pattern that matched %{{.*}} in DFS post-order prevents vectorizing %{{.*}}
      affine.for %i18 = 0 to %M { // vectorized due to scalar -> vector
        %a18 = affine.load %A[%c0, %c0] : memref<?x?xf32>
      }
    }
    return
// CHECK-DAG: [[ARG_P:%[0-9]+]] = dim %{{.*}}, %[[C2]] : memref<?x?x?xf32>
    %c1 = constant 1 : index
    %c2 = constant 2 : index
    %M = dim %A, %c0 : memref<?x?xf32>
    %N = dim %A, %c1 : memref<?x?xf32>
    %P = dim %B, %c2 : memref<?x?x?xf32>
 
 // CHECK: affine.for %{{.*}}{{[0-9]*}} = 0 to %{{[0-9]*}} {
 // CHECK:   for [[IV18:%[a-zA-Z0-9]+]] = 0 to [[ARG_M]] step 128
-// CHECK:      %{{.*}} = affine.apply #map0(%{{.*}})
-// CHECK-NEXT: %{{.*}} = affine.apply #map0(%{{.*}})
+// CHECK:      %{{.*}} = affine.apply #[[$map_id1]](%{{.*}})
+// CHECK-NEXT: %{{.*}} = affine.apply #[[$map_id1]](%{{.*}})
 // CHECK-NEXT: %{{.*}} = constant 0.0{{.*}}: f32
 // CHECK-NEXT: {{.*}} = vector.transfer_read %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}} {permutation_map = #[[$map_proj_d0d1_0]]} : memref<?x?xf32>, vector<128xf32>
    affine.for %i17 = 0 to %M { // not vectorized, the 1-D pattern that matched %i18 in DFS post-order prevents vectorizing %{{.*}}
      affine.for %i18 = 0 to %M { // vectorized due to scalar -> vector
        %a18 = affine.load %A[%c0, %c0] : memref<?x?xf32>
      }
    }
    return

mlir/test/Dialect/Affine/SuperVectorize/vectorize_2d.mlir

func @vectorize_matmul(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
   affine.for %i0 = affine_map<(d0) -> (d0)>(%c0) to affine_map<(d0) -> (d0)>(%M) {
     affine.for %i1 = affine_map<(d0) -> (d0)>(%c0) to affine_map<(d0) -> (d0)>(%N) {
       %cst = constant 0.000000e+00 : f32
       affine.store %cst, %arg2[%i0, %i1] : memref<?x?xf32>
     }
   }
   //      VECT:  affine.for %[[I2:.*]] = #[[$map_id1]](%[[C0]]) to #[[$map_id1]](%[[M]]) step 4 {
   // VECT-NEXT:    affine.for %[[I3:.*]] = #[[$map_id1]](%[[C0]]) to #[[$map_id1]](%[[N]]) step 8 {
-  // VECT-NEXT:      affine.for %[[I4:.*]] = #map5(%[[C0]]) to #[[$map_id1]](%[[K]]) {
+  // VECT-NEXT:      affine.for %[[I4:.*]] = #[[$map_id1]](%[[C0]]) to #[[$map_id1]](%[[K]]) {
   //      VECT:        %[[A:.*]] = vector.transfer_read %{{.*}}[%[[I4]], %[[I3]]], %{{.*}} {permutation_map = #[[$map_proj_d0d1_zerod1]]} : memref<?x?xf32>, vector<4x8xf32>
   //      VECT:        %[[B:.*]] = vector.transfer_read %{{.*}}[%[[I2]], %[[I4]]], %{{.*}} {permutation_map = #[[$map_proj_d0d1_d0zero]]} : memref<?x?xf32>, vector<4x8xf32>
   // VECT-NEXT:        %[[C:.*]] = mulf %[[B]], %[[A]] : vector<4x8xf32>
   //      VECT:        %[[D:.*]] = vector.transfer_read %{{.*}}[%[[I2]], %[[I3]]], %{{.*}} : memref<?x?xf32>, vector<4x8xf32>
   // VECT-NEXT:        %[[E:.*]] = addf %[[D]], %[[C]] : vector<4x8xf32>
   //      VECT:        vector.transfer_write %[[E]], %{{.*}}[%[[I2]], %[[I3]]] : vector<4x8xf32>, memref<?x?xf32>
   affine.for %i2 = affine_map<(d0) -> (d0)>(%c0) to affine_map<(d0) -> (d0)>(%M) {
     affine.for %i3 = affine_map<(d0) -> (d0)>(%c0) to affine_map<(d0) -> (d0)>(%N) {