[mlir][Affine][VectorOps] Fix super vectorizer utility (D85869)

Adding missing code that should have been part of "D85869: Utility to
vectorize loop nest using strategy."

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D88346

Author: dcaballe

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

@@ -122,7 +122,7 @@ void vectorizeAffineLoops(
 /// 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,
+vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
                         const VectorizationStrategy &strategy);
 
 /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1.

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

@@ -1338,7 +1338,7 @@ void Vectorize::runOnFunction() {
 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");
+  assert(loops[0].size() == 1 && "Expected only one root loop");
 
   // Traverse loops outer-to-inner to check some invariants.
   for (int i = 1, end = loops.size(); i < end; ++i) {

mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir

@@ -1,5 +1,6 @@
 // RUN: mlir-opt %s -affine-super-vectorizer-test -vector-shape-ratio 4 -vector-shape-ratio 8 2>&1 | FileCheck %s
 // RUN: mlir-opt %s -affine-super-vectorizer-test -vector-shape-ratio 2 -vector-shape-ratio 5 -vector-shape-ratio 2 2>&1 | FileCheck %s -check-prefix=TEST-3x4x5x8
+// RUN: mlir-opt %s -affine-super-vectorizer-test -vectorize-affine-loop-nest 2>&1 | FileCheck %s -check-prefix=VECNEST
 
 func @vector_add_2d(%arg0: index, %arg1: index) -> f32 {
   // Nothing should be matched in this first block.
@@ -35,3 +36,27 @@ func @vector_add_2d(%arg0: index, %arg1: index) -> f32 {
   %9 = load %2[%c7, %c42] : memref<?x?xf32>
   return %9 : f32
 }
+
+// VECNEST-LABEL: func @double_loop_nest
+func @double_loop_nest(%a: memref<20x30xf32>, %b: memref<20xf32>) {
+
+  affine.for %i = 0 to 20 {
+    %b_ld = affine.load %b[%i] : memref<20xf32>
+    affine.for %j = 0 to 30 {
+      %a_ld = affine.load %a[%i, %j] : memref<20x30xf32>
+      affine.store %a_ld, %a[%i, %j] : memref<20x30xf32>
+    }
+    affine.store %b_ld, %b[%i] : memref<20xf32>
+  }
+
+  return
+}
+
+// VECNEST:       affine.for %{{.*}} = 0 to 20 step 4 {
+// VECNEST:         vector.transfer_read
+// VECNEST-NEXT:    affine.for %{{.*}} = 0 to 30 {
+// VECNEST:           vector.transfer_read
+// VECNEST-NEXT:      vector.transfer_write
+// VECNEST-NEXT:    }
+// VECNEST-NEXT:    vector.transfer_write
+// VECNEST:       }

mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

@@ -14,12 +14,14 @@
 #include "mlir/Analysis/NestedMatcher.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/LoopUtils.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/STLExtras.h"
@@ -67,6 +69,12 @@ static llvm::cl::opt<bool> clTestNormalizeMaps(
         "where each AffineAffineApplyOp in the composition is a single output "
         "operation."),
     llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestVecAffineLoopNest(
+    "vectorize-affine-loop-nest",
+    llvm::cl::desc(
+        "Enable testing for the 'vectorizeAffineLoopNest' utility by "
+        "vectorizing the outermost loops found"),
+    llvm::cl::cat(clOptionsCategory));
 
 namespace {
 struct VectorizerTestPass
@@ -84,6 +92,9 @@ struct VectorizerTestPass
   void testSlicing(llvm::raw_ostream &outs);
   void testComposeMaps(llvm::raw_ostream &outs);
   void testNormalizeMaps();
+
+  /// Test for 'vectorizeAffineLoopNest' utility.
+  void testVecAffineLoopNest();
 };
 
 } // end anonymous namespace
@@ -246,10 +257,26 @@ void VectorizerTestPass::testNormalizeMaps() {
   }
 }
 
-void VectorizerTestPass::runOnFunction() {
-  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
-  NestedPatternContext mlContext;
+/// Test for 'vectorizeAffineLoopNest' utility.
+void VectorizerTestPass::testVecAffineLoopNest() {
+  std::vector<SmallVector<AffineForOp, 2>> loops;
+  gatherLoops(getFunction(), loops);
+
+  // Expected only one loop nest.
+  if (loops.empty() || loops[0].size() != 1)
+    return;
 
+  // We vectorize the outermost loop found with VF=4.
+  AffineForOp outermostLoop = loops[0][0];
+  VectorizationStrategy strategy;
+  strategy.vectorSizes.push_back(4 /*vectorization factor*/);
+  strategy.loopToVectorDim[outermostLoop] = 0;
+  std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize;
+  loopsToVectorize.push_back({outermostLoop});
+  vectorizeAffineLoopNest(loopsToVectorize, strategy);
+}
+
+void VectorizerTestPass::runOnFunction() {
   // Only support single block functions at this point.
   FuncOp f = getFunction();
   if (!llvm::hasSingleElement(f))
@@ -258,23 +285,30 @@ void VectorizerTestPass::runOnFunction() {
   std::string str;
   llvm::raw_string_ostream outs(str);
 
-  if (!clTestVectorShapeRatio.empty())
-    testVectorShapeRatio(outs);
+  { // Tests that expect a NestedPatternContext to be allocated externally.
+    NestedPatternContext mlContext;
 
-  if (clTestForwardSlicingAnalysis)
-    testForwardSlicing(outs);
+    if (!clTestVectorShapeRatio.empty())
+      testVectorShapeRatio(outs);
 
-  if (clTestBackwardSlicingAnalysis)
-    testBackwardSlicing(outs);
+    if (clTestForwardSlicingAnalysis)
+      testForwardSlicing(outs);
 
-  if (clTestSlicingAnalysis)
-    testSlicing(outs);
+    if (clTestBackwardSlicingAnalysis)
+      testBackwardSlicing(outs);
 
-  if (clTestComposeMaps)
-    testComposeMaps(outs);
+    if (clTestSlicingAnalysis)
+      testSlicing(outs);
+
+    if (clTestComposeMaps)
+      testComposeMaps(outs);
+
+    if (clTestNormalizeMaps)
+      testNormalizeMaps();
+  }
 
-  if (clTestNormalizeMaps)
-    testNormalizeMaps();
+  if (clTestVecAffineLoopNest)
+    testVecAffineLoopNest();
 
   if (!outs.str().empty()) {
     emitRemark(UnknownLoc::get(&getContext()), outs.str());