// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "third_party/llvm/llvm/include/llvm/ADT/DenseSet.h"
#include "third_party/llvm/llvm/include/llvm/ADT/SmallVector.h"
#include "third_party/llvm/llvm/include/llvm/Support/Allocator.h"
#include "third_party/llvm/llvm/include/llvm/Support/Casting.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Dialect/StandardOps/Ops.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Attributes.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/BlockAndValueMapping.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/Location.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/MLIRContext.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/PatternMatch.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/IR/StandardTypes.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/Pass.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Pass/PassRegistry.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/DialectConversion.h"
#include "third_party/llvm/llvm/projects/google_mlir/include/mlir/Transforms/Utils.h"
#include "third_party/mlir_edge/iree/compiler/IR/Dialect.h"
#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLDialect.h"
#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/HLOps.h"
#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/LLDialect.h"
#include "third_party/mlir_edge/iree/compiler/IR/Interpreter/LLOps.h"
#include "third_party/mlir_edge/iree/compiler/IR/Ops.h"
#include "third_party/mlir_edge/iree/compiler/Serialization/BytecodeTables.h"
#include "third_party/mlir_edge/iree/schemas/bytecode/interpreter_bytecode_v0.h"

namespace mlir {
namespace iree_compiler {

namespace {

struct LowerBranchOpPattern
    : public OpRewritePattern<IREEInterp::HL::BranchOp> {
  using OpRewritePattern<IREEInterp::HL::BranchOp>::OpRewritePattern;

  PatternMatchResult matchAndRewrite(IREEInterp::HL::BranchOp op,
                                     PatternRewriter &rewriter) const {
    SmallVector<Value *, 8> operands{op.getOperation()->getOperands()};

    rewriter.replaceOpWithNewOp<IREEInterp::LL::BranchOp>(op, op.getDest(),
                                                          operands);
    return matchSuccess();
  }
};

struct LowerCondCondBranchOpPattern
    : public OpRewritePattern<IREEInterp::HL::CondBranchOp> {
  using OpRewritePattern<IREEInterp::HL::CondBranchOp>::OpRewritePattern;

  PatternMatchResult matchAndRewrite(IREEInterp::HL::CondBranchOp op,
                                     PatternRewriter &rewriter) const {
    SmallVector<Value *, 8> trueOperands{op.getTrueOperands()};
    SmallVector<Value *, 8> falseOperands{op.getFalseOperands()};

    rewriter.replaceOpWithNewOp<IREEInterp::LL::CondBranchOp>(
        op, op.getCondition(), op.getTrueDest(), trueOperands,
        op.getFalseDest(), falseOperands);
    return matchSuccess();
  }
};

// Returns true if the op defined by |opName| (like 'iree_ll_interp.reshape')
// uses output operands for results (like iree_ll_interp.add_i) or returns real
// results.
bool opTakesOutputOperands(llvm::StringRef opName) {
  if (!opName.consume_front("iree_ll_interp.")) {
    assert(false && "op not part of IREE LL Interpreter dialect");
    return false;
  }
  auto opcode = GetInterpreterOpcodeByName(opName.str());
  assert(opcode.hasValue() && "op has no corresponding opcode");
  const auto &info = GetInterpreterOpcodeInfo(opcode.getValue());
  for (auto &operand : info.operands) {
    if (operand == iree::OperandEncoding::kOutputSlot ||
        operand == iree::OperandEncoding::kVariadicOutputSlots) {
      return true;
    }
  }
  return false;
}

template <typename SrcOp, typename DstOp>
class SimpleOpLowering : public OpRewritePattern<SrcOp> {
  using OpRewritePattern<SrcOp>::OpRewritePattern;

  PatternMatchResult matchAndRewrite(SrcOp op,
                                     PatternRewriter &rewriter) const {
    SmallVector<Value *, 8> operands{op.getOperation()->getOperands()};

    // Most ops take results as output operands to populate during execution.
    // Certain ops, like reshape, return references to existing memrefs and
    // should still retain their results.
    if (!opTakesOutputOperands(DstOp::getOperationName())) {
      SmallVector<Type, 8> resultTypes{op.getOperation()->getResultTypes()};

      rewriter.replaceOpWithNewOp<DstOp>(op, resultTypes, operands,
                                         op.getAttrs());
      return this->matchSuccess();
    }

    for (Value *result : op.getOperation()->getResults()) {
      auto memRefType = result->getType().cast<MemRefType>();
      if (!memRefType.hasStaticShape()) {
        // TODO(benvanik): real thing here - dynamic shaping required.
        // This should emit a shape calculation based on the operation. Most
        // are likely simple and by running DCE after this we can clean up
        // parts that are static or unused.
        op.emitOpError() << "uses unsupported dynamic shapes";
        return this->matchFailure();
      }
      ArrayRef<Value *> dim_pieces;
      auto allocOp = rewriter.create<IREEInterp::LL::AllocHeapOp>(
          op.getLoc(), memRefType, dim_pieces);
      operands.push_back(allocOp);
      result->replaceAllUsesWith(allocOp);
    }
    ArrayRef<Type> resultTypes;
    rewriter.create<DstOp>(op.getLoc(), resultTypes, operands, op.getAttrs());
    op.erase();
    return this->matchSuccess();
  }
};

}  // namespace

class LowerInterpreterDialectPass
    : public FunctionPass<LowerInterpreterDialectPass> {
 public:
  void runOnFunction() override {
    OwningRewritePatternList patterns;
    patterns.insert<LowerBranchOpPattern, LowerCondCondBranchOpPattern>(
        &getContext());
    patterns.insert<
        SimpleOpLowering<IREE::ConstantOp, IREEInterp::LL::ConstantOp>,
        SimpleOpLowering<IREEInterp::HL::CopyOp, IREEInterp::LL::DynamicCopyOp>,
        SimpleOpLowering<IREEInterp::HL::SliceOp,
                         IREEInterp::LL::DynamicSliceOp>>(&getContext());
#define SAME_NAME_SIMPLE_PATTERN(op_name) \
  SimpleOpLowering<IREEInterp::HL::op_name, IREEInterp::LL::op_name>
    // clang-format off
    patterns.insert<
        SAME_NAME_SIMPLE_PATTERN(AssignOp),
        SAME_NAME_SIMPLE_PATTERN(AbsFOp),
        SAME_NAME_SIMPLE_PATTERN(AbsIOp),
        SAME_NAME_SIMPLE_PATTERN(AddFOp),
        SAME_NAME_SIMPLE_PATTERN(AddIOp),
        SAME_NAME_SIMPLE_PATTERN(AllocHeapOp),
        SAME_NAME_SIMPLE_PATTERN(AndOp),
        SAME_NAME_SIMPLE_PATTERN(Atan2FOp),
        SAME_NAME_SIMPLE_PATTERN(BreakOp),
        SAME_NAME_SIMPLE_PATTERN(BroadcastOp),
        SAME_NAME_SIMPLE_PATTERN(CallOp),
        SAME_NAME_SIMPLE_PATTERN(CallIndirectOp),
        SAME_NAME_SIMPLE_PATTERN(CeilFOp),
        SAME_NAME_SIMPLE_PATTERN(ClampFOp),
        SAME_NAME_SIMPLE_PATTERN(CloneOp),
        SAME_NAME_SIMPLE_PATTERN(CmpFOp),
        SAME_NAME_SIMPLE_PATTERN(CmpIOp),
        SAME_NAME_SIMPLE_PATTERN(CondAssignOp),
        SAME_NAME_SIMPLE_PATTERN(ConvertSSOp),
        SAME_NAME_SIMPLE_PATTERN(ConvertUUOp),
        SAME_NAME_SIMPLE_PATTERN(ConvertSUOp),
        SAME_NAME_SIMPLE_PATTERN(ConvertUSOp),
        SAME_NAME_SIMPLE_PATTERN(CondBreakOp),
        SAME_NAME_SIMPLE_PATTERN(CosFOp),
        SAME_NAME_SIMPLE_PATTERN(DimOp),
        SAME_NAME_SIMPLE_PATTERN(DivFOp),
        SAME_NAME_SIMPLE_PATTERN(DivISOp),
        SAME_NAME_SIMPLE_PATTERN(DivIUOp),
        SAME_NAME_SIMPLE_PATTERN(ExpFOp),
        SAME_NAME_SIMPLE_PATTERN(LogFOp),
        SAME_NAME_SIMPLE_PATTERN(RsqrtFOp),
        SAME_NAME_SIMPLE_PATTERN(FloorFOp),
        SAME_NAME_SIMPLE_PATTERN(LengthOp),
        SAME_NAME_SIMPLE_PATTERN(MatMulFOp),
        SAME_NAME_SIMPLE_PATTERN(MatMulIOp),
        SAME_NAME_SIMPLE_PATTERN(MaxFOp),
        SAME_NAME_SIMPLE_PATTERN(MaxISOp),
        SAME_NAME_SIMPLE_PATTERN(MaxIUOp),
        SAME_NAME_SIMPLE_PATTERN(MinFOp),
        SAME_NAME_SIMPLE_PATTERN(MinISOp),
        SAME_NAME_SIMPLE_PATTERN(MinIUOp),
        SAME_NAME_SIMPLE_PATTERN(MulAddFOp),
        SAME_NAME_SIMPLE_PATTERN(MulAddIOp),
        SAME_NAME_SIMPLE_PATTERN(MulFOp),
        SAME_NAME_SIMPLE_PATTERN(MulIOp),
        SAME_NAME_SIMPLE_PATTERN(NotOp),
        SAME_NAME_SIMPLE_PATTERN(OrOp),
        SAME_NAME_SIMPLE_PATTERN(PadOp),
        SAME_NAME_SIMPLE_PATTERN(RankOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceSumIOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceSumFOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceMinIOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceMinFOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceMaxIOp),
        SAME_NAME_SIMPLE_PATTERN(ReduceMaxFOp),
        SAME_NAME_SIMPLE_PATTERN(ReshapeOp),
        SAME_NAME_SIMPLE_PATTERN(ReturnOp),
        SAME_NAME_SIMPLE_PATTERN(SelectOp),
        SAME_NAME_SIMPLE_PATTERN(ShapeOp),
        SAME_NAME_SIMPLE_PATTERN(ShiftLeftOp),
        SAME_NAME_SIMPLE_PATTERN(ShiftRightArithmeticOp),
        SAME_NAME_SIMPLE_PATTERN(ShiftRightLogicalOp),
        SAME_NAME_SIMPLE_PATTERN(SinFOp),
        SAME_NAME_SIMPLE_PATTERN(SplitOp),
        SAME_NAME_SIMPLE_PATTERN(SubFOp),
        SAME_NAME_SIMPLE_PATTERN(SubIOp),
        SAME_NAME_SIMPLE_PATTERN(TanhFOp),
        SAME_NAME_SIMPLE_PATTERN(TileOp),
        SAME_NAME_SIMPLE_PATTERN(TraceOp),
        SAME_NAME_SIMPLE_PATTERN(TransposeOp),
        SAME_NAME_SIMPLE_PATTERN(ReverseOp),
        SAME_NAME_SIMPLE_PATTERN(XorOp)>(&getContext());
    // clang-format on
#undef SAME_NAME_SIMPLE_PATTERN

    ConversionTarget target(getContext());
    target.addLegalDialect<IREELLInterpreterDialect>();
    target.addLegalOp<FuncOp, IREE::ReturnOp>();
    if (failed(applyFullConversion(getFunction(), target, patterns))) {
      return signalPassFailure();
    }
  }
};

std::unique_ptr<OpPassBase<FuncOp>> createLowerInterpreterDialectPass() {
  return std::make_unique<LowerInterpreterDialectPass>();
}

static PassRegistration<LowerInterpreterDialectPass> pass(
    "lower-iree-interpreter-hl-to-ll", "Lowers IREE HL ops to IREE LL ops");

}  // namespace iree_compiler
}  // namespace mlir