[mlir][Affine] Add pass options to supervectorizer

[jsetoain]

Summary

The only current options to create a supervectorization pass from an
external dialect is to use createSuperVectorizePass with the virtual
vector dimensions as a parameter, but the pass accepts other parameters.

This patch makes use of the autogenerated pass creation options to enable
pass configuration through AffineVectorizeOptions.

[jsetoain]

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

1720–1727

Note to reviewers: While working on this patch I noticed this "duplicated" code up here. It doesn't seem to be doing anything, and I've checked it can be safely removed. I really don't understand why it was left here, if it was an accident or something I'm missing, but I wanted to point it out. If it's unnecessary, we might as well get rid of it with this patch.

[jsetoain]

clang-format

[dcaballe]

Thanks, Javier! LGTM

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

1720–1727

It looks like it was accidentally left there.

[rriddle]

Pass generation was updated such that this should already be exposed via AffineVectorizeOptions. See https://reviews.llvm.org/D134607 as an example for updating to enable auto generation of the create methods.

[jsetoain]

Move on to the new autogenerated pass builders

[jsetoain]

In D134632#3817111, @rriddle wrote:

Pass generation was updated such that this should already be exposed via AffineVectorizeOptions. See https://reviews.llvm.org/D134607 as an example for updating to enable auto generation of the create methods.

Ah! I somehow missed this. Thanks for pointer, River! It looks much better this way. On the flip side, this might break things downstream if somebody was using the old method with the "virtualVectorSize" parameter. I think this is cleaner and the potential minor issue acceptable, but if it's not, I'm happy to re-introduce it for "backwards compatibility". In any case, I will let it sit here for a few days so people can see it coming.

[jsetoain]

Oops...

[dcaballe]

Awesome! I was not aware of that either. Thanks, River!
Yeah, this pass is pretty old. I wouldn't be surprised if more opportunities like this one come up.

On the flip side, this might break things downstream if somebody was using the old method with the "virtualVectorSize" parameter.

That sounds like a simple update. I would go with it.

mlir/include/mlir/Dialect/Affine/Passes.h

 /// line if provided.
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLoopUnrollAndJamPass(int unrollJamFactor = -1);
 
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
 std::unique_ptr<OperationPass<func::FuncOp>> createPipelineDataTransferPass();
 
-/// Creates a pass to vectorize loops, operations and data types using a
-/// target-independent, n-D super-vector abstraction.
-std::unique_ptr<OperationPass<func::FuncOp>>
-createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize);
-/// Overload relying on pass options for initialization.
-std::unique_ptr<OperationPass<func::FuncOp>> createSuperVectorizePass();
-
 /// Populate patterns that expand affine index operations into more fundamental
 /// operations (not necessarily restricted to Affine dialect).
 void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns);
 
 /// Creates a pass to expand affine index operations into more fundamental
 /// operations (not necessarily restricted to Affine dialect).
 std::unique_ptr<Pass> createAffineExpandIndexOpsPass();
 

mlir/include/mlir/Dialect/Affine/Passes.td

let description = [{
     }
     ```
   }];
   let constructor = "mlir::createAffineScalarReplacementPass()";
 }
 
 def AffineVectorize : Pass<"affine-super-vectorize", "func::FuncOp"> {
   let summary = "Vectorize to a target independent n-D vector abstraction";
-  let constructor = "mlir::createSuperVectorizePass()";
   let dependentDialects = ["vector::VectorDialect"];
   let options = [
     ListOption<"vectorSizes", "virtual-vector-size", "int64_t",
                "Specify an n-D virtual vector size for vectorization">,
     // Optionally, the fixed mapping from loop to fastest varying MemRef
     // dimension for all the MemRefs within a loop pattern:
     //   the index represents the loop depth, the value represents the k^th
     //   fastest varying memory dimension.

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

static NestedPattern &vectorTransferPattern() {
   return pattern;
 }
 
 namespace {
 
 /// Base state for the vectorize pass.
 /// Command line arguments are preempted by non-empty pass arguments.
 struct Vectorize : public impl::AffineVectorizeBase<Vectorize> {
-  Vectorize() = default;
-  Vectorize(ArrayRef<int64_t> virtualVectorSize);
+  using Base::Base;
+
   void runOnOperation() override;
 };
 
 } // namespace
 
-Vectorize::Vectorize(ArrayRef<int64_t> virtualVectorSize) {
-  vectorSizes = virtualVectorSize;
-}
-
 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;
for (NestedMatch &match : intersectingMatches) {
       if (succeeded(vectorizeRootMatch(match, strategy)))
         break;
     }
   }
 
   LLVM_DEBUG(dbgs() << "\n");
 }
 
-std::unique_ptr<OperationPass<func::FuncOp>>
-createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
-  return std::make_unique<Vectorize>(virtualVectorSize);
-}
-std::unique_ptr<OperationPass<func::FuncOp>> createSuperVectorizePass() {
-  return std::make_unique<Vectorize>();
-}
-
 /// Applies vectorization to the current function by searching over a bunch of
 /// predetermined patterns.
 void Vectorize::runOnOperation() {
   func::FuncOp f = getOperation();
   if (!fastestVaryingPattern.empty() &&
       fastestVaryingPattern.size() != vectorSizes.size()) {
     f.emitRemark("Fastest varying pattern specified with different size than "
                  "the vector size.");
vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
                         const VectorizationStrategy &strategy) {
   // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
   NestedPatternContext mlContext;
   if (failed(verifyLoopNesting(loops)))
     return failure();
   return vectorizeLoopNest(loops, strategy);
 }
 
-std::unique_ptr<OperationPass<func::FuncOp>>
-createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
-  return std::make_unique<Vectorize>(virtualVectorSize);
-}
-std::unique_ptr<OperationPass<func::FuncOp>> createSuperVectorizePass() {
-  return std::make_unique<Vectorize>();
-}
-
 } // namespace mlir