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

#include "vpux/compiler/core/feasible_memory_scheduler_control_edges.hpp"
#include "vpux/compiler/core/cost_model_utils.hpp"
#include "vpux/compiler/core/feasible_scheduler_utils.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/ops.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/ops_interfaces.hpp"
#include "vpux/compiler/dialect/VPUIP/utils/utils.hpp"
#include "vpux/compiler/utils/analysis.hpp"
#include "vpux/compiler/utils/attributes.hpp"

using namespace vpux;

//
// Feasible Memory Scheduler Control Edges support
//

FeasibleMemorySchedulerControlEdges::FeasibleMemorySchedulerControlEdges(
        VPU::MemoryKind memKind, AsyncDepsInfo& depsInfo, AliasesInfo& aliasInfo, Logger log,
        LinearScan<mlir::Value, LinearScanHandler>& scan)
        : _log(log), _memKind(memKind), _depsInfo(depsInfo), _aliasInfo(aliasInfo), _scan(scan) {
    _log.setName("feasible-memory-scheduler-control-edges");
}

// This method will update all AsyncExecOp token dependencies so that resulting
// execution is aligned with order generated by list-scheduler
void FeasibleMemorySchedulerControlEdges::insertDependenciesBasic(
        ArrayRef<FeasibleMemoryScheduler::ScheduledOpInfo> scheduledOps) {
    // Go through all the tasks and add token dependencies between
    // all tasks with start time t to all tasks with time t+1
    _log.trace("Get dependencies based on scheduler time decisions");
    _log = _log.nest();
    for (auto opIt = scheduledOps.begin(); opIt != scheduledOps.end(); opIt++) {
        if (!opIt->isOriginalOp()) {
            continue;
        }

        int64_t nextTimeDiff = 0;
        for (auto nextTimeOpIt = opIt; nextTimeOpIt != scheduledOps.end(); nextTimeOpIt++) {
            if (!nextTimeOpIt->isOriginalOp()) {
                continue;
            } else if (nextTimeDiff == 0 && nextTimeOpIt->cycleBegin_ > opIt->cycleBegin_) {
                nextTimeDiff = nextTimeOpIt->cycleBegin_ - opIt->cycleBegin_;
            }

            if (nextTimeDiff != 0) {
                if (nextTimeOpIt->cycleBegin_ == opIt->cycleBegin_ + nextTimeDiff) {
                    // Insert dependency between op at time t to op at
                    // time t+1
                    auto srcAsyncOp = _depsInfo.getExecuteOpAtIndex(opIt->op_);
                    auto dstAsyncOp = _depsInfo.getExecuteOpAtIndex(nextTimeOpIt->op_);
                    _log.trace("Dep: {0} -> {1}", opIt->op_, nextTimeOpIt->op_);
                    VPUX_THROW_UNLESS((srcAsyncOp != nullptr) && (dstAsyncOp != nullptr),
                                      "srcAsyncOp/dstAsyncOp not located based on index");
                    _depsInfo.addDependency(srcAsyncOp, dstAsyncOp);
                } else if (nextTimeOpIt->cycleBegin_ > (opIt->cycleBegin_ + nextTimeDiff)) {
                    break;
                }
            }
        }
    }
    _log = _log.unnest();
}

// This method will update all AsyncExecOp token dependencies for a given executor type
// so that resulting execution is aligned with order generated by list-scheduler
void FeasibleMemorySchedulerControlEdges::insertScheduleOrderDepsForExecutor(
        ArrayRef<FeasibleMemoryScheduler::ScheduledOpInfo> scheduledOps, config::ExecutorKind executorKind,
        uint32_t executorInstance) {
    // Go through all the tasks with given executor type
    _log.trace("Insert control edges aligned with schedule order for provided executor");
    _log = _log.nest();
    for (auto opIt = scheduledOps.begin(); opIt != scheduledOps.end(); opIt++) {
        if (!opIt->isOriginalOp()) {
            continue;
        }

        if (opIt->queueType.execKind != executorKind) {
            continue;
        }

        if (executorInstance >= opIt->executorInstanceMask.size() || !opIt->executorInstanceMask[executorInstance]) {
            continue;
        }

        auto srcAsyncOp = _depsInfo.getExecuteOpAtIndex(opIt->op_);

        int64_t nextTimeDiff = 0;
        for (auto nextTimeOpIt = opIt; nextTimeOpIt != scheduledOps.end(); nextTimeOpIt++) {
            if (!nextTimeOpIt->isOriginalOp()) {
                continue;
            }

            if (nextTimeOpIt->queueType.execKind != opIt->queueType.execKind ||
                nextTimeOpIt->queueType.id != opIt->queueType.id) {
                continue;
            }

            if (executorInstance >= nextTimeOpIt->executorInstanceMask.size() ||
                !nextTimeOpIt->executorInstanceMask[executorInstance]) {
                continue;
            }

            if (nextTimeDiff == 0 && nextTimeOpIt->cycleBegin_ > opIt->cycleBegin_) {
                nextTimeDiff = nextTimeOpIt->cycleBegin_ - opIt->cycleBegin_;
            }

            if (nextTimeDiff != 0) {
                if (nextTimeOpIt->cycleBegin_ == opIt->cycleBegin_ + nextTimeDiff) {
                    // Insert dependency between op at time t to op at
                    // time t+1
                    auto dstAsyncOp = _depsInfo.getExecuteOpAtIndex(nextTimeOpIt->op_);
                    _log.trace("Dep: {0} -> {1}", opIt->op_, nextTimeOpIt->op_);
                    VPUX_THROW_UNLESS((srcAsyncOp != nullptr) && (dstAsyncOp != nullptr),
                                      "srcAsyncOp/dstAsyncOp not located based on index");
                    auto srcCycleEnd = getAsyncExecuteCycleEnd(srcAsyncOp);
                    auto dstCycleBegin = getAsyncExecuteCycleBegin(dstAsyncOp);
                    VPUX_THROW_UNLESS(srcCycleEnd <= dstCycleBegin,
                                      "Order of execution not preserved from {0} with cycleEnd = {1} to {2} with "
                                      "cycleBegin = {3}",
                                      opIt->op_, srcCycleEnd, nextTimeOpIt->op_, dstCycleBegin);
                    _depsInfo.addDependency(srcAsyncOp, dstAsyncOp);
                } else if (nextTimeOpIt->cycleBegin_ > (opIt->cycleBegin_ + nextTimeDiff)) {
                    break;
                }
            }
        }
    }
    _log = _log.unnest();
}

// For given operation input and output root buffers update ScheduledOpOneResource structure for
// produced or consumed memory range so that control edge algorithm can later track all ranges
// throughout the schedule and insert dependencies for overlaps.
void vpux::updateScheduledOpsResourcesForControlEdgeBasic(std::list<ScheduledOpOneResource>& scheduledOpsResources,
                                                          LinearScan<mlir::Value, LinearScanHandler>& scan,
                                                          size_t opIndex, const ValueOrderedSet& inputBuffers,
                                                          const ValueOrderedSet& outputBuffers, Logger& log) {
    // For all identified buffers used by operation create separate entries with information
    // about memory ranges to properly identify range producer and consumers at a given time
    auto updateResources = [&](const ValueOrderedSet& bufs, ScheduledOpOneResource::EResRelation relType) {
        for (auto& buf : bufs) {
            if (!isBufAllocOp(buf.getDefiningOp())) {
                continue;
            }
            auto addressStart = scan.handler().getAddress(buf);
            auto addressEnd = addressStart + scan.handler().getSize(buf) - 1;
            log.trace("op = '{0}'\t {1} = [{2} - {3}]", opIndex,
                      (relType == ScheduledOpOneResource::EResRelation::CONSUMER) ? "input" : "output", addressStart,
                      addressEnd);
            scheduledOpsResources.push_back(ScheduledOpOneResource(opIndex, addressStart, addressEnd, relType));
        }
    };

    updateResources(inputBuffers, ScheduledOpOneResource::EResRelation::CONSUMER);
    updateResources(outputBuffers, ScheduledOpOneResource::EResRelation::PRODUCER);
}

// For provided operation (execOp) identify operands and update ScheduledOpOneResource structure for
// each produced or consumed memory range so that control edge algorithm can later track all ranges
// throughout the schedule and insert dependencies for overlaps.
// This method is capable of identifying SubViews on the chain from operand to root buffer to understand
// if given operation uses just part of some root buffer.
void vpux::updateScheduledOpsResourcesForControlEdge(std::list<ScheduledOpOneResource>& scheduledOpsResources,
                                                     AliasesInfo& aliasInfo,
                                                     LinearScan<mlir::Value, LinearScanHandler>& scan,
                                                     VPU::MemoryKind memKind, size_t opIndex,
                                                     mlir::async::ExecuteOp execOp, Logger& log) {
    SmallVector<mlir::Value> inputOperands;
    SmallVector<mlir::Value> outputOperands;

    auto populateTargetBuffers = [&](const mlir::ValueRange& buffers, SmallVector<mlir::Value>& targetVec) {
        for (const auto& buffer : buffers) {
            const auto type = mlir::dyn_cast<vpux::NDTypeInterface>(buffer.getType());
            if (type == nullptr || type.getMemoryKind() != memKind) {
                continue;
            }
            targetVec.push_back(buffer);
        }
    };

    // Get operation buffers for all operands. Go through each layer op and
    // store in a set all root buffers
    auto* bodyBlock = execOp.getBody();
    for (auto& innerOp : bodyBlock->getOperations()) {
        auto layerOp = mlir::dyn_cast<VPUIP::LayerOpInterface>(innerOp);
        if (layerOp == nullptr) {
            continue;
        }

        const auto inputs = VPUIP::getInputsSanitized(layerOp);
        populateTargetBuffers(inputs, inputOperands);

        const auto outputs = layerOp.getOutputs();
        populateTargetBuffers(outputs, outputOperands);
    }

    // For all identified buffers used by operation create separate entries with information
    // about memory ranges to properly identify range producer and consumers at a given time
    auto updateResources = [&](SmallVector<mlir::Value>& operands, ScheduledOpOneResource::EResRelation relType) {
        for (auto& operand : operands) {
            auto buf = aliasInfo.getRoot(operand);
            if (!isBufAllocOp(buf.getDefiningOp())) {
                continue;
            }

            std::optional<ScheduledOpOneResource::ResourceView> resView = std::nullopt;

            auto operandSize = static_cast<size_t>(
                    mlir::cast<vpux::NDTypeInterface>(operand.getType()).getCompactAllocSize().count());
            auto allocSize = scan.handler().getSize(buf);
            if (operandSize < allocSize) {
                // Check chain of connections from operand to root buffer and look
                // for SubViewOp that will define chunk of buffer used by operation
                VPUIP::SubViewOp subViewOp = nullptr;
                mlir::Value sourceAlias = operand;

                do {
                    if (auto tmpSubViewOp = sourceAlias.getDefiningOp<VPUIP::SubViewOp>()) {
                        // If subViewOp has already been encountered treat this as an unsupported scenario and limit
                        // cases only with single SubView on the chain between operand to root buffer
                        // TODO: Extend support for multiple subviews (E#106837)
                        if (subViewOp) {
                            log.trace("More than 1 SubView for op - {0} identified in operand->root chain. Skip "
                                      "optimization",
                                      opIndex);
                            subViewOp = nullptr;
                            break;
                        }

                        subViewOp = tmpSubViewOp;
                    }
                } while ((sourceAlias = aliasInfo.getSource(sourceAlias)));

                if (subViewOp) {
                    auto subViewShape = parseIntArrayAttr<int64_t>(subViewOp.getStaticSizes());
                    auto subViewOffsets = parseIntArrayAttr<int64_t>(subViewOp.getStaticOffsets());
                    auto subViewStrides =
                            subViewOp.getStaticStrides().has_value()
                                    ? parseIntArrayAttr<int64_t>(subViewOp.getStaticStrides().value())
                                    : SmallVector<int64_t>(
                                              mlir::cast<vpux::NDTypeInterface>(subViewOp.getSource().getType())
                                                      .getRank(),
                                              1);

                    resView = ScheduledOpOneResource::ResourceView(
                            {buf, std::move(subViewOffsets), std::move(subViewShape), std::move(subViewStrides)});
                }
            }

            auto addressStart = scan.handler().getAddress(buf);
            auto addressEnd = addressStart + allocSize - 1;
            log.trace("op = '{0}'\t {1} = [{2} - {3}]", opIndex,
                      (relType == ScheduledOpOneResource::EResRelation::CONSUMER) ? "input" : "output", addressStart,
                      addressEnd);
            scheduledOpsResources.push_back(
                    ScheduledOpOneResource(opIndex, addressStart, addressEnd, relType, resView));
        }
    };

    updateResources(inputOperands, ScheduledOpOneResource::EResRelation::CONSUMER);
    updateResources(outputOperands, ScheduledOpOneResource::EResRelation::PRODUCER);
}

// Insert control flow for overlapping memory regions
void FeasibleMemorySchedulerControlEdges::insertMemoryControlEdges(
        ArrayRef<FeasibleMemoryScheduler::ScheduledOpInfo> scheduledOps) {
    std::list<ScheduledOpOneResource> scheduledOpsResources;
    // set of modified operations, e.g.
    // DMA         : input  [0-100]   : cycles [180-220]
    // DPU         : output [500-600] : cycles [200-250]
    // SPILL-WRITE : input  [500-600] : cycles [250-280]
    // SPILL-READ  : output [0-100]   : cycles [280-310]
    // can be optimized to:
    // DMA         : input  [0-100]   : cycles [180-220]
    // DPU         : output [0-100]   : cycles [200-250]
    // now DMA with DPU can not coexist
    std::unordered_set<size_t> modifiedOps;

    _log.trace("Insert control edges for overlapping memory resources");

    // Analyze output from feasible scheduler and prepare list of scheduled
    // operations with their resource and time as needed by control edge
    // generation algorithm
    for (auto& scheduledOp : scheduledOps) {
        VPUX_THROW_UNLESS(scheduledOp.isOriginalOp(), "Invalid operation identified for control edge insertion");

        auto opIndex = scheduledOp.op_;
        auto execOp = _depsInfo.getExecuteOpAtIndex(opIndex);

        // store operations modified by spilling optimizations
        if (scheduledOp.modifiedMemoryRange_) {
            modifiedOps.insert(opIndex);
        }

        // Check all operands of operation and prepare entries in scheduledOpsResources that will be used
        // by control edge algorithm to generate new dependencies
        updateScheduledOpsResourcesForControlEdge(scheduledOpsResources, _aliasInfo, _scan, _memKind, opIndex, execOp,
                                                  _log);
    }

    ControlEdgeSet controlEdges;
    ControlEdgeGenerator controlEdgeGenerator;
    // Generate control edges for overlapping memory regions
    controlEdgeGenerator.generateControlEdges(scheduledOpsResources.begin(), scheduledOpsResources.end(), controlEdges);

    _log = _log.nest();

    // Apply dependencies from controlEdges set into depsInfo.
    // Later they should be transfered to token based dependencies between AsyncExecuteOps
    updateControlEdgesInDepsInfo(_depsInfo, controlEdges, _log, modifiedOps);

    _log = _log.unnest();
}

// After all new dependencies have been prepared call this function to make actual changes in IR
void FeasibleMemorySchedulerControlEdges::updateDependenciesInIR() {
    _log.trace("Update token dependencies in IR");
    _depsInfo.updateTokenDependencies();
}

void vpux::updateControlEdgesInDepsInfo(AsyncDepsInfo& depsInfo, ControlEdgeSet& controlEdges, Logger& log,
                                        const std::unordered_set<size_t>& modifiedOps) {
    // Store information about optimized pairs which can coexist based on scheduler decision
    // and assignment of overlapping cycles
    // Map represents all control edges for a given sink node
    //  key - sink node
    //  vector of values - source nodes
    // TODO: This is to be removed when control edge subview awareness is fully completed (E#106837)
    std::map<size_t, SmallVector<size_t>> optimizedEdgesBasedOnSinkOp;
    // Map represents all control edges for a given  source node
    //  key - source node
    //  vector of values - sink nodes
    // TODO: This is to be removed when control edge subview awareness is fully completed (E#106837)
    std::map<size_t, SmallVector<size_t>> optimizedEdgesBasedOnSourceOp;

    for (auto itr = controlEdges.begin(); itr != controlEdges.end(); ++itr) {
        if (itr->_source == itr->_sink) {
            continue;
        }

        auto sourceOp = depsInfo.getExecuteOpAtIndex(itr->_source);
        auto sinkOp = depsInfo.getExecuteOpAtIndex(itr->_sink);

        if (!modifiedOps.count(itr->_source) && !modifiedOps.count(itr->_sink)) {
            // Coexisting pairs may be modified by spilling optimizations
            auto sourceOpCycleStart = getAsyncExecuteCycleBegin(sourceOp);
            auto sourceOpCycleEnd = getAsyncExecuteCycleEnd(sourceOp);
            auto sinkOpCycleStart = getAsyncExecuteCycleBegin(sinkOp);
            auto sinkOpCycleEnd = getAsyncExecuteCycleEnd(sinkOp);

            // If there is cycle overlap then scheduler assumed operations can coexist and there
            // is no need for memory control edge
            if (sinkOpCycleStart < sourceOpCycleEnd && sinkOpCycleEnd > sourceOpCycleStart) {
                optimizedEdgesBasedOnSinkOp[itr->_sink].push_back(itr->_source);
                optimizedEdgesBasedOnSourceOp[itr->_source].push_back(itr->_sink);
                continue;
            }
        }

        log.trace("Dep: {0} -> {1}", itr->_source, itr->_sink);
        depsInfo.addDependency(sourceOp, sinkOp);
    }

    if (optimizedEdgesBasedOnSinkOp.empty() && optimizedEdgesBasedOnSourceOp.empty()) {
        return;
    }

    log.trace("Identified edges to optimize due to overlapping cycles - {0}", optimizedEdgesBasedOnSourceOp);

    // Traverse again to update dependencies from optimized deps sources
    // to destinations of optimized control flow sinks
    for (auto itr = controlEdges.begin(); itr != controlEdges.end(); ++itr) {
        if (itr->_source == itr->_sink) {
            continue;
        }

        auto thisOpOptimEdgesItr = optimizedEdgesBasedOnSinkOp.find(itr->_source);
        if (thisOpOptimEdgesItr != optimizedEdgesBasedOnSinkOp.end()) {
            for (auto& src : thisOpOptimEdgesItr->second) {
                auto sourceOp = depsInfo.getExecuteOpAtIndex(src);
                auto sinkOp = depsInfo.getExecuteOpAtIndex(itr->_sink);

                log.trace("Dep: {0} -> {1}", src, itr->_sink);
                depsInfo.addDependency(sourceOp, sinkOp);
            }
        }

        thisOpOptimEdgesItr = optimizedEdgesBasedOnSourceOp.find(itr->_sink);
        if (thisOpOptimEdgesItr != optimizedEdgesBasedOnSourceOp.end()) {
            for (auto& dst : thisOpOptimEdgesItr->second) {
                auto sourceOp = depsInfo.getExecuteOpAtIndex(itr->_source);
                auto sinkOp = depsInfo.getExecuteOpAtIndex(dst);

                log.trace("Dep: {0} -> {1}", itr->_source, dst);
                depsInfo.addDependency(sourceOp, sinkOp);
            }
        }
    }
}
