//
// Copyright (C) 2024-2025 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#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/sibling_ops_analysis.hpp"
#include "vpux/compiler/dialect/VPU/utils/tile_utils.hpp"
#include "vpux/compiler/dialect/config/IR/utils.hpp"

#include <mlir/IR/PatternMatch.h>
#include <mlir/Transforms/GreedyPatternRewriteDriver.h>

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

using namespace vpux;

namespace {

//
// OutputPipelineTilingPass
//

class OutputPipelineTilingPass final : public VPU::impl::OutputPipelineTilingBase<OutputPipelineTilingPass> {
public:
    explicit OutputPipelineTilingPass(bool enablePrefetchTiling, Logger log)
            : _enablePrefetchTiling(enablePrefetchTiling) {
        Base::initLogger(log, Base::getArgumentName());
    }

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

private:
    bool isTilingAdjustmentBeneficial(VPU::NCEOpInterface nceOp, const OutputTiling& origTiling,
                                      const OutputTiling& newTiling, std::shared_ptr<VPU::LayerCostModel> costModel);
    void safeRunOnFunc() final;
    bool _enablePrefetchTiling = true;
};

mlir::LogicalResult OutputPipelineTilingPass::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";
    }
    return mlir::success();
}

bool OutputPipelineTilingPass::isTilingAdjustmentBeneficial(VPU::NCEOpInterface nceOp, const OutputTiling& origTiling,
                                                            const OutputTiling& newTiling,
                                                            std::shared_ptr<VPU::LayerCostModel> costModel) {
    auto origOp = nceOp.getOperation();

    // 1.Check if new output tiles can all be evenly unrolled

    // For example, input distributed type:
    // shape = [1, 48, 10, 771]
    // num_tiles = [1, 1, 2, 1]
    // alignment = [1, 1, 4, 1]
    // would be unrolled into shape [1, 48, 8, 771] and [1, 48, 2, 771], such uneven unrolling is not beneficial
    auto strategy = mlir::cast<VPU::ClusteredOpInterface>(origOp).getMultiClusterStrategy();
    auto canTileBeEvenlyUnrolled = [&](auto& tileInfo) {
        auto curTileTypes = VPU::getTileTypes(origOp, tileInfo, strategy);
        VPUX_THROW_WHEN(curTileTypes.empty(), "Tile types vector is empty");
        auto curTileInputType = curTileTypes[0];
        auto inputDistType = mlir::dyn_cast<VPU::DistributedTensorType>(curTileInputType);
        if (inputDistType == nullptr) {
            return true;
        }

        auto distributionAttr = inputDistType.getDistribution();
        auto perClusterShape = inputDistType.getPerClusterComputeShapes();
        if (distributionAttr.getNumTiles() == nullptr) {
            return true;
        }

        const auto tilingScheme = parseIntArrayAttr<int64_t>(distributionAttr.getNumTiles());
        const auto axisDim = Dim(vpux::VPU::getDistributedTilingAxis(tilingScheme));
        // First cluster has the largest per-cluster shape
        // Compare other cluster shapes with the largest per-cluster shape
        // If any cluster has a 'too' smaller cluster shape, it's considered to be unevenly unrolled
        const auto firstClusterShape = perClusterShape.front();
        for (const auto& shape : perClusterShape) {
            if (shape[axisDim] <= firstClusterShape[axisDim] / 2) {
                _log.nest().debug("Input tile type {0} is not evenly unrolled at {1}", inputDistType, origOp->getLoc());
                return false;
            }
        }

        return true;
    };
    if (!llvm::all_of(newTiling, canTileBeEvenlyUnrolled)) {
        return false;
    }

    // 2.Check if new tiling strategy has smaller cost

    // If the op does not have MC strategy, use Clustering by default
    auto mcStrategy = VPU::MultiClusterStrategy::Clustering;
    auto clusteredNCEOp = mlir::dyn_cast<VPU::ClusteredOpInterface>(origOp);
    if (clusteredNCEOp != nullptr) {
        auto strategy = clusteredNCEOp.getMultiClusterStrategy();
        if (strategy.has_value()) {
            mcStrategy = strategy.value();
        }
    }

    // TODO: calculate layer cost with considering output pipelining to benefit more cases, see E#110180
    // origCost is calculated accurately, but newCost is calculated too big because layer cost value is calculated
    // without considering output pipelining right now.
    // However, if the 'big' newCost is still smaller than the accurate origCost, new strategy must be better than the
    // old one. As a result, the costs checking still can ensure new tiling strategy is optimal even though newCost is
    // not so accurate.
    auto [newCostWithoutPrefetching, newCostWithPrefetching] =
            costModel->getDPUandDMATimeCostWithCustomTiling(nceOp, mcStrategy, newTiling);
    auto [origCostWithoutPrefetching, origCostWithPrefetching] =
            costModel->getDPUandDMATimeCostWithCustomTiling(nceOp, mcStrategy, origTiling);

    auto tilingInfoOp = mlir::dyn_cast<VPU::TilingInfoOpInterface>(nceOp.getOperation());
    double newCost = 0;
    double origCost = 0;
    if (_enablePrefetchTiling && tilingInfoOp != nullptr &&
        tilingInfoOp.isSupportedTiling(newTiling, vpux::TilingMode::PREFETCHING, _log)) {
        newCost = newCostWithPrefetching;
    } else {
        newCost = newCostWithoutPrefetching;
    }
    if (_enablePrefetchTiling && tilingInfoOp != nullptr &&
        tilingInfoOp.isSupportedTiling(origTiling, vpux::TilingMode::PREFETCHING, _log)) {
        origCost = origCostWithPrefetching;
    } else {
        origCost = origCostWithoutPrefetching;
    }

    _log.nest().trace("Original cost: {0} , Current cost: {1}", origCost, newCost);

    return newCost < origCost;
}

void OutputPipelineTilingPass::safeRunOnFunc() {
    if (!_enablePrefetchTiling) {
        return;
    }

    auto func = getOperation();
    auto module = func->getParentOfType<mlir::ModuleOp>();
    auto maybeLayerCostModelAnalysis = getCachedParentAnalysis<VPU::LayerCostModelAnalysis>(module);
    auto layerCostModel =
            VPU::LayerCostModelAnalysis::getOrCreateLayerCostModel(maybeLayerCostModelAnalysis, &getContext(), _log);

    auto siblingsOpsAnalysis = getAnalysis<VPU::SiblingOpsAnalysis>();
    auto costModel = std::make_shared<VPU::LayerCostModel>(func, _enablePrefetchTiling, siblingsOpsAnalysis,
                                                           layerCostModel, _log);

    func->walk([&](VPU::TilingBuilderOpInterface tilingBuilderOp) {
        auto origOp = tilingBuilderOp.getOperation();

        auto nceOp = mlir::dyn_cast<VPU::NCEOpInterface>(origOp);
        if (nceOp == nullptr) {
            return;
        }

        if (mlir::isa<VPU::NCEMatMulOp>(origOp)) {
            // E126102: Skip MatMulOp since it does not support cost, but this pass is cost based.
            return;
        }

        // Don't increase tiles number when layer doesn't have tilingStrategy
        // This prevents changing tiling strategy for VF tiled operation:
        // At this stage, VF tiling has been applied and VF tiled operations do not have tilingStrategy, so VF tiled
        // operations are not impacted
        if (!origOp->hasAttr(tilingStrategy)) {
            return;
        }

        const auto origTilingStrategy =
                Shape(parseIntArrayAttr<int64_t>(mlir::cast<mlir::ArrayAttr>(origOp->getAttr(tilingStrategy))));
        const auto nonOneDims = getNonOneDim(origTilingStrategy);
        if (nonOneDims.size() != 1) {
            _log.nest().trace("Operation {0} has nested tiling strategy {1}", origOp->getLoc(), origTilingStrategy);
            return;
        }

        const auto tilingModeOpt =
                VPU::getTilingMode(tilingBuilderOp.getOperation(), _enablePrefetchTiling, nullptr, _log);
        if (!tilingModeOpt.has_value() || tilingModeOpt.value().first != TilingMode::PIPELINING) {
            return;
        }
        const auto tilingMode = tilingModeOpt.value().first;

        const auto outputShape = getShape(origOp->getResult(0));
        auto origTiles = fillDividedTiles(origOp, origTilingStrategy, outputShape);
        if (mlir::failed(origTiles)) {
            return;
        }

        _log.trace("Attempting to generate tiling strategy for output pipelining at {0}", origOp->getLoc());
        const auto alignRequirement = getAlignDimAndSize(origOp);
        const auto dimToAlign = alignRequirement.first;
        const auto dimAlignment = alignRequirement.second;
        const auto& maxNumTiles = tilingBuilderOp.getMaxNumTiles();

        // Set attribute to find a valid tiling strategy for output pipelining
        origOp->setAttr(outputPipelining, mlir::BoolAttr::get(origOp->getContext(), true));
        const auto adjustTiling = [&]() {
            mlir::FailureOr<OutputTiling> findSupportedTileSize;
            auto prefetchableTilesOnDim = origTilingStrategy;
            auto targetDim = *nonOneDims.begin();
            _log.trace(" Target dimension for output pipelining: {0} from {1}", targetDim, origTilingStrategy);
            do {
                findSupportedTileSize = isSupportedTileSize(origOp, prefetchableTilesOnDim, tilingMode, _log);
                if (!mlir::failed(findSupportedTileSize)) {
                    _log.trace(" Skip as current strategy is already pipelining");
                    break;
                }

                if (!isDimLeftToTile(prefetchableTilesOnDim, maxNumTiles, targetDim)) {
                    break;
                }

                auto nextTileSearchResult = getNextTiling(targetDim, dimToAlign, dimAlignment, prefetchableTilesOnDim,
                                                          maxNumTiles, outputShape);
                if (mlir::failed(nextTileSearchResult)) {
                    break;
                }
                prefetchableTilesOnDim = nextTileSearchResult.value();
            } while (mlir::failed(findSupportedTileSize) &&
                     (prefetchableTilesOnDim[targetDim] <=
                      MAX_OUTPUT_PIPELINE_TILING_TIME * origTilingStrategy[targetDim]));

            _log.trace("New candidate tiling strategy for output pipelining: {0} at {1}", prefetchableTilesOnDim,
                       origOp->getLoc());
            if (!mlir::failed(findSupportedTileSize) && prefetchableTilesOnDim != origTilingStrategy) {
                // find an available tiling strategy for output pipelining
                auto curTiles = findSupportedTileSize.value();
                _log.debug("Found output pipelining tiling strategy: {0} at {1}", curTiles, origOp->getLoc());
                if (isTilingAdjustmentBeneficial(nceOp, origTiles.value(), curTiles, costModel)) {
                    origOp->setAttr(tilingStrategy, getIntArrayAttr(origOp->getContext(), curTiles[0].axis));
                    _log.debug("Overwrite original tiling strategy: {0} with output pipelining tiling strategy: {1} at "
                               "{2}",
                               origTilingStrategy, prefetchableTilesOnDim, origOp->getLoc());
                    return true;
                }
            }
            return false;
        };

        auto tilingAdjusted = adjustTiling();
        origOp->removeAttr(outputPipelining);
        // Due to some scheduling problem, we need to set a CMX fragment threshold for better pipelining. Currently if
        // it can not be pipelined, we can adjust the threshold to try it again. We only support tiling on H or W now.
        // #E187660 remove the attribute and pass the fragment threshold as a parameter
        if (!tilingAdjusted) {
            origOp->setAttr(outputPipeliningMinFragmentation, mlir::BoolAttr::get(origOp->getContext(), true));
            adjustTiling();
            origOp->removeAttr(outputPipeliningMinFragmentation);
        }
    });
}

}  // namespace

std::unique_ptr<mlir::Pass> vpux::VPU::createOutputPipelineTilingPass(bool enablePrefetchTiling, Logger log) {
    return std::make_unique<OutputPipelineTilingPass>(enablePrefetchTiling, log);
}
