// 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> 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 &registry) {
    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)
      .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);
                         })
      .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
