| // 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/GlobalOptimization/Passes.h" |
| #include "iree/compiler/Dialect/Flow/Transforms/Passes.h" |
| #include "iree/compiler/Dialect/TensorExt/IR/TensorExtDialect.h" |
| #include "iree/compiler/Dialect/Util/Transforms/Passes.h" |
| #include "iree/compiler/DispatchCreation/Passes.h" |
| #include "iree/compiler/Modules/IO/Parameters/Transforms/Passes.h" |
| #include "iree/compiler/Utils/PassUtils.h" |
| #include "mlir/Dialect/Linalg/Passes.h" |
| #include "mlir/Dialect/MemRef/Transforms/Passes.h" |
| #include "mlir/Transforms/Passes.h" |
| |
| namespace mlir::iree_compiler::GlobalOptimization { |
| |
| using FunctionLikeNest = |
| MultiOpNest<IREE::Util::InitializerOp, IREE::Util::FuncOp>; |
| |
| static llvm::cl::opt<bool> clEnableQuantizedMatmulReassociation( |
| "iree-global-opt-enable-quantized-matmul-reassociation", |
| llvm::cl::desc( |
| "Enables reassociation of quantized matmul ops (experimental)."), |
| llvm::cl::init(false)); |
| static llvm::cl::opt<bool> clConvertQDQToIntegerMath( |
| "iree-global-opt-enable-qdq-to-integer-math", |
| llvm::cl::desc("Rewrites contractions over dequantized operands into " |
| "integer contractions with zero point corrections (" |
| "changes floating-point rounding and exceptional-value " |
| "behavior)."), |
| llvm::cl::init(true)); |
| static llvm::cl::opt<bool> clEnableTransposePropagation( |
| "iree-global-opt-propagate-transposes", |
| llvm::cl::desc( |
| "Enables propagation of transpose ops to improve fusion chances."), |
| llvm::cl::init(true)); |
| static llvm::cl::opt<bool> clEnableAttentionVTranspose( |
| "iree-global-opt-enable-attention-v-transpose", |
| llvm::cl::desc("Enables transposition of v operand of attention ops,"), |
| llvm::cl::init(true)); |
| |
| // TODO(hanchung): Remove the flag. We don't want to do early materialization by |
| // default. Because it won't work for heterogeneous computing. This is not the |
| // right layer for handling such information. |
| static llvm::cl::opt<bool> clEnableEarlyMaterialization( |
| "iree-global-opt-enable-early-materialization", |
| llvm::cl::desc( |
| "Enables early materialization on encodings. Note, this flag should be " |
| "false eventually. This does not work for heterogeneous computing."), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<DemoteType> clDemoteContractionInputsType( |
| "iree-global-opt-demote-contraction-inputs-type", |
| llvm::cl::desc( |
| "Demotes inputs (LHS, RHS) of contraction ops to a narrow type."), |
| llvm::cl::values(clEnumValN(DemoteType::F16, "f16", "Demote to f16."), |
| clEnumValN(DemoteType::BF16, "bf16", "Demote to bf16."))); |
| |
| static llvm::cl::opt<DemoteOperation> clDemoteContractionInputsOperations( |
| "iree-global-opt-demote-contraction-inputs-operations", |
| llvm::cl::desc("Select the type of contraction ops to demote."), |
| llvm::cl::values(clEnumValN(DemoteOperation::All, "all", |
| "Demote all contraction ops."), |
| clEnumValN(DemoteOperation::Conv, "conv", |
| "Only demote convolution ops."), |
| clEnumValN(DemoteOperation::Matmul, "matmul", |
| "Only demote matmul ops."), |
| clEnumValN(DemoteOperation::None, "none", |
| "Demote no contraction ops.")), |
| llvm::cl::init(DemoteOperation::None)); |
| |
| static llvm::cl::opt<DispatchCreation::EncodingOptions> clSetEncodingStrategy( |
| "iree-global-opt-set-encoding-strategy", |
| llvm::cl::desc("Set the encoding strategy for operations."), |
| llvm::cl::values(clEnumValN( |
| DispatchCreation::EncodingOptions::Generic, "generic", |
| "Using EncodingAttr which encodes as much information as possible")), |
| llvm::cl::init(DispatchCreation::EncodingOptions::Generic)); |
| |
| static llvm::cl::opt<bool> clWarnOnUninitializedValues( |
| "iree-global-opt-enable-warn-on-uninitialized-values", |
| llvm::cl::desc("Warn on some classes of uses of uninitialized values."), |
| llvm::cl::init(true)); |
| |
| static llvm::cl::opt<bool> clEnableEdgeReshapePropagation( |
| "iree-global-opt-experimental-enable-edge-reshape-propagation", |
| llvm::cl::desc( |
| "Enables propagation of reshapes on the edges of the program " |
| "in transpose propagation. This workaround for better performance and " |
| "will be removed soon."), |
| llvm::cl::init(false)); |
| |
| void buildGlobalOptExprHoistingPassPipeline( |
| OpPassManager &passManager, const TransformOptions &transformOptions) { |
| IREE::Util::ExprHoistingOptions options; |
| options.maxSizeIncreaseThreshold = |
| transformOptions.constExprMaxSizeIncreaseThreshold; |
| options.registerDependentDialectsFn = [](DialectRegistry ®istry) { |
| registry.insert<IREE::TensorExt::IREETensorExtDialect>(); |
| }; |
| passManager.addPass(IREE::Util::createHoistIntoGlobalsPass(options)); |
| } |
| |
| void buildGlobalOptimizationPassPipeline( |
| OpPassManager &mainPassManager, const TransformOptions &transformOptions) { |
| // Import parameters before any global optimization passes so that the inlined |
| // parameters are available for folding. |
| if (!transformOptions.parameterImportPaths.empty()) { |
| IREE::IO::Parameters::ImportParametersPassOptions importParametersOptions; |
| importParametersOptions.scopePaths.assign( |
| transformOptions.parameterImportPaths.begin(), |
| transformOptions.parameterImportPaths.end()); |
| importParametersOptions.keys.assign( |
| transformOptions.parameterImportKeys.begin(), |
| transformOptions.parameterImportKeys.end()); |
| importParametersOptions.maximumSize = |
| transformOptions.parameterImportMaximumSize; |
| mainPassManager.addPass(IREE::IO::Parameters::createImportParametersPass( |
| importParametersOptions)); |
| } |
| |
| // Preprocessing passes to get the program into a canonical state. |
| FunctionLikeNest(mainPassManager) |
| .addPredicatedPass(clWarnOnUninitializedValues, |
| createWarnOnUninitializedValuesPass) |
| .addPredicatedPass(transformOptions.stripAssertions, |
| IREE::Util::createStripDebugOpsPass) |
| .addPass(IREE::Util::createOptimizeIntArithmeticPass) |
| .addPass(createLinalgQuantizedConvToConvPass) |
| .addPass(createLinalgQuantizedMatmulToMatmulPass) |
| // Match direct dequantize producers before convolution lowering and |
| // unit-dimension folding insert reshapes between them and the |
| // contraction. |
| .addPredicatedPass(clConvertQDQToIntegerMath, |
| createConvertQDQToIntegerMathPass) |
| .addPredicatedPass(transformOptions.useIm2colForConvs, |
| createConvertConv2DToImg2ColPass) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(createRemoveZeroExtentTensorsPass) |
| .addPass(createDetachElementwiseFromNamedOpsPass) |
| .addPass(mlir::createSimplifyDepthwiseConvPass); |
| mainPassManager.addPass(createEraseUnusedLinalgOperandsPass()); |
| |
| // Expand tensor shapes into SSA values and optimize the whole program. |
| // The more we are able to equate shape dimensions at this level the |
| // better our fusions will be. |
| mainPassManager.addPass(createExpandTensorShapesPass()); |
| |
| FunctionLikeNest(mainPassManager) |
| // Preprocess the input to a form more amenable for fusion |
| // - Convert all elementwise ops to Linalg |
| // - Remove unit-extent dimensions. |
| .addPass(mlir::createConvertElementwiseToLinalgPass) |
| // RaiseSpecialOps, by virtue of implementing various peephole |
| // optimizations, is sensitive to surrounding IR structure. Thus we run |
| // this pass both before unit dim folding + consteval, as well as after. |
| .addPass(createRaiseSpecialOpsPass) |
| // We decompose and transpose concatenations immediately before folding |
| // unit extent dims because this allows decoupling unit dims in the |
| // concatenation from the transposes that are introduced. |
| .addPass([&]() { |
| return createDecomposeConcatPass(transformOptions.outerDimConcat); |
| }) |
| // We generalize certain named ops immediately before folding unit extent |
| // dims as the unit dim folding pass updates indexing maps and is better |
| // at working with generics. By this point we have already done any |
| // specialized raising and the op names are no longer useful. |
| .addPass([&]() { |
| GeneralizeLinalgNamedOpsPassOptions opt; |
| opt.enableGeneralizeMatmul = transformOptions.generalizeMatmul; |
| return createGeneralizeLinalgNamedOpsPass(opt); |
| }) |
| .addPredicatedPass(!clEnableEdgeReshapePropagation, |
| DispatchCreation::createInsertTensorBarriersPass); |
| mainPassManager.addPass(DispatchCreation::createFoldUnitExtentDimsPass()); |
| FunctionLikeNest(mainPassManager) |
| .addPass(DispatchCreation::createFoldReshapesIntoTensorBarriersPass) |
| .addPass([&]() { |
| return createDemoteContractionInputsPass( |
| clDemoteContractionInputsType, clDemoteContractionInputsOperations); |
| }) |
| .addPredicatedPass(clEnableQuantizedMatmulReassociation, |
| createFuseDequantizationMatmulPass) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(mlir::createCSEPass) |
| // Convert broadcast + batch_matmul before transpose propagation so we |
| // can match the pure broadcast patterns before they get fused with |
| // transposes. |
| .addPass(GlobalOptimization::createConvertBatchMatmulToMatmulPass); |
| |
| FunctionLikeNest(mainPassManager) |
| // Propagate transposes immediately before set encoding/data tiling |
| // because transpose propagation cannot take an opinion on the preferred |
| // layout of various operations. This simplifies local propagation |
| // decisions as SetEncoding is expected to pick the ideal layout for |
| // that operation anyway, and this way we only need to make such a |
| // decision once. |
| .addPredicatedPass(clEnableTransposePropagation, |
| [&]() { |
| PropagateLinalgTransposePassOptions options; |
| options.enableAggressivePropagation = |
| transformOptions.aggressiveTransposePropagation; |
| options.enableConvolutionPropagation = |
| transformOptions.propagateTransposesThroughConv; |
| options.enableSinkTransposeThroughPad = |
| transformOptions.sinkTransposeThroughPad; |
| options.enableAttentionVTranspose = |
| clEnableAttentionVTranspose; |
| options.enableEdgeReshapePropagation = |
| clEnableEdgeReshapePropagation; |
| return createPropagateLinalgTransposePass(options); |
| }) |
| // Retry for contractions whose dequantize producers were exposed by |
| // reshape/transpose propagation, before selecting integer layouts. |
| .addPredicatedPass(clConvertQDQToIntegerMath, |
| createConvertQDQToIntegerMathPass) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(mlir::createCSEPass); |
| mainPassManager.addPass( |
| GlobalOptimization::createConvertStridedContractionToContractionPass()); |
| |
| // Enable data tiling after they are in a canonical form. |
| if (transformOptions.dataTiling) { |
| FunctionLikeNest(mainPassManager) |
| .addPass(DispatchCreation::createAnnotateDataTilingHintsPass) |
| .addPass([&]() { |
| return DispatchCreation::createSetEncodingPass( |
| DispatchCreation::SetEncodingPassOptions{clSetEncodingStrategy}); |
| }); |
| if (clEnableEarlyMaterialization) { |
| mainPassManager.addPass(createMaterializeHomogeneousEncodingsPass()); |
| } |
| mainPassManager.addPass(IREE::Flow::createCanonicalizePass()); |
| mainPassManager.addPass(createCSEPass()); |
| mainPassManager.addPass(createSimplifyPackUnpackPass()); |
| FunctionLikeNest(mainPassManager).addPass(createDataLayoutPropagationPass); |
| } |
| |
| FunctionLikeNest(mainPassManager) |
| // Generalize transposes and any other remaining named linalg ops that can |
| // now be represented as generics. |
| .addPass(createGeneralizeLinalgNamedOpsPass) |
| // Hoist loop invariants (e.g. from scf loops) with zero-trip-check. |
| .addPass(createGlobalLoopInvariantCodeMotionPass) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(mlir::createCSEPass) |
| |
| // Simplify util.global accesses early on; this can help with dispatch |
| // region formation as redundant store-loads are removed. |
| .addPass(IREE::Util::createSimplifyGlobalAccessesPass) |
| |
| // Aggressive cleanup. |
| .addPass(IREE::Util::createApplyPatternsPass); |
| |
| // Module level cleanup and canonicalization of util.global (and other |
| // util ops). |
| mainPassManager.addPass(IREE::Util::createFoldGlobalsPass()); |
| mainPassManager.addPass(IREE::Util::createIPOPass()); |
| |
| FunctionLikeNest(mainPassManager) |
| .addPass(IREE::Util::createOptimizeIntArithmeticPass) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(createCSEPass); |
| |
| if (transformOptions.constExprHoisting) { |
| buildGlobalOptExprHoistingPassPipeline(mainPassManager, transformOptions); |
| } |
| |
| if (transformOptions.buildConstEvalPassPipeline) { |
| transformOptions.buildConstEvalPassPipeline(mainPassManager); |
| } |
| |
| if (transformOptions.numericPrecisionReduction) { |
| mainPassManager.addPass(createInferNumericNarrowingPass()); |
| mainPassManager.addPass(createOptimizeNumericsPass()); |
| mainPassManager.addPass(createCleanupNumericNarrowingPass()); |
| } |
| |
| FunctionLikeNest(mainPassManager) |
| .addPass(IREE::Flow::createCanonicalizePass) |
| .addPass(mlir::createCSEPass) |
| // After running const-eval to a fixed point and folding unit extent dims, |
| // try any new raising opportunities. |
| .addPass(createRaiseSpecialOpsPass); |
| |
| // Export after const-eval. If the user wants to keep the input constants |
| // as is in the final parameter archive, they will probably want to disable |
| // const-eval, or could run this pass as preprocessing. There might be a |
| // configuration in the future where users want to limit const-eval to smaller |
| // constants that aren't exported and skip it for larger parameters, but this |
| // is a sensible place for the common case of wanting const-eval in the final |
| // artifact + archive. |
| if (!transformOptions.parameterExportPath.empty()) { |
| IREE::IO::Parameters::ExportParametersPassOptions exportParametersOptions; |
| exportParametersOptions.scopePath = transformOptions.parameterExportPath; |
| exportParametersOptions.minimumSize = |
| transformOptions.parameterExportMinimumSize; |
| mainPassManager.addPass(IREE::IO::Parameters::createExportParametersPass( |
| exportParametersOptions)); |
| } |
| |
| if (!transformOptions.parameterSplatPath.empty()) { |
| IREE::IO::Parameters::GenerateSplatParameterArchivePassOptions |
| generateSplatOptions; |
| generateSplatOptions.filePath = transformOptions.parameterSplatPath; |
| mainPassManager.addPass( |
| IREE::IO::Parameters::createGenerateSplatParameterArchivePass( |
| generateSplatOptions)); |
| } |
| } |
| |
| namespace { |
| #define GEN_PASS_REGISTRATION |
| #include "iree/compiler/GlobalOptimization/Passes.h.inc" // IWYU pragma: export |
| } // namespace |
| |
| void registerGlobalOptimizationPipeline() { |
| registerPasses(); |
| |
| PassPipelineRegistration<TransformOptions> |
| globalOptimizationTransformPassPipeline( |
| "iree-global-optimization-transformation-pipeline", |
| "Runs the IREE global optimization transformation pipeline", |
| [](OpPassManager &passManager, |
| const TransformOptions &transformOptions) { |
| buildGlobalOptimizationPassPipeline(passManager, transformOptions); |
| }); |
| PassPipelineRegistration<TransformOptions> |
| globalOptimizationConstantHoistingPassPipeline( |
| "iree-global-optimization-hoist-constant-expressions", |
| "Hoists constant expressions with the preferred storage types for " |
| "global optimization", |
| [](OpPassManager &passManager, |
| const TransformOptions &transformOptions) { |
| buildGlobalOptExprHoistingPassPipeline(passManager, |
| transformOptions); |
| }); |
| } |
| |
| } // namespace mlir::iree_compiler::GlobalOptimization |
