| // Copyright 2020 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-dialects/Dialect/LinalgTransform/Passes.h" |
| #include "iree/compiler/Codegen/Common/CPU/Passes.h" |
| #include "iree/compiler/Codegen/Common/PassUtils.h" |
| #include "iree/compiler/Codegen/Common/Passes.h" |
| #include "iree/compiler/Codegen/Common/TileSizeSelection.h" |
| #include "iree/compiler/Codegen/LLVMCPU/Passes.h" |
| #include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h" |
| #include "iree/compiler/Dialect/Util/Transforms/Passes.h" |
| #include "iree/compiler/Utils/PassUtils.h" |
| #include "llvm/ADT/TypeSwitch.h" |
| #include "llvm/Support/CommandLine.h" |
| #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" |
| #include "mlir/Conversion/ArithToArmSME/ArithToArmSME.h" |
| #include "mlir/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.h" |
| #include "mlir/Conversion/ArmSMEToSCF/ArmSMEToSCF.h" |
| #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h" |
| #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" |
| #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" |
| #include "mlir/Conversion/VectorToArmSME/VectorToArmSME.h" |
| #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.h" |
| #include "mlir/Dialect/Arith/Transforms/Passes.h" |
| #include "mlir/Dialect/ArmSME/Transforms/Passes.h" |
| #include "mlir/Dialect/Func/IR/FuncOps.h" |
| #include "mlir/Dialect/Linalg/Passes.h" |
| #include "mlir/Dialect/MemRef/Transforms/Passes.h" |
| #include "mlir/Pass/PassManager.h" |
| #include "mlir/Transforms/Passes.h" |
| |
| #define DEBUG_TYPE "iree-llvmcpu-pass-pipelines" |
| |
| namespace mlir::iree_compiler { |
| |
| /// Command line options used purely for development purposes. Not to be relied |
| /// on in any way. |
| static llvm::cl::opt<bool> clFailOnOutOfBoundsStackAllocation( |
| "iree-llvmcpu-fail-on-out-of-bounds-stack-allocation", |
| llvm::cl::desc("fail if the upper bound of dynamic stack allocation cannot " |
| "be solved"), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clFailOnLargeVector( |
| "iree-llvmcpu-fail-on-large-vector", |
| llvm::cl::desc("fail if there are operations with large vectors"), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clCheckLinalgVectorization( |
| "iree-llvmcpu-check-linalg-vectorization", |
| llvm::cl::desc( |
| "Runs the pass to check if all the Linalg ops are vectorized"), |
| llvm::cl::init(false)); |
| |
| static llvm::cl::opt<bool> clUseFastMinMaxOps( |
| "iree-llvmcpu-use-fast-min-max-ops", |
| llvm::cl::desc( |
| "Use `arith.minf/maxf` instead of `arith.minimumf/maximumf` ops"), |
| llvm::cl::init(false)); |
| |
| static llvm::cl::opt<bool> clEnableReassociateFpReductions( |
| "iree-llvmcpu-reassociate-fp-reductions", |
| llvm::cl::desc("Enables reassociation for FP reductions"), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clSkipIntermediateRoundings( |
| "iree-llvmcpu-skip-intermediate-roundings", |
| llvm::cl::desc( |
| "Allow skipping intermediate roundings. For example, in f16 matmul " |
| "kernels on targets with only f32 arithmetic, we have to perform each " |
| "multiply-accumulate in f32, and if this flag is false, then we have " |
| "to round those f32 accumulators to the nearest f16 every time, which " |
| "is slow."), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clInstrumentMemoryAccesses{ |
| "iree-llvmcpu-instrument-memory-accesses", |
| llvm::cl::desc("Instruments memory accesses in dispatches when dispatch " |
| "instrumentation is enabled."), |
| llvm::cl::init(false)}; |
| |
| static llvm::cl::opt<bool> clUseSoftmaxInterFusion( |
| "iree-llvmcpu-use-decompose-softmax-fuse", |
| llvm::cl::desc("Enables inter-pass fusion for the DecomposeSoftmax pass."), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clEnableVectorContractCustomKernels( |
| "iree-llvmcpu-enable-vector-contract-custom-kernels", |
| llvm::cl::desc("Enables vector contract custom kernels for " |
| "LLVMCPUMmt4dVectorLowering pass."), |
| llvm::cl::init(false)); |
| |
| static llvm::cl::opt<bool> clTileDispatchUsingForall( |
| "iree-llvmcpu-tile-dispatch-using-forall", |
| llvm::cl::desc("Enable tile and distribute to workgroups using scf.forall"), |
| llvm::cl::init(true)); |
| |
| // By default, IREE does not enable the Armv9-A streaming SVE mode in the |
| // presence of scalable vectors (even when using `+sme`), as currently there's |
| // no cost model of when it could be beneficial. This flag will effectively make |
| // IREE/LLVM switch from SVE to SSVE in dispatch regions with supported |
| // scalable vector operations. |
| static llvm::cl::opt<bool> clForceArmStreaming( |
| "iree-llvmcpu-force-arm-streaming", |
| llvm::cl::desc( |
| "Enables Armv9-A streaming SVE mode for any dispatch region that " |
| "contains supported scalable vector operations (i.e., use SSVE rather " |
| "than SVE). Requires the +sme feature flag."), |
| llvm::cl::init(false)); |
| |
| // TODO: Enable `TileDispatchUsingForall` for every pipeline. |
| static void addTileAndDistributePasses(OpPassManager &funcPassManager) { |
| if (clTileDispatchUsingForall) { |
| funcPassManager.addPass( |
| createTileAndDistributeToWorkgroupsUsingForallOpPass()); |
| } else { |
| funcPassManager.addPass(createTileAndDistributeToWorkgroupsPass()); |
| funcPassManager.addPass(createCSEPass()); |
| funcPassManager.addPass(createConvertToDestinationPassingStylePass()); |
| funcPassManager.addPass(createFoldAffineMinInDistributedLoopsPass()); |
| } |
| funcPassManager.addPass(createConfigTrackingCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| funcPassManager.addPass(createFuseTensorPadWithConsumerPass()); |
| funcPassManager.addPass(createConcretizePadResultShapePass()); |
| } |
| |
| //===---------------------------------------------------------------------===// |
| // Codegen configuration verifications. |
| //===---------------------------------------------------------------------===// |
| |
| static bool isValidInterchange(ArrayRef<int64_t> interchange, int numLoops) { |
| if (interchange.empty()) |
| return true; |
| llvm::SmallDenseSet<int64_t> s; |
| s.insert(interchange.begin(), interchange.end()); |
| for (int i = 0; i < numLoops; ++i) { |
| if (!s.contains(i)) |
| return false; |
| } |
| return true; |
| } |
| |
| LogicalResult verifyDoubleTilingExpertPassPipelineConfig( |
| Operation *op, TilingConfig &tilingConfig, |
| IREE::Codegen::TranslationInfoAttr translationInfo, |
| ArrayRef<int64_t> workgroupSize) { |
| if (!workgroupSize.empty()) { |
| return op->emitOpError( |
| "expected workgroup size to be empty for CPU pipelines"); |
| } |
| |
| // Verify that the translation info is using the right pipeline. |
| if (translationInfo.getDispatchLoweringPassPipeline() != |
| IREE::Codegen::DispatchLoweringPassPipeline::CPUDoubleTilingExpert) { |
| return op->emitOpError("expected pipeline in translation_info to be ") |
| << stringifyEnum(IREE::Codegen::DispatchLoweringPassPipeline:: |
| CPUDoubleTilingExpert); |
| } |
| |
| if (tilingConfig.getNumTilingLevels() == 6) { |
| // TODO: update verification. |
| return success(); |
| } |
| |
| if (tilingConfig.getNumTilingLevels() != 4) { |
| return op->emitOpError("expected four tiling levels, got ") |
| << tilingConfig.getNumTilingLevels(); |
| } |
| |
| auto interfaceOp = dyn_cast_or_null<TilingInterface>(op); |
| if (interfaceOp) { |
| llvm::SmallDenseSet<unsigned> pLoopsSet; |
| for (auto [index, iteratorType] : |
| llvm::enumerate(interfaceOp.getLoopIteratorTypes())) { |
| if (iteratorType == utils::IteratorType::parallel) { |
| pLoopsSet.insert(index); |
| } |
| } |
| |
| SmallVector<int64_t> secondLevelTileSizes; |
| std::tie(secondLevelTileSizes, std::ignore) = |
| tilingConfig.getVectorCommonParallelSizes(); |
| for (auto [index, tileSize] : llvm::enumerate(secondLevelTileSizes)) { |
| if (tileSize != 0 && !pLoopsSet.contains(index)) { |
| return op->emitOpError( |
| "expected only parallel dims to be set in the second tiling " |
| "level, got ") |
| << index << "-th tile size set"; |
| } |
| } |
| |
| SmallVector<int64_t> thirdLevelTileSizes; |
| std::tie(thirdLevelTileSizes, std::ignore) = |
| tilingConfig.getVectorReductionSizes(); |
| for (auto [index, tileSize] : llvm::enumerate(thirdLevelTileSizes)) { |
| if (tileSize != 0 && pLoopsSet.contains(index)) { |
| return op->emitOpError( |
| "expected only reduction dims to be set in the third tiling " |
| "level, got ") |
| << index << "-th tile size set"; |
| } |
| } |
| } |
| |
| // Verify interchange |
| auto tileSizesForLevel = tilingConfig.getTileSizes(); |
| for (int level = 0; level < tilingConfig.getNumTilingLevels(); level++) { |
| auto interchange = tilingConfig.getTileInterchangeSizes(level); |
| auto &tileSizes = tileSizesForLevel[level]; |
| if (!isValidInterchange(interchange, tileSizes.size())) { |
| return op->emitOpError("expected [0, ") |
| << tileSizes.size() << ") to be set exactly once in interchange #" |
| << level; |
| } |
| } |
| |
| // Verify that native vector size is empty. |
| SmallVector<int64_t> nativeVectorSize = tilingConfig.getNativeVectorSizes(); |
| if (!nativeVectorSize.empty()) { |
| return op->emitOpError("native_vector_size must be empty"); |
| } |
| return success(); |
| } |
| |
| LogicalResult verifyConvTileAndDecomposeExpertConfig( |
| Operation *op, TilingConfig &tilingConfig, |
| IREE::Codegen::TranslationInfoAttr translationInfo, |
| ArrayRef<int64_t> workgroupSize) { |
| if (!isa<linalg::ConvolutionOpInterface>(op)) |
| return success(); |
| |
| if (tilingConfig.getNumTilingLevels() == 6) { |
| // TODO: update verification. |
| return success(); |
| } |
| |
| if (tilingConfig.getNumTilingLevels() != 4) { |
| return op->emitOpError("expected four tiling levels, got ") |
| << tilingConfig.getNumTilingLevels(); |
| } |
| |
| linalg::LinalgOp linalgOp = cast<linalg::LinalgOp>(op); |
| SmallVector<int64_t> shape = linalgOp.getStaticLoopRanges(); |
| for (auto sizes : tilingConfig.getTileSizes()) { |
| for (auto [i, size] : llvm::enumerate(sizes)) { |
| if (size == 1) |
| shape[i] = 1; |
| if (shape[i] == -1 || size == 0) |
| continue; |
| if (shape[i] % size != 0) { |
| shape[i] = -1; |
| } else { |
| shape[i] = size; |
| } |
| } |
| } |
| |
| int64_t khSize, kwSize, ohSize, owSize; |
| auto isSizeExtracted = |
| TypeSwitch<Operation *, LogicalResult>(op) |
| .Case<linalg::Conv2DNhwcHwcfOp, linalg::DepthwiseConv2DNhwcHwcOp, |
| linalg::PoolingNhwcSumOp, linalg::PoolingNhwcMaxOp, |
| linalg::PoolingNhwcMaxUnsignedOp, linalg::PoolingNhwcMinOp, |
| linalg::PoolingNhwcMinUnsignedOp>([&](auto) { |
| // Shape: N, OH, OW, OC, KH, KW, (IC) |
| khSize = shape[4]; |
| kwSize = shape[5]; |
| ohSize = shape[1]; |
| owSize = shape[2]; |
| return success(); |
| }) |
| .Case<linalg::Conv2DNchwFchwOp>([&](auto) { |
| // Shape: N, OC, OH, OW, (IC), KH, KW |
| khSize = shape[5]; |
| kwSize = shape[6]; |
| ohSize = shape[2]; |
| owSize = shape[3]; |
| return success(); |
| }) |
| .Case<linalg::PoolingNchwSumOp, linalg::PoolingNchwMaxOp>([&](auto) { |
| // Shape: N, OC, OH, OW, KH, KW |
| khSize = shape[4]; |
| kwSize = shape[5]; |
| ohSize = shape[2]; |
| owSize = shape[3]; |
| return success(); |
| }) |
| .Default([&](auto) { return failure(); }); |
| if (failed(isSizeExtracted)) { |
| return op->emitOpError("unsupported conv types"); |
| } |
| |
| bool removeH = (khSize == 1 && ohSize == 1); |
| bool removeW = (kwSize == 1 && owSize == 1); |
| if (!removeH && !removeW) { |
| return op->emitOpError("can't decompose the conv op"); |
| } |
| |
| return success(); |
| } |
| |
| //===---------------------------------------------------------------------===// |
| // Codegen pipelines. |
| //===---------------------------------------------------------------------===// |
| |
| void buildLLVMCPUVectorLoweringPipeline( |
| OpPassManager &funcPassManager, |
| const LLVMCPUVectorLoweringPassOptions &options) { |
| funcPassManager.addPass(createDropVectorUnitDimsPass()); |
| funcPassManager.addPass(createLLVMCPUVirtualVectorLoweringPass( |
| LLVMCPUVirtualVectorLoweringPassOptions{options.splitVectorTransfersTo, |
| options.enableArmI8mm})); |
| |
| // Make sure we remove redundant vector ops (e.g., vector tranposes) before we |
| // lower them and can't be optimized away anymore. |
| funcPassManager.addPass(createCanonicalizerPass()); |
| |
| LLVMCPUVectorTransferLoweringPassOptions transferLoweringOptions{}; |
| if (!options.enableArmSME) { |
| // The ArmSME dialect has its own (more specific) lowerings for scalable |
| // vectors that occur later in the pipeline, so only enable the general |
| // lowerings if SME is not available. |
| transferLoweringOptions.enableScalableLowerings = true; |
| } |
| funcPassManager.addPass( |
| createLLVMCPUVectorTransferLoweringPass(transferLoweringOptions)); |
| funcPassManager.addPass(createLLVMCPUVectorTransposeLoweringPass( |
| LLVMCPUVectorTransposeLoweringPassOptions{ |
| options.lowerVectorTransposeToAVX2})); |
| |
| // Potentially removes shape_cast and broadcast on unit dims before shape_cast |
| // lowering. |
| funcPassManager.addPass(createCanonicalizerPass()); |
| |
| // 'vector.shape_cast' are very expensive operations that are even generated |
| // by some of the lowerings above (e.g., transpose lowering). There are |
| // chances to cancel them out if they are not lowered too early so we lower |
| // them at the very end of the pass. |
| funcPassManager.addPass(createLLVMCPUVectorShapeCastLoweringPass()); |
| } |
| |
| void addCPUBufferOpsTileAndVectorizePipeline( |
| OpPassManager &funcPassManager, TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| |
| // Skip tiling reduction loops because this is expected to apply on copy ops |
| // only. |
| funcPassManager.addPass( |
| createLLVMCPUTilePass(tilingConfig.getVectorCommonParallelLevel())); |
| funcPassManager.addPass(createLLVMCPUPeelPass()); |
| { |
| GenericVectorizationPassOptions options; |
| options.useConfiguredVectorSizes = pipelineOpt.useConfiguredVectorSizes; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| options.vectorizeGatherAccesses = true; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| // Run IREE specific passes before vector lowering expert. |
| funcPassManager.addPass(createRemoveSingleIterationLoopPass()); |
| |
| { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| } |
| |
| void addMultiTilingExpertPassPipeline(OpPassManager &funcPassManager, |
| TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| |
| SmallVector<int64_t> allFusableLevels(tilingConfig.getFusableLevels()); |
| // Apply tile and fuse to all the non-distribution fusable levels. Skip |
| // distribution level as that level has been fused already. |
| if (allFusableLevels.size() > 1) { |
| llvm::SmallSetVector<int64_t, 4> fusableLevels(allFusableLevels.begin(), |
| allFusableLevels.end()); |
| for (int i = 0; i < tilingConfig.getNumTilingLevels(); ++i) { |
| if (i == tilingConfig.getDistributionLevel()) |
| continue; |
| if (fusableLevels.contains(i)) { |
| funcPassManager.addPass(createLLVMCPUTileAndFusePass(i)); |
| funcPassManager.addPass(createFuseTensorPadWithConsumerPass()); |
| funcPassManager.addPass(createConcretizePadResultShapePass()); |
| continue; |
| } |
| |
| if (i == tilingConfig.getVectorReductionLevel()) { |
| // Run SplitReductionPass before the final reduction Fuse pass, because |
| // SplitReductionPass takes care of banked-tiling. |
| funcPassManager.addPass( |
| createLLVMCPUSplitReductionPass(clEnableReassociateFpReductions)); |
| funcPassManager.addPass(createLLVMCPUTileRootAndFuseInputOperands(i)); |
| continue; |
| } |
| |
| funcPassManager.addPass(createLLVMCPUTileRootAndFuseInputOperands(i)); |
| } |
| } |
| |
| if (pipelineOpt.enablePeeling) { |
| funcPassManager.addPass(createLLVMCPUPeelPass()); |
| } |
| |
| if (pipelineOpt.enableAArch64SME) { |
| funcPassManager.addPass(createLLVMCPU2DScalableTo1DScalablePass()); |
| } |
| |
| { |
| funcPassManager.addPass(createTensorToVectorVectorizePadPass()); |
| if (pipelineOpt.decomposePackUnPackOps) { |
| funcPassManager.addPass(createDecomposePackUnPackOpsPass()); |
| funcPassManager.addPass(createConfigTrackingCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| } |
| |
| GenericVectorizationPassOptions options; |
| options.useConfiguredVectorSizes = pipelineOpt.useConfiguredVectorSizes; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| options.vectorizePadding = true; |
| options.vectorizeGatherAccesses = true; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| addCPUBufferizePasses(funcPassManager); |
| |
| // Run IREE specific passes before vector lowering expert. |
| funcPassManager.addPass(createRemoveSingleIterationLoopPass()); |
| |
| { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| } |
| |
| void addConvTileAndDecomposeExpertPassPipeline( |
| OpPassManager &funcPassManager, TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| |
| // Run LLVMTileAndFuse firstly in case that we have fill + conv + generic |
| // ops. At this stage, we do not apply vectorization. The reduction dim won't |
| // get tiled if the case is conv + generic op. In this case, we have to tile |
| // along reduction dim again, which needs them to be Linalg ops form. |
| |
| funcPassManager.addPass(createLLVMCPUTileAndFusePass( |
| tilingConfig.getVectorCommonParallelLevel())); |
| funcPassManager.addPass(createFuseTensorPadWithConsumerPass()); |
| funcPassManager.addPass(createConcretizePadResultShapePass()); |
| |
| funcPassManager.addPass(createLLVMCPUTileRootAndFuseInputOperands( |
| tilingConfig.getVectorReductionLevel())); |
| funcPassManager.addPass( |
| createLLVMCPUTileAndFusePass(tilingConfig.getVectorInnerParallelLevel())); |
| funcPassManager.addPass(createDecomposeConvolutionToLowerDimOpsPass()); |
| |
| funcPassManager.addPass(createFuseTensorPadWithConsumerPass()); |
| funcPassManager.addPass(createConcretizePadResultShapePass()); |
| |
| if (pipelineOpt.enablePeeling) { |
| funcPassManager.addPass(createLLVMCPUPeelPass()); |
| } |
| |
| { |
| funcPassManager.addPass(createTensorToVectorVectorizePadPass()); |
| GenericVectorizationPassOptions options; |
| options.useConfiguredVectorSizes = pipelineOpt.useConfiguredVectorSizes; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| options.vectorizePadding = true; |
| options.vectorizeGatherAccesses = true; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| // Eliminate redundant transfer_read/write to avoid stack allocations. |
| funcPassManager.addPass(createOptimizeVectorTransferPass( |
| OptimizeVectorTransferPassOptions{/*flatten=*/true})); |
| |
| addCPUBufferizePasses(funcPassManager); |
| |
| // Run IREE specific passes before vector lowering expert. |
| funcPassManager.addPass(createRemoveSingleIterationLoopPass()); |
| |
| { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "shuffle"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| } |
| |
| void addMmt4dTilingExpertPassPipeline(OpPassManager &funcPassManager, |
| TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| |
| funcPassManager.addPass(createLLVMCPUTileRootAndFuseProducerConsumer( |
| static_cast<int64_t>(tilingConfig.getVectorCommonParallelLevel()))); |
| // The below two passes are nop if the "mmt4d" is explicitly excluded in the |
| // ukernels attribute. |
| funcPassManager.addPass(createCPUPrepareUkernelsPass()); |
| funcPassManager.addPass( |
| createCPULowerToUKernelsPass(clSkipIntermediateRoundings)); |
| funcPassManager.addPass(createLLVMCPUTileRootAndFuseInputOperands( |
| static_cast<int64_t>(tilingConfig.getVectorReductionLevel()))); |
| |
| { |
| GenericVectorizationPassOptions options; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| options.vectorizePadding = true; |
| options.vectorizeGatherAccesses = true; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| |
| addCPUBufferizePasses(funcPassManager); |
| |
| // Vector lowering of Mmt4d. |
| funcPassManager.addPass(createLLVMCPUMmt4dVectorLoweringPass( |
| LLVMCPUMmt4dVectorLoweringPassOptions{ |
| clEnableVectorContractCustomKernels})); |
| |
| // Generic vector lowering. |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| |
| void addCPUDataTilingPipeline(OpPassManager &funcPassManager, |
| TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| |
| // The below two passes are nop if pack/unpack is not specified in ukernels |
| // attribute. By default, they are disabled. |
| funcPassManager.addPass(createCPUPrepareUkernelsPass()); |
| funcPassManager.addPass( |
| createCPULowerToUKernelsPass(clSkipIntermediateRoundings)); |
| |
| funcPassManager.addPass( |
| createLLVMCPUTilePass(tilingConfig.getVectorCommonParallelLevel())); |
| if (pipelineOpt.decomposePackUnPackOps) { |
| funcPassManager.addPass(createDecomposePackUnPackOpsPass()); |
| } |
| |
| { |
| GenericVectorizationPassOptions options; |
| options.useConfiguredVectorSizes = pipelineOpt.useConfiguredVectorSizes; |
| options.vectorizePadding = true; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| addCPUBufferizePasses(funcPassManager); |
| |
| { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| } |
| |
| void addCPULinalgExtTileAndVectorizePipeline( |
| OpPassManager &funcPassManager, TilingConfig &tilingConfig, |
| LLVMCPUPipelineOptions &pipelineOpt) { |
| addTileAndDistributePasses(funcPassManager); |
| funcPassManager.addPass( |
| createLLVMCPUTilePass(tilingConfig.getVectorCommonParallelLevel())); |
| // TODO: Remove the pass once we have PartialReductionOpInterface implemented |
| // for AttentionOp. |
| funcPassManager.addPass( |
| IREE::LinalgExt::createConvertAttentionToOnlineAttentionPass()); |
| funcPassManager.addPass( |
| createLLVMCPUTilePass(tilingConfig.getVectorReductionLevel())); |
| funcPassManager.addPass( |
| IREE::LinalgExt::createDecomposeWinogradTransformPass()); |
| funcPassManager.addPass(IREE::LinalgExt::createDecomposeAttentionPass()); |
| |
| { |
| GenericVectorizationPassOptions options; |
| options.useConfiguredVectorSizes = pipelineOpt.useConfiguredVectorSizes; |
| options.enableVectorMasking = pipelineOpt.enableVectorMasking; |
| funcPassManager.addPass(createGenericVectorizationPass(options)); |
| funcPassManager.addPass(createOptimizeTensorInsertExtractSlicesPass()); |
| funcPassManager.addPass(createCanonicalizerPass()); |
| funcPassManager.addPass(createCSEPass()); |
| if (clFailOnLargeVector) { |
| funcPassManager.addPass(createLLVMCPUVerifyVectorSizeLegalityPass()); |
| } |
| } |
| |
| addCPUBufferizePasses(funcPassManager); |
| |
| { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.lowerVectorTransposeToAVX2 = pipelineOpt.lowerToAVX2; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| options.enableArmI8mm = pipelineOpt.enableAArch64I8mm; |
| options.enableArmSME = pipelineOpt.enableAArch64SME; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| } |
| } |
| |
| void addCPUDefaultPassPipeline(OpPassManager &funcPassManager, |
| FailureOr<TilingConfig> &tilingConfig) { |
| if (succeeded(tilingConfig) && |
| tilingConfig.value().getNumTilingLevels() > 1) { |
| addTileAndDistributePasses(funcPassManager); |
| funcPassManager.addPass(createLLVMCPUTileAndFusePass( |
| tilingConfig.value().getVectorCommonParallelLevel())); |
| } |
| addCPUBufferizePasses(funcPassManager); |
| } |
| |
| static void addLowerToLLVMPasses(OpPassManager &modulePassManager, |
| bool enableAArch64SME) { |
| // TODO: Remove the following pass and plumb support for #hal.descriptor_type |
| // memory space through the stack. |
| FunctionLikeNest(modulePassManager) |
| .addPass(createEraseHALDescriptorTypeFromMemRefPass); |
| |
| // Lower `ukernel.*` ops to function calls |
| modulePassManager.addPass(createLowerUKernelOpsToCallsPass()); |
| |
| FunctionLikeNest(modulePassManager) |
| // LinalgExt -> SCF |
| .addPass(IREE::LinalgExt::createLinalgExtToLoopsPass) |
| // Linalg -> SCF |
| .addPass(createMemrefCopyToLinalgPass) |
| .addPredicatedPass(clCheckLinalgVectorization, |
| createLLVMCPUEmitVectorizationRemarksPass) |
| .addPass(createConvertLinalgToLoopsPass) |
| .addPass(createConvertBf16ArithToF32Pass) |
| .addPass(createConvertBf16ToUInt16BuffersPass) |
| .addPass(createCanonicalizerPass) |
| .addPass(createCSEPass); |
| |
| // Handled tensor-type constants. |
| addConstantBufferizePasses(modulePassManager); |
| |
| FunctionLikeNest(modulePassManager) |
| .addPass(createFoldTensorExtractOpPass) |
| // Handle complex operation conversion. |
| .addPass(createConvertComplexToStandardPass) |
| // math dialect elementry functions -> polynomial form. |
| .addPass(createPolynomialApproximationPass) |
| .addPass(createHoistStaticallyBoundAllocationsPass) |
| // Use `arith.minf/maxf` instead of `arith.minimumf/maximumf`. |
| .addPredicatedPass(clUseFastMinMaxOps, createReplaceSlowMinMaxOpsPass); |
| |
| if (enableAArch64SME) { |
| modulePassManager.addPass(mlir::arm_sme::createVectorLegalizationPass()); |
| FunctionLikeNest(modulePassManager) |
| .addPredicatedPass( |
| clForceArmStreaming, |
| [] { |
| // 1. Enable Armv9-A streaming mode without ZA (i.e., SSVE) for |
| // dispatch regions that contain scalable vectors when forced via |
| // the --iree-llvmcpu-force-arm-streaming flag. |
| return mlir::arm_sme::createEnableArmStreamingPass( |
| mlir::arm_sme::ArmStreamingMode::StreamingLocally, |
| mlir::arm_sme::ArmZaMode::Disabled, |
| /*ifRequiredByOps=*/false, |
| /*ifContainsScalableVectors=*/true); |
| }) |
| .addPass(createCanonicalizerPass) |
| .addPass(createCSEPass) |
| .addPass(mlir::createArithToArmSMEConversionPass) |
| .addPass(mlir::createConvertVectorToArmSMEPass) |
| .addPass([] { |
| // 2. Enable ZA for dispatch regions that contain ArmSME ops (which |
| // all make use of the ZA state). |
| return mlir::arm_sme::createEnableArmStreamingPass( |
| mlir::arm_sme::ArmStreamingMode::StreamingLocally, |
| mlir::arm_sme::ArmZaMode::NewZA, |
| /*ifRequiredByOps=*/true); |
| }) |
| .addPass(mlir::createConvertArmSMEToSCFPass); |
| } |
| |
| FunctionLikeNest(modulePassManager) |
| // Resolve get_buffer_descriptor ops. All structural buffer manipulations |
| // must conclude before this point. |
| .addPass(createIREEExpandStridedMetadataPass) |
| .addPass(createCleanupBufferAllocViewPass) |
| // Checking stack allocation before converting to CF dialect is easier. |
| .addPass([&]() { |
| return createLLVMCPUCheckIRBeforeLLVMConversionPass( |
| LLVMCPUCheckIRBeforeLLVMConversionPassOptions{ |
| clFailOnOutOfBoundsStackAllocation}); |
| }) |
| // SCF -> CF |
| .addPass(createConvertSCFToCFPass) |
| .addPass(createCanonicalizerPass) |
| .addPass(createCSEPass) |
| // (HAL, IREE, Linalg, CF) -> LLVM |
| .addPass(arith::createArithExpandOpsPass) |
| .addPass(memref::createExpandOpsPass) |
| .addPass(memref::createFoldMemRefAliasOpsPass) |
| .addPass(createEmulateNarrowTypePass) |
| .addPass(createCanonicalizerPass) |
| .addPass(createCSEPass) |
| .addPredicatedPass(clInstrumentMemoryAccesses, |
| createInstrumentMemoryAccessesPass); |
| |
| if (enableAArch64SME) { |
| FunctionLikeNest(modulePassManager).addPass([&] { |
| return createConvertArmSMEToLLVMPass(); |
| }); |
| } |
| modulePassManager.addPass( |
| createConvertToLLVMPass(clEnableReassociateFpReductions)); |
| modulePassManager.addPass(createReconcileUnrealizedCastsPass()); |
| |
| // We rely on MLIR symbol visibility being correct after this point and need |
| // to mirror the LLVM linkage that was assigned during conversion. |
| modulePassManager.addPass(createLLVMCPUSynchronizeSymbolVisibilityPass()); |
| |
| modulePassManager.addPass(createCanonicalizerPass()); |
| modulePassManager.addPass(createCSEPass()); |
| modulePassManager.addNestedPass<LLVM::LLVMFuncOp>( |
| createAddFastMathFlagsPass()); |
| } |
| |
| void buildLLVMCPUCodegenConfigurationPassPipelineImpl( |
| OpPassManager &modulePassManager) { |
| { |
| FunctionLikeNest funcPassManager(modulePassManager); |
| addCommonTargetExecutablePreprocessingPasses(funcPassManager, |
| clUseSoftmaxInterFusion); |
| } |
| modulePassManager.addPass(createMaterializeUserConfigsPass()); |
| FunctionLikeNest(modulePassManager) |
| .addPass(createRematerializeParallelOpsPass) |
| // TODO(#13888): This(createExpandF16OpToF32Pass()) pass is being added |
| // way to late and should insted be be done during lowering to LLVM. |
| .addPass(createExpandF16OpToF32Pass) |
| .addPass(createMaterializeDeviceEncodingPass) |
| // TODO: Remove the following pass the plumb support for |
| // #hal.descriptor_type memory space through the stack. |
| .addPass(createEraseHALDescriptorTypeFromMemRefPass); |
| |
| modulePassManager.addPass(createLLVMCPUSelectLoweringStrategyPass()); |
| LLVM_DEBUG({ |
| llvm::dbgs() << "LLVMCPU codegen configuration pass pipeline:\n"; |
| modulePassManager.printAsTextualPipeline(llvm::dbgs()); |
| llvm::dbgs() << "\n"; |
| }); |
| } |
| |
| void buildLLVMCPUCodegenConfigurationPassPipeline( |
| OpPassManager &variantPassManager) { |
| OpPassManager &modulePassManager = variantPassManager.nest<ModuleOp>(); |
| buildLLVMCPUCodegenConfigurationPassPipelineImpl(modulePassManager); |
| } |
| |
| void buildLLVMCPUCodegenPassPipeline(OpPassManager &variantPassManager, |
| bool enableAArch64SME) { |
| |
| { |
| OpPassManager &modulePassManager = variantPassManager.nest<ModuleOp>(); |
| modulePassManager.addPass(createLowerExecutableUsingTransformDialectPass()); |
| FunctionLikeNest(modulePassManager) |
| .addPass(createLLVMCPULowerExecutableTargetPass) |
| .addPass(createVerifyWorkgroupDistributionPass); |
| } |
| |
| variantPassManager.addPass(createReconcileTranslationInfoPass()); |
| variantPassManager.addPass(createLowerAffinePass()); |
| variantPassManager.addPass(IREE::Util::createDropCompilerHintsPass()); |
| |
| // Run conversion to LLVM at `ModuleOp` granularity. |
| { |
| OpPassManager &modulePassManager = variantPassManager.nest<ModuleOp>(); |
| addLowerToLLVMPasses(modulePassManager, enableAArch64SME); |
| } |
| LLVM_DEBUG({ |
| llvm::dbgs() << "LLVMCPU codegen pass pipeline:\n"; |
| variantPassManager.printAsTextualPipeline(llvm::dbgs()); |
| llvm::dbgs() << "\n"; |
| }); |
| } |
| |
| // NOTE: this runs on the top-level program module containing all |
| // hal.executable ops. |
| void buildLLVMCPULinkingPassPipeline(OpPassManager &modulePassManager, |
| std::optional<std::string> target) { |
| // Link together executables. This may produce some IR duplication. |
| LLVMCPULinkExecutablesPassOptions linkOptions; |
| linkOptions.target = target.value_or(""); |
| modulePassManager.addPass(createLLVMCPULinkExecutablesPass(linkOptions)); |
| |
| // Cleanup IR duplication. |
| modulePassManager.addNestedPass<IREE::HAL::ExecutableOp>( |
| mlir::createCanonicalizerPass()); |
| |
| // Assign final executable constant and import ordinals. |
| auto &variantPassManager = modulePassManager.nest<IREE::HAL::ExecutableOp>() |
| .nest<IREE::HAL::ExecutableVariantOp>(); |
| variantPassManager.addPass(createLLVMCPUAssignConstantOrdinalsPass()); |
| variantPassManager.addPass(createLLVMCPUAssignImportOrdinalsPass()); |
| } |
| |
| //===---------------------------------------------------------------------===// |
| // Register LLVMCPU Passes |
| //===---------------------------------------------------------------------===// |
| |
| namespace { |
| #define GEN_PASS_REGISTRATION |
| #include "iree/compiler/Codegen/LLVMCPU/Passes.h.inc" |
| } // namespace |
| |
| void registerCodegenLLVMCPUPasses() { |
| // Generated. |
| registerPasses(); |
| |
| static PassPipelineRegistration<> LLVMCPUConfigPipeline( |
| "iree-codegen-llvmcpu-configuration-pipeline", |
| "Runs the translation strategy configuration pipeline on Linalg for CPU", |
| [](OpPassManager &modulePassManager) { |
| buildLLVMCPUCodegenConfigurationPassPipeline(modulePassManager); |
| }); |
| |
| static PassPipelineRegistration<> LLVMCPUBufferizationPipeline( |
| "iree-codegen-llvmcpu-bufferization-pipeline", |
| "Runs the bufferization pipeline for CPU", |
| [](OpPassManager &funcPassManager) { |
| addCPUBufferizePasses(funcPassManager); |
| }); |
| |
| static PassPipelineRegistration<> LLVMCPUVectorLoweringPipeline( |
| "iree-codegen-llvmcpu-vector-lowering-pipeline", |
| "Runs the translation strategy configuration pipeline on Linalg for CPU", |
| [](OpPassManager &funcPassManager) { |
| LLVMCPUVectorLoweringPassOptions options; |
| options.splitVectorTransfersTo = "linalg-copy"; |
| buildLLVMCPUVectorLoweringPipeline(funcPassManager, options); |
| }); |
| |
| struct LinalgToLLVMPipelineOptions |
| : public PassPipelineOptions<LinalgToLLVMPipelineOptions> { |
| Option<bool> enableArmSME{ |
| *this, "enable-arm-sme", |
| llvm::cl::desc("Enable the ArmSME lowering pipeline.")}; |
| }; |
| |
| static PassPipelineRegistration<LinalgToLLVMPipelineOptions> |
| LinalgLLVMPipeline( |
| "iree-codegen-linalg-to-llvm-pipeline", |
| "Runs the progressive lowering pipeline from Linalg to LLVM", |
| [](OpPassManager &variantPassManager, |
| LinalgToLLVMPipelineOptions const &options) { |
| buildLLVMCPUCodegenPassPipeline(variantPassManager, |
| options.enableArmSME); |
| }); |
| |
| static PassPipelineRegistration<> LLVMCPULinkingPipeline( |
| "iree-codegen-llvmcpu-linking-pipeline", |
| "Runs the LLVMCPU HAL executable linking pipeline", |
| [](OpPassManager &modulePassManager) { |
| buildLLVMCPULinkingPassPipeline(modulePassManager); |
| }); |
| } |
| |
| } // namespace mlir::iree_compiler |
