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

#include "iree/compiler/Codegen/LLVMCPU/Passes.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-llvmcpu-peel"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_LLVMCPUPEELPASS
#include "iree/compiler/Codegen/LLVMCPU/Passes.h.inc"

namespace {

// Gathers tiled loops that aren't distribution loops from previous tiling
// stages.
void collectLoopsToPeel(Operation *op,
                        llvm::SmallSetVector<scf::ForOp, 8> &loopsToPeel) {
  if (!iree_compiler::getLoweringConfig(op)) {
    return;
  }

  int maxNumLoopsToPeel =
      TypeSwitch<Operation *, int>(op)
          .Case(
              [](linalg::LinalgOp linalgOp) { return linalgOp.getNumLoops(); })
          .Case([](linalg::PackOp packOp) { return packOp.getSourceRank(); })
          .Default([](auto) { return 0; });
  for (int i = 0; i < maxNumLoopsToPeel; ++i) {
    op = op->getParentOfType<scf::ForOp>();
    auto loop = cast_or_null<scf::ForOp>(op);
    if (!loop || iree_compiler::isTiledAndDistributedLoop(loop)) {
      break;
    }

    LDBG() << "Loop to peel\n  " << *op;
    loopsToPeel.insert(loop);
  }
}

class LLVMCPUPeelPass : public impl::LLVMCPUPeelPassBase<LLVMCPUPeelPass> {
public:
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<tensor::TensorDialect, linalg::LinalgDialect,
                    scf::SCFDialect>();
  }
  void runOnOperation() override;
};

void LLVMCPUPeelPass::runOnOperation() {
  MLIRContext *context = &getContext();
  mlir::FunctionOpInterface funcOp = getOperation();

  llvm::SmallSetVector<scf::ForOp, 8> uniqueLoopsToPeel;
  funcOp.walk([&](Operation *op) {
    if (isa<linalg::LinalgOp, linalg::PackOp>(op)) {
      LDBG() << "Gather loops to peel from candidate op\n  " << *op;
      collectLoopsToPeel(op, uniqueLoopsToPeel);
    }
  });

  LDBG() << "Peeling loops";
  // Visiting the loops in outer-to-inner order will prevent loops nested in
  // partial iterations to be peeled again.
  SmallVector<scf::ForOp, 8> outerToInnerLoopsToPeel(uniqueLoopsToPeel.rbegin(),
                                                     uniqueLoopsToPeel.rend());
  IRRewriter rewriter(context);
  linalg::peelLoops(rewriter, outerToInnerLoopsToPeel);

  LDBG() << "Canonicalizing loops";
  RewritePatternSet patterns(context);
  linalg::populateLinalgTilingCanonicalizationPatterns(patterns);
  scf::populateSCFForLoopCanonicalizationPatterns(patterns);
  memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
  context->getLoadedDialect<tensor::TensorDialect>()
      ->getCanonicalizationPatterns(patterns);
  if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
    LDBG() << "----- cleanup failed -----";
    return signalPassFailure();
  }
}
} // namespace
} // namespace mlir::iree_compiler
