// 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/compiler/Codegen/Interfaces/Interfaces.h"

#include "iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h"
#include "iree/compiler/Codegen/Dialect/GPU/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/Codegen/Dialect/GPU/TransformExtensions/IREEGPUExtensions.h"
#include "iree/compiler/Codegen/Dialect/Map/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/Codegen/Dialect/PCF/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/Codegen/ExternalInterfaces/Interfaces.h"
#include "iree/compiler/Codegen/Interfaces/BufferizationInterfaces.h"
#include "iree/compiler/Codegen/Interfaces/HoistableRegionOpInterface.h"
#include "iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.h"
#include "iree/compiler/Codegen/Interfaces/ProcessorOpInterfaces.h"
#include "iree/compiler/Codegen/Interfaces/TensorMaskingOpInterface.h"
#include "iree/compiler/Codegen/Interfaces/VectorizableOpInterface.h"
#include "iree/compiler/Codegen/LLVMCPU/TransformExtensions/LLVMCPUExtensions.h"
#include "iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.h"
#include "iree/compiler/Dialect/LinalgExt/TransformExtensions/LinalgExtExtensionsOps.h"
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
#include "mlir/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
#include "mlir/Dialect/GPU/Transforms/IndexedAccessOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h"
#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
#include "mlir/Dialect/Vector/Transforms/IndexedAccessOpInterfaceImpl.h"
#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

namespace mlir::iree_compiler {

namespace {

// linalg.generic operations are also loop-like, but they don't have
// LoopLikeOpInterface implemented for them.
// Ideally, we should be checking if the linalg.generic op has a trip
// count of zero, but while that is possible and can be written using
// ValueBoundsConstraintSet, it is usually not needed. Unlike loops,
// which can have arbitrary operations inside them, the loop invariant
// operations inside a linalg.generic operation are usually
// operations performed on scalars. Hoisting scalar constants does not
// have a big cost even if the trip count is zero.
struct LinalgGenericHoistableRegionModel final
    : IREE::Codegen::HoistableRegionOpInterface::ExternalModel<
          LinalgGenericHoistableRegionModel, linalg::GenericOp> {
  SmallVector<Region *> getHoistableRegions(Operation *op) const {
    return {&cast<linalg::GenericOp>(op).getBodyRegion()};
  }

  bool isHoistable(Operation *self, Operation *op) const {
    // linalg.index is marked Pure but is semantically tied to the enclosing
    // iteration space, so it must not be hoisted despite passing the
    // memory-effect and speculatability checks.
    return !isa<linalg::IndexOp>(op) && isMemoryEffectFree(op) &&
           isSpeculatable(op);
  }
};

} // namespace

void registerCodegenInterfaces(DialectRegistry &registry) {
  registerProcessorOpInterfaceExternalModels(registry);
  registerCodegenExternalInterfaces(registry);
  registerIREEGPUExternalInterfaces(registry);
  registerIREEMapExternalInterfaces(registry);
  registerPCFExternalInterfaces(registry);
  registerBufferizationInterfaces(registry);
  registerTensorMaskingOpInterface(registry);
  // TODO: when warranted, move to its own file.
  registry.addExtensions<IREE::LinalgExt::LinalgExtTransformOpsExtension>();
  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
    linalg::GenericOp::attachInterface<LinalgGenericHoistableRegionModel>(*ctx);
  });
  registerPartitionableLoopsInterfaceModels(registry);
  registerVectorizableOpInterfaceExternalModels(registry);
  registerTransformDialectCommonExtension(registry);
  registerTransformDialectIREEGPUExtension(registry);
  registerTransformDialectLLVMCPUExtension(registry);
  registerTransformDialectLLVMGPUExtension(registry);
  linalg::registerTilingInterfaceExternalModels(registry);
  affine::registerTransformDialectExtension(registry);
  affine::registerValueBoundsOpInterfaceExternalModels(registry);
  arith::registerValueBoundsOpInterfaceExternalModels(registry);
  bufferization::registerTransformDialectExtension(registry);
  gpu::registerIndexedAccessOpInterfaceExternalModels(registry);
  gpu::registerValueBoundsOpInterfaceExternalModels(registry);
  gpu::registerTransformDialectExtension(registry);
  gpu::registerValueBoundsOpInterfaceExternalModels(registry);
  linalg::registerTransformDialectExtension(registry);
  linalg::registerValueBoundsOpInterfaceExternalModels(registry);
  linalg::registerSubsetOpInterfaceExternalModels(registry);
  memref::registerTransformDialectExtension(registry);
  memref::registerValueBoundsOpInterfaceExternalModels(registry);
  scf::registerTransformDialectExtension(registry);
  scf::registerValueBoundsOpInterfaceExternalModels(registry);
  tensor::registerSubsetOpInterfaceExternalModels(registry);
  tensor::registerTransformDialectExtension(registry);
  tensor::registerValueBoundsOpInterfaceExternalModels(registry);
  transform::registerLoopExtension(registry);
  vector::registerIndexedAccessOpInterfaceExternalModels(registry);
  vector::registerSubsetOpInterfaceExternalModels(registry);
  vector::registerTransformDialectExtension(registry);
  vector::registerValueBoundsOpInterfaceExternalModels(registry);
}

} // namespace mlir::iree_compiler
