integrations/tensorflow/iree-dialects/tools/iree-dialects-opt/iree-dialects-opt.cpp

// Copyright 2021 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/Input/InputDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Async/IR/Async.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/PDL/IR/PDL.h"
#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
#include "mlir/Dialect/SCF/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/Dialect.h"
#include "mlir/Tools/mlir-opt/MlirOptMain.h"
#include "mlir/Transforms/Passes.h"

using namespace mlir;
namespace IREE = mlir::iree_compiler::IREE;

namespace mlir {
namespace test_ext {
/// Test passes, do not deserve an include.
void registerTestLinalgTransformWrapScope();
} // namespace test_ext
} // namespace mlir

int main(int argc, char **argv) {
  registerAsmPrinterCLOptions();
  registerMLIRContextCLOptions();

  DialectRegistry registry;
  registry.insert<
      // clang-format off
      // Local dialects
      mlir::iree_compiler::IREE::Input::IREEInputDialect,
      // Upstream dialects
      mlir::async::AsyncDialect,
      mlir::arith::ArithDialect,
      mlir::AffineDialect,
      mlir::cf::ControlFlowDialect,
      mlir::func::FuncDialect,
      mlir::linalg::LinalgDialect,
      mlir::memref::MemRefDialect,
      mlir::pdl::PDLDialect,
      mlir::pdl_interp::PDLInterpDialect,
      mlir::scf::SCFDialect,
      mlir::tensor::TensorDialect,
      mlir::transform::TransformDialect
      // clang-format on
      >();

  // Core dialect passes.
  memref::registerMemRefPasses();
  registerTransformsPasses();
  registerSCFPasses();

  // External models.
  mlir::linalg::registerTilingInterfaceExternalModels(registry);

  mlir::linalg::registerTransformDialectExtension(registry);
  mlir::scf::registerTransformDialectExtension(registry);

  return mlir::asMainReturnCode(
      mlir::MlirOptMain(argc, argv, "MLIR modular optimizer driver\n", registry,
                        // Note: without preloading, 3 tests fail atm.
                        /*preloadDialectsInContext=*/true));
}