//
// Copyright (C) 2023-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "vpux/compiler/core/tiling.hpp"
#include "vpux/compiler/dialect/VPU/IR/dialect.hpp"
#include "vpux/compiler/dialect/VPU/IR/ops/dpu.hpp"
#include "vpux/compiler/dialect/VPU/transforms/passes.hpp"
#include "vpux/compiler/dialect/VPU/utils/cost_model/layer_vpunn_cost.hpp"
#include "vpux/compiler/dialect/VPU/utils/generate_tiling.hpp"
#include "vpux/compiler/dialect/VPU/utils/manual_strategy_utils.hpp"
#include "vpux/compiler/dialect/VPU/utils/nce_utils.hpp"
#include "vpux/compiler/dialect/VPU/utils/op_tiling_cache.hpp"
#include "vpux/compiler/dialect/VPU/utils/sibling_ops_analysis.hpp"
#include "vpux/compiler/dialect/VPU/utils/tiling_pass_config_utils.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/dialect.hpp"
#include "vpux/compiler/dialect/config/IR/utils.hpp"
#include "vpux/compiler/dialect/config/utils/config_option_utils.hpp"
#include "vpux/compiler/utils/loop.hpp"

#include <mutex>

namespace vpux::VPU {
#define GEN_PASS_DECL_TILINGSTRATEGYASSIGNMENT
#define GEN_PASS_DEF_TILINGSTRATEGYASSIGNMENT
#include "vpux/compiler/dialect/VPU/passes.hpp.inc"
}  // namespace vpux::VPU

using namespace vpux;

namespace {

//
// TilingStrategyAssignment
//
// There are three modes of tiling strategies as defined in vpux::TilingMode:
// 1. ISOLATED tiling: Split operations with smallest tiling number to make them fit into CMX
// 2. PIPELINING tiling: Overlap the DPU time of earlier sub-tile with the DMA time of the later ones
//              Two possible scenarios where the PIPELINING could be triggered:
//              a). When ISOLATED tiling is required, the tiling number will be increased to satisfy PIPELINING
//              b). When the constant weights of an operation is larger than the threshold
//                  tiling number will be increased to satisfy PIPELINING
//                  even though the operation doesn't require ISOLATED tiling originally
//              A precondition is that a feasible tiling strategy must exist to make PIPELINING work
//              Otherwise it will fallback to ISOLATED tiling or non-tiling
// 3. PREFETCHING tiling: Overlap the DPU time of parent operation with the DMA time the child
//

//
// TilingStrategyAssignmentPass
//
class TilingStrategyAssignmentPass final :
        public VPU::impl::TilingStrategyAssignmentBase<TilingStrategyAssignmentPass> {
public:
    TilingStrategyAssignmentPass(bool enablePrefetchTiling, bool enableVpunnCostForTiling,
                                 StringRef enableShaveDDRAccessOptimization, bool enableDynamicDimAlignment, Logger log)
            : _enablePrefetchTiling(enablePrefetchTiling),
              _enableVpunnCostForTiling(enableVpunnCostForTiling),
              _shaveDDRAccessOptimizationMode(
                      vpux::VPU::getShaveDDRAccessOptimizationMode(enableShaveDDRAccessOptimization)),
              _enableDynamicDimAlignment(enableDynamicDimAlignment) {
        Base::initLogger(log, Base::getArgumentName());
    }

    mlir::LogicalResult initialize(mlir::MLIRContext* ctx) final;

private:
    void safeRunOnFunc() final;
    std::vector<vpux::VPU::StrategyWithCost> getStrategies(VPU::TilingBuilderOpInterface origOp, TilingMode tilingMode);

    bool _enablePrefetchTiling = true;
    bool _enableVpunnCostForTiling = false;
    VPU::EnableShaveDDRAccessOptimization _shaveDDRAccessOptimizationMode;
    std::shared_ptr<vpux::VPU::LayerCostModel> _costModel;
    bool _enableDynamicDimAlignment = false;
};

mlir::LogicalResult TilingStrategyAssignmentPass::initialize(mlir::MLIRContext* ctx) {
    if (mlir::failed(Base::initialize(ctx))) {
        return mlir::failure();
    }
    if (tilingMode.hasValue()) {
        _log.trace("Overloading the default value {0} of the '_enablePrefetchTiling' field to the value {1} of the "
                   "pass option 'tilingMode' generated by MLIR",
                   _enablePrefetchTiling, tilingMode.getValue());
        _enablePrefetchTiling = tilingMode.getValue() == "PREFETCH";
    }
    if (enableVpunnCostForTiling.hasValue()) {
        _log.trace("Overloading the default value {0} of the '_enableVpunnCostForTiling' field to the value {1} of the "
                   "pass option "
                   "'enableVpunnCostForTiling' generated by MLIR",
                   _enableVpunnCostForTiling, enableVpunnCostForTiling);
        _enableVpunnCostForTiling = enableVpunnCostForTiling;
    }
    if (enableShaveDDRAccessOptimization.hasValue()) {
        _shaveDDRAccessOptimizationMode =
                vpux::VPU::getShaveDDRAccessOptimizationMode(enableShaveDDRAccessOptimization);
    }
    if (enableDynamicDimAlignment.hasValue()) {
        _log.trace("Overloading the default value 'false' of the 'enableDynamicDimAlignment' field to the value "
                   "{0} of the pass option 'enableDynamicDimAlignment' generated by MLIR",
                   enableDynamicDimAlignment.getValue());
        _enableDynamicDimAlignment = enableDynamicDimAlignment.getValue();
    }
    return mlir::success();
}

std::vector<vpux::VPU::StrategyWithCost> TilingStrategyAssignmentPass::getStrategies(
        VPU::TilingBuilderOpInterface origOp, TilingMode tilingMode) {
    _log.trace("Assign: '{0}' at '{1}'", origOp->getName(), origOp->getLoc());
    auto op = origOp.getOperation();

    std::vector<vpux::VPU::StrategyWithCost> strategies{};

    // Cost based strategy assignment is not applied to NCEMatMulOp yet #E126102
    auto& costModelUtils = VPU::getICostModelUtilsInterface(op->getContext());
    if (!_enableVpunnCostForTiling || !mlir::isa<VPU::NCEOpInterface>(op) ||
        (VPU::isNCEWithInt4Weights(op) && !costModelUtils.isNCEWithInt4WeightsSupported()) ||
        mlir::isa<VPU::NCEMatMulOp>(op)) {
        auto strategy = origOp.getTilingStrategy(tilingMode, _log);
        if (mlir::succeeded(strategy)) {
            strategies.push_back({strategy.value()[0].axis, {0, 0}});
        }
    } else {
        _log.trace("Get tiling strategy based on VPUNN DPU+DMA cost for NCE ops");
        strategies = getHwLayerTilingStrategiesWithCost(origOp, tilingMode, _costModel, _log);
    }
    return strategies;
}

void TilingStrategyAssignmentPass::safeRunOnFunc() {
    auto func = getOperation();
    auto module = func->getParentOfType<mlir::ModuleOp>();

    if (_enableDynamicDimAlignment) {
        VPU::setDynamicDimAlignment(func);
    }

    auto maybeLayerCostModelAnalysis = getCachedParentAnalysis<VPU::LayerCostModelAnalysis>(module);
    auto layerCostModel =
            VPU::LayerCostModelAnalysis::getOrCreateLayerCostModel(maybeLayerCostModelAnalysis, &getContext(), _log);
    auto strategyLayerCost = std::make_unique<VPU::LayerVPUNNCost>(func);

    auto siblingOpsAnalysis = getAnalysis<VPU::EagerSiblingOpsAnalysis>();
    llvm::DenseSet<VPU::TilingBuilderOpInterface> opsToCheckForPrefetch{};
    llvm::DenseMap<VPU::TilingBuilderOpInterface, std::vector<vpux::VPU::StrategyWithCost>> strategiesForOp{};
    std::mutex prefetchSetMutex;
    std::mutex strategiesMapMutex;
    const auto opsToTile = func.getOps<VPU::TilingBuilderOpInterface>();
    const auto numOps = std::distance(opsToTile.begin(), opsToTile.end());
    _costModel = std::make_shared<vpux::VPU::LayerCostModel>(vpux::VPU::LayerCostModel(
            func, _enablePrefetchTiling, siblingOpsAnalysis, std::move(layerCostModel), _log));

    // Get cost of possible ISOLATED/PIPELINING strategies in parallel. Note that it is not possible to
    // say which strategy will be the best since whether strategy can be prefetched depends on parent op tiling.
    auto loopExecPolicy = numOps < MIN_OPS_FOR_PARALLEL_TILING ? LoopExecPolicy::Sequential : LoopExecPolicy::Parallel;
    loop_1d(loopExecPolicy, func.getContext(), numOps, [&](const size_t index) {
        auto it = opsToTile.begin();
        std::advance(it, index);
        auto tilingOp = *it;
        auto op = tilingOp.getOperation();
        if (_shaveDDRAccessOptimizationMode == VPU::EnableShaveDDRAccessOptimization::TRUE) {
            auto ddrAccessOp = mlir::dyn_cast<VPU::DDRAccessOpInterface>(op);
            if (ddrAccessOp != nullptr && ddrAccessOp.isDDRAccessNecessaryOrBeneficial(_log)) {
                return;
            }
        }

        auto tilingModeOpt = vpux::VPU::getTilingMode(op, _enablePrefetchTiling, strategyLayerCost, _log);
        if (tilingModeOpt.has_value()) {
            auto [tilingMode, prefetchingCandidate] = tilingModeOpt.value();
            if (tilingMode == TilingMode::ISOLATED || tilingMode == TilingMode::PIPELINING) {
                auto strategies = getStrategies(tilingOp, tilingMode);
                std::lock_guard<std::mutex> guard(strategiesMapMutex);
                strategiesForOp.insert(std::make_pair(tilingOp, strategies));
            }
            if (prefetchingCandidate) {
                std::lock_guard<std::mutex> guard(prefetchSetMutex);
                opsToCheckForPrefetch.insert(tilingOp);
            }
        }
    });

    auto getBestStrategy = [this](VPU::TilingBuilderOpInterface tilingOp,
                                  std::vector<vpux::VPU::StrategyWithCost>& strategies,
                                  bool assumePrefetching) -> Shape {
        double bestCost = strategies.front().cost.costWithoutPrefetching;
        Shape bestStrategy = strategies.front().strategy;
        auto tilingInfoOp = mlir::dyn_cast<VPU::TilingInfoOpInterface>(tilingOp.getOperation());
        for (auto& [strategy, costs] : strategies) {
            double currentCost = 0;
            auto [costWithoutPrefetching, costWithPrefetching] = costs;
            if (assumePrefetching) {
                currentCost = costWithPrefetching;
            } else {
                currentCost =
                        _enablePrefetchTiling && (tilingInfoOp != nullptr) &&
                                        tilingInfoOp.isSupportedTilingStrategy(strategy, TilingMode::PREFETCHING, _log)
                                ? costWithPrefetching
                                : costWithoutPrefetching;
            }
            if (currentCost < bestCost ||
                (currentCost == bestCost &&
                 isNewTileWithSameCostHasPotentialDMABenefits(tilingOp.getOperation(), bestStrategy, strategy))) {
                bestCost = currentCost;
                bestStrategy = strategy;
            }
        }
        return bestStrategy;
    };

    // Choose best strategy based on whether it can be prefetched. Also tries to generate prefetching strategies for ops
    // that can be prefetched.
    func->walk([&](VPU::TilingBuilderOpInterface tilingOp) {
        if (strategiesForOp.contains(tilingOp)) {
            auto strategies = strategiesForOp.at(tilingOp);
            if (!strategies.empty()) {
                auto bestStrategy = getBestStrategy(tilingOp, strategies, false);
                tilingOp.getOperation()->setAttr(tilingStrategy,
                                                 getIntArrayAttr(tilingOp.getOperation()->getContext(), bestStrategy));
            }
        }
        if (opsToCheckForPrefetch.contains(tilingOp) &&
            VPU::prefetchTilingConditionSatisfied(tilingOp.getOperation(), _log.nest())) {
            auto strategies = getStrategies(tilingOp, TilingMode::PREFETCHING);
            if (!strategies.empty()) {
                auto bestStrategy = getBestStrategy(tilingOp, strategies, true);
                tilingOp.getOperation()->setAttr(tilingStrategy,
                                                 getIntArrayAttr(tilingOp.getOperation()->getContext(), bestStrategy));
            }
        }
    });

    // Set minimum tiling strategy if any for unassigned ones
    func->walk([&](VPU::TilingBuilderOpInterface tilingOp) {
        auto op = tilingOp.getOperation();
        if (!op->hasAttr(tilingStrategy)) {
            const auto minNumTiles = getMinNumTiles(op);
            if (llvm::any_of(minNumTiles, [](int64_t v) {
                    return v > 1;
                })) {
                op->setAttr(tilingStrategy, getIntArrayAttr(op->getContext(), minNumTiles));
            }
        }
    });

    auto& cache = VPU::getGlobalOpTilingCache();
    cache.printStats(_log);
    cache.cleanUp();
    VPU::removeDynamicDimAlignment(func);
}
}  // namespace

std::unique_ptr<mlir::Pass> vpux::VPU::createTilingStrategyAssignmentPass(bool enablePrefetchTiling,
                                                                          bool enableVpunnCostForTiling,
                                                                          StringRef enableShaveDDRAccessOptimization,
                                                                          bool enableDynamicDimAlignment, Logger log) {
    return std::make_unique<TilingStrategyAssignmentPass>(enablePrefetchTiling, enableVpunnCostForTiling,
                                                          enableShaveDDRAccessOptimization, enableDynamicDimAlignment,
                                                          log);
}
