[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();
 
+/// Options for the supervectorizer. It should be kept in sync with the options
+/// defined for the pass.
+struct SuperVectorizeOptions {
+  SuperVectorizeOptions(ArrayRef<int64_t> virtualVectorSize,
+                        bool vectorizeReductions = false,
+                        ArrayRef<int64_t> fastestVaryingPattern = {})
+      : virtualVectorSize{virtualVectorSize},
+        vectorizeReductions{vectorizeReductions}, fastestVaryingPattern{
+                                                      fastestVaryingPattern} {}
+
+  ArrayRef<int64_t> virtualVectorSize;
+  bool vectorizeReductions{false};
+  ArrayRef<int64_t> fastestVaryingPattern;
+};
+
 /// 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);
+createSuperVectorizePass(const SuperVectorizeOptions &options);
 /// 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

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

 }
 
 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);
+  Vectorize(const SuperVectorizeOptions& options);
   void runOnOperation() override;
 };
 
 } // namespace
 
-Vectorize::Vectorize(ArrayRef<int64_t> virtualVectorSize) {
-  vectorSizes = virtualVectorSize;
+Vectorize::Vectorize(const SuperVectorizeOptions& options) {
+  vectorSizes = options.virtualVectorSize;
+  vectorizeReductions = options.vectorizeReductions;
+  fastestVaryingPattern = options.fastestVaryingPattern;
 }
 
 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()) {
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,
   // 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);
+createSuperVectorizePass(const SuperVectorizeOptions& options) {
+  return std::make_unique<Vectorize>(options);
 }
 std::unique_ptr<OperationPass<func::FuncOp>> createSuperVectorizePass() {
   return std::make_unique<Vectorize>();
 }
 
 } // namespace mlir