#ifndef MLIR_DIALECT_AFFINE_UTILS_H
#define MLIR_DIALECT_AFFINE_UTILS_H
+#include "mlir/Support/LLVM.h"
+
namespace mlir {
class AffineForOp;
class AffineIfOp;
class AffineParallelOp;
struct LogicalResult;
+class Operation;
/// Replaces parallel affine.for op with 1-d affine.parallel op.
/// mlir::isLoopParallel detect the parallel affine.for ops.
/// significant code expansion in some cases.
LogicalResult hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr);
+/// Vectorizes affine loops in 'loops' using the n-D vectorization factors in
+/// 'vectorSizes'. By default, each vectorization factor is applied
+/// inner-to-outer to the loops of each loop nest. 'fastestVaryingPattern' can
+/// be optionally used to provide a different loop vectorization order.
+void vectorizeAffineLoops(
+ Operation *parentOp,
+ llvm::DenseSet<Operation *, DenseMapInfo<Operation *>> &loops,
+ ArrayRef<int64_t> vectorSizes, ArrayRef<int64_t> fastestVaryingPattern);
+
} // namespace mlir
#endif // MLIR_DIALECT_AFFINE_UTILS_H
#include "mlir/Analysis/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Vector/VectorOps.h"
#include "mlir/Dialect/Vector/VectorUtils.h"
return signalPassFailure();
}
- // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
- NestedPatternContext mlContext;
-
DenseSet<Operation *> parallelLoops;
f.walk([¶llelLoops](AffineForOp loop) {
if (isLoopParallel(loop))
parallelLoops.insert(loop);
});
+ vectorizeAffineLoops(f, parallelLoops, vectorSizes, fastestVaryingPattern);
+}
+
+namespace mlir {
+
+/// Vectorizes affine loops in 'loops' using the n-D vectorization factors in
+/// 'vectorSizes'. By default, each vectorization factor is applied
+/// inner-to-outer to the loops of each loop nest. 'fastestVaryingPattern' can
+/// be optionally used to provide a different loop vectorization order.
+void vectorizeAffineLoops(Operation *parentOp, DenseSet<Operation *> &loops,
+ ArrayRef<int64_t> vectorSizes,
+ ArrayRef<int64_t> fastestVaryingPattern) {
+ // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
+ NestedPatternContext mlContext;
+
for (auto &pat :
- makePatterns(parallelLoops, vectorSizes.size(), fastestVaryingPattern)) {
+ makePatterns(loops, vectorSizes.size(), fastestVaryingPattern)) {
LLVM_DEBUG(dbgs() << "\n******************************************");
LLVM_DEBUG(dbgs() << "\n******************************************");
- LLVM_DEBUG(dbgs() << "\n[early-vect] new pattern on Function\n");
- LLVM_DEBUG(f.print(dbgs()));
+ LLVM_DEBUG(dbgs() << "\n[early-vect] new pattern on parent op\n");
+ LLVM_DEBUG(parentOp->print(dbgs()));
+
unsigned patternDepth = pat.getDepth();
SmallVector<NestedMatch, 8> matches;
- pat.match(f, &matches);
+ pat.match(parentOp, &matches);
// Iterate over all the top-level matches and vectorize eagerly.
// This automatically prunes intersecting matches.
for (auto m : matches) {
}
std::unique_ptr<OperationPass<FuncOp>>
-mlir::createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
+createSuperVectorizePass(ArrayRef<int64_t> virtualVectorSize) {
return std::make_unique<Vectorize>(virtualVectorSize);
}
-std::unique_ptr<OperationPass<FuncOp>> mlir::createSuperVectorizePass() {
+std::unique_ptr<OperationPass<FuncOp>> createSuperVectorizePass() {
return std::make_unique<Vectorize>();
}
+
+} // namespace mlir