// 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

//===- FoldUnitExtentDims.cpp - Pass to fold unit extent dims of tensors -===//
//
// Light weight wrapper to call the patterns to fold unit extent dims with
// IREE control.
//
//===----------------------------------------------------------------------===//

#include "iree/compiler/Dialect/Flow/IR/FlowTypes.h"
#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
#include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
#include "iree/compiler/Dialect/Util/Analysis/Explorer.h"
#include "iree/compiler/DispatchCreation/Passes.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-dispatch-creation-fold-unit-extent-dims"

namespace mlir::iree_compiler::DispatchCreation {

#define GEN_PASS_DEF_FOLDUNITEXTENTDIMSPASS
#define GEN_PASS_DEF_FOLDUNITEXTENTDIMSFORFUNCPASS
#include "iree/compiler/DispatchCreation/Passes.h.inc"

//===----------------------------------------------------------------------===//
// Attention Pattern
//===----------------------------------------------------------------------===//

namespace {
/// Simplify collapse_shape(expand_shape) by removing unneeded unit dimensions
/// that get expanded and subsequently collapsed.
///
/// TODO: move this upstream with the other reshape folding patterns. This can
/// also be generalized to fold non-unit dimensions.
///
/// For example:
/// ```
/// %0 = expand_shape ... tensor<3x3x10xf32> into tensor<3x3x5x1x2xf32>
/// %1 = collapse_shape %1 ... tensor<3x3x5x1x2xf32> into tensor<9x5x2xf32>
/// ```
/// simplifies to
/// ```
/// %0 = expand_shape ... tensor<3x3x10xf32> into tensor<3x3x5x2xf32>
/// %1 = collapse_shape %1 ... tensor<3x3x5x2xf32> into tensor<9x5x2xf32>
/// ```
struct DropUnitDimsFromCollapseOfExpand
    : OpRewritePattern<tensor::CollapseShapeOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(tensor::CollapseShapeOp collapseOp,
                                PatternRewriter &rewriter) const override {
    auto expandOp = collapseOp.getSrc().getDefiningOp<tensor::ExpandShapeOp>();
    if (!expandOp) {
      return failure();
    }

    const SmallVector<ReassociationIndices, 4> collapseReassoc =
        collapseOp.getReassociationIndices();
    ArrayRef<int64_t> interShape = expandOp.getType().getShape();
    ArrayRef<int64_t> outShape = collapseOp.getType().getShape();
    SmallVector<int64_t> interToOutMap(expandOp.getType().getRank());

    // Construct a set of dimensions (with respect to the intermediate
    // shape) that are unit length and get collapsed away by the final
    // `collapseOp` op.
    llvm::SmallDenseSet<int64_t> toDrop;
    for (const auto &[outDim, indices] : llvm::enumerate(collapseReassoc)) {
      for (auto [innerIdx, inDim] : llvm::enumerate(indices)) {
        // Can't drop this dim if it isnt statically 1 or if it isn't being
        // combined with any other dimensions.
        if (indices.size() == 1 || interShape[inDim] != 1) {
          continue;
        }

        // If outShape[outDim] == 1, we must preserve 1 unit dim,
        // so we drop the first. If the first is the only unit dim,
        // we can't drop it anyway.
        if (outShape[outDim] == 1 && innerIdx == 0) {
          continue;
        }
        toDrop.insert(inDim);
      }
    }

    // Remove dimensions from `toDrop` that weren't introduced by the
    // `expandOp` op.
    const auto expandReassoc = expandOp.getReassociationIndices();
    for (const auto &indices : expandReassoc) {
      // If all of indices are in `toDrop`, we must preserve at least one
      // to avoid an empty reassociation map during expansion.
      // This can happen when outShape does not have a unit dimension
      // corresponding to the unit dimensions being dropped here.
      if (llvm::all_of(indices,
                       [&](int64_t idx) { return toDrop.contains(idx); })) {
        toDrop.erase(indices[0]);
      }
    }
    if (toDrop.empty()) {
      return rewriter.notifyMatchFailure(collapseOp,
                                         "Didn't find any unit dims to drop");
    }

    // Construct a new intermediate shape without the foldable dims.
    SmallVector<int64_t> newInterShape;
    SmallVector<OpFoldResult> newInterSizes;
    newInterShape.reserve(interShape.size() - toDrop.size());
    newInterSizes.reserve(interShape.size() - toDrop.size());
    SmallVector<OpFoldResult> origInterSizes = getMixedValues(
        expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter);
    for (auto [idx, ofr] : llvm::enumerate(origInterSizes)) {
      if (!toDrop.contains(idx)) {
        std::optional<int64_t> staticDim = getConstantIntValue(ofr);
        newInterShape.push_back(staticDim.value_or(ShapedType::kDynamic));
        newInterSizes.push_back(ofr);
      }
    }

    /// Attempts to push back a new `ReassociationIndices` to `reassoc` that
    /// does not include the dimensions in `toDrop`.
    /// Returns true if new `ReassociationIndices` were appended to `reassoc`.
    auto pushBackReassociation =
        [&toDrop](SmallVectorImpl<ReassociationIndices> &reassoc, int64_t start,
                  int64_t origStart, int64_t count) -> bool {
      ReassociationIndices indices;
      indices.reserve(count);
      int64_t dim = start;
      for (auto idx : llvm::seq<int64_t>(origStart, origStart + count)) {
        if (!toDrop.contains(idx)) {
          indices.push_back(dim++);
        }
      }
      if (!indices.empty()) {
        reassoc.push_back(std::move(indices));
        return true;
      }
      return false;
    };

    // Construct new reassociations for the `collapse_shape` that does
    // not include the dropped dimensions.
    SmallVector<ReassociationIndices, 4> newCollapseReassoc;
    int64_t collapsedDim = 0;
    for (auto dim : llvm::seq<int64_t>(0, outShape.size())) {
      bool changed = pushBackReassociation(newCollapseReassoc, collapsedDim,
                                           collapseReassoc[dim][0],
                                           collapseReassoc[dim].size());
      if (changed) {
        collapsedDim += newCollapseReassoc.back().size();
      }
    }

    // Construct new reassociations for the `expand_shape` op that does
    // not include the dropped dimensions.
    SmallVector<ReassociationIndices, 4> newExpandReassoc;
    ArrayRef<int64_t> srcShape = expandOp.getSrcType().getShape();
    int64_t expandedDim = 0;
    for (auto dim : llvm::seq<int64_t>(0, srcShape.size())) {
      bool changed = pushBackReassociation(newExpandReassoc, expandedDim,
                                           expandReassoc[dim][0],
                                           expandReassoc[dim].size());
      if (changed) {
        expandedDim += newExpandReassoc.back().size();
      }
    }

    // Construct the new expand_shape and collapse_shape ops.
    // Note: we must handle the cases where the expand/collapse is no longer
    // needed. Both ops require a non-identity reassociation (i.e. they can't be
    // no-ops).
    Value newExpanded = expandOp.getSrc();
    // Special case: If the source has no shape (e.g., `tensor<f32>`), we
    // produce an empty reassociation, but still need to insert an
    // `expand_shape`.
    bool noSrcShape = srcShape.empty() && !newInterShape.empty();
    if (noSrcShape ||
        !llvm::all_of(newExpandReassoc,
                      llvm::hasSingleElement<ReassociationIndicesRef>)) {
      newExpanded = tensor::ExpandShapeOp::create(
          rewriter, expandOp.getLoc(),
          RankedTensorType::get(newInterShape,
                                expandOp.getType().getElementType()),
          expandOp.getSrc(), newExpandReassoc, newInterSizes);
    }

    Value result = newExpanded;
    if (!llvm::all_of(newCollapseReassoc,
                      llvm::hasSingleElement<ReassociationIndicesRef>)) {
      result = tensor::CollapseShapeOp::create(rewriter, collapseOp.getLoc(),
                                               collapseOp.getType(),
                                               newExpanded, newCollapseReassoc);
    }
    rewriter.replaceOp(collapseOp, result);
    return success();
  }
};

// Fold unit dims from `tensor.extract` ops.
struct FoldUnitDimsFromExtractOp : OpRewritePattern<tensor::ExtractOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractOp extractOp,
                                PatternRewriter &rewriter) const override {
    RankedTensorType srcType = extractOp.getTensor().getType();
    if (srcType.getShape().empty() ||
        llvm::none_of(srcType.getShape(),
                      [](int64_t size) { return size == 1; })) {
      return failure();
    }
    SmallVector<Value> oldIndices = extractOp.getIndices();

    SmallVector<int64_t> newShape;
    SmallVector<Value> newIndices;
    SmallVector<ReassociationIndices> reassoc;
    ReassociationIndices currReassoc;

    // Build reassociation groups where each non-unit dimension forms one output
    // dimension, and unit dimensions are grouped with adjacent non-unit dims.
    for (auto [idx, size] : llvm::enumerate(srcType.getShape())) {
      currReassoc.push_back(idx);

      if (size != 1) {
        // Non-unit dimension: this forms one output dimension
        // Finish current group and start a new one
        reassoc.push_back(std::move(currReassoc));
        currReassoc.clear();
        newShape.push_back(size);
        newIndices.push_back(oldIndices[idx]);
      }
    }

    // If we have trailing unit dims, merge them with the last group
    if (!currReassoc.empty() && !reassoc.empty()) {
      reassoc.back().append(currReassoc.begin(), currReassoc.end());
    }

    rewriter.setInsertionPointAfterValue(extractOp.getTensor());
    auto collapseOp = tensor::CollapseShapeOp::create(
        rewriter, extractOp.getLoc(), extractOp.getTensor(), reassoc);

    rewriter.setInsertionPointAfter(extractOp);
    auto newExtract = tensor::ExtractOp::create(
        rewriter, extractOp.getLoc(), extractOp.getResult().getType(),
        collapseOp.getResult(), newIndices);
    rewriter.replaceOp(extractOp, newExtract);
    return success();
  }
};

} // namespace

template <typename AttentionOpTy>
static SmallVector<unsigned>
getAllowedAttentionUnitDims(AttentionOpTy attentionOp) {
  SetVector<unsigned> allowedUnitDims;
  allowedUnitDims.insert_range(
      llvm::seq<unsigned>(0, attentionOp.getNumLoops()));
  FailureOr<IREE::LinalgExt::AttentionOpDetail> maybeOpInfo =
      IREE::LinalgExt::AttentionOpDetail::get(
          attentionOp.getQueryMap(), attentionOp.getKeyMap(),
          attentionOp.getValueMap(), attentionOp.getOutputMap());
  if (failed(maybeOpInfo)) {
    return {};
  }

  SmallVector<int64_t> loopRanges = attentionOp.getStaticLoopRanges();
  llvm::SmallDenseSet<unsigned> preservedUnitDims;
  auto preserveInnermostUnitDim = [&](ArrayRef<int64_t> dims) {
    if (dims.empty()) {
      return;
    }
    if (llvm::all_of(dims, [&](int64_t dim) { return loopRanges[dim] == 1; })) {
      preservedUnitDims.insert(dims.back());
    }
  };
  preserveInnermostUnitDim(maybeOpInfo->getK1Dims());
  preserveInnermostUnitDim(maybeOpInfo->getNDims());

  allowedUnitDims.set_subtract(preservedUnitDims);
  return allowedUnitDims.takeVector();
}

static linalg::ControlDropUnitDims getControlDropUnitDimsOptions() {
  linalg::ControlDropUnitDims options;
  auto defaultFn = options.controlFn;

  options.controlFn = [=](Operation *op) {
    // Ignore operations already in dispatches.
    if (!IREE::Flow::isNonNullAndOutsideDispatch(op)) {
      return SmallVector<unsigned>{};
    }
    if (auto attentionOp = dyn_cast<IREE::LinalgExt::AttentionOp>(op)) {
      return getAllowedAttentionUnitDims(attentionOp);
    }
    if (auto attentionOp = dyn_cast<IREE::LinalgExt::OnlineAttentionOp>(op)) {
      return getAllowedAttentionUnitDims(attentionOp);
    }
    if (isa<IREE::LinalgExt::LinalgExtOp>(op)) {
      return IREE::LinalgExt::defaultControlDropUnitDims(op);
    }
    return defaultFn(op);
  };

  auto getCollapsedEncoding =
      [](Attribute encoding, ArrayRef<int64_t> origShape,
         ArrayRef<int64_t> targetShape) -> FailureOr<Attribute> {
    if (!encoding) {
      return encoding;
    }
    auto serializableEncoding =
        dyn_cast<IREE::Encoding::SerializableAttr>(encoding);
    if (!serializableEncoding) {
      // If the encoding doesn't implement the interface, it can't be collapsed.
      return failure();
    }
    Attribute collapsedEncoding =
        serializableEncoding.getCollapsedEncoding(origShape, targetShape);
    if (!collapsedEncoding) {
      // If the interface function returns nullptr, the encoding can't be
      // collapsed.
      return failure();
    }
    return collapsedEncoding;
  };

  options.collapseFn =
      [=](RewriterBase &rewriter, Location loc, Value operand,
          ArrayRef<int64_t> targetShape,
          ArrayRef<ReassociationIndices> reassociation,
          const linalg::ControlDropUnitDims &control) -> FailureOr<Value> {
    auto tensorType = cast<RankedTensorType>(operand.getType());
    assert(control.rankReductionStrategy ==
               linalg::ControlDropUnitDims::RankReductionStrategy::
                   ReassociativeReshape &&
           "unexpected rank reduction strategy");
    ArrayRef<int64_t> origShape = tensorType.getShape();
    Attribute encoding = tensorType.getEncoding();
    FailureOr<Attribute> maybeCollapsedEncoding =
        getCollapsedEncoding(encoding, origShape, targetShape);
    if (failed(maybeCollapsedEncoding)) {
      // If the tensor has an encoding and that encoding does not implement the
      // interface or cannot be collapsed, return failure to abort the
      // transformation.
      return failure();
    }
    auto targetType =
        RankedTensorType::get(targetShape, tensorType.getElementType(),
                              maybeCollapsedEncoding.value());
    return tensor::CollapseShapeOp::create(rewriter, loc, targetType, operand,
                                           reassociation)
        .getResult();
  };

  options.expandFn =
      [=](RewriterBase &rewriter, Location loc, Value result, Value origDest,
          ArrayRef<ReassociationIndices> reassociation,
          const linalg::ControlDropUnitDims &control) -> FailureOr<Value> {
    auto origResultType = cast<RankedTensorType>(origDest.getType());
    assert(control.rankReductionStrategy ==
               linalg::ControlDropUnitDims::RankReductionStrategy::
                   ReassociativeReshape &&
           "unknown rank reduction strategy");
    ArrayRef<int64_t> origShape = origResultType.getShape();
    Attribute origEncoding = origResultType.getEncoding();
    auto newResultType = cast<RankedTensorType>(result.getType());
    ArrayRef<int64_t> targetShape = newResultType.getShape();
    FailureOr<Attribute> maybeCollapsedEncoding =
        getCollapsedEncoding(origEncoding, origShape, targetShape);
    if (failed(maybeCollapsedEncoding)) {
      // If the tensor has an encoding and that encoding does not implement the
      // interface or cannot be collapsed, return failure to abort the
      // transformation.
      return failure();
    }
    return tensor::ExpandShapeOp::create(rewriter, loc, origResultType, result,
                                         reassociation)
        .getResult();
  };
  return options;
}

/// Populate patterns for the core unit-extent dim folding transformations.
/// These patterns are applied with a walk-based driver.
static void populateFoldUnitDimsPatterns(RewritePatternSet &patterns,
                                         linalg::ControlDropUnitDims &options) {
  linalg::populateFoldUnitExtentDimsPatterns(patterns, options);
  IREE::LinalgExt::populateFoldUnitExtentDimsPatterns(patterns, options);
}

/// Populate canonicalization patterns that clean up after unit-extent dim
/// folding. These patterns are applied with a greedy driver.
static void populateFoldUnitDimsCanonicalizationPatterns(
    RewritePatternSet &patterns, linalg::ControlDropUnitDims &options) {
  patterns.insert<DropUnitDimsFromCollapseOfExpand, FoldUnitDimsFromExtractOp>(
      patterns.getContext());
  linalg::populateMoveInitOperandsToInputPattern(patterns);
  linalg::populateFoldUnitExtentDimsCanonicalizationPatterns(patterns, options);
}

static LogicalResult
foldUnitDimsOnGlobal(IRRewriter &rewriter, IREE::Util::GlobalOpInterface global,
                     SmallVector<IREE::Util::GlobalLoadOpInterface> loadOps,
                     SmallVector<IREE::Util::GlobalStoreOpInterface> storeOps,
                     SymbolTable moduleSymbols) {
  // Create a new transformed GlobalOp.
  SmallVector<int64_t> newShape;
  auto globalType = cast<RankedTensorType>(global.getGlobalType());
  for (auto size : globalType.getShape()) {
    if (size != 1) {
      newShape.push_back(size);
    }
  }
  auto newGlobalType = globalType.clone(newShape);
  auto initialValue = global.getGlobalInitialValue();
  if (!initialValue) {
    return success();
  }
  // TODO: Handle other cases
  auto newInitialValue =
      llvm::TypeSwitch<Attribute, Attribute>(initialValue)
          .Case<IREE::Util::UninitializedAttr>([&](Attribute) {
            return IREE::Util::UninitializedAttr::get(rewriter.getContext(),
                                                      newGlobalType);
          })
          .Case(
              // TODO: Remove this case once frontends have caught up, we should
              // not have stream.parameter.named at this level.
              [&](IREE::Stream::NamedParameterAttr attr) {
                return IREE::Stream::NamedParameterAttr::get(
                    rewriter.getContext(), newGlobalType, attr.getScope(),
                    attr.getKey(), attr.getConfig());
              })
          .Case([&](IREE::Flow::NamedParameterAttr attr) {
            return IREE::Flow::NamedParameterAttr::get(
                rewriter.getContext(), newGlobalType, attr.getScope(),
                attr.getKey(), attr.getConfig());
          })
          .Default([&](Attribute) { return nullptr; });
  if (!newInitialValue) {
    return success();
  }
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(global);
  auto newGlobal =
      clone(rewriter, global, global->getResultTypes(), global->getOperands());
  newGlobal.setGlobalType(newGlobalType);
  newGlobal.setGlobalInitialValue(newInitialValue);

  // Rewrite loads and stores to use the new global.
  auto expandShapeReInds =
      getReassociationIndicesForReshape(globalType, newGlobalType);
  if (!expandShapeReInds) {
    return failure();
  }

  for (auto load : loadOps) {
    rewriter.setInsertionPoint(load);
    auto newLoad = clone(rewriter, load, {newGlobalType}, load->getOperands());
    newLoad.setGlobalAttr(FlatSymbolRefAttr::get(newGlobal.getGlobalName()));
    rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
        load, globalType, newLoad->getResult(0), expandShapeReInds.value());
  }
  for (auto store : storeOps) {
    rewriter.setInsertionPoint(store);
    Value collapse = tensor::CollapseShapeOp::create(
        rewriter, store.getLoc(), newGlobalType, store->getOperand(0),
        expandShapeReInds.value());
    auto newStore =
        clone(rewriter, store, store->getResultTypes(), store->getOperands());
    newStore.setGlobalAttr(FlatSymbolRefAttr::get(newGlobal.getGlobalName()));
    newStore.setStoredGlobalValue(collapse);
    rewriter.eraseOp(store);
  }
  rewriter.eraseOp(global);
  return success();
}

namespace {
struct FoldUnitExtentDimsPass final
    : impl::FoldUnitExtentDimsPassBase<FoldUnitExtentDimsPass> {
  void runOnOperation() override;
};

struct FoldUnitExtentDimsForFuncPass final
    : impl::FoldUnitExtentDimsForFuncPassBase<FoldUnitExtentDimsForFuncPass> {
  void runOnOperation() override;
};

} // namespace

void FoldUnitExtentDimsPass::runOnOperation() {
  mlir::ModuleOp moduleOp = getOperation();
  MLIRContext *context = &getContext();

  SymbolTable moduleSymbols(moduleOp);
  Explorer explorer(moduleOp, TraversalAction::RECURSE);
  explorer.initialize();
  IRRewriter rewriter(context);

  // Fold unit dims of GlobalOpInterface ops.
  explorer.forEachGlobal([&](const Explorer::GlobalInfo *globalInfo) {
    IREE::Util::GlobalOpInterface global = globalInfo->op;
    auto tensorType = dyn_cast<RankedTensorType>(global.getGlobalType());
    if (!tensorType || !global.isGlobalPrivate()) {
      return;
    }
    if (llvm::none_of(tensorType.getShape(),
                      [](int64_t size) { return size == 1; })) {
      return;
    }
    SmallVector<IREE::Util::GlobalLoadOpInterface> loadOps =
        llvm::to_vector(globalInfo->getLoads());
    SmallVector<IREE::Util::GlobalStoreOpInterface> storeOps =
        llvm::to_vector(globalInfo->getStores());
    if (failed(foldUnitDimsOnGlobal(rewriter, global, loadOps, storeOps,
                                    moduleSymbols))) {
      return signalPassFailure();
    }
  });

  linalg::ControlDropUnitDims options = getControlDropUnitDimsOptions();
  // Apply fold unit extent dims patterns with walk-based driver.
  {
    RewritePatternSet patterns(context);
    populateFoldUnitDimsPatterns(patterns, options);
    walkAndApplyPatterns(moduleOp, std::move(patterns));
  }

  // Apply canonicalization patterns with greedy driver.
  {
    RewritePatternSet patterns(context);
    populateFoldUnitDimsCanonicalizationPatterns(patterns, options);
    if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
}

void FoldUnitExtentDimsForFuncPass::runOnOperation() {
  MLIRContext *context = &getContext();
  Operation *op = getOperation();
  linalg::ControlDropUnitDims options = getControlDropUnitDimsOptions();
  // Apply fold unit extent dims patterns with walk-based driver.
  {
    RewritePatternSet patterns(context);
    populateFoldUnitDimsPatterns(patterns, options);
    walkAndApplyPatterns(op, std::move(patterns));
  }

  // Apply canonicalization patterns with greedy driver.
  {
    RewritePatternSet patterns(context);
    populateFoldUnitDimsCanonicalizationPatterns(patterns, options);
    if (failed(applyPatternsGreedily(op, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
}

} // namespace mlir::iree_compiler::DispatchCreation
