[Mlir-commits] [mlir] 93936da - [mlir][Affine][VectorOps] Fix super vectorizer utility (D85869)
Diego Caballero
llvmlistbot at llvm.org
Mon Sep 28 16:33:14 PDT 2020
Author: Diego Caballero
Date: 2020-09-28T16:24:11-07:00
New Revision: 93936da90479c4ac444f989ba060eedec4133154
URL: https://github.com/llvm/llvm-project/commit/93936da90479c4ac444f989ba060eedec4133154
DIFF: https://github.com/llvm/llvm-project/commit/93936da90479c4ac444f989ba060eedec4133154.diff
LOG: [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
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Utils.h
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir
mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index c83955eb0891..4f36bb800dd1 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/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.
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 1feeecb5c400..52256deb10bf 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/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) {
diff --git a/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir b/mlir/test/Dialect/Affine/SuperVectorize/vector_utils.mlir
index 4e06028802a6..bf899cdbb975 100644
--- 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 @@ 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: }
diff --git a/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp b/mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp
index cfac2dce2300..25cc2a6e160d 100644
--- 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 @@ 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());
More information about the Mlir-commits
mailing list