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

[dcaballe]

Summary

Adding missing code that should have been part of "D85869: Utility to
vectorize loop nest using strategy." Sorry, it made it to my stash and
D85869 commit passed build and testing without it.

[nicolasvasilache]

Note it does not pass CI, did not look why

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

303

weird chars

[dcaballe]

Thanks, Nicolas. The failure doesn't seem related to this:

Failed Tests (1):
  LLVM :: Bindings/Go/go.test

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

303

I think this is the way Phabricator represents that we are just changing the indentation here.

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

 /// 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,
+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.
 /// 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

 /// 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");
+  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) {
     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]) {

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

 // 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.
   // CHECK-NOT:matched: {{.*}} = alloc{{.*}}
   // CHECK-NOT:matched: {{.*}} = constant 0{{.*}}
   // CHECK-NOT:matched: {{.*}} = constant 1{{.*}}
   %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
   %1 = alloc(%arg0, %arg1) : memref<?x?xf32>
func @vector_add_2d(%arg0: index, %arg1: index) -> f32 {
   // CHECK-NOT:matched: {{.*}} = constant 42{{.*}}
   // CHECK-NOT:matched: {{.*}} = load{{.*}}
   // CHECK-NOT:matched: return {{.*}}
   %c7 = constant 7 : index
   %c42 = constant 42 : index
   %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

 //===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl --------------===//
 //
 // 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 file implements a simple testing pass for vectorization functionality.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #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"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "affine-super-vectorizer-test"
 
static llvm::cl::opt<bool> clTestComposeMaps(
     llvm::cl::cat(clOptionsCategory));
 static llvm::cl::opt<bool> clTestNormalizeMaps(
     "normalize-maps",
     llvm::cl::desc(
         "Enable testing the normalization of AffineAffineApplyOp "
         "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
     : public PassWrapper<VectorizerTestPass, FunctionPass> {
   static constexpr auto kTestAffineMapOpName = "test_affine_map";
   static constexpr auto kTestAffineMapAttrName = "affine_map";
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
   }
 
   void runOnFunction() override;
   void testVectorShapeRatio(llvm::raw_ostream &outs);
   void testForwardSlicing(llvm::raw_ostream &outs);
   void testBackwardSlicing(llvm::raw_ostream &outs);
   void testSlicing(llvm::raw_ostream &outs);
   void testComposeMaps(llvm::raw_ostream &outs);
   void testNormalizeMaps();
+
+  /// Test for 'vectorizeAffineLoopNest' utility.
+  void testVecAffineLoopNest();
 };
 
 } // end anonymous namespace
 
 void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {
   auto f = getFunction();
   using matcher::Op;
   SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(),
pattern.match(f, &toErase);
     }
   }
   // We should now be able to erase everything in reverse order in this test.
   for (auto m : llvm::reverse(toErase)) {
     m.getMatchedOperation()->erase();
   }
 }
 
-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))
     return;
 
   std::string str;
   llvm::raw_string_ostream outs(str);
 
+  { // Tests that expect a NestedPatternContext to be allocated externally.
+    NestedPatternContext mlContext;
+
     if (!clTestVectorShapeRatio.empty())
       testVectorShapeRatio(outs);
 
     if (clTestForwardSlicingAnalysis)
       testForwardSlicing(outs);
 
     if (clTestBackwardSlicingAnalysis)
       testBackwardSlicing(outs);
 
     if (clTestSlicingAnalysis)
       testSlicing(outs);
 
     if (clTestComposeMaps)
       testComposeMaps(outs);
 
     if (clTestNormalizeMaps)
       testNormalizeMaps();
+  }
+
+  if (clTestVecAffineLoopNest)
+    testVecAffineLoopNest();
 
   if (!outs.str().empty()) {
     emitRemark(UnknownLoc::get(&getContext()), outs.str());
   }
 }
 
 namespace mlir {
 void registerVectorizerTestPass() {
   PassRegistration<VectorizerTestPass> pass(
       "affine-super-vectorizer-test",
       "Tests vectorizer standalone functionality.");
 }
 } // namespace mlir