diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h --- a/mlir/include/mlir/Dialect/Affine/Utils.h +++ b/mlir/include/mlir/Dialect/Affine/Utils.h @@ -122,7 +122,7 @@ /// loops = {{%i3}}, to vectorize only the second innermost loop; /// loops = {{%i1}}, to vectorize only the middle loop. LogicalResult -vectorizeAffineLoopNest(const std::vector> &loops, +vectorizeAffineLoopNest(std::vector> &loops, const VectorizationStrategy &strategy); /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1. diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp --- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp @@ -1338,7 +1338,7 @@ static void verifyLoopNesting(const std::vector> &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) { diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir --- a/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir +++ b/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 @@ %9 = load %2[%c7, %c42] : memref 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: } diff --git a/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp b/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp --- a/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp +++ b/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 @@ "where each AffineAffineApplyOp in the composition is a single output " "operation."), llvm::cl::cat(clOptionsCategory)); +static llvm::cl::opt 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 @@ 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::runOnFunction() { - // Thread-safe RAII local context, BumpPtrAllocator freed on exit. - NestedPatternContext mlContext; +/// Test for 'vectorizeAffineLoopNest' utility. +void VectorizerTestPass::testVecAffineLoopNest() { + std::vector> 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> 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 @@ 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());