//===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
//
// 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 implements logic for testing Linalg transformations.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Vector/VectorOps.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"

using namespace mlir;
using namespace mlir::linalg;

namespace {
struct TestLinalgTransforms
    : public PassWrapper<TestLinalgTransforms, FunctionPass> {
  TestLinalgTransforms() = default;
  TestLinalgTransforms(const TestLinalgTransforms &pass) {}

  void getDependentDialects(DialectRegistry &registry) const override {
    // clang-format off
    registry.insert<AffineDialect,
                    memref::MemRefDialect,
                    scf::SCFDialect,
                    StandardOpsDialect,
                    vector::VectorDialect,
                    gpu::GPUDialect>();
    // clang-format on
  }
  StringRef getArgument() const final {
    return "test-linalg-transform-patterns";
  }
  StringRef getDescription() const final {
    return "Test Linalg transformation patterns by applying them greedily.";
  }

  void runOnFunction() override;

  Option<bool> testPatterns{*this, "test-patterns",
                            llvm::cl::desc("Test a mixed set of patterns"),
                            llvm::cl::init(false)};
  Option<bool> testMatmulToVectorPatterns1dTiling{
      *this, "test-matmul-to-vector-patterns-tile-1d",
      llvm::cl::desc(
          "Test a fused pass that applies patterns from matmul to vectors via "
          "1-d tiling"),
      llvm::cl::init(false)};
  Option<bool> testMatmulToVectorPatterns2dTiling{
      *this, "test-matmul-to-vector-patterns-tile-2d",
      llvm::cl::desc(
          "Test a fused pass that applies patterns from matmul to vectors via "
          "2-d tiling"),
      llvm::cl::init(false)};
  Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
                                    llvm::cl::desc("Test promotion options"),
                                    llvm::cl::init(false)};
  Option<bool> testTileAndDistributionOptions{
      *this, "test-tile-and-distribute-options",
      llvm::cl::desc("Test tile and distribute options"),
      llvm::cl::init(false)};
  Option<bool> testVectorTransferForwardingPatterns{
      *this, "test-vector-transfer-forwarding-patterns",
      llvm::cl::desc(
          "Test a fused pass that forwards linalg.copy to vector.transfer"),
      llvm::cl::init(false)};
  Option<bool> testGenericToVectorPattern{
      *this, "test-linalg-to-vector-patterns",
      llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
                     "in vector.contract form"),
      llvm::cl::init(false)};
  Option<bool> testTilePattern{*this, "test-tile-pattern",
                               llvm::cl::desc("Test tile pattern"),
                               llvm::cl::init(false)};
  Option<bool> testTileScalarizeDynamicDims{
      *this, "test-tile-scalarize-dynamic-dims",
      llvm::cl::desc("Test tiling of dynamic dims by 1"),
      llvm::cl::init(false)};
  Option<bool> testTransformPadTensor{
      *this, "test-transform-pad-tensor",
      llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
      llvm::cl::init(false)};
  Option<bool> testGeneralizePadTensor{
      *this, "test-generalize-pad-tensor",
      llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
      llvm::cl::init(false)};
  Option<bool> testSwapSubTensorPadTensor{
      *this, "test-swap-subtensor-padtensor",
      llvm::cl::desc("Test rewrite of subtensor(pad_tensor) into "
                     "pad_tensor(subtensor)"),
      llvm::cl::init(false)};
  ListOption<int64_t> peeledLoops{
      *this, "peeled-loops",
      llvm::cl::desc("Loops to be peeled when test-tile-pattern"),
      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
  ListOption<int64_t> tileSizes{
      *this, "tile-sizes",
      llvm::cl::desc("Linalg tile sizes for test-tile-pattern"),
      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
  ListOption<unsigned> testTiledLoopPeeling{
      *this, "test-tiled-loop-peeling",
      llvm::cl::desc("Test peeling of linalg.tiled_loop ops"),
      llvm::cl::OneOrMore, llvm::cl::MiscFlags::CommaSeparated};
  Option<bool> skipPartial{
      *this, "skip-partial",
      llvm::cl::desc("Skip loops inside partial iterations during peeling"),
      llvm::cl::init(false)};
  Option<std::string> loopType{
      *this, "loop-type",
      llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
                     "tiled_loop"),
      llvm::cl::init("for")};
};
} // end anonymous namespace

static void applyPatterns(FuncOp funcOp) {
  MLIRContext *ctx = funcOp.getContext();
  RewritePatternSet patterns(ctx);

  //===--------------------------------------------------------------------===//
  // Linalg tiling patterns.
  //===--------------------------------------------------------------------===//
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({2000, 3000, 4000}),
      LinalgTransformationFilter(StringAttr::get(ctx, "MEM"),
                                 StringAttr::get(ctx, "L3")));
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({200, 300, 400}),
      LinalgTransformationFilter(StringAttr::get(ctx, "L3"),
                                 StringAttr::get(ctx, "L2")));
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
      LinalgTransformationFilter(StringAttr::get(ctx, "L2"),
                                 StringAttr::get(ctx, "L1")));
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({2, 3, 4}),
      LinalgTransformationFilter(StringAttr::get(ctx, "L1"),
                                 StringAttr::get(ctx, "REG")));

  patterns.add<LinalgTilingPattern<MatvecOp>>(
      ctx,
      LinalgTilingOptions().setTileSizes({5, 6}).setLoopType(
          LinalgTilingLoopType::ParallelLoops),
      LinalgTransformationFilter(ArrayRef<StringAttr>{},
                                 StringAttr::get(ctx, "L1")));

  patterns.add<LinalgTilingPattern<DotOp>>(
      ctx, LinalgTilingOptions().setTileSizes(8000),
      LinalgTransformationFilter(
          ArrayRef<StringAttr>{StringAttr::get(ctx, "MEM"),
                               StringAttr::get(ctx, "L3"),
                               StringAttr::get(ctx, "L2")},
          StringAttr::get(ctx, "REG")));

  //===--------------------------------------------------------------------===//
  // Linalg tiling and permutation patterns.
  //===--------------------------------------------------------------------===//
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx,
      LinalgTilingOptions()
          .setTileSizes({2000, 3000, 4000})
          .setInterchange({1, 2, 0}),
      LinalgTransformationFilter(StringAttr::get(ctx, "__with_perm__"),
                                 StringAttr::get(ctx, "L2__with_perm__")));
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx,
      LinalgTilingOptions()
          .setTileSizes({200, 300, 400})
          .setInterchange({1, 0, 2}),
      LinalgTransformationFilter(StringAttr::get(ctx, "L2__with_perm__"),
                                 StringAttr::get(ctx, "L1__with_perm__")));
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
      LinalgTransformationFilter(StringAttr::get(ctx, "L1__with_perm__"),
                                 StringAttr::get(ctx, "REG__with_perm__")));

  patterns.add<LinalgTilingPattern<MatvecOp>>(
      ctx, LinalgTilingOptions().setTileSizes({5, 6}).setInterchange({1, 0}),
      LinalgTransformationFilter(StringAttr::get(ctx, "__with_perm__"),
                                 StringAttr::get(ctx, "L1__with_perm__")));

  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx,
      LinalgTilingOptions()
          .setTileSizes({16, 8, 4})
          .setInterchange({1, 2, 0})
          .setLoopType(LinalgTilingLoopType::ParallelLoops),
      LinalgTransformationFilter(
          StringAttr::get(ctx, "par__with_perm__"),
          StringAttr::get(ctx, "after_par__with_perm__")));

  //===--------------------------------------------------------------------===//
  // Linalg to loops patterns.
  //===--------------------------------------------------------------------===//
  patterns.add<LinalgLoweringPattern<DotOp>>(
      ctx,
      /*loweringType=*/LinalgLoweringType::Loops,
      LinalgTransformationFilter(StringAttr::get(ctx, "REG")));

  //===--------------------------------------------------------------------===//
  // Linalg distribution patterns.
  //===--------------------------------------------------------------------===//
  LinalgLoopDistributionOptions distributionOptions;

  //===--------------------------------------------------------------------===//
  // Linalg to vector contraction patterns.
  //===--------------------------------------------------------------------===//
  patterns.add<LinalgVectorizationPattern>(
      ctx, LinalgTransformationFilter(StringAttr::get(ctx, "VECTORIZE"))
               .addOpFilter<MatmulOp, FillOp, CopyOp, GenericOp>());

  //===--------------------------------------------------------------------===//
  // Linalg generic interchange pattern.
  //===--------------------------------------------------------------------===//
  patterns.add<GenericOpInterchangePattern>(
      ctx,
      /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
      LinalgTransformationFilter(ArrayRef<StringAttr>{},
                                 StringAttr::get(ctx, "PERMUTED")));

  //===--------------------------------------------------------------------===//
  // Linalg subview operands promotion.
  //===--------------------------------------------------------------------===//
  patterns.add<LinalgPromotionPattern<MatmulOp>>(
      ctx, LinalgPromotionOptions().setUseFullTileBuffersByDefault(true),
      LinalgTransformationFilter(StringAttr::get(ctx, "_promote_views_"),
                                 StringAttr::get(ctx, "_views_promoted_")));
  patterns.add<LinalgPromotionPattern<MatmulOp>>(
      ctx,
      LinalgPromotionOptions()
          .setOperandsToPromote({0})
          .setUseFullTileBuffersByDefault(true),
      LinalgTransformationFilter(
          StringAttr::get(ctx, "_promote_first_view_"),
          StringAttr::get(ctx, "_first_view_promoted_")));
  patterns.add<LinalgPromotionPattern<FillOp>>(
      ctx,
      LinalgPromotionOptions()
          .setOperandsToPromote({1})
          .setUseFullTileBuffers({false, true})
          .setAlignment(32),
      LinalgTransformationFilter(
          StringAttr::get(ctx, "_promote_views_aligned_"),
          StringAttr::get(ctx, "_views_aligned_promoted_")));

  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));

  // Drop the marker.
  funcOp.walk([](LinalgOp op) {
    op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
  });
}

static void fillL1TilingAndMatmulToVectorPatterns(
    FuncOp funcOp, StringRef startMarker,
    SmallVectorImpl<RewritePatternSet> &patternsVector) {
  MLIRContext *ctx = funcOp.getContext();
  patternsVector.emplace_back(
      ctx, std::make_unique<LinalgTilingPattern<MatmulOp>>(
               ctx,
               LinalgTilingOptions()
                   .setTileSizes({8, 12, 16})
                   .setInterchange({1, 0, 2}),
               LinalgTransformationFilter(StringAttr::get(ctx, startMarker),
                                          StringAttr::get(ctx, "L1"))));

  patternsVector.emplace_back(
      ctx,
      std::make_unique<LinalgPromotionPattern<MatmulOp>>(
          ctx, LinalgPromotionOptions().setUseFullTileBuffersByDefault(true),
          LinalgTransformationFilter(StringAttr::get(ctx, "L1"),
                                     StringAttr::get(ctx, "VEC"))));

  patternsVector.emplace_back(
      ctx, std::make_unique<LinalgVectorizationPattern>(
               MatmulOp::getOperationName(), ctx, LinalgVectorizationOptions(),
               LinalgTransformationFilter(StringAttr::get(ctx, "VEC"))));
  patternsVector.back().add<LinalgVectorizationPattern>(
      ctx, LinalgTransformationFilter().addFilter(
               [](Operation *op) { return success(isa<FillOp, CopyOp>(op)); }));
}

//===----------------------------------------------------------------------===//
// Test promotion callbacks
//===----------------------------------------------------------------------===//

// Allocation call back
static Optional<Value> allocCallBackFn(OpBuilder &b, memref::SubViewOp subView,
                                       ArrayRef<Value> boundingSubViewSize,
                                       DataLayout &layout) {
  SmallVector<int64_t, 4> shape(boundingSubViewSize.size(), -1);
  return b
      .create<memref::AllocOp>(
          subView.getLoc(),
          MemRefType::get(shape, subView.getType().getElementType(),
                          /*affineMapComposition =*/{}, 3),
          boundingSubViewSize)
      .getResult();
}

// Deallocation callback
static LogicalResult deallocCallBackFn(OpBuilder &b, Value buffer) {
  b.create<memref::DeallocOp>(buffer.getLoc(), buffer);
  return success();
}

// Copy in call back
static LogicalResult copyCallBackFn(OpBuilder &b, Value src, Value dst,
                                    bool isOutput) {
  auto floatType = src.getType().cast<MemRefType>().getElementType();
  if (!floatType.isa<FloatType>())
    return failure();
  if (!isOutput) {
    Value cst = b.create<arith::ConstantOp>(src.getLoc(),
                                            FloatAttr::get(floatType, 42.0));
    b.create<FillOp>(src.getLoc(), cst, dst);
  }
  b.create<CopyOp>(src.getLoc(), src, dst);
  return success();
}

static void fillPromotionCallBackPatterns(MLIRContext *ctx,
                                          RewritePatternSet &patterns) {
  patterns.add<LinalgTilingPattern<MatmulOp>>(
      ctx, LinalgTilingOptions().setTileSizes({16, 16, 16}),
      LinalgTransformationFilter(StringAttr::get(ctx, "START"),
                                 StringAttr::get(ctx, "PROMOTE")));
  patterns.add<LinalgPromotionPattern<MatmulOp>>(
      ctx,
      LinalgPromotionOptions()
          .setOperandsToPromote({0, 2})
          .setUseFullTileBuffers({false, false})
          .setAllocationDeallocationFns(allocCallBackFn, deallocCallBackFn)
          .setCopyInOutFns(
              [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                return copyCallBackFn(b, src, dst, false);
              },
              [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
                return copyCallBackFn(b, src, dst, true);
              }),
      LinalgTransformationFilter(StringAttr::get(ctx, "PROMOTE")));
}

template <typename IdOp, typename NProcsOp>
static SmallVector<ProcInfo, 2>
getGpuProcIds(OpBuilder &b, Location loc, ArrayRef<Range> parallelLoopRanges) {
  size_t count = std::min<size_t>(3, parallelLoopRanges.size());
  SmallVector<ProcInfo, 2> procInfo(count);
  const char *xyz[] = {"x", "y", "z"};
  Type indexType = b.getIndexType();
  for (unsigned i = 0; i < count; ++i) {
    procInfo[count - 1 - i] = {
        b.create<IdOp>(loc, indexType, b.getStringAttr(xyz[i])),
        b.create<NProcsOp>(loc, indexType, b.getStringAttr(xyz[i]))};
  }
  return procInfo;
}

static void fillTileAndDistributePatterns(MLIRContext *context,
                                          RewritePatternSet &patterns) {
  {
    LinalgLoopDistributionOptions cyclicNprocsEqNiters;
    cyclicNprocsEqNiters.distributionMethod.resize(
        2, DistributionMethod::CyclicNumProcsEqNumIters);
    cyclicNprocsEqNiters.procInfo =
        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsEqNiters),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute1"),
            StringAttr::get(context, "after_distribute1")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsGeNiters;
    cyclicNprocsGeNiters.distributionMethod.resize(
        2, DistributionMethod::CyclicNumProcsGeNumIters);
    cyclicNprocsGeNiters.procInfo =
        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsGeNiters),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute2"),
            StringAttr::get(context, "after_distribute2")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsDefault;
    cyclicNprocsDefault.distributionMethod.resize(2,
                                                  DistributionMethod::Cyclic);
    cyclicNprocsDefault.procInfo =
        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsDefault),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute3"),
            StringAttr::get(context, "after_distribute3")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsMixed1;
    cyclicNprocsMixed1.distributionMethod = {
        DistributionMethod::CyclicNumProcsEqNumIters,
        DistributionMethod::CyclicNumProcsGeNumIters};
    cyclicNprocsMixed1.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsMixed1),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute4"),
            StringAttr::get(context, "after_distribute4")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsMixed2;
    cyclicNprocsMixed2.distributionMethod = {
        DistributionMethod::CyclicNumProcsGeNumIters,
        DistributionMethod::Cyclic};
    cyclicNprocsMixed2.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsMixed2),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute5"),
            StringAttr::get(context, "after_distribute5")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsMixed3;
    cyclicNprocsMixed3.distributionMethod = {
        DistributionMethod::Cyclic,
        DistributionMethod::CyclicNumProcsEqNumIters};
    cyclicNprocsMixed3.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;

    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::ParallelLoops)
            .setDistributionOptions(cyclicNprocsMixed3),
        LinalgTransformationFilter(
            StringAttr::get(context, "distribute6"),
            StringAttr::get(context, "after_distribute6")));
  }

  {
    LinalgLoopDistributionOptions cyclicNprocsEqNiters;
    cyclicNprocsEqNiters.distributionMethod.resize(2,
                                                   DistributionMethod::Cyclic);
    cyclicNprocsEqNiters.procInfo =
        getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
    patterns.add<LinalgTilingPattern<MatmulOp>>(
        context,
        LinalgTilingOptions()
            .setTileSizes({8, 8, 4})
            .setLoopType(LinalgTilingLoopType::Loops)
            .setDistributionOptions(cyclicNprocsEqNiters),
        LinalgTransformationFilter(
            StringAttr::get(context, "tensors_distribute1"),
            StringAttr::get(context, "tensors_after_distribute1")));
  }
}

static void
applyMatmulToVectorPatterns(FuncOp funcOp,
                            bool testMatmulToVectorPatterns1dTiling,
                            bool testMatmulToVectorPatterns2dTiling) {
  MLIRContext *ctx = funcOp.getContext();
  SmallVector<RewritePatternSet, 4> stage1Patterns;
  if (testMatmulToVectorPatterns1dTiling) {
    fillL1TilingAndMatmulToVectorPatterns(funcOp, "START", stage1Patterns);
  } else if (testMatmulToVectorPatterns2dTiling) {
    stage1Patterns.emplace_back(
        ctx, std::make_unique<LinalgTilingPattern<MatmulOp>>(
                 ctx,
                 LinalgTilingOptions()
                     .setTileSizes({768, 264, 768})
                     .setInterchange({1, 2, 0}),
                 LinalgTransformationFilter(StringAttr::get(ctx, "START"),
                                            StringAttr::get(ctx, "L2"))));
    fillL1TilingAndMatmulToVectorPatterns(funcOp, "L2", stage1Patterns);
  }
  {
    // Canonicalization patterns
    RewritePatternSet canonicalizationPatterns(funcOp.getContext());
    vector::populateVectorTransferPermutationMapLoweringPatterns(
        canonicalizationPatterns);
    vector::populateVectorReductionToContractPatterns(canonicalizationPatterns);
    stage1Patterns.push_back(std::move(canonicalizationPatterns));
  }
  SmallVector<FrozenRewritePatternSet, 4> frozenStage1Patterns;
  llvm::move(stage1Patterns, std::back_inserter(frozenStage1Patterns));
  FrozenRewritePatternSet stage2Patterns =
      getLinalgTilingCanonicalizationPatterns(ctx);
  (void)applyStagedPatterns(funcOp, frozenStage1Patterns,
                            std::move(stage2Patterns));
}

static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
  RewritePatternSet forwardPattern(funcOp.getContext());
  forwardPattern.add<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
  forwardPattern.add<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(forwardPattern));
}

static void applyLinalgToVectorPatterns(FuncOp funcOp) {
  RewritePatternSet patterns(funcOp.getContext());
  patterns.add<LinalgVectorizationPattern>(
      funcOp.getContext(),
      LinalgTransformationFilter()
          .addOpFilter<ContractionOpInterface, FillOp, CopyOp, GenericOp>());
  populatePadTensorOpVectorizationPatterns(patterns);
  populateConvolutionVectorizationPatterns(patterns);
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}

static void applyPadTensorToGenericPatterns(FuncOp funcOp) {
  RewritePatternSet patterns(funcOp.getContext());
  patterns.add<PadTensorOpTransformationPattern>(funcOp.getContext());
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}

static void applyGeneralizePadTensorPatterns(FuncOp funcOp) {
  RewritePatternSet patterns(funcOp.getContext());
  patterns.add<GeneralizePadTensorOpPattern>(funcOp.getContext());
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}

static void applyExtractSliceOfPadTensorSwapPattern(FuncOp funcOp) {
  RewritePatternSet patterns(funcOp.getContext());
  patterns.add<ExtractSliceOfPadTensorSwapPattern>(funcOp.getContext());
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}

static void applyTilePattern(FuncOp funcOp, std::string loopType,
                             ArrayRef<int64_t> tileSizes,
                             ArrayRef<int64_t> peeledLoops,
                             bool scalarizeDynamicDims) {
  MLIRContext *context = funcOp.getContext();
  RewritePatternSet tilingPattern(context);
  LinalgTilingLoopType type =
      llvm::StringSwitch<LinalgTilingLoopType>(loopType)
          .Case("for", LinalgTilingLoopType::Loops)
          .Case("affine", LinalgTilingLoopType::AffineLoops)
          .Case("parallel", LinalgTilingLoopType::ParallelLoops)
          .Case("tiled_loop", LinalgTilingLoopType::TiledLoops);
  auto linalgTilingOptions = linalg::LinalgTilingOptions()
                                 .setPeeledLoops(peeledLoops)
                                 .setLoopType(type);
  if (scalarizeDynamicDims) {
    linalgTilingOptions.scalarizeDynamicDims();
    assert(tileSizes.empty() &&
           "tileSizes and scalarizeDynamicDims is mutually exclusive");
  } else {
    linalgTilingOptions.setTileSizes(tileSizes);
  }
  tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulOp>,
                    linalg::LinalgTilingPattern<linalg::GenericOp>>(
      context, linalgTilingOptions,
      linalg::LinalgTransformationFilter(StringAttr::get(context, "tile")));
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
}

static constexpr char kPeeledLoopsLabel[] = "__peeled_loops__";
static constexpr char kPartialIterationLabel[] = "__partial_iteration__";

namespace {
/// Peel TiledLoopOps, i.e., split them into two loops: One loop where the
/// `idx`-th loop contains only "full" iterations and a second loop for the
/// remaining partial iteration (if any).
struct TiledLoopPeelingPattern : public OpRewritePattern<TiledLoopOp> {
  TiledLoopPeelingPattern(MLIRContext *ctx, int64_t idx, bool skipPartial)
      : OpRewritePattern<TiledLoopOp>(ctx), idx(idx), skipPartial(skipPartial) {
  }

  LogicalResult matchAndRewrite(TiledLoopOp loopOp,
                                PatternRewriter &rewriter) const override {
    SmallVector<int64_t> peeledLoops;
    if (loopOp->hasAttr(kPeeledLoopsLabel)) {
      auto attr = loopOp->getAttr(kPeeledLoopsLabel).cast<ArrayAttr>();
      peeledLoops =
          llvm::to_vector<4>(llvm::map_range(attr, [](Attribute attr) {
            return attr.cast<IntegerAttr>().getInt();
          }));
      // Check if the loop was already peeled.
      if (llvm::find(peeledLoops, idx) != peeledLoops.end())
        return failure();
    }
    if (skipPartial && loopOp->hasAttr(kPartialIterationLabel))
      // No peeling of loop nests with a partial iteration.
      return failure();

    if (static_cast<int64_t>(loopOp.iterator_types().size()) <= idx)
      return failure();

    // Peel loop and canonicalize.
    TiledLoopOp result;
    if (failed(linalg::peelAndCanonicalizeTiledLoop(rewriter, loopOp, idx,
                                                    result)))
      return failure();

    // Apply label, so that the same loop is not rewritten a second time.
    peeledLoops.push_back(idx);
    rewriter.updateRootInPlace(loopOp, [&]() {
      loopOp->setAttr(kPeeledLoopsLabel, rewriter.getI64ArrayAttr(peeledLoops));
    });
    result->setAttr(kPeeledLoopsLabel, rewriter.getI64ArrayAttr(peeledLoops));
    result->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());

    return success();
  }

  /// Index of loop to peel.
  int64_t idx;

  /// If set to true, do not peel TiledLoopOps with a partial iteration.
  bool skipPartial;
};
} // namespace

static void applyTiledLoopPeelingPattern(FuncOp funcOp,
                                         ArrayRef<unsigned> loops,
                                         bool skipPartial) {
  MLIRContext *ctx = funcOp.getContext();
  RewritePatternSet patterns(ctx);
  for (unsigned idx : loops)
    patterns.add<TiledLoopPeelingPattern>(ctx, idx, skipPartial);
  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));

  // Drop the markers.
  funcOp.walk([](TiledLoopOp op) {
    op->removeAttr(kPeeledLoopsLabel);
    op->removeAttr(kPartialIterationLabel);
  });
}

/// Apply transformations specified as patterns.
void TestLinalgTransforms::runOnFunction() {
  auto lambda = [&](void *) {
    getFunction().walk([](LinalgOp op) {
      op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
    });
  };
  std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};

  if (testPromotionOptions) {
    RewritePatternSet patterns(&getContext());
    fillPromotionCallBackPatterns(&getContext(), patterns);
    (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
    return;
  }
  if (testTileAndDistributionOptions) {
    RewritePatternSet patterns(&getContext());
    fillTileAndDistributePatterns(&getContext(), patterns);
    (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
    return;
  }
  if (testPatterns)
    return applyPatterns(getFunction());
  if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
    return applyMatmulToVectorPatterns(getFunction(),
                                       testMatmulToVectorPatterns1dTiling,
                                       testMatmulToVectorPatterns2dTiling);
  if (testVectorTransferForwardingPatterns)
    return applyVectorTransferForwardingPatterns(getFunction());
  if (testGenericToVectorPattern)
    return applyLinalgToVectorPatterns(getFunction());
  if (testTransformPadTensor)
    return applyPadTensorToGenericPatterns(getFunction());
  if (testGeneralizePadTensor)
    return applyGeneralizePadTensorPatterns(getFunction());
  if (testSwapSubTensorPadTensor)
    return applyExtractSliceOfPadTensorSwapPattern(getFunction());
  if (testTiledLoopPeeling.hasValue())
    return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling,
                                        skipPartial);
  if (testTilePattern)
    return applyTilePattern(getFunction(), loopType, tileSizes, peeledLoops,
                            /*scalarizeDynamicDims=*/false);
  if (testTileScalarizeDynamicDims)
    return applyTilePattern(getFunction(), loopType, tileSizes,
                            /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
}

namespace mlir {
namespace test {
void registerTestLinalgTransforms() {
  PassRegistration<TestLinalgTransforms>();
}
} // namespace test
} // namespace mlir