// Copyright 2022 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 <deque>

#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Interfaces/PartitionableLoopsInterface.h"
#include "iree/compiler/Codegen/Transforms/Transforms.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/TilingInterface.h"

#define DEBUG_TYPE "tile-dispatch-using-interface"

namespace mlir::iree_compiler {

/// Helper method to adjust the interchange vector to match the iteration
/// domain.
static SmallVector<int64_t>
fillInterchangeVector(ArrayRef<unsigned> interchangeVector,
                      size_t iterationDomainSize) {
  SmallVector<int64_t> filledVector;
  for (auto v : interchangeVector)
    filledVector.push_back(v);
  if (filledVector.size() < iterationDomainSize) {
    auto range = llvm::seq<unsigned>(filledVector.size(), iterationDomainSize);
    filledVector.append(range.begin(), range.end());
  }
  if (filledVector.size() > iterationDomainSize)
    filledVector.resize(iterationDomainSize);
  return filledVector;
}

/// Given the `lb` and `step` of a loop, return the lower bound and step to use
/// for a distributed loop. Replace the iteration domain to
/// - lb_partitioned = lb + procId * step
/// - step_partitioned = step * nprocs
static std::tuple<OpFoldResult, OpFoldResult>
getDistributeLBAndStep(OpBuilder &b, Location loc, OpFoldResult lb,
                       OpFoldResult step, OpFoldResult procId,
                       OpFoldResult nprocs) {
  AffineExpr s0, s1, s2;
  bindSymbols(b.getContext(), s0, s1, s2);
  auto offsetMap = AffineMap::get(0, 3, {s0 + s1 * s2});
  auto stepMap = AffineMap::get(0, 2, {s0 * s1});
  OpFoldResult distributeLB = affine::makeComposedFoldedAffineApply(
      b, loc, offsetMap, ArrayRef<OpFoldResult>{lb, procId, step});
  OpFoldResult distributeStep = affine::makeComposedFoldedAffineApply(
      b, loc, stepMap, ArrayRef<OpFoldResult>{step, nprocs});
  return {distributeLB, distributeStep};
}

// Helper function to change arith.constant to i64 attribute.
static void changeArithCstToI64Attr(OpBuilder &b,
                                    MutableArrayRef<OpFoldResult> constants) {
  for (OpFoldResult &val : constants) {
    if (auto dyn_cast = llvm::dyn_cast_if_present<Value>(val)) {
      APInt intVal;
      if (matchPattern(dyn_cast, m_ConstantInt(&intVal))) {
        val = b.getI64IntegerAttr(intVal.getSExtValue());
      }
    }
  }
}

//===----------------------------------------------------------------------===//
// TileDispatchUsingSCFForOp implementation.
//===----------------------------------------------------------------------===//

/// This implementation mirrors the implementation in
/// include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h and
/// lib/Dialect/SCF/Transforms/TileUsingInterface.cpp. It is adapted to do
/// distribution and also use `flow.dispatch.tensor.load/store` instead of
/// `tensor.extract_slice/insert_slice`.

/// Generate an empty loop nest that represents the tiled loop nest shell.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `tileSizeVals` is the tile sizes to use. Zero represent untiled loops.
/// - In `offsets` and `sizes` return the multi-dimensional offset and size of
/// the
///   tile processed within the inner most loop.
static SmallVector<scf::ForOp> generateTileLoopNest(
    OpBuilder &builder, Location loc, ArrayRef<Range> loopRanges,
    ArrayRef<OpFoldResult> tileSizeVals,
    ArrayRef<linalg::DistributionMethod> distributionMethod,
    ArrayRef<linalg::ProcInfo> procInfo, SmallVector<OpFoldResult> &offsets,
    SmallVector<OpFoldResult> &sizes, SmallVector<Value> &workgroupCount) {
  assert(!loopRanges.empty() && "expected at least one loop range");
  assert(loopRanges.size() == tileSizeVals.size() &&
         "expected as many tile sizes as loop ranges");
  assert(loopRanges.size() == distributionMethod.size() &&
         "expected as many entries in distribution method list as number of "
         "loop ranges");
  OpBuilder::InsertionGuard guard(builder);
  SmallVector<scf::ForOp> loops;
  offsets.resize(loopRanges.size());
  sizes.resize(loopRanges.size());

  // The tile size to use (to avoid out of bounds access) is  minimum of
  // `tileSize` and `ub - iv`, where `iv` is the induction variable
  // of the tiled loop.
  AffineExpr s0, s1, d0, s2;
  bindDims(builder.getContext(), d0);
  bindSymbols(builder.getContext(), s0, s1, s2);
  AffineMap minMap = AffineMap::get(1, 2, {s0, s1 - d0}, builder.getContext());
  auto createBoundedTileSize = [&](OpFoldResult iv, OpFoldResult tileSize,
                                   OpFoldResult size) -> OpFoldResult {
    if (isConstantIntValue(tileSize, 1)) {
      return tileSize;
    }
    return affine::makeComposedFoldedAffineMin(
        builder, loc, minMap, ArrayRef<OpFoldResult>{iv, tileSize, size});
  };

  unsigned procDim = 0;
  AffineMap numIterationsMap =
      AffineMap::get(0, 3, (s1 - s0).ceilDiv(s2), builder.getContext());
  for (auto [idx, loopRange] : llvm::enumerate(loopRanges)) {
    // Capturing structured bindings in lambdas is a c++20 feature, so we have
    // to declare a local variable for it.
    int index = idx;
    OpFoldResult lb = loopRange.offset;
    OpFoldResult ub = loopRange.size;
    OpFoldResult step = tileSizeVals[index];
    // No loops if tile size is zero. Set offset and size to the loop
    // offset and size.
    if (isConstantIntValue(tileSizeVals[index], 0)) {
      offsets[index] = lb;
      sizes[index] = ub;
      continue;
    }

    auto method = distributionMethod[index];
    if (method != linalg::DistributionMethod::None) {
      Value numWorkgroups = getValueOrCreateConstantIndexOp(
          builder, loc,
          affine::makeComposedFoldedAffineApply(
              builder, loc, numIterationsMap,
              {loopRange.offset, loopRange.size, tileSizeVals[index]}));
      workgroupCount.push_back(numWorkgroups);
      std::tie(lb, step) = getDistributeLBAndStep(builder, loc, lb, step,
                                                  procInfo[procDim].procId,
                                                  procInfo[procDim].nprocs);
      procDim++;
    }

    if (method == linalg::DistributionMethod::CyclicNumProcsEqNumIters) {
      offsets[index] = lb;
      sizes[index] = createBoundedTileSize(lb, tileSizeVals[index], ub);
      continue;
    }

    Value lbVal = getValueOrCreateConstantIndexOp(builder, loc, lb);
    Value ubVal = getValueOrCreateConstantIndexOp(builder, loc, ub);
    Value stepVal = getValueOrCreateConstantIndexOp(builder, loc, step);
    auto loop = builder.create<scf::ForOp>(
        loc, lbVal, ubVal, stepVal, ValueRange{},
        [&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
            ValueRange /*iterArgs*/) {
          sizes[index] = createBoundedTileSize(iv, tileSizeVals[index], ub);
          builder.create<scf::YieldOp>(loc);
        });
    offsets[index] = loop.getInductionVar();
    loops.push_back(loop);
    builder.setInsertionPoint(loop.getBody()->getTerminator());
  }

  // Update the sizes if it contains arith.index with i64 attrs.
  // TODO: tensor.extract_slice is unable to determine the
  // result type if arith.constant is present. This is a workaround
  // to ensure that the result type is determined.
  changeArithCstToI64Attr(builder, sizes);

  return loops;
}

/// Replace the `flow.dispatch.tensor.store` of the `untiledValue` with a tiled
/// `flow.dispatch.tensor.store` that writes only a tile of the result at
/// offsets given by `tiledOffsets` and sizes given by `tiledSizes`, using
/// `tiledValue` as the source.
static LogicalResult replaceStoresWithTiledVersion(
    RewriterBase &rewriter, OpResult untiledValue, Value tiledValue,
    ArrayRef<OpFoldResult> tileOffsets, ArrayRef<OpFoldResult> tileSizes,
    Block *innerLoopBody) {
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(innerLoopBody->getTerminator());
  SmallVector<IREE::Flow::DispatchTensorStoreOp> storeOps;
  for (OpOperand &use : untiledValue.getUses()) {
    auto storeOp = dyn_cast<IREE::Flow::DispatchTensorStoreOp>(use.getOwner());
    if (storeOp && storeOp.getValue() == use.get()) {
      storeOps.push_back(storeOp);
    }
  }
  if (storeOps.empty())
    return success();
  if (storeOps.size() != 1) {
    return rewriter.notifyMatchFailure(untiledValue.getOwner(),
                                       "expected a single store for the op");
  }
  auto storeOp = storeOps.front();

  SmallVector<OpFoldResult> tileStrides(tileOffsets.size(),
                                        rewriter.getIndexAttr(1));
  SmallVector<OpFoldResult> combinedOffsets, combinedSizes, combinedStrides;
  SliceAndDynamicDims clonedSliceAndVals =
      cloneOffsetsSizesAndStrides(rewriter, storeOp);

  if (failed(affine::mergeOffsetsSizesAndStrides(
          rewriter, storeOp.getLoc(), clonedSliceAndVals.offsets,
          clonedSliceAndVals.sizes, clonedSliceAndVals.strides,
          storeOp.getDroppedDims(), tileOffsets, tileSizes, tileStrides,
          combinedOffsets, combinedSizes, combinedStrides))) {
    return rewriter.notifyMatchFailure(
        storeOp, "failed to create tiled flow.dispatch.tensor.store op");
  }

  rewriter.create<IREE::Flow::DispatchTensorStoreOp>(
      storeOp.getLoc(), tiledValue, storeOp.getTarget(),
      clonedSliceAndVals.dynamicDims, combinedOffsets, combinedSizes,
      combinedStrides);
  rewriter.eraseOp(storeOp);
  return success();
}

/// Replace all `flow.dispatch.tensor.store` operations that use values produced
/// by `untiledOp` as source, with tiled stores, with tiled values produced by
/// `tiledOp`.
static LogicalResult replaceAllStoresWithTiledVersion(
    RewriterBase &rewriter, TilingInterface untiledOp,
    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
    ValueRange tiledValues, Block *innerLoopBody) {
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(innerLoopBody->getTerminator());
  for (auto [index, result] : llvm::enumerate(untiledOp->getResults())) {
    SmallVector<OpFoldResult> resultOffsets, resultSizes;
    if (failed(untiledOp.getResultTilePosition(rewriter, index, offsets, sizes,
                                               resultOffsets, resultSizes))) {
      return rewriter.notifyMatchFailure(
          untiledOp, "failed to rewrite destructive update");
    }
    if (failed(replaceStoresWithTiledVersion(
            rewriter, llvm::cast<OpResult>(result), tiledValues[index],
            resultOffsets, resultSizes, innerLoopBody))) {
      return failure();
    }
  }
  return success();
}

FailureOr<IREETilingResult>
tileDispatchUsingSCFFopOp(RewriterBase &rewriter, TilingInterface op,
                          linalg::LinalgTilingOptions options) {
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointAfter(op);

  if (!options.tileSizeComputationFunction) {
    return rewriter.notifyMatchFailure(
        op, "missing tile size computation function");
  }

  // 1. Get the range of the loops that are represented by the operation.
  SmallVector<Range> iterationDomain = op.getIterationDomain(rewriter);
  Location loc = op.getLoc();
  size_t numLoops = iterationDomain.size();
  if (numLoops == 0) {
    return IREETilingResult();
  }

  // 2. Materialize the tile sizes. Enforce the convention that "tiling by zero"
  // skips tiling a particular dimension. This convention is significantly
  // simpler to handle instead of adjusting affine maps to account for missing
  // dimensions.
  SmallVector<OpFoldResult> tileSizes =
      getAsOpFoldResult(options.tileSizeComputationFunction(rewriter, op));
  tileSizes.resize(numLoops, rewriter.getIndexAttr(0));

  IREETilingResult tilingResult;
  tilingResult.tiledLoops.resize(numLoops, false);
  AffineExpr s0, s1, s2, s3; // lb, ub, step, tileSize
  bindSymbols(rewriter.getContext(), s0, s1, s2, s3);
  AffineExpr numTilesExprs = (s1 - s0).ceilDiv(s2 * s3);
  for (auto [index, iteratorType, range, tileSize] :
       llvm::enumerate(op.getLoopIteratorTypes(), iterationDomain, tileSizes)) {
    // If distribution is specified, only parallel loops are tiled.
    if (options.distribution && iteratorType != utils::IteratorType::parallel) {
      continue;
    }
    // If tile size is 0, it isnt tiled.
    if (isConstantIntValue(tileSize, 0)) {
      continue;
    }
    // If number of tiles is statically know to be 1, the loop isnt tiled.
    OpFoldResult numTiles = affine::makeComposedFoldedAffineApply(
        rewriter, loc, numTilesExprs,
        {range.offset, range.size, range.stride, tileSize});
    if (isConstantIntValue(numTiles, 1)) {
      continue;
    }

    tilingResult.tiledLoops.set(index);
  }

  if (!tilingResult.tiledLoops.any()) {
    return IREETilingResult();
  }

  {
    SmallVector<OpFoldResult> offsets, sizes;
    SmallVector<Value> workgroupCount;
    // If there is an interchange specified, permute the iteration domain and
    // the tile sizes.
    SmallVector<int64_t> interchangeVector;
    if (!options.interchangeVector.empty()) {
      interchangeVector = fillInterchangeVector(options.interchangeVector,
                                                iterationDomain.size());
    }
    if (!interchangeVector.empty()) {
      if (!isPermutationVector(interchangeVector)) {
        return rewriter.notifyMatchFailure(
            op, "invalid intechange vector, not a permutation of the entire "
                "iteration space");
      }

      applyPermutationToVector(iterationDomain, interchangeVector);
      applyPermutationToVector(tileSizes, interchangeVector);
    }

    // If there is distribution specified, adjust the loop ranges. Note that
    // permutation happens *before* distribution.
    SmallVector<linalg::DistributionMethod> distributionMethods(
        iterationDomain.size(), linalg::DistributionMethod::None);
    SmallVector<linalg::ProcInfo> procInfo;
    if (options.distribution) {
      SmallVector<Range> parallelLoopRanges;
      for (auto loopIdx : llvm::seq<unsigned>(0, numLoops)) {
        if (tilingResult.tiledLoops.test(loopIdx)) {
          AffineExpr s0, s1;
          bindSymbols(rewriter.getContext(), s0, s1);
          OpFoldResult parallelLoopStep = affine::makeComposedFoldedAffineApply(
              rewriter, loc, s0 * s1,
              {iterationDomain[loopIdx].stride, tileSizes[loopIdx]});
          Range r = {iterationDomain[loopIdx].offset,
                     iterationDomain[loopIdx].size, parallelLoopStep};
          parallelLoopRanges.emplace_back(std::move(r));
        }
      }

      procInfo =
          options.distribution->procInfo(rewriter, loc, parallelLoopRanges);

      unsigned partitionedLoopIdx = 0;
      for (auto loopIdx : llvm::seq<unsigned>(0, numLoops)) {
        if (!tilingResult.tiledLoops.test(loopIdx)) {
          continue;
        }
        distributionMethods[loopIdx] =
            procInfo[partitionedLoopIdx++].distributionMethod;
      }
    }

    // 3. Materialize an empty loop nest that iterates over the tiles. These
    // loops for now do not return any values even if the original operation has
    // results.
    tilingResult.loops = generateTileLoopNest(
        rewriter, loc, iterationDomain, tileSizes, distributionMethods,
        procInfo, offsets, sizes, workgroupCount);

    if (!interchangeVector.empty()) {
      auto inversePermutation = invertPermutationVector(interchangeVector);
      applyPermutationToVector(offsets, inversePermutation);
      applyPermutationToVector(sizes, inversePermutation);
    }

    LLVM_DEBUG({
      if (!tilingResult.loops.empty()) {
        llvm::errs() << "LoopNest shell :\n";
        tilingResult.loops.front().dump();
        llvm::errs() << "\n";
      }
    });

    // 4. Generate the tiled implementation within the inner most loop.
    if (!tilingResult.loops.empty())
      rewriter.setInsertionPoint(
          tilingResult.loops.back().getBody()->getTerminator());
    FailureOr<TilingResult> tiledImplementation =
        op.getTiledImplementation(rewriter, offsets, sizes);
    if (failed(tiledImplementation)) {
      return rewriter.notifyMatchFailure(
          op, "failed to generate tiled implementation");
    }
    std::swap(tilingResult.tiledOps, tiledImplementation->tiledOps);
    std::swap(tilingResult.tiledValues, tiledImplementation->tiledValues);

    LLVM_DEBUG({
      if (!tilingResult.loops.empty()) {
        llvm::errs() << "After tiled implementation :\n";
        tilingResult.loops.front().dump();
        llvm::errs() << "\n";
      }
    });
    std::swap(tilingResult.tileOffsets, offsets);
    std::swap(tilingResult.tileSizes, sizes);
    std::swap(tilingResult.workgroupCount, workgroupCount);
  }

  if (op->getNumResults() == 0) {
    rewriter.eraseOp(op);
    return tilingResult;
  }
  Block *body = tilingResult.loops.empty()
                    ? op->getBlock()
                    : tilingResult.loops.back().getBody();
  // Rewrite all `flow.dispatch.tensor.store` operation with tiled version
  // of the store. Its valid to this for all stores of the root untiled op.
  if (failed(replaceAllStoresWithTiledVersion(
          rewriter, op, tilingResult.tileOffsets, tilingResult.tileSizes,
          tilingResult.tiledValues, body))) {
    return failure();
  }
  return tilingResult;
}

//===----------------------------------------------------------------------===//
// tileAndFuseDispatchUsingSCFForOp implementation.
//===----------------------------------------------------------------------===//

/// Find all producers to fuse and return them in sorted order;
static SmallVector<Operation *> getAllFusableProducers(TilingInterface op) {
  llvm::SetVector<Operation *> producers;
  std::deque<Operation *> worklist;
  worklist.push_back(op);

  while (!worklist.empty()) {
    Operation *currOp = worklist.front();
    worklist.pop_front();
    for (OpOperand &operand : currOp->getOpOperands()) {
      Operation *definingOp = operand.get().getDefiningOp();
      auto tilingInterfaceProducer =
          dyn_cast_or_null<TilingInterface>(definingOp);
      if (!tilingInterfaceProducer || isa<tensor::PadOp>(definingOp) ||
          producers.count(tilingInterfaceProducer)) {
        continue;
      }
      worklist.push_back(tilingInterfaceProducer);
      producers.insert(tilingInterfaceProducer);
    }
  }

  SmallVector<Operation *> sortedOps(producers.begin(), producers.end());
  mlir::computeTopologicalSorting(sortedOps);
  return sortedOps;
}

/// Return all slices that are used to access a tile of the producer. Assume
/// that `tiledOps` are in "reverse" order of their appearance in the IR.
static SmallVector<tensor::ExtractSliceOp>
getAllFusableProducerUses(Operation *untiledOp,
                          ArrayRef<Operation *> tiledOps) {
  SmallVector<tensor::ExtractSliceOp> sliceOps;
  for (auto tiledOp : llvm::reverse(tiledOps)) {
    for (OpOperand &operand : llvm::reverse(tiledOp->getOpOperands())) {
      auto sliceOp = operand.get().getDefiningOp<tensor::ExtractSliceOp>();
      if (!sliceOp || sliceOp.getSource().getDefiningOp() != untiledOp)
        continue;
      sliceOps.push_back(sliceOp);
    }
  }
  return sliceOps;
}

FailureOr<IREETileAndFuseResult>
tileAndFuseDispatchUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
                                 linalg::LinalgTilingOptions tilingOptions) {
  IREETileAndFuseResult tileAndFuseResult;
  auto fusableProducers = getAllFusableProducers(op);
  // Apply the tiling pattern.
  FailureOr<IREETilingResult> tilingResult =
      tileDispatchUsingSCFFopOp(rewriter, op, tilingOptions);
  if (failed(tilingResult)) {
    return failure();
  }
  tileAndFuseResult.tiledAndFusedOps = tilingResult->tiledOps;
  tileAndFuseResult.loops = tilingResult->loops;
  tileAndFuseResult.workgroupCount = tilingResult->workgroupCount;

  // If there is no tiling then there is nothing to do for fusion.
  if (!tilingResult->tiledLoops.any()) {
    return tileAndFuseResult;
  }

  // Get all ops to tile and fuse.
  auto fusableProducersRef = llvm::ArrayRef(fusableProducers);
  while (!fusableProducersRef.empty()) {
    auto fusableProducer = cast<TilingInterface>(fusableProducersRef.back());
    fusableProducersRef = fusableProducersRef.drop_back();

    // Find a slice that is used to access the producer. Get all the slice ops.
    // It is assumed that the slice ops are returned in-order, so that you
    // could use the first slice as the insertion point.
    auto sliceOps = getAllFusableProducerUses(
        fusableProducer, tileAndFuseResult.tiledAndFusedOps);

    for (auto sliceOp : sliceOps) {
      OpBuilder::InsertionGuard g(rewriter);
      rewriter.setInsertionPoint(sliceOp);

      // Generate the tiled implementation of the producer.
      OpResult untiledValue = llvm::cast<OpResult>(sliceOp.getSource());
      FailureOr<TilingResult> swapSliceResult =
          tensor::replaceExtractSliceWithTiledProducer(rewriter, sliceOp,
                                                       untiledValue);
      if (failed(swapSliceResult) || swapSliceResult->tiledValues.size() != 1) {
        return rewriter.notifyMatchFailure(sliceOp,
                                           "fusion along slice op failed");
      }

      Block *body = tileAndFuseResult.loops.empty()
                        ? op->getBlock()
                        : tileAndFuseResult.loops.back().getBody();
      // 2c. Finally replace any `flow.dispatch.tensor.store` operation with
      //     tiled version of the operation. It is only valid to do this under
      //     the above assumption that the producer and consumer share the loops
      //     that can be tiled.
      if (failed(replaceStoresWithTiledVersion(
              rewriter, untiledValue, swapSliceResult->tiledValues[0],
              sliceOp.getMixedOffsets(), sliceOp.getMixedSizes(), body))) {
        return failure();
      }
      rewriter.replaceOp(sliceOp, swapSliceResult->tiledValues[0]);
      tileAndFuseResult.tiledAndFusedOps.append(swapSliceResult->tiledOps);
    }
  }

  return tileAndFuseResult;
}

namespace {

//===----------------------------------------------------------------------===//
// SwapExtractSliceWithDispatchTensorLoad
//===----------------------------------------------------------------------===//

/// Pattern to swap `flow.dispatch.tensor.load` -> `tensor.extract_slice` with
/// `flow.dispatch.tensor.load` of the slice.
struct SwapExtractSliceWithDispatchTensorLoad
    : public OpRewritePattern<tensor::ExtractSliceOp> {
  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    auto loadOp =
        sliceOp.getSource().getDefiningOp<IREE::Flow::DispatchTensorLoadOp>();
    if (!loadOp)
      return failure();

    SmallVector<OpFoldResult> combinedOffsets, combinedSizes, combinedStrides;
    if (failed(affine::mergeOffsetsSizesAndStrides(
            rewriter, loadOp.getLoc(), loadOp, sliceOp, loadOp.getDroppedDims(),
            combinedOffsets, combinedSizes, combinedStrides))) {
      return rewriter.notifyMatchFailure(
          sliceOp, "failed to fold offsets, sizes and strides with load op");
    }
    rewriter.replaceOpWithNewOp<IREE::Flow::DispatchTensorLoadOp>(
        sliceOp, sliceOp.getType(), loadOp.getSource(), loadOp.getSourceDims(),
        combinedOffsets, combinedSizes, combinedStrides);
    return success();
  }
};

//===----------------------------------------------------------------------===//
// SwapExtractSliceWithTensorEmpty
//===----------------------------------------------------------------------===//

/// Pattern to swap `empty` -> `tensor.extract_slice` with
/// `empty` of the slice.
struct SwapExtractSliceWithTensorEmpty
    : public OpRewritePattern<tensor::ExtractSliceOp> {
  using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    auto emptyTensorOp = sliceOp.getSource().getDefiningOp<tensor::EmptyOp>();
    if (!emptyTensorOp)
      return failure();

    SmallVector<OpFoldResult> mixedSizes = sliceOp.getMixedSizes();
    if (mixedSizes.size() != sliceOp.getType().getRank()) {
      SmallVector<OpFoldResult> rankReducedMixedSizes;
      rankReducedMixedSizes.reserve(sliceOp.getType().getRank());
      auto droppedDims = sliceOp.getDroppedDims();
      for (auto [index, size] : llvm::enumerate(mixedSizes)) {
        if (droppedDims.test(index))
          continue;
        rankReducedMixedSizes.push_back(size);
      }
      std::swap(mixedSizes, rankReducedMixedSizes);
    }
    rewriter.replaceOpWithNewOp<tensor::EmptyOp>(
        sliceOp, mixedSizes, sliceOp.getType().getElementType());
    return success();
  }
};

} // namespace

void populateTileAndDistributeToWorkgroupsCleanupPatterns(
    RewritePatternSet &patterns) {
  MLIRContext *context = patterns.getContext();
  patterns.insert<SwapExtractSliceWithDispatchTensorLoad,
                  SwapExtractSliceWithTensorEmpty>(context);
}

} // namespace mlir::iree_compiler
