// 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/Dialect/Stream/IR/StreamDialect.h"

#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/InliningUtils.h"

namespace mlir::iree_compiler::IREE::Stream {

namespace {

// Used to control inlining behavior.
struct StreamInlinerInterface : DialectInlinerInterface {
  using DialectInlinerInterface::DialectInlinerInterface;

  bool isLegalToInline(Operation *call, Operation *callable,
                       bool wouldBeCloned) const final {
    // Sure!
    return true;
  }
  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                       IRMapping &valueMapping) const final {
    // Sure!
    return true;
  }

  bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                       IRMapping &valueMapping) const final {
    // Sure!
    return true;
  }
};

struct StreamFolderInterface : DialectFoldInterface {
  using DialectFoldInterface::DialectFoldInterface;

  bool shouldMaterializeInto(Region *region) const override {
    // TODO(benvanik): redirect constants to the region scope when small.
    return false;
  }
};

// Tries to fold away unrealized_conversion_cast ops if the downstream consumers
// don't need the extra information. These are inserted during conversion or
// transforms that may interop with external dialects.
//
// Specifically matches:
//   %0 = builtin.unrealized_conversion_cast %arg0, %arg1 :
//        !stream.resource<transient>, index to !stream.resource<transient>
struct StripResourceConversionCastPattern
    : OpRewritePattern<UnrealizedConversionCastOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(UnrealizedConversionCastOp castOp,
                                PatternRewriter &rewriter) const override {
    auto result = castOp.getResult(0);
    if (!isa<IREE::Stream::ResourceType>(result.getType())) {
      return failure();
    }
    assert(castOp.getNumOperands() == 2 &&
           "expect resource, index -> resource");
    auto resourceValue = castOp.getOperand(0);
    auto sizeValue = castOp.getOperand(1);
    for (auto &use : llvm::make_early_inc_range(result.getUses())) {
      if (auto sizeOp =
              dyn_cast<IREE::Stream::ResourceSizeOp>(use.getOwner())) {
        rewriter.replaceOp(sizeOp, sizeValue);
      } else {
        rewriter.modifyOpInPlace(use.getOwner(),
                                 [&]() { use.set(resourceValue); });
      }
    }
    rewriter.eraseOp(castOp);
    return success();
  }
};

} // namespace

StreamDialect::StreamDialect(MLIRContext *context)
    : Dialect(getDialectNamespace(), context, TypeID::get<StreamDialect>()) {
  context->loadDialect<IREE::Util::UtilDialect>();
  context->loadDialect<mlir::complex::ComplexDialect>();

  registerAttributes();
  registerTypes();

#define GET_OP_LIST
  addOperations<
#include "iree/compiler/Dialect/Stream/IR/StreamOps.cpp.inc"
      >();

  addInterfaces<StreamInlinerInterface>();
  addInterfaces<StreamFolderInterface>();
}

void StreamDialect::getCanonicalizationPatterns(
    RewritePatternSet &results) const {
  results.insert<StripResourceConversionCastPattern>(getContext());
}

Operation *StreamDialect::materializeConstant(OpBuilder &builder,
                                              Attribute value, Type type,
                                              Location loc) {
  if (mlir::func::ConstantOp::isBuildableWith(value, type)) {
    return mlir::func::ConstantOp::create(builder, loc, type,
                                          cast<FlatSymbolRefAttr>(value));
  } else if (arith::ConstantOp::isBuildableWith(value, type)) {
    return arith::ConstantOp::create(builder, loc, type,
                                     cast<TypedAttr>(value));
  } else if (isa<IREE::Stream::TimepointAttr>(value)) {
    return IREE::Stream::TimepointImmediateOp::create(builder, loc);
  }
  return nullptr;
}

} // namespace mlir::iree_compiler::IREE::Stream
