// Copyright 2022 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/Pipelines/Options.h"
#include "llvm/Passes/OptimizationLevel.h"

IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::BindingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::InputDialectOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(
    mlir::iree_compiler::GlobalOptimizationOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::ParameterOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::SchedulingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::PreprocessingOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::GlobalPipelineOptions);
IREE_DEFINE_COMPILER_OPTION_FLAGS(mlir::iree_compiler::DispatchCreationOptions);

namespace mlir::iree_compiler {

void GlobalPipelineOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE global pipeline options controlling the entire compilation flow.");

  binder.topLevelOpt(
      "iree-opt-level", optLevel,
      llvm::cl::desc("Global optimization level to apply to the entire "
                     "compilation flow."),
      llvm::cl::cat(category));

  binder.opt<bool>(
      "iree-opt-data-tiling", dataTiling,
      llvm::cl::desc(
          "Enables data tiling optimization. There are two data-tiling "
          "paths; it is the global flag that enables data-tiling with the"
          "suggested path. You can choose one of them to enable explicitly, if "
          "you have a preference. See `iree-global-opt-data-tiling` and "
          "`iree-dispatch-creation-data-tiling` for details. Note that this "
          "flag will override the other two flags, because it has a higher "
          "priority."),
      llvm::cl::cat(category));

  binder.opt<bool>(
      "iree-opt-const-expr-hoisting", constExprHoisting,
      llvm::cl::desc(
          "Hoists the results of latent constant expressions into immutable "
          "global initializers for evaluation at program load."),
      llvm::cl::cat(category));
  binder.opt<int64_t>(
      "iree-opt-const-expr-max-size-increase-threshold",
      constExprMaxSizeIncreaseThreshold,
      llvm::cl::desc("Maximum byte size increase allowed for constant expr "
                     "hoisting policy to allow hoisting."),
      llvm::cl::cat(category));
}

void BindingOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE translation binding support options.");

  binder.opt<bool>(
      "iree-native-bindings-support", native,
      llvm::cl::desc(
          "Include runtime support for native IREE ABI-compatible bindings."),
      llvm::cl::cat(category));
  binder.opt<bool>("iree-tflite-bindings-support", tflite,
                   llvm::cl::desc("Include runtime support for the IREE TFLite "
                                  "compatibility bindings."),
                   llvm::cl::cat(category));
}

void InputDialectOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE options for controlling the input transformations to apply.");
  binder.opt<std::string>(
      "iree-input-type", inputTypeMnemonic,
      llvm::cl::desc(
          // clang-format off
          "Specifies the input program representation:\n"
          "  =none          - No input dialect transformation.\n"
          "  =auto          - Analyze the input program to choose conversion.\n"
// NOTE: The plugin system does not have a good way to populate CL help
// messages, so we err on the side of being helpful and populating plugin
// options here, even though it is a layering violation.
#ifdef IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_STABLEHLO
          "  =stablehlo     - Legalize from StableHLO ops (including VHLO deserialization).\n"
          "  =stablehlo_xla - Legalize from StableHLO ops (including VHLO deserialization and XLA de-tupling).\n"
#endif // IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_STABLEHLO
#ifdef IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_TOSA
          "  =tosa          - Legalize from TOSA ops.\n"
#endif  // IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_TOSA
#ifdef IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_TORCH
          "  =tm_tensor     - Legalize a subset of Torch input ops.\n"
          "  =torch         - Legalize from the 'torch' dialect.\n"
#endif  // IREE_COMPILER_PLUGIN_HAVE_STATIC_INPUT_TORCH
          "  =*             - An extensible input type defined in a plugin."
          // clang-format on
          ),
      llvm::cl::cat(category));

  binder.opt<bool>(
      "iree-input-demote-i64-to-i32", demoteI64ToI32,
      llvm::cl::desc("Converts all i64 ops and values into i32 counterparts "
                     "unconditionally before main global optimizations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-input-demote-f32-to-f16", demoteF32ToF16,
      llvm::cl::desc("Converts all f32 ops and values into f16 counterparts "
                     "unconditionally before main global optimizations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-input-demote-f64-to-f32", demoteF64ToF32,
      llvm::cl::desc("Converts all f64 ops and values into f32 counterparts "
                     "unconditionally before main global optimizations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-input-promote-f16-to-f32", promoteF16ToF32,
      llvm::cl::desc("Converts all f16 ops and values into f32 counterparts "
                     "unconditionally before main global optimizations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-input-promote-bf16-to-f32", promoteBF16ToF32,
      llvm::cl::desc("Converts all bf16 ops and values into f32 counterparts "
                     "unconditionally before main global optimizations."),
      llvm::cl::cat(category));
}

InputDialectOptions::Type InputDialectOptions::parseInputTypeMnemonic() {
  if (inputTypeMnemonic == "none") {
    return Type::none;
  } else if (inputTypeMnemonic == "auto") {
    return Type::auto_detect;
  } else {
    return Type::plugin;
  }
}

void PreprocessingOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "Preprocessing options",
      "IREE options for apply custom preprocessing before normal IREE "
      "compilation flow. These options are exercised in the following order:\n"
      " 1. `iree-preprocessing-pass-pipeline`,\n"
      " 2. `iree-preprocessing-transform-spec-filename`,\n"
      " 3. `iree-preprocessing-pdl-spec-filename`");

  binder.opt<std::string>(
      "iree-preprocessing-pass-pipeline", preprocessingPassPipeline,
      llvm::cl::desc("Textual description of the pass pipeline to run before "
                     "running normal IREE compilation pipelines."),
      llvm::cl::cat(category));
  binder.opt<std::string>(
      "iree-preprocessing-transform-spec-filename",
      preprocessingTransformSpecFilename,
      llvm::cl::desc(
          "File name of a transform dialect spec to use for preprocessing."),
      llvm::cl::cat(category));
  binder.opt<std::string>(
      "iree-preprocessing-pdl-spec-filename", preprocessingPDLSpecFilename,
      llvm::cl::desc("File name of a PDL spec to use for preprocessing."),
      llvm::cl::cat(category));
}

void ParameterOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category("IREE Parameter Options");

  // Parameter import/export options.
  binder.list<std::string>(
      "iree-parameter-import", importPaths,
      llvm::cl::desc("File paths to archives to import parameters from with an "
                     "optional `scope=` prefix."),
      llvm::cl::cat(category));
  binder.list<std::string>(
      "iree-parameter-import-keys", importKeys,
      llvm::cl::desc("List of parameter keys to import. Any matching keys from "
                     "any scope will be imported."),
      llvm::cl::cat(category));
  binder.opt<int64_t>(
      "iree-parameter-import-maximum-size", importMaximumSize,
      llvm::cl::desc("Maximum size of parameters to import or 0 to disable "
                     "automatic import."),
      llvm::cl::cat(category));

  binder.opt<std::string>(
      "iree-parameter-export", exportPath,
      llvm::cl::desc("File path to an archive to export parameters to with an "
                     "optional `scope=` prefix."),
      llvm::cl::cat(category));
  binder.opt<int64_t>(
      "iree-parameter-export-minimum-size", exportMinimumSize,
      llvm::cl::desc("Minimum size of constants to export to the parameter "
                     "archive."),
      llvm::cl::cat(category));

  binder.opt<std::string>(
      "iree-parameter-splat", splatPath,
      llvm::cl::desc("File path to create a parameter archive of splat values "
                     "from all parameter backed globals."),
      llvm::cl::cat(category));

  // Parameter encoder options.
  binder.opt<ParameterEncoderMode>(
      "iree-parameter-encoder-mode", encoderMode,
      llvm::cl::desc("Controls how the encoder manages parameters."),
      llvm::cl::values(
          clEnumValN(ParameterEncoderMode::Consolidate, "consolidate",
                     "Merge all encoded and original parameters into a single "
                     "consolidated scope."),
          clEnumValN(ParameterEncoderMode::Overlay, "overlay",
                     "Only produce encoded parameters and leave original "
                     "parameters untouched.")),
      llvm::cl::cat(category));

  binder.opt<std::string>(
      "iree-parameter-encoder-output-file", encoderOutputFile,
      llvm::cl::desc(".mlir/.mlirbc file path to write the split parameter "
                     "encoder module to (empty = disabled)."),
      llvm::cl::cat(category));

  binder.opt<std::string>(
      "iree-parameter-encoder-output-scope", encoderOutputScope,
      llvm::cl::desc("Parameter scope for the encoder output parameters."),
      llvm::cl::cat(category));

  // Deprecated flags aliasing the new ones above.
  binder.opt<std::string>(
      "iree-opt-export-parameters", exportPath,
      Deprecated("use --iree-parameter-export=<path> instead"),
      llvm::cl::Hidden,
      llvm::cl::desc("File path to an archive to export parameters to with an "
                     "optional `scope=` prefix."),
      llvm::cl::cat(category));
  binder.opt<std::string>(
      "iree-opt-splat-parameters", splatPath,
      Deprecated("use --iree-parameter-splat=<path> instead"), llvm::cl::Hidden,
      llvm::cl::desc(
          "File path to create a parameter archive of splat values out of all "
          "parameter backed globals."),
      llvm::cl::cat(category));
}

void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE options for controlling global optimizations.");
  auto initAtOpt = binder.optimizationLevel(
      "iree-global-optimization-opt-level", optLevel,
      llvm::cl::desc("Optimization level for the this pipeline"),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-opt-aggressively-propagate-transposes",
      aggressiveTransposePropagation,
      {initAtOpt(llvm::OptimizationLevel::O0, false),
       initAtOpt(llvm::OptimizationLevel::O3, true)},
      llvm::cl::desc(
          "Propagates transposes to named ops even when the resulting op will "
          "be a linalg.generic"),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-global-opt-propagate-transposes-through-conv",
      propagateTransposesThroughConv,
      llvm::cl::desc(
          "Enables propagation of transpose ops through convolutions."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-global-opt-enable-sink-transpose-through-pad",
      sinkTransposeThroughPad,
      llvm::cl::desc("Enables sinking transpose through pad operations."),
      llvm::cl::cat(category));
  binder.opt<bool>("iree-opt-outer-dim-concat", outerDimConcat,
                   {initAtOpt(llvm::OptimizationLevel::O0, false),
                    initAtOpt(llvm::OptimizationLevel::O1, true)},
                   llvm::cl::desc("Transposes all concatenations to happen"
                                  "along the outer most dimension."),
                   llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-global-opt-data-tiling", dataTiling,
      llvm::cl::desc(
          "Enables data tiling path starting from GlobalOptimization phase."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-opt-const-eval", constEval,
      llvm::cl::desc("Enables eager evaluation of constants using the full "
                     "compiler and runtime (on by default)."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-opt-numeric-precision-reduction", numericPrecisionReduction,
      llvm::cl::desc(
          "Reduces numeric precision to lower bit depths where possible."),
      llvm::cl::cat(category));
  binder.opt<bool>("iree-opt-strip-assertions", stripAssertions,
                   {initAtOpt(llvm::OptimizationLevel::O0, false),
                    initAtOpt(llvm::OptimizationLevel::O1, true)},
                   llvm::cl::desc("Strips debug assertions after any useful "
                                  "information has been extracted."),
                   llvm::cl::cat(category));

  binder.opt<bool>(
      "iree-opt-generalize-matmul", generalizeMatmul,
      {initAtOpt(llvm::OptimizationLevel::O0, false),
       initAtOpt(llvm::OptimizationLevel::O2, true)},
      llvm::cl::desc("Convert named matmul ops to linalg generic ops during "
                     "global optimization to enable better fusion."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-global-opt-use-im2col-for-convs", useIm2colForConvs,
      llvm::cl::desc(
          "Enables converting convolution operations to im2col + "
          "matmul form to leverage optimized matmul implementations."),
      llvm::cl::cat(category));
}

void SchedulingOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE options for controlling host/device scheduling.");

  binder.opt<ExecutionModel>(
      "iree-execution-model", executionModel,
      llvm::cl::desc("Specifies the execution model used for scheduling tensor "
                     "compute operations."),
      llvm::cl::values(
          clEnumValN(
              ExecutionModel::HostOnly, "host-only",
              "Host-local code only that does not need execution scheduling."),
          clEnumValN(ExecutionModel::AsyncInternal, "async-internal",
                     "Full HAL using asynchronous host/device execution "
                     "internally but exporting functions as if synchronous."),
          clEnumValN(ExecutionModel::AsyncExternal, "async-external",
                     "Full HAL using asynchronous host/device execution both "
                     "internally and externally."),
          clEnumValN(ExecutionModel::InlineStatic, "inline-static",
                     "Inline host-local in-process execution with executable "
                     "code statically linked into the host program."),
          clEnumValN(ExecutionModel::InlineDynamic, "inline-dynamic",
                     "Inline host-local in-process execution using dynamic "
                     "executables.")),
      llvm::cl::cat(category));

  binder.opt<InitializationMode>(
      "iree-scheduling-initialization-mode", initializationMode,
      llvm::cl::desc(
          "Specifies the initialization mode for parameters and globals."),
      llvm::cl::values(
          clEnumValN(InitializationMode::Synchronous, "sync",
                     "Synchronously initialize all parameters and globals "
                     "prior to returning from the module initializer."),
          clEnumValN(InitializationMode::Asynchronous, "async",
                     "Asynchronously initialize all parameters and globals and "
                     "return immediately from the module initializer without "
                     "waiting for them to complete.")),
      llvm::cl::cat(category));

  binder.opt<bool>(
      "iree-scheduling-optimize-bindings", optimizeBindings,
      llvm::cl::desc(
          "Enables binding fusion and dispatch site specialization."),
      llvm::cl::cat(category));

  binder.opt<DumpOutputFormat>(
      "iree-scheduling-dump-statistics-format", dumpStatisticsFormat,
      llvm::cl::desc("Dumps statistics in the specified output format."),
      llvm::cl::cat(category),
      llvm::cl::values(
          clEnumValN(DumpOutputFormat::Pretty, "pretty",
                     "Human-readable pretty printed output."),
          clEnumValN(DumpOutputFormat::Verbose, "verbose",
                     "Pretty printed output with additional IR."),
          clEnumValN(DumpOutputFormat::CSV, "csv", "Comma separated values."),
          clEnumValN(DumpOutputFormat::JSON, "json",
                     "JSON output with structures for data exchange")));
  binder.opt<std::string>("iree-scheduling-dump-statistics-file",
                          dumpStatisticsFile,
                          llvm::cl::desc("File path to write statistics to; or "
                                         "`` for stderr or `-` for stdout."),
                          llvm::cl::cat(category));
}

void DispatchCreationOptions::bindOptions(OptionsBinder &binder) {
  static llvm::cl::OptionCategory category(
      "IREE options for controlling dispatch region creation.");
  auto initAtOpt = binder.optimizationLevel(
      "iree-dispatch-creation-opt-level", optLevel,
      llvm::cl::desc("Optimization level for the this pipeline"),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-enable-aggressive-fusion", enableAggressiveFusion,
      {initAtOpt(llvm::OptimizationLevel::O0, false),
       initAtOpt(llvm::OptimizationLevel::O2, true)},
      llvm::cl::desc("Aggressive fusion opportunities that are behind a flag "
                     "since all backends dont support it yet"),
      llvm::cl::cat(category));
  binder.opt<bool>("iree-dispatch-creation-fuse-multi-use", enableFuseMultiUse,
                   llvm::cl::desc("Fuse operations with multiple uses."),
                   llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-data-tiling", dataTiling,
      llvm::cl::desc(
          "Enables data tiling path starting from DispatchCreation phase."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-enable-split-reduction", enableSplitReduction,
      llvm::cl::desc(
          "Enable split-reduction for certain reduction operations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-split-reduction-low-parallelism",
      splitReductionLowParallelism,
      llvm::cl::desc("Use the low-parallelism split-reduction parameter set "
                     "(RDNA-class targets) instead of the high-parallelism "
                     "default (CDNA-class targets)."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-enable-aggressive-reshape-movement",
      enableAggressiveReshapeMovement,
      llvm::cl::desc(
          "Enable aggressive reshape movement (bubbling expand/collapse "
          "shapes across reduction ops)."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-flow-enable-pad-handling", enablePadHandling,
      llvm::cl::desc("Enable native handling of tensor.pad operations."),
      llvm::cl::cat(category));
  binder.opt<bool>(
      "iree-dispatch-creation-enable-fuse-padding-into-linalg-consumer-ops",
      enableFusePaddingIntoLinalgConsumerOps,
      llvm::cl::desc("Enable fusing tensor.pad ops into Linalg consumer ops."),
      llvm::cl::cat(category));
}

} // namespace mlir::iree_compiler
