Migrate LLVMCPU vectorization to its own pass. (#12837)
The below related code are deprecated because they are no longer used by
anyone.
- Deprecate SingleTilingExpert
- Deprecate StrategyVectorize pass
The VectorizationPattern and LinalgVectorizationPattern are not deleted
because they are used by LLVMGPU.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
index 57b2f91..0a8cc1f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel
@@ -34,6 +34,7 @@
"LLVMCPUTile.cpp",
"LLVMCPUTileAndFuse.cpp",
"LLVMCPUUnfuseFMAOps.cpp",
+ "LLVMCPUVectorization.cpp",
"Passes.cpp",
"TargetMLTransformInfo.cpp",
"Utils.cpp",
@@ -74,6 +75,7 @@
"@llvm-project//llvm:TargetParser",
"@llvm-project//mlir:AffineDialect",
"@llvm-project//mlir:AffineToStandard",
+ "@llvm-project//mlir:AffineUtils",
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:ArithToLLVM",
@@ -108,6 +110,7 @@
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:SCFToControlFlow",
"@llvm-project//mlir:SCFTransforms",
+ "@llvm-project//mlir:SCFUtils",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TensorTransforms",
"@llvm-project//mlir:TosaDialect",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
index acac23d..da0f65b 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt
@@ -38,6 +38,7 @@
"LLVMCPUTile.cpp"
"LLVMCPUTileAndFuse.cpp"
"LLVMCPUUnfuseFMAOps.cpp"
+ "LLVMCPUVectorization.cpp"
"Passes.cpp"
"TargetMLTransformInfo.cpp"
"Utils.cpp"
@@ -53,6 +54,7 @@
LLVMTargetParser
MLIRAffineDialect
MLIRAffineToStandard
+ MLIRAffineUtils
MLIRAnalysis
MLIRArithDialect
MLIRArithToLLVM
@@ -86,6 +88,7 @@
MLIRSCFDialect
MLIRSCFToControlFlow
MLIRSCFTransforms
+ MLIRSCFUtils
MLIRTensorDialect
MLIRTensorTransforms
MLIRTosaDialect
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorization.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorization.cpp
new file mode 100644
index 0000000..47440d0
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorization.cpp
@@ -0,0 +1,230 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed 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
+
+#include "iree/compiler/Codegen/PassDetail.h"
+#include "iree/compiler/Codegen/Passes.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#define DEBUG_TYPE "iree-llvmcpu-vectorization"
+
+namespace mlir {
+namespace iree_compiler {
+namespace {
+/// Returns the op that contains lowering config. Checks whether the provided op
+/// contains the lowering config and returns it. Otherwise, tries to find the
+/// lowering config across the function. If there are multiple ops with the same
+/// lowering configs, returns the first one found. Returns failure if there are
+/// multiple op with different lowering config.
+/// TODO(hanchung): This is copied from LinalgTensorCodegenDriver.cpp. We should
+/// refactor it to Utils.h.
+static FailureOr<Operation *> getRootOp(Operation *op) {
+ // Check for self first.
+ if (iree_compiler::getLoweringConfig(op)) {
+ return op;
+ }
+
+ // Get the function op.
+ auto funcOp = dyn_cast<func::FuncOp>(op);
+ if (!funcOp) {
+ funcOp = op->getParentOfType<func::FuncOp>();
+ }
+
+ assert(funcOp && "Missing funcOp");
+
+ Operation *rootOp = nullptr;
+ mlir::iree_compiler::IREE::Codegen::LoweringConfigAttr rootLoweringConfig;
+ auto result = funcOp.walk([&](Operation *op) -> WalkResult {
+ auto loweringConfig = iree_compiler::getLoweringConfig(op);
+ if (!loweringConfig) {
+ return WalkResult::advance();
+ }
+ if (rootLoweringConfig) {
+ if (rootLoweringConfig != loweringConfig) {
+ return WalkResult::interrupt();
+ }
+ } else {
+ rootOp = op;
+ rootLoweringConfig = loweringConfig;
+ }
+ return WalkResult::advance();
+ });
+
+ if (!rootOp || result.wasInterrupted()) {
+ return failure();
+ }
+ return rootOp;
+}
+
+/// Computes the canonical shape used to vectorize this dispatch. Retrieves
+/// the vectorization tile sizes (parallel and reduction levels) out of the
+/// lowering config and adjusts them to the format expected by the Linalg
+/// vectorizer.
+static SmallVector<int64_t> getCanonicalVectorShape(func::FuncOp funcOp) {
+ FailureOr<Operation *> rootOp = getRootOp(funcOp);
+ if (failed(rootOp)) {
+ return {};
+ }
+
+ unsigned numTileLevels =
+ mlir::iree_compiler::getNumTileLevels(rootOp.value());
+ if (numTileLevels < 3) {
+ return {};
+ }
+
+ // Retrieve the tile sizes from the last two tiling levels (parallel and
+ // reduction) used for vectorization.
+ SmallVector<int64_t> canonicalVectorShape =
+ mlir::iree_compiler::getTileSizes(rootOp.value(), numTileLevels - 2);
+ SmallVector<int64_t> reductionTileSizes =
+ mlir::iree_compiler::getTileSizes(rootOp.value(), numTileLevels - 1);
+
+ if (!reductionTileSizes.empty()) {
+ assert(canonicalVectorShape.size() == reductionTileSizes.size() &&
+ "Unexpected tile sizes");
+
+ // Combine the reduction tile sizes with the parallel tile sizes already in
+ // the canonical vector shape.
+ for (int i = 0, end = canonicalVectorShape.size(); i < end; ++i) {
+ if (reductionTileSizes[i] > 0)
+ canonicalVectorShape[i] = reductionTileSizes[i];
+ }
+ }
+
+ // Replace zeros in canonical vector shape to turn it into a valid shape.
+ std::replace(canonicalVectorShape.begin(), canonicalVectorShape.end(), 0, 1);
+ return canonicalVectorShape;
+}
+
+// Give the canonical vector shape of a dispatch, returns the vector sizes for a
+// particular linalg op within that dispatch.
+static SmallVector<int64_t> getVectorSizes(
+ linalg::LinalgOp linalgOp, ArrayRef<int64_t> canonicalVectorShape) {
+ FailureOr<Operation *> rootOp = getRootOp(linalgOp);
+ if (failed(rootOp)) {
+ return {};
+ }
+
+ // TODO: Infer the tiles sizes for an op that is not the root op.
+ if (*rootOp != linalgOp.getOperation()) {
+ return {};
+ }
+
+ if (canonicalVectorShape.empty()) {
+ return {};
+ }
+
+ assert(canonicalVectorShape.size() >= linalgOp.getNumLoops() &&
+ "Unexpected canonical vector shape or number of loops");
+
+ // Return the valid canonical vector shape subset based on the number of loops
+ // of the linalg op.
+ SmallVector<int64_t> vecSize(
+ canonicalVectorShape.take_front(linalgOp.getNumLoops()));
+ for (auto [idx, val] : llvm::enumerate(linalgOp.getStaticLoopRanges())) {
+ if (ShapedType::isDynamic(val)) continue;
+ vecSize[idx] = std::max(vecSize[idx], val);
+ }
+
+ return vecSize;
+}
+
+class LLVMCPUVectorizationPass
+ : public LLVMCPUVectorizationBase<LLVMCPUVectorizationPass> {
+ public:
+ using LLVMCPUVectorizationBase::LLVMCPUVectorizationBase;
+ LLVMCPUVectorizationPass(const LLVMCPUVectorizationPassOptions &options) {
+ this->enableVectorMasking.setValue(options.enableVectorMasking);
+ this->vectorizePadding.setValue(options.vectorizePadding);
+ this->vectorizeGatherAccesses.setValue(options.vectorizeGatherAccesses);
+ }
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<tensor::TensorDialect, linalg::LinalgDialect,
+ vector::VectorDialect>();
+ }
+ void runOnOperation() override;
+};
+
+void LLVMCPUVectorizationPass::runOnOperation() {
+ MLIRContext *context = &getContext();
+ auto funcOp = getOperation();
+ SmallVector<int64_t> canonicalVectorShape;
+ if (enableVectorMasking) {
+ canonicalVectorShape = getCanonicalVectorShape(funcOp);
+ }
+
+ IRRewriter rewriter(context);
+ SmallVector<linalg::LinalgOp> candidates;
+ funcOp.walk(
+ [&](linalg::LinalgOp linalgOp) { candidates.push_back(linalgOp); });
+ for (auto linalgOp : candidates) {
+ SmallVector<int64_t> vectorSizes;
+ if (enableVectorMasking) {
+ vectorSizes.append(getVectorSizes(linalgOp, canonicalVectorShape));
+ }
+ (void)linalg::vectorize(rewriter, linalgOp, vectorSizes,
+ vectorizeGatherAccesses);
+ };
+
+ // TODO: Move this down the pipeline once we have the ODM-based masking
+ // representation.
+ RewritePatternSet vectorizationPatterns(funcOp.getContext());
+ vector::populateVectorMaskLoweringPatternsForSideEffectingOps(
+ vectorizationPatterns);
+ vector::populateVectorTransferPermutationMapLoweringPatterns(
+ vectorizationPatterns);
+ vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
+ vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
+ linalg::LinalgCopyVTWForwardingPattern>(
+ funcOp.getContext(), /*benefit=*/2);
+ vector::TransferReadOp::getCanonicalizationPatterns(vectorizationPatterns,
+ funcOp.getContext());
+ vector::TransferWriteOp::getCanonicalizationPatterns(vectorizationPatterns,
+ funcOp.getContext());
+ (void)applyPatternsAndFoldGreedily(funcOp, std::move(vectorizationPatterns));
+
+ // Apply the pad tensor op vectorization separately to avoid running the
+ // GenericPadOpVectorizationPattern too early.
+ // TODO: Improve once we have better infrastructure to control pattern
+ // application.
+ if (vectorizePadding) {
+ RewritePatternSet patterns(funcOp.getContext());
+ linalg::populatePadOpVectorizationPatterns(patterns);
+ (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+ }
+
+ // Gathers all innermost loops through a post order pruned walk.
+ funcOp.walk([](Operation *op) {
+ if (auto forOp = dyn_cast<AffineForOp>(op))
+ (void)promoteIfSingleIteration(forOp);
+ else if (auto forOp = dyn_cast<scf::ForOp>(op))
+ (void)promoteIfSingleIteration(forOp);
+ });
+ linalg::hoistRedundantVectorTransfers(funcOp);
+ linalg::hoistRedundantVectorTransfersOnTensor(funcOp);
+}
+} // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createLLVMCPUVectorizationPass() {
+ return std::make_unique<LLVMCPUVectorizationPass>();
+}
+std::unique_ptr<OperationPass<func::FuncOp>> createLLVMCPUVectorizationPass(
+ const LLVMCPUVectorizationPassOptions &options) {
+ return std::make_unique<LLVMCPUVectorizationPass>(options);
+}
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
index f7b0e01..021b61f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/Passes.cpp
@@ -339,12 +339,11 @@
static_cast<int64_t>(StrategyTilingLevel::ParallelTiles)));
nestedModulePM.addNestedPass<func::FuncOp>(createLLVMCPUPeelPass());
{
- LinalgSingleTilingExpertPassOptions options;
- options.vectorize = true;
+ LLVMCPUVectorizationPassOptions options;
options.enableVectorMasking = enableVectorMasking;
options.vectorizeGatherAccesses = true;
nestedModulePM.addNestedPass<func::FuncOp>(
- createLinalgSingleTilingExpertPass(options));
+ createLLVMCPUVectorizationPass(options));
nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
}
@@ -378,13 +377,12 @@
createLLVMCPUTensorPadPass(LLVMCPUTensorPadOption::ReductionDims));
{
- LinalgSingleTilingExpertPassOptions options;
- options.vectorize = true;
+ LLVMCPUVectorizationPassOptions options;
options.enableVectorMasking = enableVectorMasking;
options.vectorizePadding = true;
options.vectorizeGatherAccesses = true;
nestedModulePM.addNestedPass<func::FuncOp>(
- createLinalgSingleTilingExpertPass(options));
+ createLLVMCPUVectorizationPass(options));
nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
}
@@ -478,17 +476,13 @@
createDecomposePackUnPackOpsPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
- LinalgSingleTilingExpertPassOptions options;
- options.vectorize = true;
+ LLVMCPUVectorizationPassOptions options;
options.enableVectorMasking = enableVectorMasking;
options.vectorizeGatherAccesses = true;
nestedModulePM.addNestedPass<func::FuncOp>(
- createLinalgSingleTilingExpertPass(options));
+ createLLVMCPUVectorizationPass(options));
nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
- // TODO(hanchung): Merge two vectorization passes into a pass. All the ops
- // should be vectorized altogether. Otherwise, there would be tensor.empty
- // ops which becomes a stack allocation in bufferization.
}
addBufferizePasses(nestedModulePM);
@@ -538,18 +532,13 @@
nestedModulePM.addNestedPass<func::FuncOp>(createVectorizePadPass());
}
- // Add the sandbox single tiling expert to vectorize.
- // We can't do the vectorization in the tiling expert above due to an issue in
- // codegen strategy pipeline. Since we are moving to the transform dialect, we
- // choose to have a workaround here by splitting them into two stages.
{
- LinalgSingleTilingExpertPassOptions options;
- options.vectorize = true;
+ LLVMCPUVectorizationPassOptions options;
options.enableVectorMasking = enableVectorMasking;
options.vectorizePadding = true;
options.vectorizeGatherAccesses = true;
nestedModulePM.addNestedPass<func::FuncOp>(
- createLinalgSingleTilingExpertPass(options));
+ createLLVMCPUVectorizationPass(options));
nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
}
@@ -582,12 +571,9 @@
nestedModulePM.addNestedPass<func::FuncOp>(createLLVMCPUTilePass(
static_cast<int64_t>(StrategyTilingLevel::ReductionTiles)));
- {
- LinalgSingleTilingExpertPassOptions options;
- options.vectorize = true;
- nestedModulePM.addNestedPass<func::FuncOp>(
- createLinalgSingleTilingExpertPass(options));
- }
+ nestedModulePM.addNestedPass<func::FuncOp>(createLLVMCPUVectorizationPass());
+ nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+ nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
addBufferizePasses(nestedModulePM);
diff --git a/compiler/src/iree/compiler/Codegen/Passes.h b/compiler/src/iree/compiler/Codegen/Passes.h
index 5011d07..489d956 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Passes.h
@@ -333,6 +333,15 @@
std::unique_ptr<OperationPass<func::FuncOp>> createLLVMCPUSplitReductionPass(
bool enableReassociateFpReductions = false);
+struct LLVMCPUVectorizationPassOptions {
+ bool enableVectorMasking = false;
+ bool vectorizePadding = false;
+ bool vectorizeGatherAccesses = false;
+};
+std::unique_ptr<OperationPass<func::FuncOp>> createLLVMCPUVectorizationPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createLLVMCPUVectorizationPass(
+ const LLVMCPUVectorizationPassOptions &options);
+
/// Performs the final conversion to LLVM dialect.
std::unique_ptr<OperationPass<ModuleOp>> createConvertToLLVMPass(
bool reassociateFpReordering = false);
diff --git a/compiler/src/iree/compiler/Codegen/Passes.td b/compiler/src/iree/compiler/Codegen/Passes.td
index 573df9e..84e553a 100644
--- a/compiler/src/iree/compiler/Codegen/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Passes.td
@@ -526,6 +526,22 @@
];
}
+def LLVMCPUVectorization :
+ Pass<"iree-llvmcpu-vectorization", "func::FuncOp"> {
+ let summary = "Pass to perform vectorization on tensor/linalg ops.";
+ let constructor = "mlir::createLinalgVectorizationExpertPass()";
+ let options = [
+ Option<"enableVectorMasking", "enable-vector-masking", "bool",/*default=*/"false",
+ "Enable vector masking during vectorization.">,
+ Option<"vectorizePadding", "vectorize-padding", "bool", /*default=*/"false",
+ "Rewrite all tensor.pad ops in the function to vector form.">,
+ Option<"vectorizeGatherAccesses", "vectorize-gather-accesses", "bool", /*default=*/"false",
+ "Enable vectorizaiton of operations that may generate vector.gather operations.">,
+ ];
+ let constructor =
+ "mlir::iree_compiler::createLLVMCPUVectorizationPass()";
+}
+
def ConvertToLLVM :
Pass<"iree-convert-to-llvm", "ModuleOp"> {
let summary =
diff --git a/compiler/src/iree/compiler/Codegen/Sandbox/LinalgTensorCodegenDriver.cpp b/compiler/src/iree/compiler/Codegen/Sandbox/LinalgTensorCodegenDriver.cpp
index a86526b..10dbcb9 100644
--- a/compiler/src/iree/compiler/Codegen/Sandbox/LinalgTensorCodegenDriver.cpp
+++ b/compiler/src/iree/compiler/Codegen/Sandbox/LinalgTensorCodegenDriver.cpp
@@ -31,194 +31,15 @@
using mlir::iree_compiler::IREE::LinalgExt::CodegenStrategy;
using mlir::iree_compiler::IREE::LinalgExt::LinalgTransformationFilter;
using mlir::iree_compiler::IREE::LinalgExt::LinalgTransforms;
-using mlir::iree_compiler::IREE::LinalgExt::LinalgVectorizationOptions;
using mlir::iree_compiler::IREE::LinalgExt::LinalgVectorLoweringOptions;
#define DEBUG_TYPE "iree-linalg-tensor-codegen-driver"
//===----------------------------------------------------------------------===//
-// IREE specific functions
-//===----------------------------------------------------------------------===//
-
-/// Returns the op that contains lowering config. Checks whether the provided op
-/// contains the lowering config and returns it. Otherwise, tries to find the
-/// lowering config across the function. If there are multiple ops with the same
-/// lowering configs, returns the first one found. Returns failure if there are
-/// multiple op with different lowering config.
-static FailureOr<Operation *> getRootOp(Operation *op) {
- // Check for self first.
- if (iree_compiler::getLoweringConfig(op)) {
- return op;
- }
-
- // Get the function op.
- auto funcOp = dyn_cast<func::FuncOp>(op);
- if (!funcOp) {
- funcOp = op->getParentOfType<func::FuncOp>();
- }
-
- assert(funcOp && "Missing funcOp");
-
- Operation *rootOp = nullptr;
- mlir::iree_compiler::IREE::Codegen::LoweringConfigAttr rootLoweringConfig;
- auto result = funcOp.walk([&](Operation *op) -> WalkResult {
- auto loweringConfig = iree_compiler::getLoweringConfig(op);
- if (!loweringConfig) {
- return WalkResult::advance();
- }
- if (rootLoweringConfig) {
- if (rootLoweringConfig != loweringConfig) {
- return WalkResult::interrupt();
- }
- } else {
- rootOp = op;
- rootLoweringConfig = loweringConfig;
- }
- return WalkResult::advance();
- });
-
- if (!rootOp || result.wasInterrupted()) {
- return failure();
- }
- return rootOp;
-}
-
-/// Computes the canonical shape used to vectorize this dispatch. Retrieves
-/// the vectorization tile sizes (parallel and reduction levels) out of the
-/// lowering config and adjusts them to the format expected by the Linalg
-/// vectorizer.
-static SmallVector<int64_t> getCanonicalVectorShape(func::FuncOp funcOp) {
- FailureOr<Operation *> rootOp = getRootOp(funcOp);
- if (failed(rootOp)) {
- return {};
- }
-
- unsigned numTileLevels =
- mlir::iree_compiler::getNumTileLevels(rootOp.value());
- if (numTileLevels < 3) {
- return {};
- }
-
- // Retrieve the tile sizes from the last two tiling levels (parallel and
- // reduction) used for vectorization.
- SmallVector<int64_t> canonicalVectorShape =
- mlir::iree_compiler::getTileSizes(rootOp.value(), numTileLevels - 2);
- SmallVector<int64_t> reductionTileSizes =
- mlir::iree_compiler::getTileSizes(rootOp.value(), numTileLevels - 1);
-
- if (!reductionTileSizes.empty()) {
- assert(canonicalVectorShape.size() == reductionTileSizes.size() &&
- "Unexpected tile sizes");
-
- // Combine the reduction tile sizes with the parallel tile sizes already in
- // the canonical vector shape.
- for (int i = 0, end = canonicalVectorShape.size(); i < end; ++i) {
- if (reductionTileSizes[i] > 0)
- canonicalVectorShape[i] = reductionTileSizes[i];
- }
- }
-
- // Replace zeros in canonical vector shape to turn it into a valid shape.
- std::replace(canonicalVectorShape.begin(), canonicalVectorShape.end(), 0, 1);
- return canonicalVectorShape;
-}
-
-// Give the canonical vector shape of a dispatch, returns the vector sizes for a
-// particular linalg op within that dispatch.
-static SmallVector<int64_t> getVectorSizes(
- linalg::LinalgOp linalgOp, ArrayRef<int64_t> canonicalVectorShape) {
- FailureOr<Operation *> rootOp = getRootOp(linalgOp);
- if (failed(rootOp)) {
- return {};
- }
-
- // TODO: Infer the tiles sizes for an op that is not the root op.
- if (*rootOp != linalgOp.getOperation()) {
- return {};
- }
-
- if (canonicalVectorShape.empty()) {
- return {};
- }
-
- assert(canonicalVectorShape.size() >= linalgOp.getNumLoops() &&
- "Unexpected canonical vector shape or number of loops");
-
- // Return the valid canonical vector shape subset based on the number of loops
- // of the linalg op.
- SmallVector<int64_t> vecSize(
- canonicalVectorShape.take_front(linalgOp.getNumLoops()));
- for (auto [idx, val] : llvm::enumerate(linalgOp.getStaticLoopRanges())) {
- if (ShapedType::isDynamic(val)) continue;
- vecSize[idx] = std::max(vecSize[idx], val);
- }
-
- return vecSize;
-}
-
-/// Default method to get tile sizes for tile-and-fuse in IREE. These could be
-/// ovveridden by the command line options if specified.
-static LogicalResult getTileAndFuseOptionsFromConfig(
- func::FuncOp funcOp, int64_t tilingLevel,
- SmallVector<int64_t> &tileAndFuseSizes, SmallVector<int64_t> &tileOnlySizes,
- SmallVector<int64_t> &tileInterchange) {
- if (tilingLevel == -1) {
- return success();
- }
-
- FailureOr<Operation *> rootOp = getRootOp(funcOp);
- if (failed(rootOp) || !rootOp.value()) return failure();
- auto tilingOp = cast<TilingInterface>(rootOp.value());
-
- iree_compiler::IREE::Codegen::LoweringConfigAttr loweringConfig =
- iree_compiler::getLoweringConfig(tilingOp);
- SmallVector<int64_t> tileSizes = loweringConfig.getTileSizeVals(tilingLevel);
-
- auto iteratorTypes = tilingOp.getLoopIteratorTypes();
- tileAndFuseSizes = SmallVector<int64_t>(iteratorTypes.size(), /*default=*/0);
- tileOnlySizes = SmallVector<int64_t>(iteratorTypes.size(), /*default=*/0);
- // Splits the tileSizes into two sizes vectors. We want to tile-and-fuse on
- // parallel dims and tile-only on reduction dims.
- for (size_t i = 0; i < tileSizes.size() && i < iteratorTypes.size(); ++i) {
- if (iteratorTypes[i] == utils::IteratorType::parallel) {
- tileAndFuseSizes[i] = tileSizes[i];
- } else {
- tileOnlySizes[i] = tileSizes[i];
- }
- }
-
- tileInterchange = loweringConfig.getTileInterchangeVals(tilingLevel);
-
- return success();
-}
-
-//===----------------------------------------------------------------------===//
// From Sandbox
//===----------------------------------------------------------------------===//
namespace {
-
-struct LinalgSingleTilingExpertPass
- : public LinalgSingleTilingExpertBase<LinalgSingleTilingExpertPass> {
- LinalgSingleTilingExpertPass() = default;
- LinalgSingleTilingExpertPass(
- const LinalgSingleTilingExpertPassOptions &options) {
- this->anchorFuncOpName = options.anchorFuncOpName;
- this->anchorOpName = options.anchorOpName;
- this->generalize = options.generalize;
- this->iteratorInterchange = options.iteratorInterchange;
- this->vectorize = options.vectorize;
- this->enableVectorMasking = options.enableVectorMasking;
- this->vectorizePadding = options.vectorizePadding;
- this->vectorizeGatherAccesses = options.vectorizeGatherAccesses;
- this->tilingLevel = options.tilingLevel;
- }
- LinalgSingleTilingExpertPass(const LinalgSingleTilingExpertPass &pass) {}
-
- /// Function pass entry point.
- void runOnOperation() override;
-};
-
struct LinalgVectorLoweringPass
: public LinalgVectorLoweringBase<LinalgVectorLoweringPass> {
LinalgVectorLoweringPass(int64_t vectorLoweringStage = 0) {
@@ -240,34 +61,6 @@
};
} // namespace
-void LinalgSingleTilingExpertPass::runOnOperation() {
- func::FuncOp funcOp = getOperation();
-
- LinalgVectorizationOptions vectorizationOptions;
- vectorizationOptions.setVectorizePadding(vectorizePadding);
- vectorizationOptions.setVectorizeGatherAccesses(vectorizeGatherAccesses);
- vectorizationOptions.setEnableVectorMasking(enableVectorMasking);
- if (enableVectorMasking) {
- vectorizationOptions.setCanonicalVectorSizes(
- getCanonicalVectorShape(funcOp));
- vectorizationOptions.setVectorSizeComputationFunction(getVectorSizes);
- }
-
- CodegenStrategy strategy;
- StringRef genericOpName = linalg::GenericOp::getOperationName();
- strategy.vectorizeIf(vectorize, generalize ? genericOpName : anchorOpName,
- vectorizationOptions);
-
- // Created a nested OpPassManager and run.
- OpPassManager dynamicPM(func::FuncOp::getOperationName());
- strategy.configurePassPipeline(dynamicPM, funcOp.getContext());
- dynamicPM.addPass(
- iree_compiler::IREE::LinalgExt::createLinalgStrategyEnablePass());
- if (failed(runPipeline(dynamicPM, funcOp))) {
- return signalPassFailure();
- }
-}
-
void LinalgVectorLoweringPass::runOnOperation() {
LLVM_DEBUG(llvm::dbgs() << "\n ---- Stage : " << vectorLoweringStage;);
vector::VectorTransposeLowering vectorTransposeLowering =
@@ -352,16 +145,6 @@
}
std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::createLinalgSingleTilingExpertPass() {
- return std::make_unique<LinalgSingleTilingExpertPass>();
-}
-std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::createLinalgSingleTilingExpertPass(
- const LinalgSingleTilingExpertPassOptions &passOptions) {
- return std::make_unique<LinalgSingleTilingExpertPass>(passOptions);
-}
-
-std::unique_ptr<OperationPass<func::FuncOp>>
mlir::createLinalgVectorLoweringPass(int64_t vectorLoweringStage) {
return std::make_unique<LinalgVectorLoweringPass>(vectorLoweringStage);
}
diff --git a/compiler/src/iree/compiler/Codegen/Sandbox/Passes.h b/compiler/src/iree/compiler/Codegen/Sandbox/Passes.h
index 6cd9d85..f309365 100644
--- a/compiler/src/iree/compiler/Codegen/Sandbox/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Sandbox/Passes.h
@@ -12,25 +12,6 @@
namespace mlir {
-/// Struct to control pass options for `LinalgSingleTilingExpert` pass.
-struct LinalgSingleTilingExpertPassOptions {
- std::string anchorFuncOpName = "";
- std::string anchorOpName = "";
- bool generalize = false;
- SmallVector<int64_t> iteratorInterchange = {};
- bool vectorize = false;
- bool enableVectorMasking = false;
- bool vectorizePadding = false;
- bool vectorizeGatherAccesses = false;
- int64_t tilingLevel = -1;
-};
-
-/// Creates a pass to drive one level tiling + vectorization of `LinalgOp`s.
-std::unique_ptr<OperationPass<func::FuncOp>>
-createLinalgSingleTilingExpertPass();
-std::unique_ptr<OperationPass<func::FuncOp>> createLinalgSingleTilingExpertPass(
- const LinalgSingleTilingExpertPassOptions &passOptions);
-
/// Struct to control pass options for `LinalgVectorLoweringPass` pass.
struct LinalgVectorLoweringPassOptions {
int vectorLoweringStage = 0;
@@ -65,12 +46,6 @@
//===----------------------------------------------------------------------===//
namespace iree_compiler {
-/// Creates a pass to drive one level tiling + vectorization of `LinalgOp`s with
-/// additional parameters that allow controlling the value of the pass options
-/// when building the pipeline.
-std::unique_ptr<OperationPass<func::FuncOp>> createLinalgSingleTilingExpertPass(
- int64_t tilingLevel, bool vectorize);
-
//===----------------------------------------------------------------------===//
// Registration
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Sandbox/Passes.td b/compiler/src/iree/compiler/Codegen/Sandbox/Passes.td
index 6601c58..54e618f 100644
--- a/compiler/src/iree/compiler/Codegen/Sandbox/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Sandbox/Passes.td
@@ -9,46 +9,6 @@
include "mlir/Pass/PassBase.td"
-def LinalgSingleTilingExpert
- : Pass<"linalg-single-tiling-expert-driver", "func::FuncOp"> {
- let summary = "Pass to drive transformations on Linalg on tensors.";
- let constructor = "mlir::createLinalgSingleTilingExpertPass()";
- let options = [
- // Func / op targeting options.
- Option<"anchorFuncOpName", "anchor-func", "std::string", /*default=*/"",
- "Which func op is the anchor to latch on.">,
- Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
- "Which linalg op within the func is the anchor to latch on.">,
-
- // Generalization options.
- Option<"generalize", "generalize", "bool", /*default=*/"false",
- "Convert named operations to their generic form.">,
- ListOption<"iteratorInterchange", "iterator-interchange", "int64_t",
- "Interator interchange.",
- "llvm::cl::ZeroOrMore">,
-
- // Vectorization options.
- Option<"vectorize", "vectorize", "bool", /*default=*/"false",
- "Rewrite the linalg op as a vector operation.">,
- Option<"enableVectorMasking", "enableVectorMasking", "bool", /*default=*/"false",
- "Enable vector masking during vectorization.">,
- Option<"vectorizePadding", "vectorize-padding", "bool", /*default=*/"false",
- "Rewrite all tensor.pad ops in the function to vector form.">,
- Option<"vectorizeGatherAccesses", "vectorize-gather-accesses", "bool", /*default=*/"false",
- "Enable vectorizaiton of operations that may generate vector.gather operations.">,
-
- // IREE specific options
- Option<"tilingLevel", "tiling-level", "int64_t", /*default=*/"-1",
- "Use default tiling level used to retrieve the configuration from lowering_config">
- ];
- let dependentDialects = [
- "::mlir::arith::ArithDialect", "::mlir::AffineDialect",
- "::mlir::linalg::LinalgDialect", "::mlir::scf::SCFDialect",
- "::mlir::cf::ControlFlowDialect", "::mlir::tensor::TensorDialect",
- "::mlir::vector::VectorDialect"
- ];
-}
-
def LinalgVectorLowering : Pass<"linalg-vector-lowering", "func::FuncOp"> {
let summary = "Run transformations that lower high-level vectors.";
let constructor = "mlir::createLinalgVectorLoweringPass()";
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
index 1067e5b..2f41d05 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.h
@@ -210,16 +210,8 @@
}
};
-/// Create a LinalgStrategyDecomposePass.
-// TODO: if/when we need finer control add an `opName` parameter.
-std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyDecomposePass(
- const LinalgExt::LinalgTransformationFilter &filter =
- LinalgExt::LinalgTransformationFilter());
-
-/// Create a LinalgStrategyVectorizePass.
using VectorSizeComputationFunction =
std::function<SmallVector<int64_t>(linalg::LinalgOp, ArrayRef<int64_t>)>;
-
struct LinalgVectorizationOptions {
/// Enable vector masking during vectorization.
bool enableVectorMasking = false;
@@ -268,12 +260,6 @@
}
};
-std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyVectorizePass(
- StringRef opName = "",
- const LinalgVectorizationOptions &options = LinalgVectorizationOptions(),
- const LinalgExt::LinalgTransformationFilter &filter =
- LinalgExt::LinalgTransformationFilter());
-
/// Create a LinalgStrategyEnablePass.
std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyEnablePass(
LinalgEnablingOptions opt = LinalgEnablingOptions(),
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td
index 71fd062..6f2fada 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Passes/Passes.td
@@ -100,21 +100,6 @@
// TODO: Deprecate all this.
//===---------------------------------------------------------------------====//
-def LinalgStrategyVectorizePass
- : Pass<"iree-linalg-strategy-vectorize-pass", "func::FuncOp"> {
- let summary = "Configurable pass to apply pattern-based linalg vectorization.";
- let constructor = "createLinalgStrategyVectorizePass()";
- let dependentDialects = ["linalg::LinalgDialect", "vector::VectorDialect"];
- let options = [
- Option<"anchorFuncName", "anchor-func", "std::string", /*default=*/"",
- "Which func op is the anchor to latch on.">,
- Option<"anchorOpName", "anchor-op", "std::string", /*default=*/"",
- "Which linalg op within the func is the anchor to latch on.">,
- Option<"vectorizePadding", "vectorize-padding", "bool", "false",
- "Enable vectorization of padding ops.">,
- ];
-}
-
def LinalgStrategyEnablePass
: Pass<"iree-linalg-strategy-enable-pass", "func::FuncOp"> {
let summary = "Configurable pass to enable the application of other "
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CodegenStrategy.h b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CodegenStrategy.h
index d19b3d4..35a8ac1 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CodegenStrategy.h
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgExt/Transforms/CodegenStrategy.h
@@ -37,28 +37,6 @@
LinalgExt::LinalgTransformationFilter::FilterFunction filter = nullptr;
};
-/// Represent one application of createLinalgStrategyVectorizePass.
-struct Vectorize : public Transformation {
- explicit Vectorize(
- LinalgVectorizationOptions opts = LinalgVectorizationOptions(),
- LinalgExt::LinalgTransformationFilter::FilterFunction f = nullptr)
- : Transformation(std::move(f)), options(std::move(opts)) {}
-
- Vectorize(StringRef name, LinalgVectorizationOptions opts,
- LinalgExt::LinalgTransformationFilter::FilterFunction f = nullptr)
- : Transformation(std::move(f)), opName(name), options(std::move(opts)) {}
-
- void
- addToPassPipeline(OpPassManager &pm,
- LinalgExt::LinalgTransformationFilter m) const override {
- pm.addPass(createLinalgStrategyVectorizePass(opName, options, m));
- }
-
-private:
- std::string opName;
- LinalgVectorizationOptions options;
-};
-
/// Represent one application of createLinalgStrategyLowerVectorsPass.
struct VectorLowering : public Transformation {
explicit VectorLowering(
@@ -78,24 +56,6 @@
/// Codegen strategy controls how a Linalg op is progressively lowered.
struct CodegenStrategy {
- /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
- CodegenStrategy &vectorize(
- StringRef opName,
- const LinalgVectorizationOptions &options = LinalgVectorizationOptions(),
- const LinalgExt::LinalgTransformationFilter::FilterFunction &f =
- nullptr) {
- transformationSequence.emplace_back(
- std::make_unique<Vectorize>(opName, options, f));
- return *this;
- }
- /// Conditionally append a pattern to rewrite `LinalgOpType` as a vector
- /// operation.
- CodegenStrategy &vectorizeIf(
- bool b, StringRef opName,
- LinalgVectorizationOptions options = LinalgVectorizationOptions(),
- LinalgExt::LinalgTransformationFilter::FilterFunction f = nullptr) {
- return b ? vectorize(opName, std::move(options), std::move(f)) : *this;
- }
/// Append a pattern to lower all vector operations.
CodegenStrategy &vectorLowering(LinalgVectorLoweringOptions options) {
transformationSequence.emplace_back(
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Transforms.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Transforms.cpp
index ea7d0d2..c8b0510 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Transforms.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgExt/Transforms/Transforms.cpp
@@ -126,65 +126,6 @@
}
namespace {
-/// Configurable pass to apply pattern-based linalg vectorization.
-struct LinalgStrategyVectorizePass
- : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
-
- LinalgStrategyVectorizePass() = default;
-
- LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opts,
- LinalgExt::LinalgTransformationFilter filt)
- : options(std::move(opts)), filter(std::move(filt)) {
- this->vectorizePadding = opts.vectorizePadding;
- };
-
- void runOnOperation() override {
- auto funcOp = getOperation();
- if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
- return;
-
- RewritePatternSet vectorizationPatterns(funcOp.getContext());
- if (!anchorOpName.empty()) {
- vectorizationPatterns.add<LinalgVectorizationPattern>(
- anchorOpName, funcOp.getContext(), options, filter);
- } else {
- vectorizationPatterns.add<LinalgVectorizationPattern>(funcOp.getContext(),
- options, filter);
- }
-
- // TODO: Move this down the pipeline once we have the ODM-based masking
- // representation.
- vector::populateVectorMaskLoweringPatternsForSideEffectingOps(
- vectorizationPatterns);
-
- vector::populateVectorTransferPermutationMapLoweringPatterns(
- vectorizationPatterns);
- vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
- vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
- linalg::LinalgCopyVTWForwardingPattern>(
- funcOp.getContext(), /*benefit=*/2);
- vector::TransferReadOp::getCanonicalizationPatterns(vectorizationPatterns,
- funcOp.getContext());
- vector::TransferWriteOp::getCanonicalizationPatterns(vectorizationPatterns,
- funcOp.getContext());
- (void)applyPatternsAndFoldGreedily(funcOp,
- std::move(vectorizationPatterns));
-
- // Apply the pad tensor op vectorization separately to avoid running the
- // GenericPadOpVectorizationPattern too early.
- // TODO: Improve once we have better infrastructure to control pattern
- // application.
- if (vectorizePadding) {
- RewritePatternSet patterns(funcOp.getContext());
- linalg::populatePadOpVectorizationPatterns(patterns);
- (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
- }
- }
-
- LinalgVectorizationOptions options;
- LinalgExt::LinalgTransformationFilter filter;
-};
-
/// Configurable pass to enable the application of other pattern-based linalg
/// passes.
struct LinalgStrategyEnablePass
@@ -324,13 +265,6 @@
};
} // namespace
-/// Create a LinalgStrategyVectorizePass.
-std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyVectorizePass(
- StringRef opName, const LinalgVectorizationOptions &options,
- const LinalgExt::LinalgTransformationFilter &filter) {
- return std::make_unique<LinalgStrategyVectorizePass>(opName, options, filter);
-}
-
/// Create a LinalgStrategyEnablePass.
std::unique_ptr<OperationPass<func::FuncOp>> createLinalgStrategyEnablePass(
LinalgEnablingOptions opt,
