// Copyright 2024 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/ExternalInterfaces/LinalgExtExternalModels.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"

namespace mlir::iree_compiler {
namespace {
// Used to register the LinalgFusionOpInterface with the linalg ops.
template <typename ConcreteType>
struct LinalgFusionOpInterfaceAdapter
    : IREE::LinalgExt::LinalgFusionOpInterface::ExternalModel<
          LinalgFusionOpInterfaceAdapter<ConcreteType>, ConcreteType> {
public:
  SmallVector<AffineMap> getIndexingMapsForOperands(mlir::Operation *op) const {
    return cast<ConcreteType>(op).getIndexingMapsArray();
  }

  SmallVector<AffineMap> getIndexingMapsForResults(mlir::Operation *op) const {
    auto maps = cast<ConcreteType>(op)
                    .getIndexingMaps()
                    .template getAsValueRange<AffineMapAttr>();
    return {maps.end() - cast<ConcreteType>(op).getNumResults(), maps.end()};
  }

  // Forward all the interface methods to the corresponding linalg op.
  unsigned getNumParallelLoops(mlir::Operation *op) const {
    return (cast<ConcreteType>(op).getNumParallelLoops());
  }

  unsigned getNumLoops(mlir::Operation *op) const {
    return (cast<ConcreteType>(op).getNumLoops());
  }

  SmallVector<int64_t> getStaticLoopRanges(mlir::Operation *op) const {
    return SmallVector<int64_t>(cast<ConcreteType>(op).getStaticLoopRanges());
  }

  AffineMap getIndexingMapMatchingResult(mlir::Operation *op,
                                         OpResult result) const {
    return (cast<ConcreteType>(op).getIndexingMapMatchingResult(result));
  }

  AffineMap getMatchingIndexingMap(mlir::Operation *op,
                                   OpOperand *operand) const {
    return getIndexingMapsForOperands(op)[operand->getOperandNumber()];
  }

  SmallVector<AffineMap> getIndexingMapsArray(mlir::Operation *op) const {
    auto inputMaps = getIndexingMapsForOperands(op);
    llvm::append_range(inputMaps, getIndexingMapsForResults(op));
    return inputMaps;
  }
};

struct SoftmaxFusionOpInterfaceAdapter
    : IREE::LinalgExt::LinalgFusionOpInterface::ExternalModel<
          SoftmaxFusionOpInterfaceAdapter, linalg::SoftmaxOp> {
public:
  SmallVector<AffineMap> getIndexingMapsForOperands(mlir::Operation *op) const {
    Builder b(op->getContext());
    return llvm::map_to_vector(op->getOperands(), [&b](Value operand) {
      auto rank = cast<ShapedType>(operand.getType()).getRank();
      return b.getMultiDimIdentityMap(rank);
    });
  }

  SmallVector<AffineMap> getIndexingMapsForResults(mlir::Operation *op) const {
    Builder b(op->getContext());
    return llvm::map_to_vector(
        cast<linalg::SoftmaxOp>(op).getDpsInits(),
        [&b](Value operand) -> AffineMap {
          auto rank = cast<ShapedType>(operand.getType()).getRank();
          return b.getMultiDimIdentityMap(rank);
        });
  }

  SmallVector<int64_t> getStaticLoopRanges(Operation *op) const {
    auto softmaxOp = cast<linalg::SoftmaxOp>(op);
    // Softmax loop range is the input shape.
    return SmallVector<int64_t>(softmaxOp.getInputOperandType().getShape());
  }

  AffineMap getIndexingMapMatchingResult(mlir::Operation *op,
                                         OpResult result) const {
    return getIndexingMapsForResults(op)[result.getResultNumber()];
  }

  AffineMap getMatchingIndexingMap(mlir::Operation *op,
                                   OpOperand *operand) const {
    return getIndexingMapsForOperands(op)[operand->getOperandNumber()];
  }

  SmallVector<AffineMap> getIndexingMapsArray(mlir::Operation *op) const {
    auto inputMaps = getIndexingMapsForOperands(op);
    llvm::append_range(inputMaps, getIndexingMapsForResults(op));
    return inputMaps;
  }
};

template <typename... Args>
void registerOpsWithLinalgExtOpInterface(mlir::MLIRContext *context) {
  (Args::template attachInterface<LinalgFusionOpInterfaceAdapter<Args>>(
       *context),
   ...);
}

} // namespace

void registerLinalgExtExternalModels(DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
#define GET_OP_LIST
    registerOpsWithLinalgExtOpInterface<
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
        >(ctx);
    linalg::SoftmaxOp::attachInterface<SoftmaxFusionOpInterfaceAdapter>(*ctx);
  });
  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
    IREE::LinalgExt::registerConcatOpTilingInterfaceExternalModel(ctx);
  });
}

} // namespace mlir::iree_compiler
