//===- TestTilingInterfaceTransformOps.cpp - Test `TilingInterface` ------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This file defines transform dialect operations used for testing
// TilingInterface
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Transform/IR/TransformAttrs.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "test-tiling-interface"

#define GET_OP_CLASSES
#include "TestTilingInterfaceTransformOps.h.inc"

using namespace mlir;
using namespace mlir::transform;

//===----------------------------------------------------------------------===//
// TestFuseAndYieldOp
//===----------------------------------------------------------------------===//

static llvm::SmallDenseSet<Operation *> collectTiledAndFusedOps(Operation *op) {
  SmallVector<Operation *> worklist;
  llvm::SmallDenseSet<Operation *> producers;
  worklist.push_back(op);
  producers.insert(op);
  while (!worklist.empty()) {
    Operation *current = worklist.pop_back_val();
    for (OpOperand &operand : current->getOpOperands()) {
      Operation *producer = operand.get().getDefiningOp();
      if (!producer || !isa<TilingInterface>(producer) ||
          producers.contains(producer))
        continue;
      worklist.push_back(producer);
      producers.insert(producer);
    }
  }
  return producers;
}

/// Apply a tile and fuse transformation to all payload ops and store both the
/// tiled operation as well as the created tile loops.
template <typename Range>
static LogicalResult
applyTileAndFuseToAll(RewriterBase &rewriter, Operation *transformOp,
                      Range &&payloadOps, unsigned numLoops,
                      scf::SCFTilingOptions tilingOptions,
                      TransformResults &transformResults) {
  SmallVector<Operation *> tiledOps;
  SmallVector<SmallVector<Operation *>> loopOps(numLoops);

  for (Operation *target : payloadOps) {
    auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
    if (!tilingInterfaceOp)
      return transformOp->emitError("only TilingInterface ops are supported");
    DominanceInfo dominanceInfo(tilingInterfaceOp);

    llvm::SmallDenseSet<Operation *> tiledAndFusedOps =
        collectTiledAndFusedOps(tilingInterfaceOp);
    llvm::DenseSet<Operation *> yieldReplacementsFor;
    for (auto *op : tiledAndFusedOps) {
      if (llvm::any_of(op->getUsers(), [&](Operation *user) {
            return dominanceInfo.properlyDominates(tilingInterfaceOp, user);
          })) {
        yieldReplacementsFor.insert(op);
      }
    }

    scf::SCFTileAndFuseOptions tileAndFuseOptions;
    tileAndFuseOptions.setTilingOptions(tilingOptions);

    scf::SCFTileAndFuseOptions::ControlFnTy controlFn =
        [&](tensor::ExtractSliceOp candidateSliceOp, OpResult originalProducer,
            bool isDestinationOperand)
        -> std::optional<scf::SCFTileAndFuseOptions::ControlFnResult> {
      Operation *owner = originalProducer.getOwner();
      bool yieldProducerReplacement = yieldReplacementsFor.contains(owner);
      return scf::SCFTileAndFuseOptions::ControlFnResult{
          yieldProducerReplacement};
    };
    tileAndFuseOptions.setFusionControlFn(controlFn);

    rewriter.setInsertionPoint(target);
    FailureOr<scf::SCFTileAndFuseResult> tiledResults =
        scf::tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,
                                                  tileAndFuseOptions);
    if (failed(tiledResults))
      return failure();

    // Perform the replacement of tiled and fused values.
    SmallVector<Operation *> opsToReplace{target};
    llvm::append_range(opsToReplace, tiledResults->fusedProducers);
    for (Operation *toReplace : opsToReplace) {
      for (OpResult res : toReplace->getResults())
        if (auto replacement = tiledResults->replacements.lookup(res)) {
          Operation *replacementOp = replacement.getDefiningOp();
          rewriter.replaceUsesWithIf(res, replacement, [&](OpOperand &use) {
            Operation *user = use.getOwner();
            return dominanceInfo.properlyDominates(replacementOp, user) &&
                   user->getParentOp() == replacementOp->getParentOp();
          });
        }

      if (toReplace->use_empty()) {
        rewriter.eraseOp(toReplace);
      }
    }

    // Report back the relevant handles to the transform op.
    tiledOps.push_back(tiledResults->tiledAndFusedOps.front());
    assert(tiledResults->loops.size() == numLoops &&
           "Mismatched number of loops, tile and fuse transform should have "
           "failed");
    for (unsigned int i = 0; i < numLoops; ++i)
      loopOps[i].push_back(tiledResults->loops[i]);
  }

  transformResults.set(transformOp->getOpResult(0), tiledOps);
  for (unsigned int i = 0; i < numLoops; ++i)
    transformResults.set(transformOp->getOpResult(i + 1), loopOps[i]);

  return success();
}

DiagnosedSilenceableFailure
transform::TestFuseAndYieldOp::apply(TransformRewriter &rewriter,
                                     TransformResults &transformResults,
                                     TransformState &state) {
  SmallVector<int64_t> tileSizes =
      extractFromIntegerArrayAttr<int64_t>(getTileSizes());
  SmallVector<int64_t> tileInterchange =
      extractFromIntegerArrayAttr<int64_t>(getTileInterchange());

  SmallVector<OpFoldResult> tileSizesOfr =
      getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);

  scf::SCFTilingOptions tilingOptions;
  tilingOptions.setTileSizes(tileSizesOfr).setInterchange(tileInterchange);
  if (getUseForall()) {
    tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);
  }

  LogicalResult result = applyTileAndFuseToAll(
      rewriter, getOperation(), state.getPayloadOps(getTarget()),
      tileSizes.size() - llvm::count(tileSizes, 0), tilingOptions,
      transformResults);
  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

//===----------------------------------------------------------------------===//
// TestFuseConsumerOp
//===----------------------------------------------------------------------===//

/// Fuse the consumer and store both the original consumer operation as well as
/// the fused consumer operation.
static LogicalResult
applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,
                  Operation *consumer,
                  MutableArrayRef<LoopLikeOpInterface> loops,
                  TransformResults &transformResults) {
  SmallVector<Operation *> fusedConsumerOps;
  rewriter.setInsertionPoint(consumer);

  FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =
      scf::tileAndFuseConsumer(rewriter, consumer, loops);
  if (failed(fuseConsumerResults))
    return consumer->emitOpError("failed to fuse consumer of slice");

  // Report back the relevant handles to the transform op.
  for (OpOperand *tiledAndFusedConsumerOperand :
       fuseConsumerResults->tiledAndFusedConsumerOperands) {
    fusedConsumerOps.push_back(tiledAndFusedConsumerOperand->getOwner());
  }
  transformResults.set(transformOp->getOpResult(0), fusedConsumerOps);
  for (auto [index, loop] : llvm::enumerate(loops)) {
    transformResults.set(transformOp->getOpResult(index + 1), {loop});
  }
  return success();
}

DiagnosedSilenceableFailure
transform::TestFuseConsumerOp::apply(TransformRewriter &rewriter,
                                     TransformResults &transformResults,
                                     TransformState &state) {
  Operation *consumer = *state.getPayloadOps(getConsumer()).begin();

  SmallVector<LoopLikeOpInterface> loops;
  // Since the matcher works inside-out, we need to iterate the loops in
  // reverse.
  for (auto loop : llvm::reverse(getLoops())) {
    auto loopLikeOp =
        dyn_cast<LoopLikeOpInterface>(*state.getPayloadOps(loop).begin());
    if (!loopLikeOp) {
      return DiagnosedSilenceableFailure::definiteFailure();
    }
    loops.push_back(loopLikeOp);
  }
  LogicalResult result = applyFuseConsumer(rewriter, getOperation(), consumer,
                                           loops, transformResults);
  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

void transform::TestFuseConsumerOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  consumesHandle(getConsumerMutable(), effects);
  consumesHandle(getLoopsMutable(), effects);
  producesHandle(getOperation()->getOpResults(), effects);
  modifiesPayload(effects);
}

//===----------------------------------------------------------------------===//
// TestFuseConsumerUsingSliceOp
//===----------------------------------------------------------------------===//

/// Apply fusing of consumer transformation to all payload ops and store both
/// the original consumer operation as well as the fused consumer operation.
static LogicalResult applyFuseConsumerUsingSlices(
    RewriterBase &rewriter, Operation *transformOp,
    ArrayRef<Operation *> slices, MutableArrayRef<LoopLikeOpInterface> loops,
    uint32_t numConsumerToFuse, TransformResults &transformResults) {
  SmallVector<Operation *> originalConsumerOps;
  SmallVector<Operation *> fusedConsumerOps;

  rewriter.setInsertionPoint(slices.front());

  while (numConsumerToFuse--) {
    FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =
        scf::tileAndFuseConsumerOfSlices(rewriter, slices, loops);

    if (failed(fuseConsumerResults))
      return slices.front()->emitOpError("failed to fuse consumer of slice");

    // Report back the relevant handles to the transform op.
    for (OpOperand *origConsumerOperand :
         fuseConsumerResults->origConsumerOperands) {
      originalConsumerOps.push_back(origConsumerOperand->getOwner());
    }
    for (OpOperand *tiledAndFusedConsumerOperand :
         fuseConsumerResults->tiledAndFusedConsumerOperands) {
      fusedConsumerOps.push_back(tiledAndFusedConsumerOperand->getOwner());
    }
  }

  transformResults.set(transformOp->getOpResult(0), originalConsumerOps);
  transformResults.set(transformOp->getOpResult(1), fusedConsumerOps);
  return success();
}

DiagnosedSilenceableFailure transform::TestFuseConsumerUsingSliceOp::apply(
    TransformRewriter &rewriter, TransformResults &transformResults,
    TransformState &state) {
  SmallVector<Operation *> slices;
  for (auto op : getTargets()) {
    auto *sliceOp = *state.getPayloadOps(op).begin();
    slices.push_back(sliceOp);
  }

  SmallVector<LoopLikeOpInterface> loops;
  for (auto op : llvm::reverse(getLoops())) {
    auto loopLikeOp =
        dyn_cast<LoopLikeOpInterface>(*state.getPayloadOps(op).begin());
    if (!loopLikeOp) {
      return DiagnosedSilenceableFailure::definiteFailure();
    }
    loops.push_back(loopLikeOp);
  }
  LogicalResult result =
      applyFuseConsumerUsingSlices(rewriter, getOperation(), slices, loops,
                                   getNumConsumerToFuse(), transformResults);
  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

void transform::TestFuseConsumerUsingSliceOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  consumesHandle(getTargetsMutable(), effects);
  consumesHandle(getLoopsMutable(), effects);
  producesHandle(getOperation()->getOpResults(), effects);
  modifiesPayload(effects);
}

//===----------------------------------------------------------------------===//
// TestTileUsingForallOp
//===----------------------------------------------------------------------===//

/// Apply a tiling transformation to all payload ops and store both the
/// tiled operation as well as the created tile loops.
template <typename Range>
static LogicalResult
applyTileToAll(RewriterBase &rewriter, Operation *transformOp,
               Range &&payloadOps, ArrayRef<OpFoldResult> tileSizes,
               ArrayRef<int64_t> interchange, std::optional<ArrayAttr> mapping,
               TransformResults &transformResults) {
  SmallVector<Operation *> tiledOps;
  SmallVector<Operation *> loopOps;

  for (Operation *target : payloadOps) {
    auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
    if (!tilingInterfaceOp)
      return transformOp->emitError("only TilingInterface ops are supported");
    scf::SCFTilingOptions tilingOptions;
    tilingOptions.setTileSizes(tileSizes).setInterchange(interchange);
    if (mapping) {
      tilingOptions.setMapping(mapping.value().getValue());
    }
    tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);

    rewriter.setInsertionPoint(target);
    FailureOr<scf::SCFTilingResult> tiledResults =
        scf::tileUsingSCF(rewriter, tilingInterfaceOp, tilingOptions);
    if (failed(tiledResults))
      return failure();

    // Perform the replacement of tiled and fused values.
    rewriter.replaceOp(tilingInterfaceOp, tiledResults->replacements);

    // Report back the relevant handles to the transform op.
    tiledOps.push_back(tiledResults->tiledOps.front());
    for (Operation *loop : tiledResults->loops)
      loopOps.push_back(loop);
  }

  transformResults.set(transformOp->getOpResult(0), tiledOps);
  for (auto [index, loop] : llvm::enumerate(loopOps))
    transformResults.set(transformOp->getOpResult(index + 1), {loop});

  return success();
}

DiagnosedSilenceableFailure
transform::TestTileUsingForallOp::apply(TransformRewriter &rewriter,
                                        TransformResults &transformResults,
                                        TransformState &state) {
  SmallVector<int64_t> tileSizes =
      extractFromIntegerArrayAttr<int64_t>(getTileSizes());
  SmallVector<int64_t> interchange =
      extractFromIntegerArrayAttr<int64_t>(getInterchange());
  SmallVector<OpFoldResult> tileSizesOfr =
      getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);

  LogicalResult result =
      applyTileToAll(rewriter, getOperation(), state.getPayloadOps(getTarget()),
                     tileSizesOfr, interchange, getMapping(), transformResults);
  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

void transform::TestTileUsingForallOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  consumesHandle(getTargetMutable(), effects);
  producesHandle(getOperation()->getOpResults(), effects);
  modifiesPayload(effects);
}

//===----------------------------------------------------------------------===//
// TestFuseUsingForallOp
//===----------------------------------------------------------------------===//

/// Apply a tiling transformation to all payload ops and store both the
/// tiled operation as well as the created tile loops.
template <typename Range>
static LogicalResult applyTilingToAll(
    RewriterBase &rewriter, Operation *transformOp, Range &&payloadOps,
    unsigned numLoops, TransformResults &transformResults,
    function_ref<FailureOr<scf::SCFTileAndFuseResult>(TilingInterface)>
        applyFn) {
  SmallVector<Operation *> tiledLinalgOps;
  SmallVector<SmallVector<Operation *>> loopOps(1);

  for (Operation *target : payloadOps) {
    auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
    if (!tilingInterfaceOp)
      return transformOp->emitError("only TilingInterface ops are supported");

    rewriter.setInsertionPoint(target);
    FailureOr<scf::SCFTileAndFuseResult> tiledResults =
        applyFn(tilingInterfaceOp);
    if (failed(tiledResults))
      return failure();

    // Perform the replacement of tiled and fused values.
    SmallVector<Operation *> opsToReplace{target};
    llvm::append_range(opsToReplace, tiledResults->fusedProducers);
    for (Operation *toReplace : opsToReplace) {
      for (OpResult res : toReplace->getResults())
        if (auto replacement = tiledResults->replacements.lookup(res))
          rewriter.replaceAllUsesWith(res, replacement);
      if (toReplace->use_empty())
        rewriter.eraseOp(toReplace);
    }

    // Report back the relevant handles to the transform op.
    tiledLinalgOps.push_back(tiledResults->tiledAndFusedOps.front());
    assert(tiledResults->loops.size() == 1 &&
           cast<scf::ForallOp>(tiledResults->loops[0]).getRank() == numLoops &&
           "Mismatched number of loops, tile and fuse transform should have "
           "failed");
    loopOps[0] = {tiledResults->loops[0]};
  }

  transformResults.set(transformOp->getOpResult(0), tiledLinalgOps);
  if (!loopOps.empty())
    transformResults.set(transformOp->getOpResult(1), loopOps[0]);

  return success();
}

DiagnosedSilenceableFailure
transform::TestFuseUsingForallOp::apply(TransformRewriter &rewriter,
                                        TransformResults &transformResults,
                                        TransformState &state) {
  SmallVector<int64_t> tileSizes =
      extractFromIntegerArrayAttr<int64_t>(getTileSizes());
  SmallVector<int64_t> tileInterchange =
      extractFromIntegerArrayAttr<int64_t>(getInterchange());

  scf::SCFTilingOptions tilingOptions;
  tilingOptions.interchangeVector = tileInterchange;
  SmallVector<OpFoldResult> tileSizesOfr =
      getAsIndexOpFoldResult(rewriter.getContext(), tileSizes);
  tilingOptions = tilingOptions.setTileSizes(tileSizesOfr);
  tilingOptions.setLoopType(scf::SCFTilingOptions::LoopType::ForallOp);
  scf::SCFTileAndFuseOptions tileAndFuseOptions;
  tileAndFuseOptions.tilingOptions = tilingOptions;
  LogicalResult result = applyTilingToAll(
      rewriter, getOperation(), state.getPayloadOps(getRootOp()),
      tileSizes.size() - llvm::count(tileSizes, 0), transformResults,
      [&](TilingInterface tilingInterfaceOp)
          -> FailureOr<scf::SCFTileAndFuseResult> {
        return tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,
                                                    tileAndFuseOptions);
      });
  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

void transform::TestFuseUsingForallOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  consumesHandle(getRootOpMutable(), effects);
  producesHandle(getOperation()->getOpResults(), effects);
  modifiesPayload(effects);
}

//===----------------------------------------------------------------------===//
// TestTileAndFuseOuterParallelPartialReduction
//===----------------------------------------------------------------------===//

DiagnosedSilenceableFailure
transform::TestTileAndFuseOuterParallelPartialReductionOp::apply(
    TransformRewriter &rewriter, TransformResults &transformResults,
    TransformState &state) {
  auto target =
      dyn_cast<TilingInterface>(*state.getPayloadOps(getRootOp()).begin());
  if (!target) {
    emitOpError("expected root operation to implement `TilingInterface`");
    return DiagnosedSilenceableFailure::definiteFailure();
  }

  SmallVector<unsigned> reductionDims =
      extractFromIntegerArrayAttr<unsigned>(getReductionDims());
  if (reductionDims.empty()) {
    for (auto [index, iterator] :
         llvm::enumerate(target.getLoopIteratorTypes()))
      if (iterator == utils::IteratorType::reduction)
        reductionDims.push_back(index);
  }

  if (reductionDims.empty()) {
    emitOpError(
        "no reduction dimension specified or found in the target operation");
    return DiagnosedSilenceableFailure::definiteFailure();
  }

  SmallVector<int64_t> reductionTileSizes =
      extractFromIntegerArrayAttr<int64_t>(getTileSizes());
  if (reductionTileSizes.size() != reductionDims.size()) {
    emitOpError(
        "missing tile sizes for reduction dimensions that are to be tiled");
    return DiagnosedSilenceableFailure::definiteFailure();
  }

  // Adjust tile sizes so that it corresponds to the reduction iterator types.
  SmallVector<OpFoldResult> tileSizes;
  int reductionTileSizeNum = 0;
  OpFoldResult zero = rewriter.getIndexAttr(0);
  for (auto iterator : target.getLoopIteratorTypes()) {
    if (iterator == utils::IteratorType::parallel) {
      tileSizes.push_back(zero);
      continue;
    }
    tileSizes.push_back(
        rewriter.getIndexAttr(reductionTileSizes[reductionTileSizeNum++]));
  }

  scf::SCFTilingOptions tilingOptions;
  tilingOptions.setTileSizes(tileSizes)
      .setLoopType(scf::SCFTilingOptions::LoopType::ForallOp)
      .setReductionTilingStrategy(
          ReductionTilingStrategy::PartialReductionOuterParallel)
      .setReductionDims(reductionDims);
  if (auto mapping = getMapping()) {
    tilingOptions.setMapping(getMapping().value());
  }

  LogicalResult result = applyTileAndFuseToAll(
      rewriter, getOperation(), state.getPayloadOps(getRootOp()),
      /*numLoops =*/1, tilingOptions, transformResults);

  return failed(result) ? DiagnosedSilenceableFailure::definiteFailure()
                        : DiagnosedSilenceableFailure::success();
}

//===----------------------------------------------------------------------===//
// TestTileAndFuseOuterParallelPartialReduction
//===----------------------------------------------------------------------===//

DiagnosedSilenceableFailure transform::TestTileUsingCustomLoopOp::apply(
    TransformRewriter &transformRewriter, TransformResults &transformResults,
    TransformState &state) {
  auto target =
      dyn_cast<TilingInterface>(*state.getPayloadOps(getRootOp()).begin());
  if (!target) {
    emitOpError("expected root operation to implement `TilingInterface`");
    return DiagnosedSilenceableFailure::definiteFailure();
  }

  OpFoldResult oneOfr = transformRewriter.getIndexAttr(1);

  scf::SCFTilingOptions::GenerateLoopHeaderFn loopHeaderFn =
      [&](RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,
          ArrayRef<OpFoldResult> givenTileSizes,
          ValueRange outerDestinationTensors)
      -> FailureOr<scf::SCFTilingOptions::CustomLoopHeaderInfo> {
    // Check that the strides are all 1 (to make it easier in the test).
    if (llvm::any_of(loopRanges, [](Range r) {
          return !isConstantIntValue(r.stride, 1);
        })) {
      return emitOpError("unable to handle loop ranges with strides != 1");
    }
    // Check number of tile sizes is equal to loop dimensions.
    if (loopRanges.size() != givenTileSizes.size()) {
      return emitOpError("expected number of tile sizes to be same as the "
                         "number of loops in the operation");
    }
    // For testing disallow any of the tile sizes being 0.
    if (llvm::any_of(givenTileSizes, isZeroInteger)) {
      return emitOpError("unhandled case of zero tile size");
    }
    // For testing, only handle tensor tiling.
    if (outerDestinationTensors.empty()) {
      return emitOpError("expected destination tensors");
    }

    // Compute the number of iterations for each of the loops.
    AffineExpr s0, s1, s2;
    bindSymbols(rewriter.getContext(), s0, s1, s2);
    AffineExpr numItersExpr = (s1 - s0).ceilDiv(s2); // (ub - lb) / tileSize

    SmallVector<OpFoldResult> allNumIters;
    allNumIters.reserve(loopRanges.size());
    for (auto [loopRange, tileSize] :
         llvm::zip_equal(loopRanges, givenTileSizes)) {
      OpFoldResult numIters = affine::makeComposedFoldedAffineApply(
          rewriter, loc, numItersExpr,
          {loopRange.offset, loopRange.size, tileSize});
      allNumIters.push_back(numIters);
    }
    if (allNumIters.empty()) {
      return emitOpError("invalid empty tile sizes and loop ranges");
    }

    AffineExpr mulExpr = s0 * s1;
    OpFoldResult cumulative = oneOfr;
    for (auto numIters : allNumIters) {
      cumulative = affine::makeComposedFoldedAffineApply(
          rewriter, loc, mulExpr, {cumulative, numIters});
    }

    Value zeroVal = arith::ConstantIndexOp::create(rewriter, loc, 0);
    Value oneVal = arith::ConstantIndexOp::create(rewriter, loc, 1);
    Value ub = getValueOrCreateConstantIndexOp(rewriter, loc, cumulative);

    SmallVector<OpFoldResult> offsets;
    SmallVector<OpFoldResult> sizes;
    SmallVector<Value> innerDestinationTensors;
    offsets.reserve(loopRanges.size());
    sizes.reserve(loopRanges.size());

    AffineExpr d0;
    bindDims(rewriter.getContext(), d0);
    AffineExpr offsetExpr = s0 + d0 * s1; // lb + iv * tileSize
    AffineMap minMap =
        AffineMap::get(1, 2, {s0 - d0, s1},
                       rewriter.getContext()); // min(ub - offset, tileSize)
    auto forOp = scf::ForOp::create(
        rewriter, loc, zeroVal, ub, oneVal, outerDestinationTensors,
        [&](OpBuilder &b, Location bodyLoc, Value linearizedIv,
            ValueRange destinations) {
          auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
              b, bodyLoc, linearizedIv, allNumIters);
          for (auto [normalizedIv, range, tileSize] : llvm::zip_equal(
                   delinearizeOp.getResults(), loopRanges, givenTileSizes)) {

            OpFoldResult normalizedIvOfr = getAsOpFoldResult(normalizedIv);
            OpFoldResult offset = affine::makeComposedFoldedAffineApply(
                b, bodyLoc, offsetExpr,
                {normalizedIvOfr, range.offset, tileSize});
            offsets.push_back(offset);

            OpFoldResult size = affine::makeComposedFoldedAffineMin(
                b, bodyLoc, minMap, {offset, range.size, tileSize});
            sizes.push_back(size);
          }
          innerDestinationTensors = llvm::to_vector(destinations);
        });
    rewriter.setInsertionPointToEnd(forOp.getBody());
    return scf::SCFTilingOptions::CustomLoopHeaderInfo{
        {cast<LoopLikeOpInterface>(forOp.getOperation())},
        offsets,
        sizes,
        innerDestinationTensors};
  };

  scf::SCFTilingOptions::GenerateLoopTerminatorFn terminatorFn =
      [&](RewriterBase &rewriter, Location loc,
          ArrayRef<LoopLikeOpInterface> loops, ValueRange tiledResults,
          ArrayRef<SmallVector<OpFoldResult>> resultOffsets,
          ArrayRef<SmallVector<OpFoldResult>> resultSizes,
          ValueRange destinationTensors) -> LogicalResult {
    SmallVector<Value> yieldValues;
    yieldValues.reserve(destinationTensors.size());
    for (auto [tiledResult, offsets, sizes, destination] : llvm::zip_equal(
             tiledResults, resultOffsets, resultSizes, destinationTensors)) {
      SmallVector<OpFoldResult> strides(offsets.size(), oneOfr);
      Value insertedVal = tensor::InsertSliceOp::create(
          rewriter, loc, tiledResult, destination, offsets, sizes, strides);
      yieldValues.push_back(insertedVal);
    }
    scf::YieldOp::create(rewriter, loc, yieldValues);
    return success();
  };

  scf::SCFTilingOptions tilingOptions;
  SmallVector<int64_t> staticTileSizes =
      extractFromIntegerArrayAttr<int64_t>(getTileSizes());
  SmallVector<OpFoldResult> tileSizes =
      getAsIndexOpFoldResult(transformRewriter.getContext(), staticTileSizes);
  tilingOptions.setTileSizes(tileSizes)
      .setLoopType(scf::SCFTilingOptions::LoopType::CustomOp)
      .setCustomLoopGenerationFns(loopHeaderFn, terminatorFn);

  OpBuilder::InsertionGuard g(transformRewriter);
  transformRewriter.setInsertionPoint(target);
  FailureOr<scf::SCFTilingResult> tiledResults =
      scf::tileUsingSCF(transformRewriter, target, tilingOptions);
  if (failed(tiledResults)) {
    return DiagnosedSilenceableFailure::definiteFailure();
  }
  transformRewriter.replaceOp(target, tiledResults->replacements);
  transformResults.set(getOperation()->getResult(0), tiledResults->tiledOps);
  transformResults.set(getOperation()->getResult(1), tiledResults->loops);

  return DiagnosedSilenceableFailure::success();
}

//===----------------------------------------------------------------------===//
// TestQueryProducerFusability
//===----------------------------------------------------------------------===//

DiagnosedSilenceableFailure transform::TestQueryProducerFusability::apply(
    TransformRewriter &rewriter, TransformResults &transformResults,
    TransformState &state) {
  for (Operation *target : state.getPayloadOps(getTarget())) {
    auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
    if (!tilingInterfaceOp) {
      return emitSilenceableError()
             << "target operation does not implement TilingInterface";
    }

    // Collect operand numbers and their corresponding producer insert_slice
    // offsets and sizes.
    SmallVector<unsigned> operandNumbers;
    SmallVector<SmallVector<OpFoldResult>> allOffsets;
    SmallVector<SmallVector<OpFoldResult>> allSizes;

    for (OpOperand &operand : target->getOpOperands()) {
      Value operandValue = operand.get();
      Operation *definingOp = operandValue.getDefiningOp();

      // Look for a producer tensor.insert_slice. This is only for testing
      // purposes and otherwise is not a useful transformation.
      if (auto insertSliceOp =
              dyn_cast_or_null<tensor::InsertSliceOp>(definingOp)) {
        operandNumbers.push_back(operand.getOperandNumber());
        allOffsets.push_back(insertSliceOp.getMixedOffsets());
        allSizes.push_back(insertSliceOp.getMixedSizes());
      }
    }

    if (!operandNumbers.empty()) {
      bool isFusable = tilingInterfaceOp.isOpFusableWithProducerSlices(
          operandNumbers, allOffsets, allSizes);

      if (isFusable) {
        target->emitRemark()
            << "can be fused with producer tensor.insert_slice ops";
      } else {
        target->emitRemark()
            << "cannot be fused with producer tensor.insert_slice ops";
      }
    }
  }

  return DiagnosedSilenceableFailure::success();
}

void transform::TestQueryProducerFusability::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  onlyReadsHandle(getTargetMutable(), effects);
  onlyReadsPayload(effects);
}

//===----------------------------------------------------------------------===//
// TestQueryConsumerFusability
//===----------------------------------------------------------------------===//

DiagnosedSilenceableFailure transform::TestQueryConsumerFusability::apply(
    TransformRewriter &rewriter, TransformResults &transformResults,
    TransformState &state) {
  for (Operation *target : state.getPayloadOps(getTarget())) {
    auto tilingInterfaceOp = dyn_cast<TilingInterface>(target);
    if (!tilingInterfaceOp) {
      return emitSilenceableError()
             << "target operation does not implement TilingInterface";
    }

    // Look for tensor.extract_slice ops that consume results of the tilable op.
    for (OpResult result : target->getResults()) {
      for (OpOperand &use : result.getUses()) {
        Operation *user = use.getOwner();

        // Look for a consumer tensor.extract_slice. This is only for testing
        // purposes and otherwise is not a useful transformation.
        if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(user)) {
          bool isFusable = tilingInterfaceOp.isOpFusableWithConsumerSlice(
              result.getResultNumber(), extractSliceOp.getMixedOffsets(),
              extractSliceOp.getMixedSizes());

          if (isFusable) {
            target->emitRemark()
                << "can be fused with consumer tensor.extract_slice op";
          } else {
            target->emitRemark()
                << "cannot be fused with consumer tensor.extract_slice op";
          }
        }
      }
    }
  }

  return DiagnosedSilenceableFailure::success();
}

void transform::TestQueryConsumerFusability::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  onlyReadsHandle(getTargetMutable(), effects);
  onlyReadsPayload(effects);
}

#define GET_OP_CLASSES
#include "TestTilingInterfaceTransformOps.cpp.inc"

namespace {
class TestTilingInterfaceDialectExtension
    : public transform::TransformDialectExtension<
          TestTilingInterfaceDialectExtension> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      TestTilingInterfaceDialectExtension)

  using Base::Base;

  void init() {
    declareDependentDialect<affine::AffineDialect>();
    declareDependentDialect<index::IndexDialect>();
    declareDependentDialect<scf::SCFDialect>();
    declareDependentDialect<tensor::TensorDialect>();

    registerTransformOps<
#define GET_OP_LIST
#include "TestTilingInterfaceTransformOps.cpp.inc"
        >();
  }
};
} // namespace

namespace test {
void registerTestTilingInterfaceTransformDialectExtension(
    DialectRegistry &registry) {
  registry.addExtensions<TestTilingInterfaceDialectExtension>();
}
} // namespace test