// 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 "compiler/IR/Dialect.h"
#include "compiler/IR/Ops.h"
#include "compiler/IR/Sequencer/HLDialect.h"
#include "compiler/IR/Sequencer/HLOps.h"
#include "compiler/IR/Sequencer/LLDialect.h"
#include "compiler/IR/Sequencer/LLOps.h"
#include "compiler/IR/StructureOps.h"
#include "compiler/Utils/TypeConversionUtils.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Module.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/StandardTypes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/Utils.h"

namespace mlir {
namespace iree_compiler {

namespace {

template <typename SrcOp>
class SequencerLoweringPattern : public OpConversionPattern<SrcOp> {
 public:
  SequencerLoweringPattern(MLIRContext *context, TypeConverter &typeConverter)
      : OpConversionPattern<SrcOp>(context), typeConverter_(typeConverter) {}

 protected:
  TypeConverter &typeConverter_;
};

// Returns an integer scalar memref containing the offset specified by |indices|
// within |type|.
Value *computeOffset(Location loc, Value *reference, Value *indices,
                     OpBuilder &builder) {
  auto referenceType = reference->getType().cast<ShapedType>();
  auto *shapeMemRef = builder
                          .create<IREESeq::LL::AllocHeapOp>(
                              loc,
                              MemRefType::get({referenceType.getRank()},
                                              builder.getIntegerType(32)),
                              ArrayRef<Value *>{})
                          .getResult();
  builder.create<IREESeq::LL::ShapeOp>(loc, reference, shapeMemRef);
  auto *resultMemRef =
      builder
          .create<IREESeq::LL::AllocHeapOp>(
              loc, MemRefType::get({}, builder.getIntegerType(32)),
              ArrayRef<Value *>{})
          .getResult();
  auto elementSizeAttr = builder.getIntegerAttr(
      builder.getIntegerType(8), referenceType.getElementTypeBitWidth() / 8);
  builder.create<IREESeq::LL::ComputeOffsetOp>(
      loc, shapeMemRef, elementSizeAttr, indices, resultMemRef);
  return resultMemRef;
}

// Returns a tuple of (offset, length) integer scalar memrefs with the range
// specified by |indices| and |lengths| within |type|.
std::pair<Value *, Value *> computeRange(Location loc, Value *reference,
                                         Value *indices, Value *lengths,
                                         OpBuilder &builder) {
  auto referenceType = reference->getType().cast<ShapedType>();
  auto *shapeMemRef = builder
                          .create<IREESeq::LL::AllocHeapOp>(
                              loc,
                              MemRefType::get({referenceType.getRank()},
                                              builder.getIntegerType(32)),
                              ArrayRef<Value *>{})
                          .getResult();
  builder.create<IREESeq::LL::ShapeOp>(loc, reference, shapeMemRef);
  auto *offsetMemRef =
      builder
          .create<IREESeq::LL::AllocHeapOp>(
              loc, MemRefType::get({}, builder.getIntegerType(32)),
              ArrayRef<Value *>{})
          .getResult();
  auto *lengthMemRef =
      builder
          .create<IREESeq::LL::AllocHeapOp>(
              loc, MemRefType::get({}, builder.getIntegerType(32)),
              ArrayRef<Value *>{})
          .getResult();
  auto elementSizeAttr = builder.getIntegerAttr(
      builder.getIntegerType(8), referenceType.getElementTypeBitWidth() / 8);
  builder.create<IREESeq::LL::ComputeRangeOp>(loc, shapeMemRef, elementSizeAttr,
                                              indices, lengths, offsetMemRef,
                                              lengthMemRef);
  return {offsetMemRef, lengthMemRef};
}

struct LowerSliceOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::SliceOp> {
  using SequencerLoweringPattern::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::SliceOp op, ArrayRef<Value *> operands,
      ConversionPatternRewriter &rewriter) const override {
    OperandAdaptor<IREESeq::HL::SliceOp> operandAdaptor(operands);
    auto range = computeRange(op.getLoc(), operandAdaptor.src(),
                              operandAdaptor.indices(),
                              operandAdaptor.lengths(), rewriter);
    rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicSliceOp>(
        op, typeConverter_.convertType(op.getType()),
        ArrayRef<Value *>{operandAdaptor.src(), range.first, range.second},
        op.getAttrs());
    return matchSuccess();
  }
};

struct LowerShapeOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::ShapeOp> {
  using SequencerLoweringPattern::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::ShapeOp op, ArrayRef<Value *> operands,
      ConversionPatternRewriter &rewriter) const override {
    auto *shapeMemRef =
        rewriter
            .create<IREESeq::LL::AllocHeapOp>(
                op.getLoc(),
                MemRefType::get({op.getType().cast<ShapedType>().getRank()},
                                rewriter.getIntegerType(64)),
                ArrayRef<Value *>{})
            .getResult();
    op.replaceAllUsesWith(shapeMemRef);
    rewriter.replaceOpWithNewOp<IREESeq::LL::ShapeOp>(op, operands[0],
                                                      shapeMemRef);
    return matchSuccess();
  }
};

struct LowerCopyOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::CopyOp> {
  using SequencerLoweringPattern::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::CopyOp op, ArrayRef<Value *> operands,
      ConversionPatternRewriter &rewriter) const override {
    OperandAdaptor<IREESeq::HL::CopyOp> operandAdaptor(operands);
    auto *srcOffsetMemRef =
        computeOffset(op.getLoc(), operandAdaptor.src(),
                      operandAdaptor.srcIndices(), rewriter);
    auto dstRange = computeRange(op.getLoc(), operandAdaptor.dst(),
                                 operandAdaptor.dstIndices(),
                                 operandAdaptor.lengths(), rewriter);
    rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicCopyOp>(
        op, operandAdaptor.src(), srcOffsetMemRef, operandAdaptor.dst(),
        dstRange.first, dstRange.second);
    return matchSuccess();
  }
};

struct LowerFillOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::FillOp> {
  using SequencerLoweringPattern::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::FillOp op, ArrayRef<Value *> operands,
      ConversionPatternRewriter &rewriter) const override {
    OperandAdaptor<IREESeq::HL::FillOp> operandAdaptor(operands);
    auto dstRange = computeRange(op.getLoc(), operandAdaptor.dst(),
                                 operandAdaptor.dstIndices(),
                                 operandAdaptor.lengths(), rewriter);
    rewriter.replaceOpWithNewOp<IREESeq::LL::DynamicFillOp>(
        op, operandAdaptor.value(), operandAdaptor.dst(), dstRange.first,
        dstRange.second);
    return matchSuccess();
  }
};

struct LowerBranchOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::BranchOp> {
  using SequencerLoweringPattern<
      IREESeq::HL::BranchOp>::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::BranchOp op, ArrayRef<Value *> properOperands,
      ArrayRef<Block *> destinations, ArrayRef<ArrayRef<Value *>> operands,
      ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREESeq::LL::BranchOp>(op, destinations[0],
                                                       operands[0]);
    return matchSuccess();
  }
};

struct LowerCondCondBranchOpPattern
    : public SequencerLoweringPattern<IREESeq::HL::CondBranchOp> {
  using SequencerLoweringPattern<
      IREESeq::HL::CondBranchOp>::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      IREESeq::HL::CondBranchOp op, ArrayRef<Value *> properOperands,
      ArrayRef<Block *> destinations, ArrayRef<ArrayRef<Value *>> operands,
      ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREESeq::LL::CondBranchOp>(
        op, properOperands[0],
        destinations[IREESeq::HL::CondBranchOp::trueIndex],
        operands[IREESeq::HL::CondBranchOp::trueIndex],
        destinations[IREESeq::HL::CondBranchOp::falseIndex],
        operands[IREESeq::HL::CondBranchOp::falseIndex]);
    return matchSuccess();
  }
};

// Rewrites an op into one with all the same operands, results, and attributes.
// Operands and results in the ops must have the same order and attributes must
// have the same name. They must also be constructed properly by the default
// builders.
template <typename SRC, typename DST>
struct LowerIdenticalOpPattern : public SequencerLoweringPattern<SRC> {
  using SequencerLoweringPattern<SRC>::SequencerLoweringPattern;

  PatternMatchResult matchAndRewrite(
      SRC op, ArrayRef<Value *> operands,
      ConversionPatternRewriter &rewriter) const override {
    SmallVector<Type, 8> originalResultTypes{
        op.getOperation()->getResultTypes()};
    SmallVector<Type, 8> resultTypes;
    if (failed(this->typeConverter_.convertTypes(originalResultTypes,
                                                 resultTypes))) {
      op.emitOpError() << "Failed to convert result types";
      return this->matchFailure();
    }
    rewriter.replaceOpWithNewOp<DST>(op, resultTypes, operands, op.getAttrs());
    return this->matchSuccess();
  }
};

}  // namespace

class LowerSequencerDialectPass : public ModulePass<LowerSequencerDialectPass> {
 public:
  void runOnModule() override {
    auto *ctx = &getContext();
    LLTypeConverter typeConverter(ctx);
    OwningRewritePatternList patterns;
    patterns.insert<
        LowerIdenticalOpPattern<IREE::ConstantOp, IREESeq::LL::ConstantOp>,
        LowerIdenticalOpPattern<IREESeq::HL::DispatchOp,
                                IREESeq::LL::DynamicDispatchOp>,
        LowerShapeOpPattern, LowerCopyOpPattern, LowerSliceOpPattern,
        LowerBranchOpPattern, LowerCondCondBranchOpPattern>(ctx, typeConverter);
#define IDENTICAL_OP_LOWERING(op_name) \
  LowerIdenticalOpPattern<IREESeq::HL::op_name, IREESeq::LL::op_name>
    patterns.insert<
        IDENTICAL_OP_LOWERING(AllocHeapOp), IDENTICAL_OP_LOWERING(CloneOp),
        IDENTICAL_OP_LOWERING(ReshapeOp), IDENTICAL_OP_LOWERING(CallOp),
        IDENTICAL_OP_LOWERING(ReturnOp)>(ctx, typeConverter);
#undef IDENTICAL_OP_LOWERING

    mlir::populateFuncOpTypeConversionPattern(patterns, ctx, typeConverter);
    ConversionTarget target(*ctx);
    target.addLegalDialect<IREELLSequencerDialect>();
    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
      return typeConverter.isSignatureLegal(op.getType());
    });

    // TODO(b/142791494): The conversion framework will recurse into the
    // executable if we just call it on the top-level module. This can't be a
    // function pass because type conversion replaces the original functions.
    auto funcsIt = getModule().getOps<FuncOp>();
    SmallVector<Operation *, 4> funcs(funcsIt.begin(), funcsIt.end());

    if (failed(applyFullConversion(funcs, target, patterns, &typeConverter))) {
      return signalPassFailure();
    }
  }
};

std::unique_ptr<OpPassBase<ModuleOp>> createLowerSequencerDialectPass() {
  return std::make_unique<LowerSequencerDialectPass>();
}

static PassRegistration<LowerSequencerDialectPass> pass(
    "iree-lower-sequencer-dialect",
    "Lowers the IREE HL sequencer dialect to the LL sequencer dialect.");

}  // namespace iree_compiler
}  // namespace mlir