// Copyright 2019 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 <algorithm>
#include <numeric>
#include <optional>

#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
#include "iree/compiler/Dialect/Util/IR/ClosureOpUtils.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "iree/compiler/Utils/ShapeUtils.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LogicalResult.h"

namespace mlir::iree_compiler::IREE::Flow {

//===----------------------------------------------------------------------===//
// Folding utilities
//===----------------------------------------------------------------------===//

namespace {

// Erases an op if it has no uses.
// This is to support ops that are "pure" but can't be marked as such because
// the MLIR CSE pass would deduplicate them.
template <typename Op>
struct ElideUnusedOp : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    if (!op.use_empty()) {
      return failure();
    }
    rewriter.eraseOp(op);
    return success();
  }
};

// Returns true if |value| is definitely empty at runtime.
static bool isTensorZeroElements(Value value) {
  auto type = dyn_cast<ShapedType>(value.getType());
  if (!type) {
    return false;
  }
  // Any static dimension being zero is definitely empty.
  for (int64_t i = 0; i < type.getRank(); ++i) {
    int64_t dim = type.getDimSize(i);
    if (dim == 0) {
      return true;
    }
  }
  return false; // may still be dynamically empty
}

// Returns true if |value| is definitely empty at runtime.
// Returns false if the value is definitely not empty or may be empty at runtime
// (one or more dynamic dimensions).
static bool isTensorOperandZeroElements(Value value) {
  return isTensorZeroElements(value);
}

// Returns true if |value| is definitely empty at runtime.
// Returns false if the value is definitely not empty or may be empty at runtime
// (one or more dynamic dimensions).
static bool isTensorResultZeroElements(Value value) {
  return isTensorZeroElements(value);
}

template <typename Op, int OperandIdx, int ResultIdx = 0>
struct ReplaceOpIfTensorOperandZeroElements : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    auto operand = op->getOperand(OperandIdx);
    if (!isTensorOperandZeroElements(operand)) {
      return failure();
    }
    auto result = op->getResult(ResultIdx);
    auto dynamicDims = op.getResultDynamicDims(result.getResultNumber());
    rewriter.replaceOpWithNewOp<IREE::Flow::TensorEmptyOp>(op, result.getType(),
                                                           dynamicDims);
    return success();
  }
};

template <typename Op, int ResultIdx>
struct ReplaceOpIfTensorResultZeroElements : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    auto result = op->getResult(ResultIdx);
    if (!isTensorResultZeroElements(result)) {
      return failure();
    }
    auto dynamicDims = op.getResultDynamicDims(result.getResultNumber());
    rewriter.replaceOpWithNewOp<IREE::Flow::TensorEmptyOp>(op, result.getType(),
                                                           dynamicDims);
    return success();
  }
};

template <typename Op, int OperandIdx, int ResultIdx = 0>
struct ReplaceOpIfTensorOperandEmpty : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    auto operand = op->getOperand(OperandIdx);
    auto emptyOp = dyn_cast_if_present<TensorEmptyOp>(operand.getDefiningOp());
    if (!emptyOp) {
      return failure();
    }
    auto result = op->getResult(ResultIdx);
    auto dynamicDims = op.getResultDynamicDims(result.getResultNumber());
    rewriter.replaceOpWithNewOp<IREE::Flow::TensorEmptyOp>(op, result.getType(),
                                                           dynamicDims);
    return success();
  }
};

// Returns a new set of dynamic dimensions for a shape carrying op when a type
// is being changed. This attempts to reuse the existing dimension values if
// they are available and will drop/insert new ones as required.
static SmallVector<Value> refreshDimsOnTypeChange(Operation *op, Type oldType,
                                                  Type newType,
                                                  ValueRange oldDims,
                                                  PatternRewriter &rewriter) {
  if (oldType == newType) {
    return llvm::to_vector(oldDims);
  }

  // Build an expanded list of all the dims - constants will be nullptr.
  // This lets us map back the new types without worrying about whether some
  // subset become static or dynamic.
  auto oldShapedType = cast<ShapedType>(oldType);
  SmallVector<Value> allOldDims(oldShapedType.getRank());
  for (unsigned i = 0; i < oldShapedType.getRank(); ++i) {
    if (oldShapedType.isDynamicDim(i)) {
      allOldDims[i] = oldDims.front();
      oldDims = oldDims.drop_front();
    }
  }

  auto newShapedType = cast<ShapedType>(newType);
  SmallVector<Value> newDims;
  for (unsigned i = 0; i < newShapedType.getRank(); ++i) {
    if (newShapedType.isDynamicDim(i)) {
      auto oldValue = allOldDims[i];
      if (oldValue) {
        // Old value valid; reuse.
        newDims.push_back(oldValue);
      } else {
        // Dimension has changed to be dynamic; insert a constant to use.
        // This sometimes happens during folding of casts and usually is cleaned
        // up pretty quickly.
        newDims.push_back(rewriter.createOrFold<arith::ConstantIndexOp>(
            op->getLoc(), oldShapedType.getDimSize(i)));
      }
    }
  }
  return newDims;
}

} // namespace

//===----------------------------------------------------------------------===//
// flow.dispatch.workgroups
//===----------------------------------------------------------------------===//

/// Helper method to take a list of values to be deduped and returns
/// - list of deduped values.
/// - mapping for a value from its position in the original list to
///   the deduped list.
static std::tuple<SmallVector<Value>, llvm::MapVector<int, int>>
dedupAndGetOldToNewPosMapping(ValueRange values) {
  llvm::MapVector<int, int> oldPosToNewPos;
  SmallVector<Value> uniquedList;
  int numUnique = 0;
  llvm::MapVector<Value, int> oldValueToNewPos;
  for (auto [index, val] : llvm::enumerate(values)) {
    if (oldValueToNewPos.count(val)) {
      oldPosToNewPos[index] = oldValueToNewPos[val];
      continue;
    }
    oldPosToNewPos[index] = numUnique;
    oldValueToNewPos[val] = numUnique;
    uniquedList.push_back(val);
    numUnique++;
  }
  return {uniquedList, oldPosToNewPos};
}

struct ReplaceDispatchResultIfZeroElements
    : OpRewritePattern<DispatchWorkgroupsOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DispatchWorkgroupsOp op,
                                PatternRewriter &rewriter) const override {
    // NOTE: we only look at used results; if unused then closure optimization
    // will drop it.
    bool didReplaceAny = false;
    for (auto result : op.getResults()) {
      if (result.use_empty()) {
        continue;
      }
      if (isTensorResultZeroElements(result)) {
        auto dynamicDims = op.getResultDynamicDims(result.getResultNumber());
        auto emptyOp = IREE::Flow::TensorEmptyOp::create(
            rewriter, result.getLoc(), result.getType(), dynamicDims);
        rewriter.replaceAllUsesWith(result, emptyOp);
        didReplaceAny = true;
      }
    }
    return didReplaceAny ? success() : failure();
  }
};

/// Deduplicate redundant workload values of a dispatch.workgroups op. This
/// requires modifying the `count` region of the op to match the new workloads.
struct ElideRedundantWorkloadValues : OpRewritePattern<DispatchWorkgroupsOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DispatchWorkgroupsOp op,
                                PatternRewriter &rewriter) const override {
    ValueRange workload = op.getWorkload();
    auto [newWorkload, oldWorkloadPosToNewWorkloadPos] =
        dedupAndGetOldToNewPosMapping(workload);
    if (newWorkload.size() == workload.size()) {
      // Nothing to do.
      return failure();
    }

    // Create a new flow.dispatch.workgroup op with new workloads.
    Location loc = op.getLoc();
    auto newWorkgroupsOp = DispatchWorkgroupsOp::create(
        rewriter, loc, newWorkload, op.getResultTypes(), op.getResultDims(),
        op.getArguments(), op.getArgumentDims(),
        op.getTiedOperandsAsIntegerList(),
        getPrunedAttributeList(op, /*elidedAttrs=*/{}));

    // Move the body over.
    Region &body = op.getWorkgroupBody();
    if (!body.empty()) {
      Region &newBody = newWorkgroupsOp.getWorkgroupBody();
      rewriter.inlineRegionBefore(body, newBody, newBody.begin());
    }

    // Move the workgroup count region over.
    Region &count = op.getWorkgroupCount();
    if (!count.empty()) {
      Region &newCount = newWorkgroupsOp.getWorkgroupCount();
      rewriter.inlineRegionBefore(count, newCount, newCount.begin());

      // Create a new entry basic block with as many arguments as the workload
      // and then merge this block with the original entry block.
      auto newWorkloadTypes =
          llvm::map_to_vector(newWorkload, [](Value v) { return v.getType(); });
      auto newWorkloadLocs =
          llvm::map_to_vector(newWorkload, [](Value v) { return v.getLoc(); });
      Block *oldCountBlock = &newCount.front();
      Block *newCountBlock = rewriter.createBlock(
          &newCount.front(), newWorkloadTypes, newWorkloadLocs);
      auto newCountBlockArgs = newCountBlock->getArguments();
      SmallVector<Value> replacements;
      replacements.resize(oldCountBlock->getNumArguments());
      for (auto [index, val] : llvm::enumerate(oldCountBlock->getArguments())) {
        replacements[index] =
            newCountBlockArgs[oldWorkloadPosToNewWorkloadPos.lookup(index)];
      }
      rewriter.mergeBlocks(oldCountBlock, newCountBlock, replacements);
    }

    // Replace the old workgroups op with the new workgroups op.
    rewriter.replaceOp(op, newWorkgroupsOp.getResults());
    return success();
  }
};

/// Deduplicate operands of the `dispatch.workgroup_count_from_slice` op. This
/// requires updating the `iree_tensor_ext.dispatch.workload.ordinal` operation
/// in the body of the `dispatch.workgroups` op to match the new positions of
/// the operands in the `dispatch.workgroup_count_from_slice`.
struct ElideRedundantOperandsOfWorkgroupCountFromSliceOp
    : OpRewritePattern<DispatchWorkgroupsOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DispatchWorkgroupsOp op,
                                PatternRewriter &rewriter) const override {
    Region &count = op.getWorkgroupCount();
    if (count.empty()) {
      return failure();
    }

    assert(
        llvm::hasSingleElement(count) &&
        "expected dispatch.workgroup op count region to have a single block");

    // Check for `dispatch.workgroup_count_from_slice` operations in the count
    // region.
    Block &countBody = count.front();
    auto countFromSliceOps =
        countBody.getOps<IREE::TensorExt::DispatchWorkgroupCountFromSliceOp>();
    if (countFromSliceOps.empty()) {
      return failure();
    }
    assert(llvm::hasSingleElement(countFromSliceOps) &&
           "expected only one dispatch.workgroup_count_from_slice op in count "
           "region");
    auto countFromSliceOp = *countFromSliceOps.begin();

    // Deduplicate the operands and get a mapping from old position to new
    // position.
    auto [newOrdinals, oldOrdinalPosToNewOrdinalPos] =
        dedupAndGetOldToNewPosMapping(countFromSliceOp.getOperands());
    if (newOrdinals.size() == countFromSliceOp.getNumOperands()) {
      return failure();
    }

    // Replace the old `dispatch.workgroup_count_from_slice` with a new op
    // with deduped operands.
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPoint(countFromSliceOp);
    rewriter
        .replaceOpWithNewOp<IREE::TensorExt::DispatchWorkgroupCountFromSliceOp>(
            countFromSliceOp, newOrdinals);

    // Adjust the iree_tensor_ext.dispatch.workload.ordinal ops in the body to
    // use the new ordinal numbers.
    Region &body = op.getWorkgroupBody();
    SmallVector<IREE::TensorExt::DispatchWorkloadOrdinalOp> ordinalOps;
    body.walk([&](IREE::TensorExt::DispatchWorkloadOrdinalOp ordinalOp) {
      ordinalOps.push_back(ordinalOp);
    });

    for (auto ordinalOp : ordinalOps) {
      int oldOrdinalPos = ordinalOp.getOrdinal().getSExtValue();
      rewriter.setInsertionPoint(ordinalOp);
      rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchWorkloadOrdinalOp>(
          ordinalOp, ordinalOp.getOperand(),
          rewriter.getIndexAttr(
              oldOrdinalPosToNewOrdinalPos.lookup(oldOrdinalPos)));
    }
    rewriter.modifyOpInPlace(op, []() {});
    return success();
  }
};

void DispatchWorkgroupsOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  // Disable constant inlining as we have done it during dispatch region
  // formation.
  IREE::Util::ClosureOptimizationOptions closureOptions;
  closureOptions.maxInlinedConstantBytes = 0;
  results.insert<IREE::Util::ClosureOptimizationPattern<DispatchWorkgroupsOp>>(
      context, closureOptions);
  results.insert<ElideRedundantWorkloadValues,
                 ElideRedundantOperandsOfWorkgroupCountFromSliceOp,
                 ReplaceDispatchResultIfZeroElements>(context);
}

//===----------------------------------------------------------------------===//
// flow.dispatch.tie_shape
//===----------------------------------------------------------------------===//

OpFoldResult DispatchTieShapeOp::fold(FoldAdaptor operands) {
  if (getDynamicDims().empty()) {
    return getOperand();
  }
  return {};
}

//===----------------------------------------------------------------------===//
// flow.dispatch
//===----------------------------------------------------------------------===//

namespace {

struct DeduplicateDispatchEntryRefs final : OpRewritePattern<DispatchOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DispatchOp dispatchOp,
                                PatternRewriter &rewriter) const override {
    auto originalAttr = dispatchOp.getEntryPointsAttr();
    auto newAttr = deduplicateArrayElements(originalAttr);
    if (newAttr == originalAttr) {
      return failure();
    }
    rewriter.modifyOpInPlace(dispatchOp,
                             [&]() { dispatchOp.setEntryPointsAttr(newAttr); });
    return success();
  }
};

} // namespace

void DispatchOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
  results.insert<DeduplicateDispatchEntryRefs>(context);
}

//===----------------------------------------------------------------------===//
// Tensor ops
//===----------------------------------------------------------------------===//

/// Reduces the provided multidimensional index into a flattended 1D row-major
/// index. The |type| is expected to be statically shaped (as all constants
/// are).
static uint64_t getFlattenedIndex(ShapedType type, ArrayRef<uint64_t> index) {
  assert(type.hasStaticShape() && "for use on statically shaped types only");
  auto rank = type.getRank();
  auto shape = type.getShape();
  uint64_t valueIndex = 0;
  uint64_t dimMultiplier = 1;
  for (int i = rank - 1; i >= 0; --i) {
    valueIndex += index[i] * dimMultiplier;
    dimMultiplier *= shape[i];
  }
  return valueIndex;
}

//===----------------------------------------------------------------------===//
// flow.tensor.constant
//===----------------------------------------------------------------------===//

OpFoldResult TensorConstantOp::fold(FoldAdaptor operands) { return getValue(); }

//===----------------------------------------------------------------------===//
// flow.tensor.dynamic_constant
//===----------------------------------------------------------------------===//

OpFoldResult TensorDynamicConstantOp::fold(FoldAdaptor operands) {
  auto dynamicType = getType();
  if (dynamicType.getNumDynamicDims() == 0) {
    return getValue();
  }
  return {};
}

namespace {

struct ExpandDynamicShapeConstant : OpRewritePattern<TensorDynamicConstantOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorDynamicConstantOp op,
                                PatternRewriter &rewriter) const override {
    auto constantOp = IREE::Flow::TensorConstantOp::create(
        rewriter, op.getLoc(), op.getValue());
    auto dynamicType = op.getType();
    auto staticType = cast<ShapedType>(op.getValue().getType());
    SmallVector<Value> dynamicDims;
    for (int64_t i = 0; i < dynamicType.getRank(); ++i) {
      if (dynamicType.isDynamicDim(i)) {
        auto dimValue = arith::ConstantIndexOp::create(rewriter, op.getLoc(),
                                                       staticType.getDimSize(i))
                            .getResult();
        dynamicDims.push_back(IREE::Util::OptimizationBarrierOp::create(
                                  rewriter, op.getLoc(), dimValue)
                                  .getResult(0));
      }
    }
    rewriter.replaceOpWithNewOp<IREE::Flow::TensorReshapeOp>(
        op, dynamicType, constantOp.getResult(), dynamicDims);
    return success();
  }
};

} // namespace

void TensorDynamicConstantOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.insert<ExpandDynamicShapeConstant>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.tie_shape
//===----------------------------------------------------------------------===//

OpFoldResult TensorTieShapeOp::fold(FoldAdaptor operands) {
  if (getDynamicDims().empty()) {
    return getOperand();
  }
  return {};
}

void TensorTieShapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorTieShapeOp, 0>>(
      context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.reshape
//===----------------------------------------------------------------------===//

OpFoldResult TensorReshapeOp::fold(FoldAdaptor operands) {
  auto sourceType = cast<ShapedType>(getSource().getType());
  auto resultType = cast<ShapedType>(getResult().getType());
  if (sourceType.getElementType() != resultType.getElementType()) {
    // Element type mismatch, this is a bitcast.
    return {};
  }
  if (compareShapesEqual(sourceType, getSourceDims(), resultType,
                         getResultDims())) {
    // Shapes match and this is a no-op so just fold to the source.
    return getSource();
  }
  return {};
}

namespace {

// Flatten a chain of reshapes or bitcasts (reshape/bitcast feeding into
// reshape or bitcast) such that a reshape only ever pulls from a non-reshape
// source. This prevents big useless chains and makes it easier to track the
// original storage for the tensor.
template <typename CastOpTy>
struct FlattenTensorCastLikeChain : OpRewritePattern<CastOpTy> {
  using OpRewritePattern<CastOpTy>::OpRewritePattern;
  LogicalResult matchAndRewrite(CastOpTy reshapeOp,
                                PatternRewriter &rewriter) const override {
    // We want the same result value/shape but to source from the ancestor. We
    // need to pull any dynamic dims from that as we don't care about the
    // intermediate reshapes.
    Value source;
    ValueRange sourceDims;
    if (auto sourceOp = dyn_cast_if_present<TensorReshapeOp>(
            reshapeOp.getSource().getDefiningOp())) {
      source = sourceOp.getSource();
      sourceDims = sourceOp.getSourceDims();
    } else if (auto sourceOp = dyn_cast_if_present<TensorBitCastOp>(
                   reshapeOp.getSource().getDefiningOp())) {
      source = sourceOp.getSource();
      sourceDims = sourceOp.getSourceDims();
    }
    if (!source) {
      return failure();
    }

    auto sourceType = cast<ShapedType>(source.getType());
    auto resultType = cast<ShapedType>(reshapeOp.getResult().getType());

    // If the element types don't match, this is a bitcast, else we can use
    // reshape.
    if (sourceType.getElementType() != resultType.getElementType()) {
      rewriter.replaceOpWithNewOp<TensorBitCastOp>(
          reshapeOp, reshapeOp.getResult().getType(), source, sourceDims,
          reshapeOp.getResultDims());
    } else {
      rewriter.replaceOpWithNewOp<TensorReshapeOp>(
          reshapeOp, reshapeOp.getResult().getType(), source, sourceDims,
          reshapeOp.getResultDims());
    }
    return success();
  }
};

struct ResolveShapedRank : OpRewritePattern<tensor::RankOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::RankOp op,
                                PatternRewriter &rewriter) const override {
    auto shapedType = cast<ShapedType>(op.getTensor().getType());
    rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op,
                                                        shapedType.getRank());
    return success();
  }
};

struct ResolveShapedDim : OpRewritePattern<tensor::DimOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::DimOp op,
                                PatternRewriter &rewriter) const override {
    if (!op.getConstantIndex().has_value()) {
      return rewriter.notifyMatchFailure(
          op, "non-constant index dim ops are unsupported");
    }
    auto idx = op.getConstantIndex().value();

    // Fold static dims from the type.
    auto shapedType = cast<ShapedType>(op.getSource().getType());
    if (!shapedType.isDynamicDim(idx)) {
      rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(
          op, shapedType.getDimSize(idx));
      return success();
    }

    // Find dims captured on shape-aware ops.
    auto dynamicDims = IREE::Util::findDynamicDims(
        op.getSource(), op->getBlock(), Block::iterator(op.getOperation()));
    if (dynamicDims.has_value()) {
      unsigned dimOffset = 0;
      for (unsigned i = 0; i < idx; ++i) {
        if (shapedType.isDynamicDim(i)) {
          ++dimOffset;
        }
      }
      rewriter.replaceOp(op, dynamicDims.value()[dimOffset]);
      return success();
    }

    return rewriter.notifyMatchFailure(op, "no dynamic dims found/usable");
  }
};

} // namespace

void TensorReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorReshapeOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorResultZeroElements<TensorReshapeOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorOperandEmpty<TensorReshapeOp, 0, 0>>(context);
  results.insert<FlattenTensorCastLikeChain<TensorReshapeOp>>(context);
  results.insert<ResolveShapedRank>(context);
  results.insert<ResolveShapedDim>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.bitcast
//===----------------------------------------------------------------------===//

OpFoldResult TensorBitCastOp::fold(FoldAdaptor operands) {
  auto sourceType = cast<ShapedType>(getSource().getType());
  auto resultType = cast<ShapedType>(getResult().getType());
  if (sourceType.getElementType() != resultType.getElementType()) {
    // Element type mismatch, this is a bitcast.
    return {};
  }
  if (compareShapesEqual(sourceType, getSourceDims(), resultType,
                         getResultDims())) {
    // Shapes match and this is a no-op so just fold to the source.
    return getSource();
  }
  return {};
}

void TensorBitCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorBitCastOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorResultZeroElements<TensorBitCastOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorOperandEmpty<TensorBitCastOp, 0, 0>>(context);
  results.insert<FlattenTensorCastLikeChain<TensorBitCastOp>>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.load
//===----------------------------------------------------------------------===//

OpFoldResult TensorLoadOp::fold(FoldAdaptor operands) {
  if (auto source = dyn_cast_if_present<ElementsAttr>(operands.getSource())) {
    // Load directly from the constant source tensor.
    if (llvm::count(operands.getIndices(), nullptr) == 0) {
      return source.getValues<Attribute>()[llvm::map_to_vector(
          operands.getIndices(), [](Attribute value) {
            return cast<IntegerAttr>(value).getValue().getZExtValue();
          })];
    }
  }
  return {};
}

namespace {

// Replace `flow.tensor.splat`-`flow.tensor.load` op-pairs by the input
// primitive value for the splat op.
struct FoldSplatLoadIntoPrimitive : OpRewritePattern<TensorLoadOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorLoadOp loadOp,
                                PatternRewriter &rewriter) const override {
    auto sourceOp =
        dyn_cast_if_present<TensorSplatOp>(loadOp.getSource().getDefiningOp());
    if (!sourceOp) {
      return failure();
    }
    rewriter.replaceOp(loadOp, sourceOp.getValue());
    return success();
  }
};

} // namespace

void TensorLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.insert<FoldSplatLoadIntoPrimitive>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.store
//===----------------------------------------------------------------------===//

OpFoldResult TensorStoreOp::fold(FoldAdaptor operands) {
  auto value = operands.getValue();
  if (!value) {
    return {};
  }
  if (auto target = dyn_cast_if_present<ElementsAttr>(operands.getTarget())) {
    // Store into the constant target tensor.
    auto targetType = cast<ShapedType>(target.getType());
    if (targetType.getRank() == 0) {
      return DenseElementsAttr::get(targetType, {value});
    }
    if (llvm::count(operands.getIndices(), nullptr) == 0) {
      uint64_t offset = getFlattenedIndex(
          targetType,
          llvm::map_to_vector(operands.getIndices(), [](Attribute value) {
            return cast<IntegerAttr>(value).getValue().getZExtValue();
          }));
      SmallVector<Attribute, 16> newContents(target.getValues<Attribute>());
      newContents[offset] = value;
      return DenseElementsAttr::get(targetType, newContents);
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// flow.tensor.alloca
//===----------------------------------------------------------------------===//

void TensorAllocaOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
  results.insert<ElideUnusedOp<TensorAllocaOp>>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.empty
//===----------------------------------------------------------------------===//

namespace {

// Folds constant dynamic dimension operands of a flow.tensor.empty into the
// static result shape, e.g. `flow.tensor.empty : tensor<?xi32>{%c4}` becomes
// `flow.tensor.empty : tensor<4xi32>`. A cast is inserted to preserve the
// original (dynamic) result type for existing users; it folds away once
// those users can accept the refined static type.
struct FoldTensorEmptyConstantDims : public OpRewritePattern<TensorEmptyOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorEmptyOp op,
                                PatternRewriter &rewriter) const override {
    auto resultType = cast<RankedTensorType>(op.getResult().getType());
    ValueRange dynamicDims = op.getResultDims();

    SmallVector<int64_t> newShape(resultType.getShape());
    // Collect only dimensions that remain dynamic after folding constants.
    SmallVector<Value> newDynamicDims;
    unsigned dimIdx = 0;
    bool didFold = false;
    for (unsigned i = 0, e = resultType.getRank(); i < e; ++i) {
      // Skip dimensions that are already static.
      if (!resultType.isDynamicDim(i)) {
        continue;
      }
      // Get the SSA operand corresponding to dynamic dimension.
      Value dim = dynamicDims[dimIdx++];
      std::optional<int64_t> constantDim = getConstantIntValue(dim);
      if (constantDim && *constantDim >= 0) {
        // Fold the constant value into the tensor's static shape.
        newShape[i] = *constantDim;
        didFold = true;
      } else {
        newDynamicDims.push_back(
            dim); // This dim stays dynamic, so preserve its SSA operand.
      }
    }
    if (!didFold) {
      return rewriter.notifyMatchFailure(op, "no constant dynamic dims");
    }

    auto newType = RankedTensorType::get(newShape, resultType.getElementType(),
                                         resultType.getEncoding());
    Value newEmpty =
        TensorEmptyOp::create(rewriter, op.getLoc(), newType, newDynamicDims);
    // Preserve the original (dynamic) result type for existing users via a
    // cast; it folds away once those users accept the refined static type.
    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, newEmpty);
    return success();
  }
};

} // namespace

void TensorEmptyOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
  results.insert<FoldTensorEmptyConstantDims>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.splat
//===----------------------------------------------------------------------===//

namespace {

struct FoldSplatReshapeIntoSplat : OpRewritePattern<TensorReshapeOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                PatternRewriter &rewriter) const override {
    auto splatOp = dyn_cast_if_present<TensorSplatOp>(
        reshapeOp.getSource().getDefiningOp());
    if (!splatOp) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<TensorSplatOp>(
        reshapeOp, reshapeOp.getResult().getType(), splatOp.getValue(),
        reshapeOp.getResultDims());
    return success();
  }
};

} // namespace

void TensorSplatOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
  // TODO(benvanik): canonicalize splat+slice to smaller splat.
  results.insert<ReplaceOpIfTensorResultZeroElements<TensorSplatOp, 0>>(
      context);
  results.insert<FoldSplatReshapeIntoSplat>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.clone
//===----------------------------------------------------------------------===//

OpFoldResult TensorCloneOp::fold(FoldAdaptor operands) {
  if (auto operand = operands.getOperand()) {
    // Constants always fold.
    return operand;
  }

  // TODO(benvanik): elide clones when safe to do so. Right now clone is
  // load-bearing to work around our lack of cross-stream scheduling. Clones are
  // inserted to avoid mutating function arguments and any logic we perform here
  // (without *also* checking all the conditions that may insert a clone) will
  // just fight.
  //
  // Once the clones are not load-bearing we can remove them in all the normal
  // cases (one user, no intervening uses between clone and consumers of
  // operands, etc).

  return {};
}

void TensorCloneOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorCloneOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorOperandEmpty<TensorCloneOp, 0, 0>>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.barrier
//===----------------------------------------------------------------------===//

namespace {

// Moves (or clones) cast-like ops after a transfer operation and possibly
// updates the transfer op to match the new type.
//
// The cast-like behavior makes more sense near consumers: the producer does not
// need the information as it has already extracted it if required before
// converting into flow ops and if there are no consumers of the transfer that
// cast-like op would have been removed anyway.
//
// Example:
//  %reshape = flow.tensor.reshape %source : tensor<1x2xf32> -> tensor<2x1xf32>
//  %target = flow.tensor.transfer %reshape : tensor<2x1xf32> to "foo"
// ->
//  %target = flow.tensor.transfer %source : tensor<1x2xf32> to "foo"
//  %reshape = flow.tensor.reshape %target : tensor<1x2xf32> -> tensor<2x1xf32>
template <typename TransferOpT, typename CastOpT>
struct SinkCastLikeOpAcrossTransfer final : OpRewritePattern<TransferOpT> {
  using OpRewritePattern<TransferOpT>::OpRewritePattern;
  LogicalResult matchAndRewrite(TransferOpT transferOp,
                                PatternRewriter &rewriter) const override {
    auto sourceOp =
        dyn_cast_if_present<CastOpT>(transferOp.getOperand().getDefiningOp());
    if (!sourceOp) {
      return rewriter.notifyMatchFailure(
          transferOp, "source op not available or does not match");
    }
    Value originValue = sourceOp.getSource();
    ValueRange originDims = sourceOp.getSourceDims();
    auto newOp = TransferOpT::create(rewriter, transferOp.getLoc(), originValue,
                                     originDims, transferOp.getTargetAttr());
    IRMapping mapper;
    mapper.map(originValue, newOp.getResult());
    rewriter.setInsertionPointAfter(newOp);
    auto clonedOp =
        cast<CastOpT>(rewriter.clone(*sourceOp.getOperation(), mapper));
    rewriter.replaceAllUsesExcept(transferOp.getResult(), clonedOp.getResult(),
                                  newOp);
    rewriter.eraseOp(transferOp);
    return success();
  }
};

} // namespace

void TensorBarrierOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
  results.add<SinkCastLikeOpAcrossTransfer<IREE::Flow::TensorBarrierOp,
                                           IREE::Flow::TensorBitCastOp>>(
      context);
  results.add<SinkCastLikeOpAcrossTransfer<IREE::Flow::TensorBarrierOp,
                                           IREE::Flow::TensorReshapeOp>>(
      context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.transfer
//===----------------------------------------------------------------------===//

namespace {

// Attempts to identify trivial cases where we locally recognize that a tensor
// is transferred to the same context it's already on. This does not look across
// control flow edges or globals and is mostly for simplifying IR that may come
// in with a transfer on every single tensor.
struct ElideRedundantTransfer : OpRewritePattern<TensorTransferOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorTransferOp op,
                                PatternRewriter &rewriter) const override {
    auto baseValue =
        IREE::Util::TiedOpInterface::findTiedBaseValue(op.getOperand());
    if (auto transferOp = dyn_cast_if_present<IREE::Flow::TensorTransferOp>(
            baseValue.getDefiningOp())) {
      if (transferOp.getTarget() == op.getTarget()) {
        rewriter.replaceOp(op, op.getOperand());
        return success();
      }
    }
    return failure();
  }
};

// Attempts to identify trivial case of chained transfer ops (A -> B -> C) and
// rewrite it as (A -> C). Writes it as A -> B and A -> C relying on dead code
// elimination to remove the unused A -> B transfer.
struct ElideIntermediateTransfer final : OpRewritePattern<TensorTransferOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorTransferOp targetTransferOp,
                                PatternRewriter &rewriter) const override {
    auto sourceTransferOp = dyn_cast_if_present<IREE::Flow::TensorTransferOp>(
        targetTransferOp.getOperand().getDefiningOp());
    if (!sourceTransferOp) {
      return failure();
    }
    if (sourceTransferOp.getTarget() == targetTransferOp.getTarget()) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<IREE::Flow::TensorTransferOp>(
        targetTransferOp, targetTransferOp->getResultTypes(),
        sourceTransferOp.getOperand(), targetTransferOp.getOperandDims(),
        targetTransferOp.getTarget());
    return success();
  }
};

} // namespace

void TensorTransferOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                   MLIRContext *context) {
  results.add<ElideRedundantTransfer>(context);
  results.add<ElideIntermediateTransfer>(context);
  results.add<SinkCastLikeOpAcrossTransfer<IREE::Flow::TensorTransferOp,
                                           IREE::Flow::TensorBitCastOp>>(
      context);
  results.add<SinkCastLikeOpAcrossTransfer<IREE::Flow::TensorTransferOp,
                                           IREE::Flow::TensorReshapeOp>>(
      context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.slice
//===----------------------------------------------------------------------===//

// Slices tensor from start to (start + length) exclusively at dim.
static ElementsAttr tensorSlice(ElementsAttr tensor, uint64_t dim,
                                uint64_t start, uint64_t length) {
  auto tensorType = cast<ShapedType>(tensor.getType());
  auto shape = llvm::to_vector(tensorType.getShape());
  if (length == shape[dim]) {
    // No need to slice.
    return tensor;
  }
  auto outputShape = shape;
  outputShape[dim] = length;
  auto outputType =
      RankedTensorType::get(outputShape, getElementTypeOrSelf(tensor));
  llvm::SmallVector<Attribute> newContents;
  newContents.reserve(outputType.getNumElements());
  auto valuesBegin = tensor.getValues<Attribute>().begin();
  int64_t step = llvm::product_of(llvm::drop_begin(shape, dim));
  int64_t num = length * step / shape[dim];
  for (int64_t offset = step / shape[dim] * start,
               numElements = tensorType.getNumElements();
       offset < numElements; offset += step) {
    newContents.append(valuesBegin + offset, valuesBegin + offset + num);
  }
  return DenseElementsAttr::get(outputType, newContents);
}

OpFoldResult TensorSliceOp::fold(FoldAdaptor operands) {
  if (llvm::count(operands.getOperands(), nullptr) == 0) {
    // Ignore DenseResources for now and do not perform folding on them.
    if (isa<DenseResourceElementsAttr>(operands.getSource())) {
      return {};
    }
    // Fully constant arguments so we can perform the slice here.
    auto tensor = dyn_cast<ElementsAttr>(operands.getSource());
    if (!tensor) {
      return {};
    }
    int64_t rank = cast<ShapedType>(getSource().getType()).getRank();
    auto start =
        llvm::map_to_vector(operands.getStartIndices(), [](Attribute value) {
          return cast<IntegerAttr>(value).getValue().getZExtValue();
        });
    auto length =
        llvm::map_to_vector(operands.getLengths(), [](Attribute value) {
          return cast<IntegerAttr>(value).getValue().getZExtValue();
        });
    for (int64_t dim = 0; dim < rank; ++dim) {
      tensor = tensorSlice(tensor, dim, start[dim], length[dim]);
    }
    return tensor;
  }
  return {};
}

namespace {

// Folds consecutive flow.tensor.slice ops by composing their offsets.
//
// Example:
//   %src -> slice[1,1] -> slice[0,1]
//      becomes
//   %src -> slice[1,2]
//
// fused_start = inner_start + outer_start
// fused_length = outer_length
struct FoldTensorSliceOfSlice : OpRewritePattern<TensorSliceOp> {
  using OpRewritePattern<TensorSliceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TensorSliceOp consumerSlice,
                                PatternRewriter &rewriter) const override {
    auto producerSlice =
        consumerSlice.getSource().getDefiningOp<TensorSliceOp>();
    if (!producerSlice) {
      return failure();
    }
    if (!producerSlice->hasOneUse()) {
      return failure();
    }

    auto loc =
        rewriter.getFusedLoc({producerSlice.getLoc(), consumerSlice.getLoc()});

    SmallVector<Value> fusedStartIndices;
    for (auto [producerStart, consumerStart] :
         llvm::zip(producerSlice.getStartIndices(),
                   consumerSlice.getStartIndices())) {
      fusedStartIndices.push_back(rewriter.createOrFold<arith::AddIOp>(
          loc, producerStart, consumerStart));
    }

    rewriter.replaceOpWithNewOp<TensorSliceOp>(
        consumerSlice, consumerSlice.getResult().getType(),
        producerSlice.getSource(), producerSlice.getSourceDims(),
        fusedStartIndices, consumerSlice.getLengths(),
        consumerSlice.getResultDims());

    return success();
  }
};

} // namespace

void TensorSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
  results.insert<FoldTensorSliceOfSlice>(context);
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorSliceOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorResultZeroElements<TensorSliceOp, 0>>(
      context);
  results.insert<ReplaceOpIfTensorOperandEmpty<TensorSliceOp, 0, 0>>(context);
}

//===----------------------------------------------------------------------===//
// flow.tensor.update
//===----------------------------------------------------------------------===//

static ElementsAttr tensorUpdate(ElementsAttr update, ElementsAttr target,
                                 ArrayRef<Attribute> startIndicesAttrs) {
  auto updateType = cast<ShapedType>(update.getType());
  auto targetType = cast<ShapedType>(target.getType());
  // If either target or update has zero element, then no update happens.
  if (updateType.getNumElements() == 0 || targetType.getNumElements() == 0) {
    return target;
  }

  int64_t rank = targetType.getRank();
  // If target is scalar, update is also scalar and is the new content.
  if (rank == 0) {
    return update;
  }

  auto startIndex = llvm::map_to_vector(startIndicesAttrs, [](Attribute value) {
    return cast<IntegerAttr>(value).getValue().getZExtValue();
  });
  auto targetValues = llvm::to_vector(target.getValues<Attribute>());
  // target indices start from startIndicesAttrs and update indices start from
  // all zeros.
  llvm::SmallVector<uint64_t> targetIndex(startIndex);
  llvm::SmallVector<uint64_t> updateIndex(rank, 0);
  int64_t numElements = updateType.getNumElements();
  while (numElements--) {
    targetValues[getFlattenedIndex(targetType, targetIndex)] =
        update.getValues<Attribute>()[updateIndex];
    // Increment the index at last dim.
    ++updateIndex.back();
    ++targetIndex.back();
    // If the index in dim j exceeds dim size, reset dim j and
    // increment dim (j-1).
    for (int64_t j = rank - 1;
         j >= 0 && updateIndex[j] >= updateType.getDimSize(j); --j) {
      updateIndex[j] = 0;
      targetIndex[j] = startIndex[j];
      if (j - 1 >= 0) {
        ++updateIndex[j - 1];
        ++targetIndex[j - 1];
      }
    }
  }
  return DenseElementsAttr::get(targetType, targetValues);
}

OpFoldResult TensorUpdateOp::fold(FoldAdaptor operands) {
  bool allIndicesConstant =
      llvm::count(operands.getStartIndices(), nullptr) == 0;
  if (operands.getUpdate() && operands.getTarget() && allIndicesConstant) {
    // Fully constant arguments so we can perform the update here.
    auto update = dyn_cast<ElementsAttr>(operands.getUpdate());
    auto target = dyn_cast<ElementsAttr>(operands.getTarget());
    if (!update || !target) {
      return {};
    }
    return tensorUpdate(update, target, operands.getStartIndices());
  } else {
    // Replace the entire tensor when the sizes match.
    auto updateType = cast<ShapedType>(getUpdate().getType());
    auto targetType = cast<ShapedType>(getTarget().getType());
    if (updateType.hasStaticShape() && targetType.hasStaticShape() &&
        updateType == targetType) {
      return getUpdate();
    }
  }
  return {};
}

namespace {

// When the target tensor is a result of a tensor.cast operation, the op needs
// to be updated to use the source of the cast as the target tensor.
struct FoldTensorUpdateOpWithCasts : OpRewritePattern<TensorUpdateOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorUpdateOp updateOp,
                                PatternRewriter &rewriter) const override {
    auto targetCastOp = updateOp.getTarget().getDefiningOp<tensor::CastOp>();
    auto updateCastOp = updateOp.getUpdate().getDefiningOp<tensor::CastOp>();
    if (!targetCastOp && !updateCastOp) {
      return failure();
    }
    Value target = (targetCastOp ? cast<Value>(targetCastOp.getSource())
                                 : cast<Value>(updateOp.getTarget()));
    Value update = (updateCastOp ? cast<Value>(updateCastOp.getSource())
                                 : cast<Value>(updateOp.getUpdate()));
    auto newOp = TensorUpdateOp::create(
        rewriter, updateOp.getLoc(), target.getType(), target,
        refreshDimsOnTypeChange(updateOp, updateOp.getTarget().getType(),
                                target.getType(), updateOp.getTargetDims(),
                                rewriter),
        updateOp.getStartIndices(), update,
        refreshDimsOnTypeChange(updateOp, updateOp.getUpdate().getType(),
                                update.getType(), updateOp.getUpdateDims(),
                                rewriter));
    rewriter.replaceOpWithNewOp<tensor::CastOp>(
        updateOp, updateOp.getResult().getType(), newOp.getResult());
    return success();
  }
};

struct ReplaceOpIfTensorUpdateOperandZeroElements
    : OpRewritePattern<TensorUpdateOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorUpdateOp op,
                                PatternRewriter &rewriter) const override {
    auto operand = op.getUpdate();
    if (!isTensorOperandZeroElements(operand)) {
      return failure();
    }
    rewriter.replaceOp(op, op.getTarget());
    return success();
  }
};

} // namespace

void TensorUpdateOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
  results.insert<FoldTensorUpdateOpWithCasts>(context);
  // target:
  results.insert<ReplaceOpIfTensorOperandZeroElements<TensorUpdateOp, 0>>(
      context);
  // update:
  results.insert<ReplaceOpIfTensorUpdateOperandZeroElements>(context);
}

//===----------------------------------------------------------------------===//
// flow.channel.split
//===----------------------------------------------------------------------===//

void ChannelSplitOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
  results.insert<ElideUnusedOp<ChannelSplitOp>>(context);
}

} // namespace mlir::iree_compiler::IREE::Flow
