// 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/Codegen/Dialect/Codegen/IR/UKernelOps.h"

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/Support/LLVM.h"

// clang-format off
#define GET_OP_CLASSES
#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp.inc" // IWYU pragma: keep
// clang-format on

namespace mlir::iree_compiler::IREE::Codegen {

//===---------------------------------------------------------------------===//
// Helpers
//===---------------------------------------------------------------------===//

/// Helper method to generate a function declaration at a module scope,
/// and a call to that function
static FailureOr<func::CallOp>
createFunctionCall(RewriterBase &rewriter, Operation *op, StringRef fnName,
                   TypeRange callArgumentTypes, TypeRange callReturnTypes,
                   ValueRange callOperands,
                   ArrayRef<NamedAttribute> fnDefAttrs) {
  FunctionType functionType =
      rewriter.getFunctionType(callArgumentTypes, callReturnTypes);

  // Create a declaration for the function type.
  Location loc = op->getLoc();
  auto moduleOp = SymbolTable::getNearestSymbolTable(op);
  // Check for duplicates.
  auto fnDecl = dyn_cast_if_present<func::FuncOp>(
      SymbolTable::lookupSymbolIn(moduleOp, fnName));
  if (!fnDecl) {
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPointToStart(&moduleOp->getRegion(0).front());
    fnDecl = func::FuncOp::create(rewriter, loc, fnName, functionType);
    SymbolTable::setSymbolVisibility(fnDecl, SymbolTable::Visibility::Private);
    for (auto attr : fnDefAttrs) {
      fnDecl->setAttr(attr.getName(), attr.getValue());
    }
    fnDecl->setAttr("llvm.bareptr", rewriter.getBoolAttr(true));
  } else if (fnDecl.getFunctionType() != functionType) {
    return rewriter.notifyMatchFailure(
        op, llvm::formatv("mismatch in function type computed during lowering "
                          "({0}) and already declared function ({1})",
                          functionType, fnDecl.getFunctionType()));
  }

  // Insert the function call.
  auto callOp = func::CallOp::create(rewriter, loc, fnDecl, callOperands);
  if (op->hasAttr("hal.executable.objects")) {
    callOp->setAttr("hal.executable.objects",
                    op->getAttr("hal.executable.objects"));
  }
  return callOp;
}

//===---------------------------------------------------------------------===//
// UKernelGenericOp
//===---------------------------------------------------------------------===//

/// Map type of operand of a `iree_codegen.ukernel.generic` operation to
/// the type(s) of the function call arguments(s) it lowers to.
static LogicalResult getCallOpType(MLIRContext *context,
                                   Type microKernelOpOperandType,
                                   ArrayRef<int64_t> stridedDims,
                                   SmallVector<Type> &callOperandTypes) {
  return TypeSwitch<Type, LogicalResult>(microKernelOpOperandType)
      .Case<FloatType, IndexType, IntegerType>([&](auto scalarType) {
        callOperandTypes.push_back(scalarType);
        return success();
      })
      .Case([&](MemRefType memrefType) {
        // Base ptr.
        callOperandTypes.push_back(MemRefType::get(
            ArrayRef<int64_t>{}, memrefType.getElementType(),
            MemRefLayoutAttrInterface{}, memrefType.getMemorySpace()));
        // Offset
        auto indexType = IndexType::get(context);
        callOperandTypes.push_back(indexType);
        // Strides.
        callOperandTypes.resize(callOperandTypes.size() + stridedDims.size(),
                                indexType);
        return success();
      })
      .Case([&](NullPointerType nullPointerType) {
        callOperandTypes.push_back(nullPointerType);
        callOperandTypes.push_back(IndexType::get(context));
        return success();
      })
      .Default(failure());
}

/// Map `operand` of a `ukernel.generic` operation to the operand(s) of
/// the function call it lowers to.
static LogicalResult lowerToCallOperands(Location loc, RewriterBase &rewriter,
                                         Value operand,
                                         ArrayRef<int64_t> stridedDims,
                                         SmallVector<Value> &callOperands) {
  return TypeSwitch<Type, LogicalResult>(operand.getType())
      .Case<FloatType, IndexType, IntegerType>([&](auto scalarType) {
        callOperands.push_back(operand);
        return success();
      })
      .Case([&](MemRefType memrefType) {
        auto extractStridedMetadataOp =
            memref::ExtractStridedMetadataOp::create(rewriter, loc, operand);
        // Base ptr.
        callOperands.push_back(extractStridedMetadataOp.getBaseBuffer());
        // Offset.
        callOperands.push_back(extractStridedMetadataOp.getOffset());
        // Strides.
        auto strides = extractStridedMetadataOp.getStrides();
        for (int64_t dim : stridedDims) {
          callOperands.push_back(strides[dim]);
        }
        return success();
      })
      .Case([&](NullPointerType /*unused*/) {
        callOperands.push_back(operand);
        callOperands.push_back(
            arith::ConstantIndexOp::create(rewriter, loc, 0));
        return success();
      })
      .Default(failure());
}

static FailureOr<func::CallOp>
lowerUKernelGenericToFunctionCall(RewriterBase &rewriter,
                                  IREE::Codegen::UKernelGenericOp op,
                                  StringRef fnName) {
  // Create the function type based on the operands and results.
  SmallVector<Type> callArgumentTypes;
  for (auto [idx, microKernelOpOperandType] :
       llvm::enumerate(op->getOperandTypes())) {
    if (failed(getCallOpType(rewriter.getContext(), microKernelOpOperandType,
                             op.getOperandStridedDims(idx),
                             callArgumentTypes))) {
      return rewriter.notifyMatchFailure(
          op, llvm::formatv("failed to lower operand type {}",
                            microKernelOpOperandType));
    }
  }
  SmallVector<Type> callResultTypes;
  for (auto [idx, resultType] : llvm::enumerate(op->getResultTypes())) {
    if (isa<ShapedType>(resultType)) {
      return rewriter.notifyMatchFailure(
          op, "cannot lower a `ShapedType` return value to function call");
    }
    if (failed(getCallOpType(rewriter.getContext(), resultType,
                             /*stridedDims=*/{}, callResultTypes))) {
      return rewriter.notifyMatchFailure(
          op, llvm::formatv("failed to lower result type {}", resultType));
    }
  }

  // Get the operands for the function call.
  SmallVector<Value> callOperands;
  for (auto [idx, operand] : llvm::enumerate(op->getOperands())) {
    if (failed(lowerToCallOperands(op->getLoc(), rewriter, operand,
                                   op.getOperandStridedDims(idx),
                                   callOperands))) {
      return rewriter.notifyMatchFailure(
          op, "failed to lower operands to function call operands");
    }
  }
  ArrayRef<NamedAttribute> fnDefAttrs = {};
  if (auto specifiedfnDefAttrs = op.getFnDefAttrs()) {
    fnDefAttrs = specifiedfnDefAttrs->getValue();
  }
  return createFunctionCall(rewriter, op, fnName, callArgumentTypes,
                            callResultTypes, callOperands, fnDefAttrs);
}

void UKernelGenericOp::build(OpBuilder &builder, OperationState &result,
                             TypeRange resultTypes, StringRef uKernelFnName,
                             ValueRange inputs, ValueRange outputs,
                             ValueRange otherOperands,
                             DictionaryAttr fnDefAttrs,
                             int64_t numStridedOuterDims) {
  SmallVector<Attribute> stridedOuterDimsAttrs;
  auto appendStridedOuterDims = [&](ValueRange operands) {
    for (Value operand : operands) {
      if (!isa<ShapedType>(operand.getType())) {
        continue;
      }
      auto outerDims = builder.getI64ArrayAttr(
          llvm::to_vector(llvm::seq<int64_t>(numStridedOuterDims)));
      stridedOuterDimsAttrs.push_back(outerDims);
    }
  };
  appendStridedOuterDims(inputs);
  appendStridedOuterDims(outputs);
  appendStridedOuterDims(otherOperands);
  auto stridedDimsArrayAttr =
      ArrayAttr::get(builder.getContext(), stridedOuterDimsAttrs);
  build(builder, result, resultTypes, uKernelFnName, inputs, outputs,
        otherOperands, fnDefAttrs, stridedDimsArrayAttr);
}

void UKernelGenericOp::build(OpBuilder &builder, OperationState &result,
                             TypeRange resultTypes, StringRef uKernelFnName,
                             ValueRange inputs, ValueRange outputs,
                             ValueRange otherOperands,
                             DictionaryAttr fnDefAttrs,
                             ArrayRef<SmallVector<int64_t>> stridedDims) {
  SmallVector<Attribute> stridedDimsAttrs = llvm::map_to_vector(
      stridedDims, [&](ArrayRef<int64_t> dims) -> Attribute {
        return builder.getI64ArrayAttr(dims);
      });
  auto stridedDimsArrayAttr =
      ArrayAttr::get(builder.getContext(), stridedDimsAttrs);
  build(builder, result, resultTypes, uKernelFnName, inputs, outputs,
        otherOperands, fnDefAttrs, stridedDimsArrayAttr);
}

SmallVector<int64_t>
UKernelGenericOp::getOperandStridedDims(int64_t operandIdx) {
  auto operandType = dyn_cast<ShapedType>(getOperandTypes()[operandIdx]);
  if (!operandType) {
    return {};
  }
  ArrayAttr dimStrides = getStridedDimsAttr();
  if (!dimStrides) {
    return llvm::to_vector(llvm::seq<int64_t>(operandType.getRank()));
  }
  int64_t dimStridesListIdx = 0;
  for (int i = 0; i < operandIdx; ++i) {
    if (isa<ShapedType>(getOperandTypes()[i])) {
      ++dimStridesListIdx;
    }
  }
  SmallVector<ArrayAttr> strideLists(dimStrides.getAsRange<ArrayAttr>());
  return llvm::map_to_vector(
      strideLists[dimStridesListIdx].getAsRange<IntegerAttr>(),
      [](IntegerAttr dim) { return dim.getInt(); });
}

MutableOperandRange UKernelGenericOp::getDpsInitsMutable() {
  return getOutputsMutable();
}

FailureOr<mlir::CallOpInterface>
UKernelGenericOp::lowerToFunctionCall(RewriterBase &rewriter) {
  return lowerUKernelGenericToFunctionCall(rewriter, *this, getUKernelFnName());
}

void UKernelGenericOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  SmallVector<OpOperand *> readOnlyOperands = getDpsInputOperands();
  llvm::append_range(readOnlyOperands,
                     llvm::make_pointer_range(getOtherOperandsMutable()));
  for (OpOperand *operand : readOnlyOperands) {
    if (!isa<MemRefType>(operand->get().getType())) {
      continue;
    }
    effects.emplace_back(MemoryEffects::Read::get(), operand,
                         SideEffects::DefaultResource::get());
  }
  for (OpOperand &operand : getDpsInitsMutable()) {
    if (!isa<MemRefType>(operand.get().getType())) {
      continue;
    }
    effects.emplace_back(MemoryEffects::Read::get(), &operand,
                         SideEffects::DefaultResource::get());
    effects.emplace_back(MemoryEffects::Write::get(), &operand,
                         SideEffects::DefaultResource::get());
  }
}

} // namespace mlir::iree_compiler::IREE::Codegen

namespace mlir::iree_compiler {

//===---------------------------------------------------------------------===//
// Register bufferization interface.
//===---------------------------------------------------------------------===//

namespace {

template <typename OpTy>
struct UKernelOpsBufferizationInterface
    : bufferization::DstBufferizableOpInterfaceExternalModel<
          UKernelOpsBufferizationInterface<OpTy>, OpTy> {
  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                          const bufferization::BufferizationOptions &options,
                          bufferization::BufferizationState &state) const {
    // TODO: Handle operations with regions if needed.
    if (op->getNumRegions() != 0) {
      op->emitOpError(
          "unhandled bufferization of micro kernel op with regions");
    }
    SmallVector<Value> bufferOpOperands;

    // Replace all `tensor` operands with corresponding `memref` operands.
    for (auto [index, operand] : llvm::enumerate(op->getOperands())) {
      // For `tensor` type operands, replace with `memref` type operand.
      if (isa<RankedTensorType>(operand.getType())) {
        FailureOr<Value> memrefOperand =
            getBuffer(rewriter, operand, options, state);
        if (failed(memrefOperand)) {
          return op->emitOpError(
              llvm::formatv("failed to bufferize operand {} ", index));
        }
        bufferOpOperands.push_back(memrefOperand.value());
        continue;
      }

      // For all other operand types, just use the same value.
      bufferOpOperands.push_back(operand);
    }

    SmallVector<Type> nonTensorResultTypes;
    SmallVector<Value> nonTensorResultValues;
    for (OpResult result : op->getResults()) {
      Type resultType = result.getType();
      if (isa<RankedTensorType>(resultType)) {
        continue;
      }
      nonTensorResultTypes.push_back(resultType);
      nonTensorResultValues.push_back(result);
    }

    auto bufferOp = OpTy::create(rewriter, op->getLoc(), nonTensorResultTypes,
                                 bufferOpOperands, op->getAttrs());
    SmallVector<Value> bufferizedResults =
        cast<DestinationStyleOpInterface>(bufferOp.getOperation())
            .getDpsInits();
    bufferizedResults.append(nonTensorResultValues);
    bufferization::replaceOpWithBufferizedValues(rewriter, op,
                                                 bufferizedResults);
    return success();
  }
};

template <typename... Ops>
struct RegisterUKernelOpsBufferizationInterface {
  static void registerOpInterface(MLIRContext *context) {
    (Ops::template attachInterface<UKernelOpsBufferizationInterface<Ops>>(
         *context),
     ...);
  }
};

} // namespace

void registerUKernelBufferizationInterface(DialectRegistry &registry) {
  registry.addExtension(
      +[](MLIRContext *context, IREE::Codegen::IREECodegenDialect *dialect) {
        RegisterUKernelOpsBufferizationInterface<
#define GET_OP_LIST
#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.cpp.inc"
            >::registerOpInterface(context);
      });
}

} // namespace mlir::iree_compiler
