| //===- LinalgTensorCodegenStrategyPass.cpp - Test Linalg codegen strategy--===// |
| // |
| // 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 logic for testing the Linalg codegen strategy. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| // TODO: avoid copy-pasta but I can't seem to be able to inherit from a pass. |
| // Will get better once upstreamed to core and it replaces the existing codegen |
| // strategy. |
| |
| #include "Transforms.h" |
| #include "llvm/ADT/SetVector.h" |
| #include "llvm/ADT/StringSwitch.h" |
| #include "mlir/Dialect/Affine/IR/AffineOps.h" |
| #include "mlir/Dialect/GPU/GPUDialect.h" |
| #include "mlir/Dialect/Linalg/IR/LinalgOps.h" |
| #include "mlir/Dialect/Linalg/Transforms/CodegenStrategy.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/StandardOps/IR/Ops.h" |
| #include "mlir/Dialect/Vector/VectorOps.h" |
| #include "mlir/IR/PatternMatch.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| |
| using namespace mlir; |
| using namespace mlir::linalg; |
| |
| namespace { |
| struct LinalgTensorCodegenStrategyPass |
| : public PassWrapper<LinalgTensorCodegenStrategyPass, FunctionPass> { |
| LinalgTensorCodegenStrategyPass() = default; |
| LinalgTensorCodegenStrategyPass(const LinalgTensorCodegenStrategyPass &pass) { |
| } |
| |
| void getDependentDialects(DialectRegistry ®istry) const override { |
| // clang-format off |
| registry.insert<AffineDialect, |
| gpu::GPUDialect, |
| linalg::LinalgDialect, |
| scf::SCFDialect, |
| StandardOpsDialect, |
| vector::VectorDialect>(); |
| // clang-format on |
| } |
| |
| template <typename LinalgNamedOp> |
| void applyStrategyToNamedLinalgOp(); |
| |
| void runOnFunction() override; |
| |
| ListOption<int64_t> tileSizes{*this, "tile-sizes", |
| llvm::cl::MiscFlags::CommaSeparated, |
| llvm::cl::desc("Specifies the tile sizes.")}; |
| Option<bool> promote{ |
| *this, "promote", |
| llvm::cl::desc("Promote the tile into a small aligned memory buffer."), |
| llvm::cl::init(false)}; |
| Option<bool> promoteFullTile{ |
| *this, "promote-full-tile-pad", |
| llvm::cl::desc("Pad the small aligned memory buffer to the tile sizes."), |
| llvm::cl::init(false)}; |
| Option<bool> vectorize{ |
| *this, "vectorize", |
| llvm::cl::desc("Rewrite the linalg op as a vector operation."), |
| llvm::cl::init(false)}; |
| Option<std::string> splitVectorTransfersTo{ |
| *this, "split-transfers", |
| llvm::cl::desc( |
| "Split vector transfers between slow (masked) and fast " |
| "(unmasked) variants. Possible options are:\n" |
| "\tnone: keep unsplit vector.transfer and pay the full price\n" |
| "\tlinalg-copy: use linalg.fill + linalg.copy for the slow path\n" |
| "\tvector-transfers: use extra small unmasked vector.transfer for" |
| " the slow path\n"), |
| llvm::cl::init("none")}; |
| Option<std::string> vectorizeContractionTo{ |
| *this, "vectorize-contraction-to", |
| llvm::cl::desc("the type of vector op to use for linalg contractions"), |
| llvm::cl::init("outerproduct")}; |
| Option<bool> unrollVectorTransfers{ |
| *this, "unroll-vector-transfers", |
| llvm::cl::desc("Enable full unrolling of vector.transfer operations"), |
| llvm::cl::init(false)}; |
| Option<bool> licm{*this, "licm", llvm::cl::desc("Enable LICM."), |
| llvm::cl::init(true)}; |
| Option<bool> hoistRedundantVectorTransfers{ |
| *this, "hoist-redundant-vector-transfers", |
| llvm::cl::desc("Enable HoistRedundantVectorTransfers"), |
| llvm::cl::init(true)}; |
| Option<bool> vectorTransferPartialRewrite{ |
| *this, "vector-transfer-partial-rewrite", |
| llvm::cl::desc("Enable rewriting of vector.transfer operations into " |
| "full/partial read/writes."), |
| llvm::cl::init(true)}; |
| Option<bool> vectorContractLowering{ |
| *this, "vector-contract-lowering", |
| llvm::cl::desc("Enable lowering of vector contractions."), |
| llvm::cl::init(true)}; |
| Option<bool> vectorToSCFConversion{ |
| *this, "vector-to-scf-conversion", |
| llvm::cl::desc("Enable vector to scf conversions."), |
| llvm::cl::init(true)}; |
| Option<std::string> anchorOpName{ |
| *this, "anchor-op", |
| llvm::cl::desc( |
| "Which single linalg op is the anchor for the codegen strategy to " |
| "latch on:\n" |
| "\tlinalg.matmul: anchor on linalg.matmul\n" |
| "\tlinalg.matmul_column_major: anchor on linalg.matmul_column_major\n" |
| "\tlinalg.copy: anchor on linalg.copy\n" |
| "\tlinalg.fill: anchor on linalg.fill\n"), |
| llvm::cl::init("")}; |
| Option<std::string> anchorFuncOpName{ |
| *this, "anchor-func", |
| llvm::cl::desc( |
| "Which single func op is the anchor for the codegen strategy to " |
| "latch on."), |
| llvm::cl::init("")}; |
| |
| Option<bool> distribute{ |
| *this, "distribute", |
| llvm::cl::desc("Distribute the linalg op into a TiledGeneric."), |
| llvm::cl::init(false)}; |
| ListOption<int64_t> distributeTileSizes{ |
| *this, "distribute-tile-sizes", llvm::cl::MiscFlags::CommaSeparated, |
| llvm::cl::desc("Specifies the tile sizes.")}; |
| Option<bool> pad{*this, "pad", llvm::cl::desc("Use padding during tiling."), |
| llvm::cl::init(false)}; |
| Option<int> hoistPadding{ |
| *this, "hoist-padding", |
| llvm::cl::desc("Hoist padding by the number of specified loops."), |
| llvm::cl::init(0)}; |
| Option<bool> vectorizePadding{ |
| *this, "vectorize-padding", |
| llvm::cl::desc("Rewrite linalg.pad_tensor in vector form."), |
| llvm::cl::init(false)}; |
| }; |
| } // end anonymous namespace |
| |
| static void runStrategy(LinalgTensorCodegenStrategyPass &pass, |
| CodegenStrategy &strategy) { |
| strategy.setEnableLICM(pass.licm) |
| .setEnableHoistRedundantVectorTransfers( |
| pass.hoistRedundantVectorTransfers) |
| .setEnableHoistRedundantVectorTransfersOnTensor( |
| pass.hoistRedundantVectorTransfers) |
| .setEnableVectorTransferPartialRewrite(pass.vectorTransferPartialRewrite) |
| .setEnableVectorContractLowering(pass.vectorContractLowering) |
| .setEnableVectorToSCFConversion(pass.vectorToSCFConversion) |
| .transform(pass.getFunction()); |
| } |
| |
| static void runGenericStrategy( |
| LinalgTensorCodegenStrategyPass &pass, LinalgTilingOptions tilingOptions, |
| vector::VectorContractLowering vectorContractLowering, |
| vector::VectorTransferSplit vectorTransferSplit) { |
| assert(!pass.anchorOpName.empty()); |
| CodegenStrategy strategy; |
| strategy |
| .tileIf<LinalgOp>(!pass.tileSizes.empty(), pass.anchorOpName, |
| tilingOptions) |
| .promoteIf<LinalgOp>( |
| pass.promote, pass.anchorOpName, |
| LinalgPromotionOptions() |
| .setAlignment(16) |
| .setUseFullTileBuffersByDefault(pass.promoteFullTile)) |
| .vectorizeIf(pass.vectorize, pass.anchorOpName) |
| .setVectorTransformsOptions( |
| vector::VectorTransformsOptions() |
| .setVectorTransformsOptions(vectorContractLowering) |
| .setVectorTransferSplit(vectorTransferSplit)) |
| .setVectorTransferToSCFOptions( |
| VectorTransferToSCFOptions().setUnroll(pass.unrollVectorTransfers)); |
| runStrategy(pass, strategy); |
| } |
| |
| template <typename OpType> |
| static void runStrategy(LinalgTensorCodegenStrategyPass &pass, |
| LinalgTilingOptions tilingOptions, |
| vector::VectorContractLowering vectorContractLowering, |
| vector::VectorTransferSplit vectorTransferSplit) { |
| CodegenStrategy strategy; |
| strategy.tileIf<OpType>(!pass.tileSizes.empty(), tilingOptions) |
| .template promoteIf<OpType>( |
| pass.promote, |
| LinalgPromotionOptions() |
| .setAlignment(16) |
| .setUseFullTileBuffersByDefault(pass.promoteFullTile)) |
| .template vectorizeIf<OpType>(pass.vectorize) |
| .setVectorTransformsOptions( |
| vector::VectorTransformsOptions() |
| .setVectorTransformsOptions(vectorContractLowering) |
| .setVectorTransferSplit(vectorTransferSplit)) |
| .setVectorTransferToSCFOptions( |
| VectorTransferToSCFOptions().setUnroll(pass.unrollVectorTransfers)); |
| runStrategy(pass, strategy); |
| } |
| |
| // For now, just assume it is the zero of type. |
| // In the future, it should be the zero of type + op. |
| static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) { |
| auto t = getElementTypeOrSelf(op.get().getType()); |
| return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t)); |
| } |
| |
| /// Apply transformations specified as patterns. |
| void LinalgTensorCodegenStrategyPass::runOnFunction() { |
| if (!anchorFuncOpName.empty() && anchorFuncOpName != getFunction().getName()) |
| return; |
| |
| if (distribute && !distributeTileSizes.empty()) { |
| LinalgTilingOptions tilingOptions; |
| tilingOptions = tilingOptions.setTileSizes(distributeTileSizes); |
| if (pad) |
| tilingOptions = tilingOptions.setPaddingValueComputationFunction( |
| getNeutralOfLinalgOp); |
| OwningRewritePatternList patterns; |
| patterns.insert<TileAndDistributePattern>( |
| TileAndDistributeOptions{tilingOptions}, |
| LinalgTransformationFilter( |
| ArrayRef<Identifier>{}, |
| {Identifier::get("distributed", getFunction().getContext())}) |
| .addFilter([](Operation *op) { |
| return success(isaContractionOpInterface(dyn_cast<LinalgOp>(op))); |
| })); |
| (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns)); |
| // Ensure we drop the marker in the end. |
| getFunction().walk([](LinalgOp op) { |
| op.removeAttr(LinalgTransforms::kLinalgTransformMarker); |
| }); |
| } |
| |
| LinalgTilingOptions tilingOptions; |
| if (!tileSizes.empty()) tilingOptions = tilingOptions.setTileSizes(tileSizes); |
| if (pad) |
| tilingOptions = |
| tilingOptions.setPaddingValueComputationFunction(getNeutralOfLinalgOp); |
| |
| vector::VectorContractLowering vectorContractLowering = |
| llvm::StringSwitch<vector::VectorContractLowering>( |
| vectorizeContractionTo.getValue()) |
| .Case("matrixintrinsics", vector::VectorContractLowering::Matmul) |
| .Case("dot", vector::VectorContractLowering::Dot) |
| .Case("outerproduct", vector::VectorContractLowering::OuterProduct) |
| .Default(vector::VectorContractLowering::OuterProduct); |
| vector::VectorTransferSplit vectorTransferSplit = |
| llvm::StringSwitch<vector::VectorTransferSplit>( |
| splitVectorTransfersTo.getValue()) |
| .Case("none", vector::VectorTransferSplit::None) |
| .Case("linalg-copy", vector::VectorTransferSplit::LinalgCopy) |
| .Case("vector-transfers", vector::VectorTransferSplit::VectorTransfer) |
| .Default(vector::VectorTransferSplit::None); |
| |
| if (!anchorOpName.empty()) { |
| using RunFnType = decltype(&runGenericStrategy); |
| RunFnType runFn = &runGenericStrategy; |
| // `linalg::PadTensorOp::getOperationName()` is not a StringLiteral, cannot |
| // use StringSwitch. |
| if (anchorOpName == CopyOp::getOperationName()) |
| runFn = &runStrategy<CopyOp>; |
| if (anchorOpName == FillOp::getOperationName()) |
| runFn = &runStrategy<FillOp>; |
| if (anchorOpName == PadTensorOp::getOperationName()) |
| runFn = &runStrategy<PadTensorOp>; |
| else if (anchorOpName == MatmulOp::getOperationName()) |
| runFn = &runStrategy<MatmulOp>; |
| else if (anchorOpName == MatmulI8I8I32Op::getOperationName()) |
| runFn = &runStrategy<MatmulI8I8I32Op>; |
| runFn(*this, tilingOptions, vectorContractLowering, vectorTransferSplit); |
| } |
| |
| // Transforms that do not require anchoring on a given op. |
| if (hoistPadding > 0) { |
| SmallVector<PadTensorOp> ops; |
| getFunction().walk([&](PadTensorOp op) { ops.push_back(op); }); |
| for (auto op : llvm::reverse(ops)) |
| (void)hoistPaddingOnTensors(op, hoistPadding); |
| } |
| if (vectorizePadding) { |
| OwningRewritePatternList extraVectorizationPatterns; |
| extraVectorizationPatterns.insert<PadTensorOpVectorizationPattern>( |
| &getContext()); |
| (void)applyPatternsAndFoldGreedily(getFunction(), |
| std::move(extraVectorizationPatterns)); |
| } |
| } |
| |
| namespace mlir { |
| namespace linalg { |
| void registerLinalgTensorCodegenStrategyPass() { |
| PassRegistration<LinalgTensorCodegenStrategyPass> |
| testLinalgCodegenStrategyPass("linalg-tensor-codegen-strategy", |
| "Linalg Tensor Codegen Strategy."); |
| } |
| } // namespace linalg |
| } // namespace mlir |
