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

#include "vpux/compiler/core/barrier_info.hpp"
#include "vpux/compiler/core/execution_group_analysis.hpp"
#include "vpux/compiler/dialect/VPUIP/utils/utils.hpp"
#include "vpux/compiler/dialect/VPURT/utils/barrier_legalization_utils.hpp"
#include "vpux/compiler/dialect/config/IR/resources.hpp"
#include "vpux/compiler/dialect/config/IR/utils.hpp"
#include "vpux/compiler/dialect/config/utils/config_option_utils.hpp"
#include "vpux/compiler/utils/attributes.hpp"
#include "vpux/compiler/utils/dma.hpp"
#include "vpux/compiler/utils/shave.hpp"
#include "vpux/utils/core/range.hpp"

#include <llvm/ADT/SetOperations.h>

#include <fstream>

using namespace vpux;

//
// Constructor
//

vpux::BarrierInfo::BarrierInfo(mlir::func::FuncOp func)
        : _log(Logger::global().nest("barrier-info", 0)),
          _func(func),
          _syncTaskAttrName(mlir::StringAttr::get(func->getContext(), "sync-task")) {
    buildBarrierMaps(func);
}

vpux::BarrierInfo::BarrierInfo(): _log(Logger::global().nest("barrier-info", 0)), _func(nullptr) {
}

//
// clearAttributes
//

void vpux::BarrierInfo::clearAttributes() {
    for (auto op : _allTaskOps) {
        VPUX_THROW_UNLESS(op.getProperties().getTaskIndex().has_value(),
                          "Remove: task index was not set for '{0}' operation at '{1}'", op->getName(), op->getLoc());
        op.getProperties().setTaskIndex(std::nullopt);
    }
    for (auto op : _allBarrierOps) {
        VPUX_THROW_UNLESS(op.getBarrierIndex().has_value(),
                          "Remove: barrier index was not set for '{0}' operation at '{1}'", op->getName(),
                          op->getLoc());
        op.setBarrierIndex(std::nullopt);
    }
}

//
// getIndex (TaskOp)
//

uint32_t vpux::BarrierInfo::getIndex(VPURT::TaskOp taskOp) const {
    const auto taskIndex = taskOp.getProperties().getTaskIndex();
    VPUX_THROW_UNLESS(taskIndex.has_value(), "Get: task index was not set for '{0}' operation at '{1}'",
                      taskOp->getName(), taskOp->getLoc());

    return checked_cast<uint32_t>(taskIndex.value());
}

//
// getIndex (BarrierOpInterface)
//

uint32_t vpux::BarrierInfo::getIndex(VPURT::BarrierOpInterface barrierOp) const {
    VPUX_THROW_UNLESS(barrierOp.getBarrierIndex().has_value(),
                      "Get: barrier index was not set for '{0}' operation at '{1}'", barrierOp->getName(),
                      barrierOp->getLoc());

    return checked_cast<uint32_t>(barrierOp.getBarrierIndex().value());
}

//
// getBarrierOpAtIndex
//

VPURT::BarrierOpInterface vpux::BarrierInfo::getBarrierOpAtIndex(size_t opIdx) const {
    VPUX_THROW_UNLESS(_allBarrierOps.size() > opIdx, "Barrier: Invalid index '{0}' for _allBarrierOps", opIdx);
    VPUX_THROW_UNLESS(getIndex(_allBarrierOps[opIdx]) == opIdx, "Barrier: Index not matching '{0}' '{1}'", opIdx,
                      _allBarrierOps[opIdx]);
    return _allBarrierOps[opIdx];
}

//
// getTaskOpAtIndex
//

VPURT::TaskOp vpux::BarrierInfo::getTaskOpAtIndex(size_t opIdx) const {
    VPUX_THROW_UNLESS(getNumOfTasks() > opIdx, "Task: Invalid index '{0}' for _allTaskOps", opIdx);
    VPUX_THROW_UNLESS(getIndex(_allTaskOps[opIdx]) == opIdx, "Task: Index not matching '{0}'", opIdx);
    return _allTaskOps[opIdx];
}

bool vpux::BarrierInfo::isTaskQueueTypeInitialized(VPURT::TaskQueueType taskQueueType) const {
    if (taskQueueType.type == config::ExecutorKind::UNKNOWN) {
        return false;
    }
    return _taskQueueTypeMap.find(taskQueueType) != _taskQueueTypeMap.end();
}

VPURT::TaskQueueType vpux::BarrierInfo::getTaskQueueType(size_t taskInd) const {
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    auto taskQueueTypeIt = llvm::find_if(_taskQueueTypeMap, [&](const auto& queueAndTaskBitVecPair) {
        return queueAndTaskBitVecPair.second.test(taskInd);
    });
    if (taskQueueTypeIt == _taskQueueTypeMap.end()) {
        return VPURT::TaskQueueType{config::ExecutorKind::UNKNOWN, 0};
    }
    return taskQueueTypeIt->first;
}

SmallVector<VPURT::TaskQueueType> vpux::BarrierInfo::getNonEmptyTaskQueueTypes() const {
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    SmallVector<VPURT::TaskQueueType> taskQueueTypes;
    for (auto taskQueueType : _taskQueueTypeMap) {
        if (!taskQueueType.second.any()) {
            continue;
        }
        taskQueueTypes.push_back(taskQueueType.first);
    }
    return taskQueueTypes;
}

size_t vpux::BarrierInfo::getLastTaskForQueueType(VPURT::TaskQueueType taskQueueType) const {
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    auto taskQueueTypeMapIt = _taskQueueTypeMap.find(taskQueueType);
    VPUX_THROW_WHEN(taskQueueTypeMapIt == _taskQueueTypeMap.end(),
                    "Task queue map not initialized for executor {0}:{1}", stringifyEnum(taskQueueType.type).data(),
                    taskQueueType.id);

    auto lastTaskInd = taskQueueTypeMapIt->second.find_last();
    VPUX_THROW_WHEN(lastTaskInd == -1, "No task for queue type {0}:{1}", stringifyEnum(taskQueueType.type).data(),
                    taskQueueType.id);

    return static_cast<size_t>(lastTaskInd);
}

std::optional<size_t> vpux::BarrierInfo::getPrevTaskOnQueue(size_t taskInd, VPURT::TaskQueueType taskQueueType) const {
    VPUX_THROW_UNLESS(getNumOfTasks() > taskInd, "Task index '{0}' out of range", taskInd);
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    VPUX_THROW_WHEN(_taskQueueTypeMap.find(taskQueueType) == _taskQueueTypeMap.end(),
                    "Task queue map not initialized for executor of task {0}", taskInd);

    const auto& taskBitVec = _taskQueueTypeMap.at(taskQueueType);

    auto prevTaskInd = taskBitVec.find_prev(taskInd);
    return (prevTaskInd == -1) ? std::nullopt : std::optional<size_t>{prevTaskInd};
}

std::optional<size_t> vpux::BarrierInfo::getNextTaskOnQueue(size_t taskInd, VPURT::TaskQueueType taskQueueType) const {
    VPUX_THROW_UNLESS(getNumOfTasks() > taskInd, "Task index '{0}' out of range", taskInd);
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    VPUX_THROW_WHEN(_taskQueueTypeMap.find(taskQueueType) == _taskQueueTypeMap.end(),
                    "Task queue map not initialized for executor of task {0}", taskInd);

    const auto& taskBitVec = _taskQueueTypeMap.at(taskQueueType);

    auto nextTaskInd = taskBitVec.find_next(taskInd);
    return (nextTaskInd == -1) ? std::nullopt : std::optional<size_t>{nextTaskInd};
}

// Return the closest previous task on the same queue that has a wait barrier
// Example:  Bar0 -> Op0 -> Op1 -> Op2
//  When called for Op2 it would return Op0 as Op0 has wait barrier
std::optional<size_t> vpux::BarrierInfo::getPrevTaskOnQueueWithWaitBar(size_t taskInd,
                                                                       VPURT::TaskQueueType taskQueueType) const {
    VPUX_THROW_UNLESS(getNumOfTasks() > taskInd, "Task index '{0}' out of range", taskInd);
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    VPUX_THROW_WHEN(_taskQueueTypeMap.find(taskQueueType) == _taskQueueTypeMap.end(),
                    "Task queue map not initialized for executor of task {0}", taskInd);

    const auto& taskBitVec = _taskQueueTypeMap.at(taskQueueType);

    auto currentTaskInd = taskInd;
    do {
        auto prevTaskInd = taskBitVec.find_prev(currentTaskInd);
        if (prevTaskInd < 0) {
            return std::nullopt;
        }
        currentTaskInd = static_cast<size_t>(prevTaskInd);
    } while (_taskWaitBarriers[currentTaskInd].empty());

    return currentTaskInd;
}

//
// getWaitBarriers
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getWaitBarriers(size_t taskInd) {
    VPUX_THROW_UNLESS(taskInd <= _taskWaitBarriers.size(), "Task not found in _taskWaitBarriers, '{0}'", taskInd);
    return _taskWaitBarriers[taskInd];
}

void vpux::BarrierInfo::dumpBarriers(std::string fileNamePrefix, std::optional<bool> wlmEnabled) {
    dumpBarrierDependencies(_taskUpdateBarriers, fileNamePrefix + ".taskUpdateBarriers");
    dumpBarrierDependencies(_taskWaitBarriers, fileNamePrefix + ".taskWaitBarriers");
    dumpBarrierDependencies(_barrierProducerMap, fileNamePrefix + ".barrierProducerMap");
    dumpBarrierDependencies(_barrierConsumerMap, fileNamePrefix + ".barrierConsumerMap");
    // Dump number of barrier slots used per task
    dumpSlots(fileNamePrefix + ".slots");

    // Dump queues
    mlir::DenseSet<vpux::config::ExecutorKind> executorKind = {
            config::ExecutorKind::DPU,
            config::ExecutorKind::DMA_NN,
            config::ExecutorKind::SHAVE_ACT,
    };
    auto taskQueueTypeMapOrig = _taskQueueTypeMap;
    clearTaskQueueTypeMap();
    initializeTaskQueueTypeMap(executorKind);
    buildTaskQueueTypeMap();
    dumpQueues(_taskQueueTypeMap, fileNamePrefix + ".taskQueueTypeMap");
    dumpTaskIdxToQueueIdxMapping(fileNamePrefix + ".taskIdxToQueueIdxMap", wlmEnabled);
    // restore original state of _taskQueueTypeMap to keep the state unchanged
    _taskQueueTypeMap = std::move(taskQueueTypeMapOrig);
}

void vpux::BarrierInfo::dumpBarrierDependencies(const SmallVector<BarrierInfo::TaskSet>& depsMap,
                                                std::string fileName) {
    VPUX_THROW_WHEN(fileName.empty(), "Output file name prefix for barriers dump was not provided");
    _log.trace("Dumping barrier dependencies to {0}", fileName);
    std::ofstream ofs(fileName);
    VPUX_THROW_UNLESS(ofs.is_open(), "Cannot open output file {0}", fileName);

    auto getTasksVec = [&](const BarrierInfo::TaskSet& deps) {
        SmallVector<size_t> v(deps.begin(), deps.end());
        llvm::sort(v.begin(), v.end());
        return v;
    };

    // SmallSet gives ordered iteration but the order is not specified - it can be either increasing or decreasing.
    // Make the order always increasing.
    for (const auto& [i, deps] : enumerate(depsMap)) {
        ofs << i << " ";
        for (auto& dep : getTasksVec(deps)) {
            ofs << dep << " ";
        }
        ofs << std::endl;
    }
    ofs.close();
}

void vpux::BarrierInfo::dumpSlots(std::string fileName) {
    VPUX_THROW_WHEN(fileName.empty(), "Output file name prefix for barriers dump was not provided");
    _log.trace("Dumping barrier slots to {0}", fileName);
    std::ofstream ofs(fileName);
    VPUX_THROW_UNLESS(ofs.is_open(), "Cannot open output file {0}", fileName);

    for (const auto& [i, taskInd] : enumerate(irange(getNumOfTasks()))) {
        ofs << i << " " << getNumOfSlotsUsedByTask(getTaskOpAtIndex(taskInd)) << std::endl;
    }
    ofs.close();
}

void vpux::BarrierInfo::dumpQueues(const std::map<VPURT::TaskQueueType, llvm::BitVector>& queueMap,
                                   std::string fileName) {
    VPUX_THROW_WHEN(fileName.empty(), "Output file name prefix for barriers dump was not provided");
    _log.trace("Dumping tasks queues to {0}", fileName);
    std::ofstream ofs(fileName);
    VPUX_THROW_UNLESS(ofs.is_open(), "Cannot open output file {0}", fileName);

    for (const auto& [queue, deps] : queueMap) {
        ofs << config::stringifyExecutorKind(queue.type).str() << "_" << static_cast<size_t>(queue.id) << "("
            << deps.count() << "): ";
        for (auto dep : deps.set_bits()) {
            ofs << dep << " ";
        }
        ofs << std::endl;
    }
    ofs.close();
}

void vpux::BarrierInfo::dumpTaskIdxToQueueIdxMapping(std::string fileName, std::optional<bool> wlmEnabled) {
    VPUX_THROW_WHEN(fileName.empty(), "Output file name prefix for barriers dump was not provided");
    _log.trace("Dumping task to invariant id mapping to {0}", fileName);
    std::ofstream ofs(fileName);
    VPUX_THROW_UNLESS(ofs.is_open(), "Cannot open output file {0}", fileName);

    mlir::DenseSet<size_t> syncDmaTaskIndices;
    _func->walk([&](VPURT::TaskOp taskOp) {
        if (taskOp.getExecutorKind() == config::ExecutorKind::DMA_NN &&
            mlir::dyn_cast<VPUIP::SyncDMAOp>(taskOp.getInnerTaskOp())) {
            syncDmaTaskIndices.insert(getIndex(taskOp));
            return;
        }
    });

    auto execGroupAnalysis =
            wlmEnabled.has_value() && wlmEnabled.value()
                    ? ExecutionGroupAnalysis(_func)
                    : ExecutionGroupAnalysis(_func, /* ignoreVariantLimit */ true, /* ignoreInvariantLimit */ false);

    ExecutionGroupListMap execGroups = execGroupAnalysis.getExecutionGroups();
    for (const auto& [queue, tasksBitVec] : _taskQueueTypeMap) {
        size_t taskIndexInQueue = 0;
        ExecutionGroupList fifoExecGroups = execGroups[queue];

        auto getExecGroupIdx = [&](size_t taskIndex) {
            for (const auto& [grpIdx, execGroup] : enumerate(fifoExecGroups)) {
                if (std::find(execGroup.begin(), execGroup.end(), taskIndex) != execGroup.end()) {
                    return std::optional<int>(static_cast<int>(grpIdx));
                }
            }
            return std::optional<int>(std::nullopt);
        };

        auto execGrpTypeStr = wlmEnabled.has_value() && wlmEnabled.value() ? "execGrp:" : "nonWLMexecGrp:";
        for (auto taskIndex : tasksBitVec.set_bits()) {
            bool isSyncDma = syncDmaTaskIndices.contains(taskIndex);
            auto blockIdx = getControlGraphBlockIndex(taskIndex);
            auto execGrp = getExecGroupIdx(taskIndex);
            std::string execGrpStr = execGrp.has_value() ? std::to_string(execGrp.value()) : "NONE";
            std::string opDescriptionStr = isSyncDma ? "SyncDMA" : "NONE";
            ofs << taskIndex << " " << taskIndexInQueue << " " << config::stringifyExecutorKind(queue.type).str() << "_"
                << static_cast<size_t>(queue.id) << " block:" << blockIdx << " " << execGrpTypeStr << execGrpStr
                << " syncPoint:" << isSyncPoint(taskIndex) << " desc:" << opDescriptionStr << std::endl;
            taskIndexInQueue++;
        }
    }
    ofs.close();
}

//
// getUpdateBarriers (by Idn)
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getUpdateBarriers(size_t taskInd) {
    VPUX_THROW_UNLESS(taskInd < _taskUpdateBarriers.size(), "Task not found in _taskUpdateBarriers, '{0}'", taskInd);
    return _taskUpdateBarriers[taskInd];
}

//
// getBarrierProducers
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getBarrierProducers(size_t barrierInd) {
    VPUX_THROW_UNLESS(barrierInd < _barrierProducerMap.size(), "Barrier not found in _barrierProducerMap, '{0}'",
                      barrierInd);
    return _barrierProducerMap[barrierInd];
}

//
// getBarrierConsumers
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getBarrierConsumers(size_t barrierInd) {
    VPUX_THROW_UNLESS(barrierInd < _barrierConsumerMap.size(), "Barrier not found in _barrierConsumerMap, '{0}'",
                      barrierInd);
    return _barrierConsumerMap[barrierInd];
}

BarrierInfo::TaskSet vpux::BarrierInfo::getBarriersUsers(const std::set<int64_t>& barrierInds) {
    BarrierInfo::TaskSet users;
    for (const auto& barrierInd : barrierInds) {
        llvm::set_union(users, getBarrierProducers(barrierInd));
        llvm::set_union(users, getBarrierConsumers(barrierInd));
    }
    return users;
}

//
// getBarrierProducers
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getBarrierProducers(VPURT::BarrierOpInterface barrierOp) {
    auto barrierInd = getIndex(barrierOp);
    return getBarrierProducers(barrierInd);
}

//
// getBarrierConsumers
//

BarrierInfo::TaskSet& vpux::BarrierInfo::getBarrierConsumers(VPURT::BarrierOpInterface barrierOp) {
    auto barrierInd = getIndex(barrierOp);
    return getBarrierConsumers(barrierInd);
}

//
// getNumOfSlotsUsed
//

size_t vpux::BarrierInfo::getNumOfSlotsUsed(VPURT::TaskOp op) {
    // This function returns the number of variants used by a task for a barrier.
    // On VPU H/W, a NCE task can be executed (across multiple DPUs) via workloads descriptors (known as variants).
    // An NCE task may have multiple workloads descriptors (which are generated in the NCE DPU workloads pass).
    // Therefore, the number of variants must be verified as they will all update a barrier and
    // contribute to the architecture specific MAX VARIANT COUNT that a barrier has.
    // Until NPU4 each variant updates the barrier to signal that it is complete but starting with NPU4 (with single DPU
    // per tile), only first/last variant of any given invariant will consume/produce a barrier in which case the
    // required slot count will be 1 as all variants of the invariant use the same FIFO. A DMA does not have variants,
    // therefore they always utilize a single barrier slot.

    if (op.getExecutorKind() == config::ExecutorKind::DPU) {
        if (VPUIP::supportsPerVariantBarrierConfiguration(op)) {
            return 1;
        }
        auto nceOp = mlir::dyn_cast<VPUIP::NCEClusterTaskOp>(op.getInnerTaskOp());
        VPUX_THROW_UNLESS(nceOp != nullptr, "Could not cast to NCE task");
        return nceOp.getNumVariants();
    }

    if (op.getExecutorKind() == config::ExecutorKind::SHAVE_ACT) {
        if (auto swKernelOp = mlir::dyn_cast<VPUIP::SwKernelOp>(op.getInnerTaskOp())) {
            auto swKernelRun = swKernelOp.getBody().getOps<VPUIP::SwKernelRun>();
            return std::distance(swKernelRun.begin(), swKernelRun.end());
        }
        return 1;
    }

    if (op.getExecutorKind() == config::ExecutorKind::DMA_NN || op.getExecutorKind() == config::ExecutorKind::M2I) {
        return 1;
    }

    // TODO: Analyze and define executor type for funcOp - E#117624
    if (op.getExecutorKind() == config::ExecutorKind::UNKNOWN) {
        return 1;
    }

    VPUX_THROW("Unsupported executor: {0}", op.getExecutorKind());
}

size_t vpux::BarrierInfo::getNumOfSlotsUsedByTask(VPURT::TaskOp op) const {
    return getNumOfSlotsUsed(op);
}

//
// getProducerSlotCount
//

size_t vpux::BarrierInfo::getProducerSlotCount(VPURT::BarrierOpInterface barrierOp) {
    size_t producerSlotCount = 0;
    for (const auto& producer : getBarrierProducers(barrierOp)) {
        producerSlotCount += getNumOfSlotsUsed(getTaskOpAtIndex(producer));
    }
    return producerSlotCount;
}

//
// getConsumerSlotCount
//

size_t vpux::BarrierInfo::getConsumerSlotCount(VPURT::BarrierOpInterface barrierOp) {
    size_t consumerSlotCount = 0;
    for (const auto& consumer : getBarrierConsumers(barrierOp)) {
        consumerSlotCount += getNumOfSlotsUsed(getTaskOpAtIndex(consumer));
    }
    return consumerSlotCount;
}

//
// getNumOfBarrierOps
//

size_t vpux::BarrierInfo::getNumOfBarrierOps() const {
    return _allBarrierOps.size();
}

size_t vpux::BarrierInfo::getBarrierMaxVariantSum() const {
    return VPUIP::getBarrierMaxVariantSum(_func);
}

//
// getNumOfTasks
//

size_t vpux::BarrierInfo::getNumOfTasks() const {
    return _allTaskOps.size();
}

size_t vpux::BarrierInfo::getNumOfTasks(vpux::config::ExecutorKind executorKind) const {
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    size_t taskCount = 0;
    for (auto& [queueType, queueTasks] : _taskQueueTypeMap) {
        if (queueType.type == executorKind) {
            taskCount += queueTasks.count();
        }
    }

    return taskCount;
}

//
// resizeBitMap
//

void vpux::BarrierInfo::resizeBitMap(SmallVector<llvm::BitVector>& bitMap, size_t length, uint32_t bits) {
    bitMap.resize(length);
    for (auto& bit : bitMap) {
        bit.resize(bits, false);
    }
}

void vpux::BarrierInfo::resetBitMap(SmallVector<llvm::BitVector>& bitMap) {
    for (auto& bit : bitMap) {
        bit.reset();
    }
}

SmallVector<size_t> vpux::BarrierInfo::getBarriersForTaskBlock(size_t blockInd, bool blockStartSyncPoint,
                                                               bool blockEndSyncPoint, bool updateBarriers) const {
    VPUX_THROW_WHEN(blockInd >= getControlGraphBlockCount(), "Invalid task block index ({0})", blockInd);
    auto [blockStartInd, blockEndInd] = getControlGraphBlockTaskRange(blockInd, blockStartSyncPoint, blockEndSyncPoint);

    llvm::BitVector barrierInd(getNumOfBarrierOps() + _fifoDependencies.size(),
                               false);  // use of BitVector here is faster than eg std::set
    if (updateBarriers) {
        for (auto taskInd = blockStartInd; taskInd <= blockEndInd; ++taskInd) {
            for (auto barIdx : _taskUpdateBarriers[taskInd]) {
                barrierInd.set(static_cast<unsigned>(barIdx));
            }
        }
    } else {
        for (auto taskInd = blockStartInd; taskInd <= blockEndInd; ++taskInd) {
            for (auto barIdx : _taskWaitBarriers[taskInd]) {
                barrierInd.set(static_cast<unsigned>(barIdx));
            }
        }
    }

    auto numOfBarriersInBlock = barrierInd.count();
    if (numOfBarriersInBlock == 0) {
        return {};
    }

    SmallVector<size_t> barriersInBlock;
    barriersInBlock.reserve(numOfBarriersInBlock);
    for (auto idx : barrierInd.set_bits()) {
        barriersInBlock.push_back(idx);
    }
    return barriersInBlock;
}

size_t vpux::BarrierInfo::getControlGraphBlockIndex(size_t taskInd) const {
    if (_syncTasksIds.empty()) {
        return 0;
    }
    VPUX_THROW_UNLESS(taskInd < _taskToBlockMap.size(), "Task index '{0}' out of range. Needs to be < '{1}", taskInd,
                      _taskToBlockMap.size());
    return _taskToBlockMap[taskInd];
}

size_t vpux::BarrierInfo::getBarrierBlockIndex(size_t barInd) {
    auto barrierLastProducer = getBarrierLatestProducer(barInd);
    if (isSyncPoint(barrierLastProducer)) {
        return getControlGraphBlockIndex(barrierLastProducer) + 1;
    }
    return getControlGraphBlockIndex(barrierLastProducer);
}

std::optional<size_t> vpux::BarrierInfo::getControlGraphSyncPoint(size_t taskInd) const {
    auto blockInd = getControlGraphBlockIndex(taskInd);
    return getControlGraphSyncPointForBlock(blockInd);
}

std::optional<size_t> vpux::BarrierInfo::getControlGraphSyncPointForBlock(size_t blockInd) const {
    if (blockInd >= _syncTasksIds.size()) {
        return std::nullopt;
    }
    return _syncTasksIds[blockInd];
}

//
//  canBarProducersControlBarConsumers
//

bool vpux::BarrierInfo::canBarProducersControlBarConsumers(const TaskSet& producers, const TaskSet& consumers,
                                                           ArrayRef<TaskSet> origWaitBarriersMap) {
    if (consumers.empty()) {
        return true;
    }
    if (producers.empty()) {
        return false;
    }

    auto haveDmaExecutorType = [&](const TaskSet& tasks) {
        auto tasksExecutorKind = haveSameExecutorKind(tasks);
        return tasksExecutorKind.has_value() && tasksExecutorKind == config::ExecutorKind::DMA_NN;
    };

    // Just because none of the producers do not share any wait barrier with some of the consumers, does not mean that
    // no producer-consumer pair cannot run in parallel. This simplification can lead to merging more barriers and
    // sometimes can cause serialization of the graph. However, the resulting schedule will have less
    // barriers and may perform faster even though DMA-DPU or DMA-SW tasks overlap can be smaller due to the
    // serialization. This likely stems from fact that DMAs can be very fast and merging barriers associated with them
    // can be beneficial even at a cost of increased serialization, thus we allow for such merges. The strategy can be
    // revisited with full-WLM enabled where there will be no overhead related with handling barriers.
    if (!haveDmaExecutorType(consumers)) {
        return false;
    }
    if (!haveDmaExecutorType(producers)) {
        return false;
    }

    if (*std::max_element(producers.begin(), producers.end()) >=
        *std::min_element(consumers.begin(), consumers.end())) {
        return false;
    }

    auto anyProducerCanRunParallelWithConsumer = llvm::any_of(producers, [&](const auto& producer) {
        return llvm::any_of(consumers, [&](const auto& consumer) {
            return origWaitBarriersMap[producer] == origWaitBarriersMap[consumer];
        });
    });
    return !anyProducerCanRunParallelWithConsumer;
}

//
// buildBarrierMaps
//

void vpux::BarrierInfo::buildBarrierMaps(mlir::func::FuncOp func) {
    _log.trace("Collect initial producer maps");

    for (auto& op : func.getOps()) {
        if (auto taskOp = mlir::dyn_cast<VPURT::TaskOp>(op)) {
            _allTaskOps.push_back(taskOp);
        } else if (auto barrierOp = mlir::dyn_cast<VPURT::BarrierOpInterface>(op)) {
            _allBarrierOps.push_back(barrierOp);
        }
    }

    _log.nest().trace("There are '{0}' VPURT::TaskOp", getNumOfTasks());
    _log.nest().trace("There are '{0}' VPURT::BarrierOpInterface", _allBarrierOps.size());

    // resize barrier maps
    _taskWaitBarriers = SmallVector<TaskSet>(getNumOfTasks(), {});
    _taskUpdateBarriers = SmallVector<TaskSet>(getNumOfTasks(), {});
    _barrierConsumerMap = SmallVector<TaskSet>(_allBarrierOps.size(), {});
    _barrierProducerMap = SmallVector<TaskSet>(_allBarrierOps.size(), {});

    // set index to task ops and barrier ops
    for (const auto& p : _allTaskOps | indexed) {
        p.value().getProperties().setTaskIndex(p.index());
        if (p.value()->hasAttr(_syncTaskAttrName)) {
            _syncTasksIds.push_back(p.index());
        }
    }

    if (!_syncTasksIds.empty()) {
        buildTaskBlockMap();
        if (_log.isActive(LogLevel::Trace)) {
            // Check block sizes
            size_t firstBlockSize = _syncTasksIds[0] + 1;
            for (size_t i = 1; i < _syncTasksIds.size(); i++) {
                size_t nextBlockSize = _syncTasksIds[i] - _syncTasksIds[i - 1];
                if (nextBlockSize != firstBlockSize) {
                    _log.trace("Inferred first block size is '{0}', block '{1}' has {2} tasks", firstBlockSize, i,
                               nextBlockSize);
                }
            }
        }
    }

    for (const auto& p : _allBarrierOps | indexed) {
        p.value().setBarrierIndex(p.index());
    }

    for (auto& op : func.getOps()) {
        if (auto taskOp = mlir::dyn_cast<VPURT::TaskOp>(op)) {
            addTaskOp(taskOp);
        } else if (auto barrierOp = mlir::dyn_cast<VPURT::BarrierOpInterface>(op)) {
            _log.nest().trace("Found 'VPURT::BarrierOpInterface' Operation at '{0}'", op.getLoc());
            VPUX_THROW_WHEN(barrierOp.getBarrier() == nullptr, "BarrierOpInterface '{0}' has no barrier.", barrierOp);

            // ensure all barrier users are TaskOps
            for (auto* user : barrierOp.getBarrier().getUsers()) {
                VPUX_THROW_WHEN(mlir::dyn_cast<VPURT::TaskOp>(user) == nullptr, "Got non-TaskOp Operation at '{0}'",
                                op.getLoc());
            }
        }
    }
}

void vpux::BarrierInfo::buildTaskBlockMap() {
    VPUX_THROW_WHEN(_syncTasksIds.empty(),
                    "Control graph split was not detected. Building task block map is not necessary.");
    size_t nTasks = getNumOfTasks();
    _taskToBlockMap = SmallVector<size_t>(nTasks);

    size_t blockStartInd = 0;
    size_t blockInd = 0;

    for (auto syncTaskInd : _syncTasksIds) {
        for (size_t taskInd = blockStartInd; taskInd <= syncTaskInd; taskInd++) {
            _taskToBlockMap[taskInd] = blockInd;
        }
        blockInd++;
        blockStartInd = syncTaskInd + 1;
    }
    // assign block index to tasks from last block
    for (size_t taskInd = blockStartInd; taskInd < nTasks; taskInd++) {
        _taskToBlockMap[taskInd] = blockInd;
    }
}

//
// addConsumer
//

void vpux::BarrierInfo::addConsumer(size_t barrierInd, size_t taskInd) {
    _log.trace("Add consumer '{0}' for barrier '{1}'", taskInd, barrierInd);
    _barrierConsumerMap[barrierInd].insert(taskInd);
    _taskWaitBarriers[taskInd].insert(barrierInd);
}

void vpux::BarrierInfo::addConsumer(VPURT::BarrierOpInterface barrierOp, size_t taskInd) {
    addConsumer(getIndex(barrierOp), taskInd);
}

//
// addConsumers
//

void vpux::BarrierInfo::addConsumers(size_t barrierInd, const TaskSet& taskInds) {
    for (const auto& taskInd : taskInds) {
        _barrierConsumerMap[barrierInd].insert(taskInd);
        _taskWaitBarriers[taskInd].insert(barrierInd);
    }
}

//
// addProducer
//

void vpux::BarrierInfo::addProducer(size_t barrierInd, size_t taskInd) {
    _log.trace("Add producer '{0}' for barrier '{1}'", taskInd, barrierInd);
    _barrierProducerMap[barrierInd].insert(taskInd);
    _taskUpdateBarriers[taskInd].insert(barrierInd);
}

void vpux::BarrierInfo::addProducer(VPURT::BarrierOpInterface barrierOp, size_t taskInd) {
    addProducer(getIndex(barrierOp), taskInd);
}

//
// addProducers
//

void vpux::BarrierInfo::addProducers(size_t barrierInd, const TaskSet& taskInds) {
    for (const auto& taskInd : taskInds) {
        _barrierProducerMap[barrierInd].insert(taskInd);
        _taskUpdateBarriers[taskInd].insert(barrierInd);
    }
}

//
// removeProducers
//
void vpux::BarrierInfo::removeProducers(size_t barrierInd, const TaskSet& taskInds) {
    for (const auto& taskInd : taskInds) {
        _barrierProducerMap[barrierInd].erase(taskInd);
        _taskUpdateBarriers[taskInd].erase(barrierInd);
    }
}

void vpux::BarrierInfo::removeProducers(VPURT::BarrierOpInterface barrierOp, const TaskSet& taskInds) {
    const auto barrierInd = getIndex(barrierOp);
    removeProducers(barrierInd, taskInds);
}

//
// removeConsumers
//
void vpux::BarrierInfo::removeConsumers(size_t barrierInd, const TaskSet& taskInds) {
    for (const auto& taskInd : taskInds) {
        _barrierConsumerMap[barrierInd].erase(taskInd);
        _taskWaitBarriers[taskInd].erase(barrierInd);
    }
}

void vpux::BarrierInfo::removeConsumers(VPURT::BarrierOpInterface barrierOp, const TaskSet& taskInds) {
    const auto barrierInd = getIndex(barrierOp);
    removeConsumers(barrierInd, taskInds);
}

//
// addTaskOp
//

void vpux::BarrierInfo::addTaskOp(VPURT::TaskOp taskOp) {
    const auto taskInd = getIndex(taskOp);
    _log.trace("Found 'TaskOp' Operation '{0}'", taskInd);

    for (const auto& bar : taskOp.getWaitBarriers()) {
        // Note: can also be VPURT::ConfigureBarrierOp
        auto barrierOp = bar.getDefiningOp<VPURT::BarrierOpInterface>();
        VPUX_THROW_WHEN(barrierOp == nullptr, "Invalid wait barrier op type {0}", bar);
        addConsumer(barrierOp, taskInd);
    }

    for (const auto& bar : taskOp.getUpdateBarriers()) {
        // Note: can also be VPURT::ConfigureBarrierOp
        auto barrierOp = bar.getDefiningOp<VPURT::BarrierOpInterface>();
        VPUX_THROW_WHEN(barrierOp == nullptr, "Invalid wait barrier op type {0}", bar);
        addProducer(barrierOp, taskInd);
    }
}

//
// addNewBarrier
//

size_t vpux::BarrierInfo::addNewBarrier(VPURT::BarrierOpInterface barrierOp) {
    size_t barrierInd = _allBarrierOps.size();
    barrierOp.setBarrierIndex(barrierInd);

    _log.trace("Add new barrier '{0}', new barrier size '{1}'", barrierInd, _allBarrierOps.size());
    _allBarrierOps.push_back(barrierOp);

    _barrierConsumerMap.push_back({});
    _barrierProducerMap.push_back({});
    return barrierInd;
}

//
// addNewTaskOp
//

size_t vpux::BarrierInfo::addNewTaskOp(VPURT::TaskOp taskOp) {
    size_t taskOpIdx = getNumOfTasks();
    taskOp.getProperties().setTaskIndex(taskOpIdx);

    _log.trace("Add new taskOp '{0}', new taskOps size '{1}'", taskOpIdx, getNumOfTasks());

    _allTaskOps.push_back(taskOp);

    _taskUpdateBarriers.push_back({});
    _taskWaitBarriers.push_back({});

    return taskOpIdx;
}

//
// setWaitBarriers
//

void vpux::BarrierInfo::setWaitBarriers(size_t taskInd, const TaskSet& barriers) {
    // remove previous wait barriers
    for (auto barrierInd : _taskWaitBarriers[taskInd]) {
        _barrierConsumerMap[static_cast<size_t>(barrierInd)].erase(taskInd);
    }

    for (auto barrierInd : barriers) {
        _barrierConsumerMap[barrierInd].insert(taskInd);
    }
    _taskWaitBarriers[taskInd] = barriers;
}

//
// setUpdateBarriers
//

void vpux::BarrierInfo::setUpdateBarriers(size_t taskInd, const TaskSet& barriers) {
    // remove previous update barriers
    for (auto barrierInd : _taskUpdateBarriers[taskInd]) {
        _barrierProducerMap[static_cast<size_t>(barrierInd)].erase(taskInd);
    }

    for (auto barrierInd : barriers) {
        _barrierProducerMap[barrierInd].insert(taskInd);
    }
    _taskUpdateBarriers[taskInd] = barriers;
}

//
// removeProducer
//

void vpux::BarrierInfo::removeProducer(size_t barrierInd, size_t taskInd) {
    _log.trace("Remove producer '{0}' for barrier '{1}'", taskInd, barrierInd);
    _barrierProducerMap[barrierInd].erase(taskInd);
    _taskUpdateBarriers[taskInd].erase(barrierInd);
}

void vpux::BarrierInfo::removeProducer(VPURT::BarrierOpInterface barrierOp, size_t taskInd) {
    removeProducer(getIndex(barrierOp), taskInd);
}

//
// removeConsumer
//
void vpux::BarrierInfo::removeConsumer(size_t barrierInd, size_t taskInd) {
    _log.trace("Remove consumer '{0}' for barrier '{1}'", taskInd, barrierInd);
    _barrierConsumerMap[barrierInd].erase(taskInd);
    _taskWaitBarriers[taskInd].erase(barrierInd);
}

void vpux::BarrierInfo::removeConsumer(VPURT::BarrierOpInterface barrierOp, size_t taskInd) {
    removeConsumer(getIndex(barrierOp), taskInd);
}

//
// resetBarrier
//

void vpux::BarrierInfo::resetBarrier(VPURT::BarrierOpInterface barrierOp) {
    const auto barrierInd = getIndex(barrierOp);
    resetBarrier(checked_cast<size_t>(barrierInd));
}

//
// resetBarrier
//

void vpux::BarrierInfo::resetBarrier(size_t barrierInd) {
    _log.trace("Reset barrier '{0}'", barrierInd);

    for (auto taskInd : _barrierProducerMap[barrierInd]) {
        _taskUpdateBarriers[static_cast<size_t>(taskInd)].erase(barrierInd);
    }
    _barrierProducerMap[barrierInd].clear();

    for (auto taskInd : _barrierConsumerMap[barrierInd]) {
        _taskWaitBarriers[static_cast<size_t>(taskInd)].erase(barrierInd);
    }
    _barrierConsumerMap[barrierInd].clear();
}

// Function that will perform split of a control graph with sync points
// present on boundaries of such split
// Example:
//  number of all task: 1200
//  requested split size: 500
//
// Graph will be split into 3 blocks with tasks indexes
// being placed in blocks in following way
//  block 0:    0 -  499
//  block 1:  500 -  999
//  block 2: 1000 - 1199
// Tasks with indexes 499 and 999 are sync points
// meaning that if there was any dependency previously present between blocks
// it will be modified to go through this sync point, example:
//  old dependency:   200 -> Bar -> 600
//  new dependencies: 200 -> newBar1 -> 499  and  499 -> newBar1 -> 600
//
// Goal of this split it to enable faster and less memory hungry optimizations
// on control graph as they can be done within single block between sync tasks
// For this example optimization can be done in 3 ranges: 0 - 499, 499 - 999, 999 - 1199
void vpux::BarrierInfo::splitControlGraphToBlocks(const size_t blockSize) {
    VPUX_THROW_WHEN(blockSize < 2, "Minimal block size is 2, requested size - {1}", blockSize);
    VPUX_THROW_WHEN(!_syncTasksIds.empty(),
                    "Partitioning on control graph was already performed, size - {0}, requested size - {1}",
                    _syncTasksIds[0] + 1, blockSize);

    auto tasksSize = getNumOfTasks();
    if (blockSize >= tasksSize) {
        _log.trace("Not enough tasks to perform split, requested split size - '{0}', num of tasks - '{1}'", blockSize,
                   tasksSize);
        return;
    }

    size_t numOfBlocks = tasksSize / blockSize;
    if (tasksSize % blockSize > 0) {
        numOfBlocks++;
    }
    auto isSyncPoint = [&](size_t taskIdx) {  // use local version of the method as the graph split is not finished.
        if (taskIdx == getNumOfTasks() - 1) {
            return false;
        }
        return (taskIdx + 1) % blockSize == 0;
    };

    _log.trace("Split control graph: num of task - '{0}', blocks size - '{1}', number of blocks - '{2}'", tasksSize,
               blockSize, numOfBlocks);

    auto barriersCount = _allBarrierOps.size();
    std::map<size_t, std::pair<size_t, size_t>> syncTaskWaitAndUpdateBarrierIndMap;

    // STEP 1:
    // Add explicit sync tasks wait and update barriers
    // which will be used to connect tasks with dependencies
    // going outside of block scope
    // If those new barriers are redundant, could have been combined with other existing barrier
    // or need a split because of large number of producer/consumers they will be optimized/processed
    // by follow-up passes
    for (size_t blockInd = 1; blockInd < numOfBlocks; blockInd++) {
        size_t syncTaskInd = blockInd * blockSize - 1;

        auto syncTaskOp = getTaskOpAtIndex(syncTaskInd);
        auto insertionOp = getTaskOpAtIndex((blockInd - 1) * blockSize);

        mlir::OpBuilder builder(insertionOp);
        //      |<-----syncWaitBar
        //  syncTask
        //      |----->syncUpdBar
        auto waitBar = builder.create<VPURT::DeclareVirtualBarrierOp>(syncTaskOp->getLoc());
        auto waitBarInd = addNewBarrier(waitBar);
        addConsumer(waitBarInd, syncTaskInd);

        builder.setInsertionPoint(syncTaskOp);
        auto updateBar = builder.create<VPURT::DeclareVirtualBarrierOp>(syncTaskOp->getLoc());
        auto updateBarInd = addNewBarrier(updateBar);

        addProducer(updateBarInd, syncTaskInd);

        _log.trace("Sync task at index '{0}', wait bar - '{1}', update bar - '{2}'", syncTaskInd, waitBarInd,
                   updateBarInd);

        syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd] = std::make_pair(waitBarInd, updateBarInd);
    }

    std::map<VPURT::TaskQueueType, size_t> firstTaskWithNoWaitBarForQueueInBlockMap;
    std::map<VPURT::TaskQueueType, size_t> lastTaskWithNoUpdateBarForQueueInBlockMap;

    // STEP 2
    // Identify operations in each block with no wait or update barriers that are first or last
    // for given queue:
    // - last op with no update barrier on each queue - link it to wait barrier of next sync task
    // - first op with no wait barrier on each queue - link it to update barrier of previous sync task
    //
    // If they will not have barriers assigned their position might be shifted
    // with respect to sync task during barrier legalization steps in follow-up passes
    for (size_t taskInd = 0; taskInd < tasksSize; taskInd++) {
        size_t taskBlockInd = taskInd / blockSize;
        auto taskUpdateBarriers = _taskUpdateBarriers[taskInd];
        size_t syncTaskInd = taskBlockInd * blockSize + blockSize - 1;

        // If task does not update any barrier store this information for each queue
        // and when moving to next block add a link to sync task. When encountering sync
        // task add connection from those ops to this sync task
        if (taskBlockInd < numOfBlocks - 1) {
            if (taskInd == syncTaskInd) {
                auto syncTaskWaitBarInd = syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd].first;
                for (const auto& queueListTaskPair : lastTaskWithNoUpdateBarForQueueInBlockMap) {
                    // Create new connection
                    // queueLastTaskInd -> syncWaitBarInd -> syncTaskInd
                    //
                    //  0(last op on queue)       0--------->|
                    //                                       |
                    //        syncWaitBar    =>         syncWaitBar
                    //             |                         |
                    //  1(sync) <--|              1(sync) <--|
                    _log.trace("New dep for task with no update barrier: task '{0}' -> sync task '{1}'",
                               queueListTaskPair.second, syncTaskInd);
                    addProducer(syncTaskWaitBarInd, queueListTaskPair.second);
                }
                lastTaskWithNoUpdateBarForQueueInBlockMap.clear();
            } else if (taskUpdateBarriers.empty()) {
                auto taskOp = getTaskOpAtIndex(taskInd);
                const auto taskQueueType = VPURT::getTaskQueueType(taskOp, false);
                lastTaskWithNoUpdateBarForQueueInBlockMap[taskQueueType] = taskInd;
            }
        }

        // If task does not wait for any barrier and is first on a given queue
        // for this block (e.g. DMA task) then add a dependency from previous syncTask
        // Otherwise after call to orderExecutionTasksAndBarriers might move it to earlier block
        // and destroy assumption about how graph got split
        if (taskBlockInd > 0) {
            if (taskInd == syncTaskInd || taskInd == tasksSize - 1) {
                size_t prevSyncTaskInd = taskBlockInd * blockSize - 1;
                auto prevSyncTaskUpdateBarInd = syncTaskWaitAndUpdateBarrierIndMap[prevSyncTaskInd].second;
                for (const auto& queueListTaskPair : firstTaskWithNoWaitBarForQueueInBlockMap) {
                    // Create new connection
                    // prevSyncTaskInd -> syncUpdBarInd -> queueFirstTaskInd
                    //
                    //  1(sync)---|               1(sync) ---|
                    //            |                          |
                    //       syncUpdBar       =>        syncUpdBar
                    //                                       |
                    //  2(first op on queue)      2<---------|
                    _log.trace("New dep for task with no wait barrier: sync task '{0}' -> task '{1}'", prevSyncTaskInd,
                               queueListTaskPair.second);
                    addConsumer(prevSyncTaskUpdateBarInd, queueListTaskPair.second);
                }
                firstTaskWithNoWaitBarForQueueInBlockMap.clear();
            } else if (_taskWaitBarriers[taskInd].empty()) {
                auto taskOp = getTaskOpAtIndex(taskInd);
                const auto taskQueueType = VPURT::getTaskQueueType(taskOp, false);
                if (firstTaskWithNoWaitBarForQueueInBlockMap.find(taskQueueType) ==
                    firstTaskWithNoWaitBarForQueueInBlockMap.end()) {
                    firstTaskWithNoWaitBarForQueueInBlockMap[taskQueueType] = taskInd;
                }
            }
        }
    }

    auto getSetWithNoSyncTasks = [&](const TaskSet& taskSet) {
        TaskSet taskSetWithNoSync;
        for (const auto& task : taskSet) {
            if (!isSyncPoint(task)) {
                taskSetWithNoSync.insert(task);
            }
        }

        return taskSetWithNoSync;
    };

    // STEP 3
    // Perform main processing to identify all ops with long dependency -
    // dependencies that go beyond the scope of existing block
    // Example:
    //  block 0:    0 -  499
    //  block 1:  500 -  999
    //  old dependency:   200 -> Bar -> 600
    //  new dependencies: 200 -> syncWaitBar -> 499(syncTask)  and  499(SyncTask) -> syncUpdBar -> 600
    // For this task scan all barriers and check such that within its producers and consumers have tasks from different
    // blocks. If yes then split those tasks and barrier so that tasks from different groups do not share barrier,
    // besides sync tasks
    for (size_t barInd = 0; barInd < barriersCount; ++barInd) {
        // Check barrier consumers and producers and check if they are from different blocks
        llvm::SmallSet<size_t, 16> blocksSet;

        VPUX_THROW_WHEN(_barrierProducerMap[barInd].empty(), "Barrier has no producers");
        VPUX_THROW_WHEN(_barrierConsumerMap[barInd].empty(), "Barrier has no consumers");

        // Gather information about barrier producers and consumers
        // and store them based on blockId
        std::map<size_t, TaskSet> barProdBlockTasksMap;
        for (const auto& taskInd : _barrierProducerMap[barInd]) {
            size_t taskBlockInd = taskInd / blockSize;
            blocksSet.insert(taskBlockInd);
            barProdBlockTasksMap[taskBlockInd].insert(taskInd);
        }

        std::map<size_t, TaskSet> barConsBlockTasksMap;
        for (const auto& taskInd : _barrierConsumerMap[barInd]) {
            size_t taskBlockInd = taskInd / blockSize;
            blocksSet.insert(taskBlockInd);
            barConsBlockTasksMap[taskBlockInd].insert(taskInd);
        }

        if (blocksSet.size() < 2) {
            // Only tasks from single block included
            // no need to split
            continue;
        }

        // Get barrier block assignment, which would correspond to earliest producer block id
        size_t barBlockId = std::min(barProdBlockTasksMap.begin()->first, barConsBlockTasksMap.begin()->first);

        // If this is last block then early exist. No need to analyze any subsequent barriers
        // as barriers are ordered and from that point no any subsequent barrier would need a split
        if (barBlockId == numOfBlocks + 1) {
            break;
        }

        _log.trace("Bar {0} needs a split, bar block id - {1}", barInd, barBlockId);

        // Remove tasks from this barrier which do not belong to this block
        for (const auto& barProdBlockTasksPair : barProdBlockTasksMap) {
            const auto blockId = barProdBlockTasksPair.first;
            if (blockId == barBlockId) {
                continue;
            }
            // Example control flow that will be shown throughout STEP 3
            // how it is modified
            // bar0 - barrier that is rossing boundary of blocks
            // 0p, 2p, 4p - barrier producer (p) tasks
            // 5c - barrier consumer (c) task
            // 1, 3 - sync tasks
            //
            //    0p -------->|         0p --------->|
            //                |                     bar0
            //    1(sync)     |         1(sync)      |
            //                |                      |
            //    2p -------->|         2p           |
            //                |    =>                |
            //    3(sync)     |         3(sync)      |
            //                |                      |
            //    4p -------->|         4p           |
            //               bar0                    |
            //    5c <--------|         5c <---------|
            _log.nest().trace("Remove producer for bar {0} from block {1}", barInd, blockId);
            removeProducers(barInd, barProdBlockTasksPair.second);

            if (blockId != barProdBlockTasksMap.rbegin()->first ||
                barConsBlockTasksMap.find(blockId) == barConsBlockTasksMap.end()) {
                // If this is not the last block produced by this barrier or there are no consumers in this block
                // add producers of this bar to now produce sync wait barrier
                //
                //    0p -------->|         0p ------------------->|
                //               bar0                             bar0
                //    1(sync)     |         1(sync)                |
                //                |                                |
                //    2p          |         2p -------->|          |
                //                |     =>          syncWaitBar    |
                //    3(sync)     |         3(sync) <---|          |
                //                |                                |
                //    4p          |         4p                     |
                //                |                                |
                //    5c <--------|         5c <-------------------|
                //
                size_t nextSyncTaskInd = blockId * blockSize - 1 + blockSize;
                if (nextSyncTaskInd < tasksSize) {
                    size_t syncTaskWaitBarInd = syncTaskWaitAndUpdateBarrierIndMap[nextSyncTaskInd].first;

                    _log.nest(2).trace("Add producers of sync bar {0} for tasks from block {1}", syncTaskWaitBarInd,
                                       blockId);
                    addProducers(syncTaskWaitBarInd, getSetWithNoSyncTasks(barProdBlockTasksPair.second));
                }
            } else {
                // If this is last block produced by this barrier and there are consumers in this block
                // create a new barrier to handle remaining deps that were produced by currently analyzed barInd
                //
                //    0p ------------------>|       0p ------------------->|
                //                         bar0                          bar0
                //    1(sync)               |       1(sync)                |
                //                          |                              |
                //    2p --------->|        |       2p -------->|          |
                //            syncWaitBar   |  =>          syncWaitBar     |
                //    3(sync) <----|        |       3(sync) <---|          |
                //                          |                              |
                //    4p                    |       4p -------->|          |
                //                          |                newBar        |
                //    5c <------------------|       5c <--------|<---------|
                //
                auto taskOp = getTaskOpAtIndex(*barProdBlockTasksPair.second.begin());
                mlir::OpBuilder builder(taskOp);
                auto newBar = builder.create<VPURT::DeclareVirtualBarrierOp>(taskOp->getLoc());
                auto newBarInd = addNewBarrier(newBar);

                _log.nest(2).trace("Add producers and consumers for new bar {0} for tasks from block {1}", newBarInd,
                                   blockId);
                addProducers(newBarInd, barProdBlockTasksPair.second);
                addConsumers(newBarInd, barConsBlockTasksMap[blockId]);
            }

            size_t syncTaskInd = blockId * blockSize - 1;
            size_t syncTaskUpdBarInd = syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd].second;

            //    0p -------------->|         0p --------------->|
            //                     bar0                         bar0
            //    1(sync)           |         1(sync)            |
            //                      |             |----->|       |
            //                      |                syncUpdBar  |
            //                      |         |<---------|       |
            //                      |         |                  |
            //    2p -------->|     |         2p ------->|       |
            //          syncWaitBar |    =>        syncWaitBar   |
            //    3(sync) <---|     |         3(sync) <--|       |
            //                      |             |              |
            //                      |             |----->|       |
            //                      |                syncUpdBar  |
            //                      |         |<---------|       |
            //                      |         |                  |
            //    4p --->|          |         4p --->|           |
            //         newBar       |              newBar        |
            //    5c <---|----------|         5c <---|-----------|
            //
            _log.nest(2).trace("Add consumers of sync bar {0} for tasks from block {1}", syncTaskUpdBarInd, blockId);
            addConsumers(syncTaskUpdBarInd, getSetWithNoSyncTasks(barProdBlockTasksPair.second));
        }
        // Remove tasks from this barrier which do not belong to this block
        for (const auto& barConsBlockTasksPair : barConsBlockTasksMap) {
            if (barConsBlockTasksPair.first == barBlockId) {
                continue;
            }
            //    0p ------------------->|        0p ----------------->|
            //                          bar0                          bar0
            //    1(sync)                |        1(sync)
            //                           |
            //    2p -------->|          |        2p ------->|
            //             syncWaitBar   |   =>          syncWaitBar
            //    3(sync) <---|          |        3(sync) <--|
            //       |                   |           |
            //       |------->|          |           |------>|
            //             syncUpdBar    |               syncUpdBar
            //    |<----------|          |        |<---------|
            //    |                      |        |
            //    4p -------->|          |        4p -------->|
            //             newBar        |                 newBar
            //    5c <--------|<---------|        5c <--------|
            //
            _log.nest().trace("Remove consumers of bar {0} from block {1}", barInd, barConsBlockTasksPair.first);
            removeConsumers(barInd, barConsBlockTasksPair.second);
        }

        size_t barSyncTaskInd = barBlockId * blockSize + blockSize - 1;

        // Sync task should be a consumer of this barrier, but only when
        // it was not the earliest producer
        if (*barProdBlockTasksMap[barBlockId].begin() != barSyncTaskInd) {
            // Remove it from producers just in case it was there
            removeProducer(barInd, barSyncTaskInd);
            //    0p ------------------>|     0p ------------------>|
            //                        bar0                        bar0
            //    1(sync)                     1(sync) <-------------|
            //
            //    2p -------->|               2p ------->|
            //            syncWaitBar                 syncWaitBar
            //    3(sync) <---|               3(sync) <--|
            //      |                     =>    |
            //      |-------->|                 |------->|
            //            syncUpdBar                  syncUpdBar
            //    |<----------|               |<---------|
            //    |                           |
            //    4p -------->|               4p ------->|
            //             newBar                     newBar
            //    5c <--------|               5c <-------|
            //
            addConsumer(barInd, barSyncTaskInd);
        }

        // Link tasks from other blocks from this barrier
        // as consumers of sync task barrier
        for (const auto& barConsBlockTasksPair : barConsBlockTasksMap) {
            if (barConsBlockTasksPair.first == barBlockId) {
                continue;
            }
            size_t syncTaskInd = barConsBlockTasksPair.first * blockSize - 1;
            size_t syncTaskUpdBarInd = syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd].second;

            //    0p ------------------->|     0p ------------------>|
            //                         bar0                        bar0
            //    1(sync) <--------------|     1(sync) <-------------|
            //      |                            |
            //      |-->syncUpdBar               |-->syncUpdBar
            //    |<---------|                 |<---------|
            //    |                            |
            //    2p -------->|                2p ------->|
            //           syncWaitBar                 syncWaitBar
            //    3(sync) <---|                3(sync) <--|
            //      |                     =>   |
            //      |-------->|                |-------------->|
            //           syncUpdBar                        syncUpdBar
            //    |<----------|                |<--------------|
            //    |                            |               |
            //    4p -------->|                4p ------->|    |
            //             newBar                      newBar  |
            //    5c <--------|                5c <-------|<---| (in this particular example this dep is redundant
            //                                                    but it will be optimized by later passes)
            //
            _log.nest().trace("Add consumers of sync bar {0} from block {1}", syncTaskUpdBarInd,
                              barConsBlockTasksPair.first);
            addConsumers(syncTaskUpdBarInd, getSetWithNoSyncTasks(barConsBlockTasksPair.second));
        }
    }

    // STEP 4
    // Add explicit dependency between sync tasks
    // If it is not needed because it was already represented implicitly by other
    // dependencies it will be removed by follow-up passes
    for (size_t blockInd = 2; blockInd < numOfBlocks; blockInd++) {
        size_t syncTaskInd = (blockInd - 1) * blockSize - 1;
        size_t nextSyncTaskInd = blockInd * blockSize - 1;

        // Create new connection
        // syncTaskInd -> syncTaskUpdateBarInd -> nextSyncTaskInd
        auto syncTaskUpdateBarInd = syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd].second;

        _log.trace("New dep: sync task '{0}' -> next sync task '{1}' using barrier '{2}'", syncTaskInd, nextSyncTaskInd,
                   syncTaskUpdateBarInd);
        addConsumer(syncTaskUpdateBarInd, nextSyncTaskInd);
    }

    // STEP 5
    // Scan all barriers and perform check and sanitization needed after
    // modifications done in STEP 3
    for (size_t barInd = 0; barInd < barriersCount; ++barInd) {
        // If after modifying control graph there is a barrier with no consumers then it
        // is redundant and related producers should be linked to next sync task
        if (_barrierConsumerMap[barInd].empty() && !_barrierProducerMap[barInd].empty()) {
            _log.trace("Empty consumers for bar - {0}", barInd);
            for (auto taskInd : _barrierProducerMap[barInd]) {
                size_t taskBlockInd = taskInd / blockSize;
                size_t syncTaskInd = taskBlockInd * blockSize + blockSize - 1;
                if (syncTaskInd >= tasksSize) {
                    continue;
                }

                if (taskInd != syncTaskInd) {
                    // Create new connection
                    // taskInd -> syncTaskWaitBarInd -> syncTaskInd
                    auto syncTaskWaitBarInd = syncTaskWaitAndUpdateBarrierIndMap[syncTaskInd].first;
                    addProducer(syncTaskWaitBarInd, taskInd);
                }
            }

            resetBarrier(barInd);
        }

        // Check if after modifying control graph there are no barriers
        // with empty producers and non empty consumers
        if (_barrierProducerMap[barInd].empty()) {
            VPUX_THROW_UNLESS(_barrierConsumerMap[barInd].empty(), "Barrier {0} has no producers but has consumers",
                              barInd);
        }
    }

    // STEP 6
    // Remove not used barriers created for sync tasks
    for (const auto& syncTaskWitAndUpdateBarrierInd : syncTaskWaitAndUpdateBarrierIndMap) {
        const auto waitBarInd = syncTaskWitAndUpdateBarrierInd.second.first;
        const auto updateBarInd = syncTaskWitAndUpdateBarrierInd.second.first;
        if (_barrierProducerMap[waitBarInd].empty()) {
            resetBarrier(waitBarInd);
        }
        if (_barrierConsumerMap[updateBarInd].empty()) {
            resetBarrier(updateBarInd);
        }
    }

    // STEP 7
    // Put dedicated attribute on sync tasks for follow-up passes so that
    // they are aware of the split and can take advantage of it (e.g. perform optimizeBarriers in parts)
    for (size_t blockInd = 1; blockInd < numOfBlocks; blockInd++) {
        size_t syncTaskInd = blockInd * blockSize - 1;
        auto syncTaskOp = getTaskOpAtIndex(syncTaskInd);
        syncTaskOp->setAttr(_syncTaskAttrName, mlir::UnitAttr::get(syncTaskOp->getContext()));
    }
}

bool vpux::BarrierInfo::verifyControlGraphSplit() {
    if (_syncTasksIds.empty()) {
        return true;
    }

    size_t firstBlockSize = _syncTasksIds[0] + 1;
    for (size_t i = 1; i < _syncTasksIds.size(); i++) {
        auto nextBlockSize = _syncTasksIds[i] - _syncTasksIds[i - 1];
        if (nextBlockSize != firstBlockSize) {
            _log.trace("Inferred first block size is '{0}', but block '{1}' has {2} tasks", firstBlockSize, i,
                       nextBlockSize);
            VPUX_THROW_WHEN(nextBlockSize == 0, "Detected 0-sized task block");
        }
    }

    auto getSyncTaskInd = [&](size_t blockInd) {
        return blockInd == _syncTasksIds.size() ? getNumOfTasks() - 1 : _syncTasksIds[blockInd];
    };

    for (size_t taskInd = 0; taskInd < getNumOfTasks(); taskInd++) {
        size_t taskBlockInd = getControlGraphBlockIndex(taskInd);
        size_t syncTaskInd = getSyncTaskInd(taskBlockInd);

        for (const auto& updateBarrierInd : _taskUpdateBarriers[taskInd]) {
            auto barrierConsumers = _barrierConsumerMap[updateBarrierInd];

            for (const auto& childTaskInd : barrierConsumers) {
                size_t childBlockInd = getControlGraphBlockIndex(childTaskInd);

                if (taskInd == syncTaskInd && taskBlockInd + 1 != childBlockInd) {
                    _log.debug("Out of block dep for sync task: {0} (block {1}) -> bar {2} -> {3} (block{4}) ", taskInd,
                               taskBlockInd, updateBarrierInd, childTaskInd, childBlockInd);
                    return false;
                } else if (taskInd != syncTaskInd && taskBlockInd != childBlockInd) {
                    _log.debug("Out of block dep for task: {0} (block {1}) -> bar {2} -> {3} (block{4}) ", taskInd,
                               taskBlockInd, updateBarrierInd, childTaskInd, childBlockInd);
                    return false;
                }
            }
        }
    }

    // Verify if tasks from different queues correctly pass through synchronization points (sync-tasks), eg.
    // queue1:  t1       t3 --|                      |-- t8 ...
    // queue2:  t2 -     |  - bW - t5 (sync-task) - bU - t7 - b ...
    // queue3:  t4 - b - |                                    |-- t6 ...
    //          ----------------------------------  --------------------
    //                       block N                    block N + 1
    //
    // - t3 must (directly or indirectly) update sync-task wait barrier bW
    // - t8 must (directly or indirectly) depend on sync-task update barrier bU
    // - t2 may update bW and t7 may wait for bU
    // - t6 must (directly or indirectly) depend on sync-task update barrier bU
    // - t2 and t5 must belong to the same block as t2 precedes sync-task t5 on single queue
    // - t4 does not need to directly update bW, because it updates a barrier that is consumed by t3 which updates bW,
    // hence there is guarantee that t4 will not execute in parallel with sync-task t5

    std::map<VPURT::TaskQueueType, llvm::BitVector> taskQueueTypeMapOrig;
    if (_func != nullptr) {
        // In units tests, _taskQueueTypeMap is not built but provided by hand. Building it here without IR is not
        // possible either due to functions that operate on ModuleOp or TaskOps which typically are not provided in unit
        // tests.
        taskQueueTypeMapOrig = _taskQueueTypeMap;
        clearTaskQueueTypeMap();
        buildTaskQueueTypeMap();
    }
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");

    size_t taskControlMapBlockInd = std::numeric_limits<size_t>::max();
    SmallVector<llvm::BitVector> taskControlMap;
    size_t taskControlMapOffset = 0;

    size_t nextTaskControlMapBlockInd = std::numeric_limits<size_t>::max();
    SmallVector<llvm::BitVector> nextTaskControlMap;
    size_t nextTaskControlMapOffset = 0;

    auto updateControlMap = [&](size_t blockInd, SmallVector<llvm::BitVector>& controlMap, size_t& controlMapOffset,
                                size_t& controlMapBlockInd) {
        if (controlMapBlockInd != blockInd) {
            if (taskControlMapBlockInd == blockInd) {
                controlMap = taskControlMap;
                controlMapOffset = taskControlMapOffset;
            } else if (nextTaskControlMapBlockInd == blockInd) {
                controlMap = nextTaskControlMap;
                controlMapOffset = nextTaskControlMapOffset;
            } else {
                std::tie(controlMap, controlMapOffset) =
                        buildTaskControlMap(blockInd, /* considerTaskFifoDependency */ true);
            }
            controlMapBlockInd = blockInd;
        }
    };

    for (size_t taskInd = 0; taskInd < getNumOfTasks(); ++taskInd) {
        if (isSyncPoint(taskInd)) {
            // there's no need to check barrier dependencies for tasks preceding and following the sync-task on the
            // sync-task's queue
            continue;
        }

        size_t taskBlockInd = getControlGraphBlockIndex(taskInd);
        auto taskSyncPoint = getControlGraphSyncPoint(taskInd);
        auto taskQueueType = getTaskQueueType(taskInd);
        auto nextTask = (taskQueueType.type != config::ExecutorKind::UNKNOWN)
                                ? getNextTaskOnQueue(taskInd, taskQueueType)
                                : std::nullopt;

        if (nextTask.has_value()) {
            size_t nextTaskBlockInd = getControlGraphBlockIndex(nextTask.value());

            if (!taskSyncPoint.has_value() || taskBlockInd == nextTaskBlockInd) {
                // No need to check dependencies to sync-tasks via update barriers for tasks from the last block (i.e.
                // when taskSyncPoint has no value), or when tasks belong to the same block
                // Separate checks for taskInd and nextTask to avoid frequent updates of taskControlMap and path checks
                // for every task.
                continue;
            }

            VPUX_THROW_WHEN(nextTask.value() == taskSyncPoint.value(),
                            "A pair of consecutive tasks on a queue, where the second one is a sync-point for the "
                            "first task, must belong to the same task block");

            // For consecutive tasks belonging to different task blocks, check if barrier exists between the last
            // task in the block and the block sync-task.
            updateControlMap(taskBlockInd, taskControlMap, taskControlMapOffset, taskControlMapBlockInd);
            if (!controlPathExistsBetweenTasksInSameBlock(taskControlMap, taskInd - taskControlMapOffset,
                                                          taskSyncPoint.value() - taskControlMapOffset, false)) {
                auto taskQueueType = getTaskQueueType(taskInd);
                VPUX_THROW_WHEN(taskQueueType.type == config::ExecutorKind::UNKNOWN,
                                "Task queue map not initialized for executor of task {0}", taskInd);
                _log.debug("Last (non-sync-point) {0}.{1} task ({2}) in block {3} does not have a barrier "
                           "dependence to block {4} sync-point {5}",
                           taskQueueType.type, taskQueueType.id, taskInd, taskBlockInd, taskBlockInd,
                           taskSyncPoint.value());
                return false;
            }

            // For consecutive tasks belonging to different task blocks, check if barrier exists between the first
            // task in the following block and sync-task from the preceding block.
            auto prevBlockSyncPoint = getControlGraphSyncPointForBlock(nextTaskBlockInd - 1);
            VPUX_THROW_UNLESS(prevBlockSyncPoint.has_value(),
                              "Could not find preceding sync-task for task ({0}) in block {1}", nextTask.value(),
                              nextTaskBlockInd - 1);

            updateControlMap(nextTaskBlockInd, nextTaskControlMap, nextTaskControlMapOffset,
                             nextTaskControlMapBlockInd);
            if (!controlPathExistsBetweenTasksInSameBlock(nextTaskControlMap,
                                                          prevBlockSyncPoint.value() - nextTaskControlMapOffset,
                                                          nextTask.value() - nextTaskControlMapOffset, false)) {
                auto taskQueueType = getTaskQueueType(nextTask.value());
                VPUX_THROW_WHEN(taskQueueType.type == config::ExecutorKind::UNKNOWN,
                                "Task queue map not initialized for executor of task {0}", nextTask.value());
                _log.debug("First {0}.{1} task ({2}) in block {3} does not have a barrier dependence to the "
                           "previous block sync-task {4}",
                           taskQueueType.type, taskQueueType.id, nextTask.value(), nextTaskBlockInd,
                           prevBlockSyncPoint.value());
                return false;
            }
        } else {
            // taskInd is the last task on the queue
            if (taskSyncPoint.has_value()) {
                // The task does not belong to the last block and a path must exist from the task to sync-task from
                // the same block
                updateControlMap(taskBlockInd, taskControlMap, taskControlMapOffset, taskControlMapBlockInd);
                if (!controlPathExistsBetweenTasksInSameBlock(taskControlMap, taskInd - taskControlMapOffset,
                                                              taskSyncPoint.value() - taskControlMapOffset, false)) {
                    auto taskQueueType = getTaskQueueType(taskInd);
                    VPUX_THROW_WHEN(taskQueueType.type == config::ExecutorKind::UNKNOWN,
                                    "Task queue map not initialized for executor of task {0}", taskInd);
                    _log.debug("Last {0}.{1} task ({2}) in block {3} does not have a barrier dependence to the "
                               "its block sync-task {4}",
                               taskQueueType.type, taskQueueType.id, taskInd, taskBlockInd, taskSyncPoint.value());
                    return false;
                }
            }
            if (taskBlockInd > 0) {
                // The task does not belong to the first block and it must depend on the previous block sync-task update
                // barrier
                auto prevBlockSyncPoint = getControlGraphSyncPointForBlock(taskBlockInd - 1);
                VPUX_THROW_UNLESS(prevBlockSyncPoint.has_value(),
                                  "Could not find preceding sync-task for task ({0}) in block {1}", taskInd,
                                  taskBlockInd - 1);

                updateControlMap(taskBlockInd, taskControlMap, taskControlMapOffset, taskControlMapBlockInd);
                if (!controlPathExistsBetweenTasksInSameBlock(taskControlMap,
                                                              prevBlockSyncPoint.value() - taskControlMapOffset,
                                                              taskInd - taskControlMapOffset, false)) {
                    auto taskQueueType = getTaskQueueType(taskInd);
                    VPUX_THROW_WHEN(taskQueueType.type == config::ExecutorKind::UNKNOWN,
                                    "Task queue map not initialized for executor of task {0}", taskInd);
                    _log.debug("First {0}.{1} task ({2}) in block {3} does not have a barrier dependence to the "
                               "previous block sync-task {4}",
                               taskQueueType.type, taskQueueType.id, taskInd, taskBlockInd, prevBlockSyncPoint.value());
                    return false;
                }
            }
        }
    }

    // restore original state of _taskQueueTypeMap to keep the state unchanged
    if (_func != nullptr) {
        _taskQueueTypeMap = std::move(taskQueueTypeMapOrig);
    }
    return true;
}

bool BarrierInfo::verifyBarriersUsersCount(size_t maxUsersCount) {
    auto checkBarrierUsersCount = [&](const auto& barrierMap, bool producers) {
        for (auto barIdx : irange(getNumOfBarrierOps())) {
            if (barrierMap[barIdx].size() > maxUsersCount) {
                _log.warning("Barrier '{0}' has {1} {2}. The maximal number of {2} is expected to be: {3}", barIdx,
                             barrierMap[barIdx].size(), producers ? "producers" : "consumers", maxUsersCount);
                return false;
            }
        }
        return true;
    };

    return checkBarrierUsersCount(_barrierProducerMap, true) && checkBarrierUsersCount(_barrierConsumerMap, false);
}

bool BarrierInfo::verifyBarriersForTaskDescriptorFetch(const ExecutionGroupListMap& executionGroups, bool wlmEnabled) {
    clearTaskQueueTypeMap();
    buildTaskQueueTypeMap();

    for (size_t blockIdx = 0; blockIdx < getControlGraphBlockCount(); ++blockIdx) {
        auto taskControlMapAndOffset = buildTaskControlMap(blockIdx, /* considerTaskFifoDependency */ true);
        std::optional<size_t> blockIdxOfTaskControlMap = blockIdx;
        std::optional<size_t> legalizationBarrierIdx = std::nullopt;
        for (auto [queue, fifoExecGroups] : executionGroups) {
            int numberOfGroups = static_cast<int>(fifoExecGroups.size());
            for (int grpIdx = 0; grpIdx < numberOfGroups - 1; grpIdx++) {
                if (!hasBarrierDependencyRequiredForDescriptorFetch(grpIdx, blockIdx, fifoExecGroups,
                                                                    taskControlMapAndOffset, blockIdxOfTaskControlMap,
                                                                    wlmEnabled, legalizationBarrierIdx)) {
                    auto lastTaskInGroup = *fifoExecGroups[grpIdx].rbegin();
                    _log.error(
                            "Task {0} from block {1} in execution group {2} ({3}.{4}) does not have an update barrier "
                            "required for task descriptor fetch",
                            lastTaskInGroup, blockIdx, grpIdx, queue.type, queue.id);
                    return false;
                }
            }
        }
    }
    return true;
}

std::optional<size_t> vpux::BarrierInfo::getFirstTaskInGroupFromBlock(int grpIdx, size_t blockIdx,
                                                                      const ExecutionGroupList& fifoExecGroups) {
    int numberOfGroups = static_cast<int>(fifoExecGroups.size());
    auto [blockStartInd, blockEndInd] =
            getControlGraphBlockTaskRange(blockIdx, /* blockStartSyncPoint */ false, /* blockEndSyncPoint */ true);

    std::optional<size_t> firstTaskInGroupInBlock = std::nullopt;
    if (grpIdx < numberOfGroups) {
        auto firstTaskInGroup = fifoExecGroups[grpIdx][0];
        if (inRange(blockStartInd, blockEndInd, firstTaskInGroup)) {
            firstTaskInGroupInBlock = firstTaskInGroup;
        }
    }
    return firstTaskInGroupInBlock;
}

bool vpux::BarrierInfo::hasBarrierDependencyRequiredForDescriptorFetch(
        int grpIdx, size_t blockIdx, const ExecutionGroupList& fifoExecGroups,
        std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
        std::optional<size_t>& blockIdxOfTaskControlMap, bool wlmEnabled,
        std::optional<size_t>& legalizationBarrierIdx) {
    int numberOfGroups = static_cast<int>(fifoExecGroups.size());
    VPUX_THROW_UNLESS(grpIdx >= 0 && grpIdx < numberOfGroups, "Invalid group index");
    legalizationBarrierIdx = std::nullopt;

    if (grpIdx + 1 == numberOfGroups) {
        // no barrier is required at the end of last group
        return true;
    }
    auto lastTaskInGroup = *fifoExecGroups[grpIdx].rbegin();
    auto firstTaskInChildGroup = *fifoExecGroups[grpIdx + 1].begin();
    if (getControlGraphBlockIndex(lastTaskInGroup) != blockIdx ||
        getControlGraphBlockIndex(firstTaskInChildGroup) != blockIdx) {
        // barrier would link tasks from different blocks, barrier is not needed
        return true;
    }

    std::optional<size_t> firstTaskInGrandChildGroup = std::nullopt;
    if (grpIdx < numberOfGroups - 2) {
        firstTaskInGrandChildGroup = fifoExecGroups[grpIdx + 2][0];
    }

    // If grand child execution group does not exist, barrier is not needed.
    if (!firstTaskInGrandChildGroup.has_value()) {
        return true;
    }

    // check if barrier between last task in execution group and first task of next (child) execution group is required
    auto endGrpUpdateBarriers = getUpdateBarriers(lastTaskInGroup);
    if (endGrpUpdateBarriers.empty()) {
        // barrier is needed when there isn't any at the end of task group
        return false;
    }

    if (wlmEnabled) {
        // For WLM case, if barrier exists at the end of execution group, new barrier is not required.
        return true;
    }

    // For non-WLM case, additional verifications are required.

    auto hasPath = [&](size_t taskA, size_t taskB) {
        return isDepFromTaskAToTaskB(taskA, taskB, taskControlMapAndOffset, blockIdxOfTaskControlMap);
    };

    // If none of the consumers of the execution group update barriers depend on grand child group,
    // then the new barrier is not needed.
    for (auto endGrpUpdateBarrier : endGrpUpdateBarriers) {
        auto barrierConsumers = getBarrierConsumers(endGrpUpdateBarrier);

        bool hasConsumersDependentOnGrandChildGrp = false;
        for (const auto& consumer : barrierConsumers) {
            if (hasPath(firstTaskInGrandChildGroup.value(), consumer)) {
                hasConsumersDependentOnGrandChildGrp = true;
                break;
            }
        }

        if (!hasConsumersDependentOnGrandChildGrp) {
            legalizationBarrierIdx = endGrpUpdateBarrier;
            return true;
        }
    }

    return false;
}

bool vpux::BarrierInfo::adjustTasksDependenciesToGraphSplitConstraints(const TaskSet& tasks) {
    if (_syncTasksIds.empty()) {
        return false;
    }

    bool dependenciesChanged = false;

    size_t controlMapBlockInd = std::numeric_limits<size_t>::max();
    SmallVector<llvm::BitVector> taskControlMap;
    size_t controlMapOffset = 0;

    auto updateControlMap = [&](size_t taskBlockInd, SmallVector<llvm::BitVector>& taskControlMap,
                                size_t& controlMapOffset) {
        if (controlMapBlockInd != taskBlockInd) {
            std::tie(taskControlMap, controlMapOffset) = buildTaskControlMap(taskBlockInd, false, true);
            controlMapBlockInd = taskBlockInd;
        }
    };

    for (auto taskInd : tasks) {
        size_t taskBlockInd = getControlGraphBlockIndex(taskInd);
        auto taskSyncPoint = getControlGraphSyncPoint(taskInd);

        // check task dependencies through its update barriers
        auto taskUpdateBarriers = getUpdateBarriers(taskInd);
        for (const auto& updateBarrierInd : taskUpdateBarriers) {
            auto barrierBlockIndex = getBarrierBlockIndex(updateBarrierInd);

            if (!taskSyncPoint.has_value()) {
                // the task belongs to last block
                VPUX_THROW_WHEN(barrierBlockIndex != taskBlockInd,
                                "taskInd '{0}' does not have sync point and should not have "
                                "out-of-block consumers",
                                taskInd);
            }

            if (taskBlockInd == barrierBlockIndex) {
                // dependence does not span across blocks
                continue;
            }

            if (isSyncPoint(taskInd) && taskBlockInd + 1 == barrierBlockIndex) {
                // sync-point can update barriers from the following block
                continue;
            }
            _log.trace("Out of block dependence for task: {0} (block {1}) -> bar {2} (block{3}) ", taskInd,
                       taskBlockInd, updateBarrierInd, barrierBlockIndex);
            dependenciesChanged = true;

            /*
              Task (taskInd) to barrier (updateBarrierInd) dependence spans across task blocks.
              Remove out-of-block connection from a producer

              /----------------------------------------------------------------------\
              0 - b0 - 1 - b1 - 2 (sync) - b2 - 3 - b3 - 4 - b4 - 5 (sync) - b5 - 6 - b6 - 7

            */

            // remove out-of-block connection from a producer (0->b6)
            removeProducer(updateBarrierInd, taskInd);

            // Add connection to nearest sync point. This is to assert that the taskInd updates a barrier after it
            // finishes.
            auto syncPointWaitBars = getWaitBarriers(taskSyncPoint.value());
            VPUX_THROW_WHEN(syncPointWaitBars.empty(), "Sync point '{0}' does not have wait barriers.", taskSyncPoint);

            updateControlMap(taskBlockInd, taskControlMap, controlMapOffset);

            if (!isSyncPoint(taskInd) &&
                !controlPathExistsBetweenTasksInSameBlock(taskControlMap, taskInd - controlMapOffset,
                                                          taskSyncPoint.value() - controlMapOffset, false)) {
                // Link out-of-block producer (0) to the nearest sync-point wait barrier (b1).
                // If the task is a sync-point, or a connection already exists, nothing needs to be done.
                for (auto syncPointWaitBar : syncPointWaitBars) {
                    addProducer(syncPointWaitBar, taskInd);
                }
            }
        }

        // check task dependencies through its wait barriers
        auto taskWaitBarriers = getWaitBarriers(taskInd);
        for (const auto& waitBarrierInd : taskWaitBarriers) {
            auto barrierBlockIndex = getBarrierBlockIndex(waitBarrierInd);
            auto barrierLastProducer = getBarrierLatestProducer(waitBarrierInd);

            if (taskBlockInd == barrierBlockIndex) {
                // dependence does not span across blocks
                continue;
            }

            if (isSyncPoint(barrierLastProducer) && taskBlockInd - 1 == barrierBlockIndex) {
                // any task from a given block can wait for update barrier of sync-point from the previous block
                continue;
            }
            _log.trace("Out of block dependence for task: {0} (block {1}) -> bar {2} (block{3}) ", taskInd,
                       taskBlockInd, waitBarrierInd, barrierBlockIndex);

            dependenciesChanged = true;

            /*
              Barrier (waitBarrierInd) to task (taskInd) dependence spans across task blocks.
              Replace out-of-block connection to a consumer

                    /----------------------------------------------------------------------\
              0 - b0 - 1 - b1 - 2 (sync) - b2 - 3 - b3 - 4 - b4 - 5 (sync) - b5 - 6        7
                                                                              \------------/
              with a connection to its preceding sync point, if no such connection exists
            */

            // Remove out-of-block connection to a consumer (b0->7)
            removeConsumer(waitBarrierInd, taskInd);
            // Remove all barrier producers if there are no more consumers for barrier waitBarrierInd.
            // (Task 0 should be removed as producer of b0 if b0 is not consumed by any task).
            if (getBarrierConsumers(waitBarrierInd).empty()) {
                removeProducers(waitBarrierInd, getBarrierProducers(waitBarrierInd));
            }

            // get sync-point preceding current task (5)
            auto prevSyncPoint = getPreviousBlockSyncPoint(taskInd);
            // if prevSyncPoint doesn't exist, task taskInd cannot have out-of-block dependency and we shouldn't be
            // here.
            VPUX_THROW_UNLESS(prevSyncPoint.has_value(), "There's not sync-point preceding task {0}", taskInd);

            updateControlMap(taskBlockInd, taskControlMap, controlMapOffset);

            if (!isSyncPoint(taskInd) &&
                !controlPathExistsBetweenTasksInSameBlock(taskControlMap, prevSyncPoint.value() - controlMapOffset,
                                                          taskInd - controlMapOffset, false)) {
                // Link out-of-block consumer (7) to the preceding sync-point update barrier (b5).
                // If the task is a sync-point, or a connection already exists, nothing needs to be done.
                auto prevSyncPointUpdateBarriers = getUpdateBarriers(prevSyncPoint.value());
                for (const auto& syncPointUpdateBarrierInd : prevSyncPointUpdateBarriers) {
                    addConsumer(syncPointUpdateBarrierInd, taskInd);
                }
            }
        }
    }

    return dependenciesChanged;
}

std::optional<size_t> vpux::BarrierInfo::getPreviousBlockSyncPoint(size_t taskInd) const {
    auto blockInd = getControlGraphBlockIndex(taskInd);
    if (blockInd == 0) {
        return std::nullopt;
    } else if (blockInd == _syncTasksIds.size()) {
        return _syncTasksIds.back();
    }
    return _syncTasksIds[blockInd - 1];
}

std::optional<size_t> vpux::BarrierInfo::getNextBlockSyncPoint(size_t taskInd) const {
    auto blockInd = getControlGraphBlockIndex(taskInd);
    if (blockInd + 1 >= _syncTasksIds.size()) {
        return std::nullopt;
    }
    return _syncTasksIds[blockInd + 1];
}

void vpux::BarrierInfo::splitBarriersWithExceedingVariantCount(size_t availableSlots, size_t maxSlotsSum,
                                                               bool maxSlotsSumLimitEnabled) {
    // split barriers if needed
    splitBarrierProducers(availableSlots, maxSlotsSum, maxSlotsSumLimitEnabled);
    // split barriers if needed
    splitBarrierConsumers(availableSlots, maxSlotsSum, maxSlotsSumLimitEnabled);
}

// split barriers if producer count > AVAILABLE SLOTS
// will produce ceil(NUM PRODUCERS / AVAILABLE SLOTS) barriers
// Note: if maxSlotsSumLimitEnabled is true, the maxAvailableSlots passed will be half of MAX_VARIANT_SUM not
// MAX_VARIANT_COUNT
/*
    x1  x2  ... xn             x1  ... x256   x257 ... xn
     \  \   /   /               \   |   /        \  |  /
         Bar0         -->          Bar0            Bar1
          |                            \          /
          u0                                u0
*/
void vpux::BarrierInfo::splitBarrierProducers(VPURT::DeclareVirtualBarrierOp barrierOp, size_t availableSlots) {
    // producers will be divided into batches
    const auto& barrierProducers = getBarrierProducers(barrierOp);
    auto barrierProducersCount = barrierProducers.size();
    _log.trace("Got barrier: '{0}' at '{1}' with barrier producer count '{2}'", barrierOp->getName(),
               barrierOp->getLoc(), barrierProducersCount);

    // consumers remain the same for all produced batched barriers
    const auto barrierConsumers = getBarrierConsumers(barrierOp);
    // create batches for producers
    auto producerBatches = createLegalVariantBatches(barrierProducers, availableSlots);

    mlir::OpBuilder builder(barrierOp);

    // create barriers based on batches
    for (const auto& producerBatch : producerBatches) {
        builder.setInsertionPoint(barrierOp);
        auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(barrierOp->getLoc());
        _log.trace("Created new barrier '{0}'", newBarrier);
        addNewBarrier(newBarrier);

        for (const auto& producer : producerBatch) {
            _log.nest().trace("Add producer '{0}' to new barrier, whose slots are {1}", producer,
                              getNumOfSlotsUsed(getTaskOpAtIndex(producer)));
            addProducer(newBarrier, producer);
        }

        for (const auto& consumer : barrierConsumers) {
            _log.nest().trace("Add consumer '{0}' to new barrier, whose slots are {1}", consumer,
                              getNumOfSlotsUsed(getTaskOpAtIndex(consumer)));
            addConsumer(newBarrier, consumer);
        }
    }

    resetBarrier(barrierOp);
    _log.trace("Barrier successfully replaced with batch of barriers");
}

// split barriers if consumer count > AVAILABLE SLOTS
// will produce ceil(NUM PRODUCERS / AVAILABLE SLOTS) barriers
// Note: if maxSlotsSumLimitEnabled is true, the maxAvailableSlots passed will be half of MAX_VARIANT_SUM not
// MAX_VARIANT_COUNT
/*
          u0                                u0
          |                            /          \
         Bar0         -->          Bar0            Bar1
     /  /   \   \               /   |   \        /  |  \
    x1  x2  ... xn             x1  ... x256   x257 ... xn
*/
void vpux::BarrierInfo::splitBarrierConsumers(VPURT::DeclareVirtualBarrierOp barrierOp, size_t availableSlots) {
    // consumers will be divided into batches
    const auto& barrierConsumers = getBarrierConsumers(barrierOp);
    auto barrierConsumersCount = barrierConsumers.size();
    _log.trace("Got barrier: '{0}' at '{1}' with barrier consumer count '{2}'", barrierOp->getName(),
               barrierOp->getLoc(), barrierConsumersCount);

    // producers remain the same for all produced batched barriers
    const auto barrierProducers = getBarrierProducers(barrierOp);
    // crete batches for consumers
    auto consumerBatches = createLegalVariantBatches(barrierConsumers, availableSlots);

    mlir::OpBuilder builder(barrierOp);

    // create barriers based on batches
    for (const auto& consumerBatch : consumerBatches) {
        builder.setInsertionPoint(barrierOp);
        auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(barrierOp->getLoc());
        _log.trace("Created new barrier '{0}'", newBarrier);
        addNewBarrier(newBarrier);

        for (const auto& producer : barrierProducers) {
            _log.nest().trace("Add producer '{0}' to new barrier, whose slots are {1}", producer,
                              getNumOfSlotsUsed(getTaskOpAtIndex(producer)));
            addProducer(newBarrier, producer);
        }

        for (const auto& consumer : consumerBatch) {
            _log.nest().trace("Add consumer '{0}' to new barrier, whose slots are {1}", consumer,
                              getNumOfSlotsUsed(getTaskOpAtIndex(consumer)));
            addConsumer(newBarrier, consumer);
        }
    }

    resetBarrier(barrierOp);
    _log.trace("Barrier successfully replaced with batch of barriers");
}

void vpux::BarrierInfo::splitBarrierProducers(size_t availableSlots, size_t maxSlotsSum, bool maxSlotsSumLimitEnabled) {
    _func->walk([&](VPURT::DeclareVirtualBarrierOp barrierOp) {
        auto producerNum = getProducerSlotCount(barrierOp);
        auto consumerNum = getConsumerSlotCount(barrierOp);
        if (maxSlotsSumLimitEnabled && (producerNum + consumerNum <= maxSlotsSum)) {
            return;
        }

        if (producerNum <= availableSlots) {
            // valid producer configuration - barrier is legal
            return;
        }

        splitBarrierProducers(barrierOp, availableSlots);
    });
}

void vpux::BarrierInfo::splitBarrierConsumers(size_t availableSlots, size_t maxSlotsSum, bool maxSlotsSumLimitEnabled) {
    _func->walk([&](VPURT::DeclareVirtualBarrierOp barrierOp) {
        auto producerNum = getProducerSlotCount(barrierOp);
        auto consumerNum = getConsumerSlotCount(barrierOp);
        if (maxSlotsSumLimitEnabled && (producerNum + consumerNum <= maxSlotsSum)) {
            return;
        }

        if (consumerNum <= availableSlots) {
            // valid consumer configuration - barrier is legal
            return;
        }

        splitBarrierConsumers(barrierOp, availableSlots);
    });
}

void vpux::BarrierInfo::removeSyncTaskAttributes() {
    if (_syncTasksIds.empty()) {
        return;
    }

    for (const auto& syncTaskInd : _syncTasksIds) {
        auto syncTaskOp = getTaskOpAtIndex(syncTaskInd);
        VPUX_THROW_UNLESS(syncTaskOp->hasAttr(_syncTaskAttrName), "Expected sync task attribute at task index '{0}'",
                          syncTaskInd);
        syncTaskOp->removeAttr(_syncTaskAttrName);
    }
}

bool vpux::BarrierInfo::isSyncPoint(size_t taskIdx) {
    auto nTasks = getNumOfTasks();
    VPUX_THROW_WHEN(taskIdx >= nTasks, "Requested task index ({0}) is too large for the number of tasks ({1})", taskIdx,
                    nTasks);

    if (_syncTasksIds.empty() || taskIdx == nTasks - 1) {
        // graph split is not done, and last task cannot be sync point
        return false;
    }

    VPUX_THROW_WHEN(_taskToBlockMap.size() != nTasks, "_taskToBlockMap was not initialized (size {0}) for {1} tasks.",
                    _taskToBlockMap.size(), nTasks);
    return _taskToBlockMap[taskIdx + 1] != _taskToBlockMap[taskIdx];
}

size_t vpux::BarrierInfo::getControlGraphBlockCount() const {
    return _syncTasksIds.size() + 1;
}

std::pair<size_t, size_t> vpux::BarrierInfo::getControlGraphBlockTaskRange(size_t blockInd, bool blockStartSyncPoint,
                                                                           bool blockEndSyncPoint) const {
    VPUX_THROW_WHEN(blockInd >= getControlGraphBlockCount(), "Invalid task block index ({0})", blockInd);
    size_t blockStartInd, blockEndInd;
    if (blockInd == 0) {
        blockStartInd = 0;  // for first block the start index of the block is the index of the first task
    } else {
        blockStartInd = blockStartSyncPoint ? _syncTasksIds[blockInd - 1] : _syncTasksIds[blockInd - 1] + 1;
    }

    if (blockInd == getControlGraphBlockCount() - 1) {
        blockEndInd = getNumOfTasks() - 1;  // for last block the end index of the block is the index of the last task
    } else {
        blockEndInd = blockEndSyncPoint ? _syncTasksIds[blockInd] : _syncTasksIds[blockInd] - 1;
    }

    return std::make_pair(blockStartInd, blockEndInd);
}

SmallVector<std::set<size_t>> vpux::BarrierInfo::findParallelTasksWithBarrierDependence(
        const BarrierInfo::TaskSet& tasks, size_t maxCount) {
    SmallVector<std::set<size_t>> parallelTasks(tasks.size());

    size_t controlMapBlockInd = std::numeric_limits<size_t>::max();
    SmallVector<llvm::BitVector> taskControlMap;
    size_t controlMapOffset = 0;

    auto updateControlMap = [&](size_t taskBlockInd, SmallVector<llvm::BitVector>& taskControlMap,
                                size_t& controlMapOffset) {
        if (controlMapBlockInd != taskBlockInd) {
            std::tie(taskControlMap, controlMapOffset) =
                    buildTaskControlMap(taskBlockInd, /* considerTaskFifoDependency */ true);
            controlMapBlockInd = taskBlockInd;
        }
    };

    buildTaskQueueTypeMap();

    // For every task from the provided set, search for parallel tasks located earlier in the IR.
    for (auto [i, taskIdx] : tasks | indexed) {
        auto taskBlockIdx = getControlGraphBlockIndex(taskIdx);
        auto [fromIdx, _] = getControlGraphBlockTaskRange(taskBlockIdx, /* blockStartSyncPoint */ false,
                                                          /* blockEndSyncPoint */ true);
        auto toIdx = taskIdx;

        // build task control map
        updateControlMap(taskBlockIdx, taskControlMap, controlMapOffset);

        // Scan paths in decreasing task order between the current task and preceding tasks up to the preceding
        // sync-point. It is expected that typically parallel tasks will be located in the close vicinity of a given
        // task.
        for (auto idx : irange(fromIdx, toIdx) | reversed) {
            if (!controlPathExistsBetweenTasksInSameBlock(taskControlMap, idx - controlMapOffset,
                                                          taskIdx - controlMapOffset, false)) {
                if (getUpdateBarriers(idx).empty() && getWaitBarriers(idx).empty()) {
                    continue;
                }
                parallelTasks[i].insert(idx);
                if (parallelTasks[i].size() >= maxCount) {
                    break;
                }
            }
        }
    }

    // If not found, search for parallel tasks located latter in the IR relative to the provided task
    for (auto [i, taskIdx] : tasks | indexed) {
        if (parallelTasks[i].size() >= maxCount) {
            continue;
        }
        auto taskBlockIdx = getControlGraphBlockIndex(taskIdx);
        auto [_, toIdx] = getControlGraphBlockTaskRange(taskBlockIdx, /* blockStartSyncPoint */ false,
                                                        /* blockEndSyncPoint */ true);
        auto fromIdx = taskIdx + 1;  // start from index of the following task

        // build task control map
        updateControlMap(taskBlockIdx, taskControlMap, controlMapOffset);

        // Scan paths in increasing task order between the current task and following tasks up to the next
        // sync-point (or the last task) inclusively. It is expected that typically parallel tasks will be located in
        // the close vicinity of a given task.
        for (auto idx : irange(fromIdx, toIdx + 1)) {
            if (!controlPathExistsBetweenTasksInSameBlock(taskControlMap, taskIdx - controlMapOffset,
                                                          idx - controlMapOffset, false)) {
                if (getUpdateBarriers(idx).empty() && getWaitBarriers(idx).empty()) {
                    continue;
                }
                parallelTasks[i].insert(idx);
                if (parallelTasks[i].size() >= maxCount) {
                    break;
                }
            }
        }
    }

    return parallelTasks;
}

//
// optimizeBarriers
//
void vpux::BarrierInfo::optimizeBarriers(bool checkValidSlotCount, bool considerTaskFifoDependency,
                                         mlir::DenseSet<vpux::config::ExecutorKind> executors) {
    if (considerTaskFifoDependency) {
        buildTaskQueueTypeMap();
    }

    // A -> B -> C
    // If B depends on A and C depends on [A, B] ==> we can remove A from C deps list,
    // since it will be implicit dependency taken from B.
    // Note: It also will merge barriers which have the same producers but different consumers

    // Barrier are optimized based on order of task ops

    _log.trace("Optimize barriers");
    _log = _log.nest();
    _log.trace("Total tasks count: {0}, total barrier count: {1}", getNumOfTasks(), _allBarrierOps.size());

    auto numBarriersBeforeOptimization = getNumOfBarrierOps();

    // Perform optimization in tasks blocks matching the distribution of synchronization points.
    // For update barriers and wait barriers we need to include sync points on both ends of the block
    // because tasks from within a given range can have updateBarriers whose consumer
    // is the upper bound sync point, and they can have waitBarriers whose producer is the lower bound sync point.
    for (size_t taskBlockIndex = 0; taskBlockIndex < getControlGraphBlockCount(); ++taskBlockIndex) {
        auto [blockStartInd, blockEndInd] = getControlGraphBlockTaskRange(taskBlockIndex);
        _log.trace("Block {0}, task range [{1}, {2}] ({3} tasks)", taskBlockIndex, blockStartInd, blockEndInd,
                   blockEndInd - blockStartInd + 1);

        if (considerTaskFifoDependency) {
            // TODO: E#126579
            auto newDepsCount = createBarrierDependenciesImpliedByFIFO(taskBlockIndex, executors);
            _log.trace("Created {0} temporary dependencies", newDepsCount);
        }

        // optimize producers
        _log.trace("Optimize producers / update barriers");
        optimizeBarrierProducers(taskBlockIndex);

        // optimize barriers which have the same producers but different consumers
        _log.trace("Optimize barriers / same producers, different consumers");
        optimizeBarriersWithSameProducers(taskBlockIndex, checkValidSlotCount);

        // optimize consumers
        _log.trace("Optimize consumers / wait barriers");
        optimizeBarrierConsumers(taskBlockIndex);

        auto numBarriersAfterOptimization = getNumOfBarrierOps();
        VPUX_THROW_WHEN(numBarriersAfterOptimization != numBarriersBeforeOptimization,
                        "Number of barriers changed during optimization.");

        if (considerTaskFifoDependency) {
            auto removedDepsCount = removeBarrierDependenciesImpliedByFIFO();
            _log.trace("Removed {0} temporary barriers", removedDepsCount);

            numBarriersAfterOptimization = getNumOfBarrierOps();
            VPUX_THROW_WHEN(numBarriersAfterOptimization != numBarriersBeforeOptimization,
                            "Incorrect number of barriers after removing barrier representation of FIFO dependencies.");
        }
    }

    _log.trace("Total tasks count: {0}, total barrier count: {1}", getNumOfTasks(), _allBarrierOps.size());
    _log = _log.unnest();
}

void vpux::BarrierInfo::optimizeBarrierProducers(size_t blockIdx) {
    // TODO: E#79318 optimize loops
    auto [blockStartInd, blockEndInd] =
            getControlGraphBlockTaskRange(blockIdx, /* blockStartSyncPoint */ true, /* blockEndSyncPoint */ true);
    const auto barriersRangeVec = getBarriersForTaskBlock(blockIdx, /* blockStartSyncPoint */ true,
                                                          /* blockEndSyncPoint */ true, /* updateBarriers */ true);
    if (barriersRangeVec.empty()) {
        return;
    }
    size_t barrierOffset = barriersRangeVec[0];

    blockEndInd = std::min(blockEndInd, getNumOfTasks() - 1);
    VPUX_THROW_WHEN(blockStartInd > blockEndInd, "Lower bound of task optimization range is greater than upper bound");

    auto blockSize = blockEndInd - blockStartInd + 1;
    auto barrierSize = barriersRangeVec.back() - barriersRangeVec.front() + 1;

    SmallVector<llvm::BitVector> updateBarriers;
    resizeBitMap(updateBarriers, blockSize, checked_cast<uint32_t>(barrierSize));

    _log.nest().trace("Collect redundant dependencies");
    for (auto taskInd = blockEndInd + 1; taskInd-- > blockStartInd;) {
        setBarrierMask(updateBarriers[taskInd - blockStartInd], _taskUpdateBarriers[taskInd], barrierOffset);
        if (taskInd == blockEndInd && isSyncPoint(taskInd)) {
            // ignore updating updateBarriers for upper bound sync point because it only has barriers whose consumers
            // are tasks from outside of the current range.
            continue;
        }
        for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
            for (auto childTaskInd : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                if (inRange(blockStartInd, blockEndInd, childTaskInd)) {
                    updateBarriers[taskInd - blockStartInd] |=
                            updateBarriers[static_cast<size_t>(childTaskInd - blockStartInd)];
                } else {
                    VPUX_THROW("Task {0} has update barriers with a consumer ({1}) from outside of the current "
                               "range [{2}, {3}].",
                               taskInd, childTaskInd, blockStartInd, blockEndInd);
                }
            }
        }
    }

    _log.nest().trace("Remove redundant dependencies");
    for (size_t taskInd = static_cast<unsigned>(blockStartInd); taskInd <= static_cast<unsigned>(blockEndInd);
         ++taskInd) {
        for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
            for (auto childTaskInd : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                if (inRange(blockStartInd, blockEndInd, childTaskInd)) {
                    updateBarriers[taskInd - blockStartInd].reset(
                            updateBarriers[static_cast<size_t>(childTaskInd - blockStartInd)]);
                }
            }
        }
        TaskSet targetUpdateBarriers;
        for (auto bar : updateBarriers[taskInd - blockStartInd].set_bits()) {
            targetUpdateBarriers.insert(bar + barrierOffset);
        }
        setUpdateBarriers(taskInd, targetUpdateBarriers);
    }
}

void vpux::BarrierInfo::optimizeBarrierConsumers(size_t blockIdx) {
    // TODO: E#79318 optimize loops
    auto [blockStartInd, blockEndInd] =
            getControlGraphBlockTaskRange(blockIdx, /* blockStartSyncPoint */ true, /* blockEndSyncPoint */ true);
    const auto barriersRangeVec = getBarriersForTaskBlock(blockIdx, /* blockStartSyncPoint  */ true,
                                                          /* blockEndSyncPoint */ true, /* updateBarriers */ false);
    if (barriersRangeVec.empty()) {
        return;
    }
    size_t barrierOffset = barriersRangeVec[0];

    blockEndInd = std::min(blockEndInd, getNumOfTasks() - 1);
    VPUX_THROW_WHEN(blockStartInd > blockEndInd, "Lower bound of task optimization range is greater than upper bound");

    auto blockSize = blockEndInd - blockStartInd + 1;
    auto barrierSize = barriersRangeVec.back() - barriersRangeVec.front() + 1;
    // optimize consumers

    SmallVector<llvm::BitVector> waitBarriers;
    resizeBitMap(waitBarriers, blockSize, checked_cast<uint32_t>(barrierSize));

    _log.nest().trace("Collect redundant dependencies");
    for (auto taskInd = blockStartInd; taskInd <= blockEndInd; ++taskInd) {
        setBarrierMask(waitBarriers[taskInd - blockStartInd], _taskWaitBarriers[taskInd], barrierOffset);
        if (taskInd == blockStartInd && isSyncPoint(taskInd)) {
            // ignore updating waitBarriers for lower bound sync point because it only has barriers whose producers
            // are tasks from outside of the current range.
            continue;
        }
        for (auto waitBarrierInd : _taskWaitBarriers[taskInd]) {
            for (auto parentTaskInd : _barrierProducerMap[waitBarrierInd]) {
                if (inRange(blockStartInd, blockEndInd, parentTaskInd)) {
                    waitBarriers[taskInd - blockStartInd] |=
                            waitBarriers[static_cast<size_t>(parentTaskInd - blockStartInd)];
                } else {
                    VPUX_THROW("Task {0} has wait barriers with a producer ({1}) from outside of the current range "
                               "[{2}, {3}].",
                               taskInd, parentTaskInd, blockStartInd, blockEndInd);
                }
            }
        }
    }

    for (auto taskInd = blockEndInd + 1; taskInd-- > blockStartInd;) {
        for (auto waitBarrierInd : _taskWaitBarriers[taskInd]) {
            for (auto parentTaskInd : _barrierProducerMap[waitBarrierInd]) {
                if (inRange(blockStartInd, blockEndInd, parentTaskInd)) {
                    waitBarriers[taskInd - blockStartInd].reset(
                            waitBarriers[static_cast<size_t>(parentTaskInd - blockStartInd)]);
                }
            }
        }

        TaskSet targetWaitBarriers;
        for (auto bar : waitBarriers[taskInd - blockStartInd].set_bits()) {
            targetWaitBarriers.insert(bar + barrierOffset);
        }
        setWaitBarriers(taskInd, targetWaitBarriers);
    }
}

void vpux::BarrierInfo::optimizeBarriersWithSameProducers(size_t blockIdx, bool checkValidSlotCount) {
    // Collect a vector of barrier indexes for current task range.
    // If a new barrier is created between subsequent calls to optimizeBarrier (eg. during looped linearization),
    // iterating over barriersRangeVec will not imply iterating over largely overlapping barrier index ranges
    // when moving from one task block to another.
    const auto barriersRangeVec =
            getBarriersForTaskBlock(blockIdx, /* blockStartSyncPoint */ true, /* blockEndSyncPoint */ false);

    if (barriersRangeVec.empty()) {
        return;
    }
    _log.nest().trace("Barrier range [{0}, {1}]", barriersRangeVec.front(), barriersRangeVec.back());
    size_t maxBarrierIndex = _allBarrierOps.size() + _fifoDependencies.size() - 1;
    VPUX_THROW_WHEN(barriersRangeVec.front() > maxBarrierIndex,
                    "Invalid barrier index for lower limit of barrier range optimization");
    VPUX_THROW_WHEN(barriersRangeVec.back() > maxBarrierIndex,
                    "Invalid barrier index for upper limit of barrier range optimization");

    const auto maxVariantCount = getBarrierMaxVariantSum();

    auto addFunc = [&](size_t sum, size_t taskInd) {
        return sum + getNumOfSlotsUsedByTask(getTaskOpAtIndex(taskInd));
    };

    // Check the slot count is valid or not if the two barriers are merged
    auto legalMergeSlotCount = [&](const size_t& bar1, const size_t& bar2) {
        // Currently used by ReduceBArrierDependencies, this is possible to be done here because
        // SplitExceedingVariantCountBarrier will legalize this in barrierLegalizationPipeline
        if (!checkValidSlotCount) {
            return true;
        }
        size_t slotCount = std::accumulate(_barrierProducerMap[bar1].begin(), _barrierProducerMap[bar1].end(),
                                           static_cast<size_t>(0), addFunc);
        BarrierInfo::TaskSet consumers = _barrierConsumerMap[bar1];
        // Use set operation to remove duplicated consumers
        llvm::set_union(consumers, _barrierConsumerMap[bar2]);
        slotCount = std::accumulate(consumers.begin(), consumers.end(), slotCount, addFunc);
        return slotCount <= maxVariantCount;
    };

    // copy barrier map onto vector for quicker comparisons between different producer sets
    SmallVector<SmallVector<size_t>> blockBarrierProducerMap(barriersRangeVec.size());
    for (size_t ind1 = 0; ind1 < barriersRangeVec.size(); ++ind1) {
        size_t barInd = barriersRangeVec[ind1];
        std::copy(_barrierProducerMap[barInd].begin(), _barrierProducerMap[barInd].end(),
                  std::back_inserter(blockBarrierProducerMap[ind1]));
    }
    _log = _log.nest();

    for (size_t ind1 = 0; ind1 < barriersRangeVec.size(); ++ind1) {
        size_t barInd = barriersRangeVec[ind1];
        for (size_t ind2 = ind1 + 1; ind2 < barriersRangeVec.size(); ++ind2) {
            size_t childBarInd = barriersRangeVec[ind2];
            if (blockBarrierProducerMap[ind1] == blockBarrierProducerMap[ind2]) {
                _log.trace("Same producers for barId '{0}' '{1}'", barInd, childBarInd);
                if (legalMergeSlotCount(barInd, childBarInd)) {
                    for (auto consumerInd : _barrierConsumerMap[childBarInd]) {
                        // move all consumers to one barrier
                        addConsumer(barInd, static_cast<size_t>(consumerInd));
                    }
                    _log.trace("New consumers number - {0}", _barrierConsumerMap[barInd].size());
                    resetBarrier(childBarInd);
                    blockBarrierProducerMap[ind2].clear();
                }
            }
        }
    }
    _log = _log.unnest();
}

bool vpux::BarrierInfo::canMergeBarriersForTasks(const BarrierInfo::TaskSet& producers, size_t availableSlots) {
    size_t producerSlotCount = 0;
    for (const auto& producerInd : producers) {
        producerSlotCount += getNumOfSlotsUsedByTask(getTaskOpAtIndex(producerInd));

        if (producerSlotCount > availableSlots) {
            // exceeding producer slot count
            return false;
        }
    }

    return true;
}

void vpux::BarrierInfo::enableUnevenVariantSplit() {
    _enableUnevenVariantSplit = true;
}

// modify barriers for task such that it is driven by a single wait barrier by merging parallel wait barriers
// or linearizing the producers of parallel wait barriers
bool vpux::BarrierInfo::eliminateParallelWaitBarriers(size_t taskInd, size_t availableSlots,
                                                      bool considerTaskExecutorType) {
    const auto waitBarriers = getWaitBarriers(taskInd);  // parallel wait barriers
    auto log = _log.nest();

    auto getBarriersProducerTasks = [&](const BarrierInfo::TaskSet& waitBarriers) {
        BarrierInfo::TaskSet newBarrierProducers;
        for (auto& waitBarrierInd : waitBarriers) {
            // merge all producers
            const auto& barrierProducers = getBarrierProducers(waitBarrierInd);
            llvm::set_union(newBarrierProducers, barrierProducers);
        }

        return newBarrierProducers;
    };

    auto getBarriersConsumerTasks = [&](const BarrierInfo::TaskSet& parallelWaitBarriers) {
        BarrierInfo::TaskSet parallelConsumers;
        for (auto& waitBarrier : parallelWaitBarriers) {
            llvm::set_union(parallelConsumers, getBarrierConsumers(waitBarrier));
        }
        return parallelConsumers;
    };

    BarrierInfo::TaskSet parallelConsumers = getBarriersConsumerTasks(waitBarriers);

    auto getBarriersConsumerTasksWithWaitBarriers = [&](const BarrierInfo::TaskSet& waitBarriers,
                                                        size_t availableSlotsForConsumer) {
        // select consumers from within parallelConsumers with given waitBarriers
        BarrierInfo::TaskSet parallelConsumersSameWaitBarriers;
        size_t parallelConsumerSlotCount = 0;

        for (auto& consumer : parallelConsumers) {
            if (waitBarriers != getWaitBarriers(consumer)) {
                continue;
            }

            parallelConsumerSlotCount += getNumOfSlotsUsedByTask(getTaskOpAtIndex(consumer));
            if (parallelConsumerSlotCount > availableSlotsForConsumer) {
                break;
            }

            parallelConsumersSameWaitBarriers.insert(consumer);
        }

        return parallelConsumersSameWaitBarriers;
    };

    auto unlinkTaskFromParallelBarriers = [&](const BarrierInfo::TaskSet& tasks,
                                              const BarrierInfo::TaskSet& waitBarriers) {
        // remove consumer task from parallel barriers
        for (auto& waitBarrierInd : waitBarriers) {
            const auto barrierConsumers = getBarrierConsumers(waitBarrierInd);
            const auto waitBarrier = getBarrierOpAtIndex(waitBarrierInd);

            // remove link from parallel barriers
            if (barrierConsumers.size() == tasks.size()) {
                // if only consumer, barrier can be reset
                resetBarrier(waitBarrier);
            } else {
                for (auto& taskInd : tasks) {
                    removeConsumer(waitBarrier, taskInd);
                }
            }
        }
    };

    log.trace("Got '{0}' parallel wait barriers for task '{1}'", waitBarriers.size(), taskInd);

    auto availableSlotsForConsumer = availableSlots;
    const auto totalSlots = 2 * availableSlots;  // by default the slots are evenly divided for producer and consumer
    auto parallelProducers = getBarriersProducerTasks(waitBarriers);
    size_t parallelProducerSlotCount = 0;
    for (auto& producer : parallelProducers) {
        parallelProducerSlotCount += getNumOfSlotsUsedByTask(getTaskOpAtIndex(producer));
    }
    // make full use of the total slots
    // originally the slots are evenly split for producers and consumers,
    // e.g., 64 slots in total, 32 for producer and 32 for consumer
    // when the producer exceeds 32 the excess producers are linearized, leading to DPU stalls
    // try to distribute more slots for producer or consumer and check the feasibility
    if (_enableUnevenVariantSplit && (parallelProducerSlotCount < totalSlots)) {
        // check if every consumer could fit consumer-slots
        // if false, use equal slots for producer and consumer (original logic)
        // else, assign slots for producer and use the remaining slots for consumer
        const bool everyConsumerFitProvidedSlots = llvm::all_of(parallelConsumers, [&](auto consumer) {
            return waitBarriers != getWaitBarriers(consumer) ||
                   getNumOfSlotsUsedByTask(getTaskOpAtIndex(consumer)) <= totalSlots - parallelProducerSlotCount;
        });
        if (everyConsumerFitProvidedSlots) {
            availableSlotsForConsumer = totalSlots - parallelProducerSlotCount;
        }
    }

    auto parallelConsumersToMerge = getBarriersConsumerTasksWithWaitBarriers(waitBarriers, availableSlotsForConsumer);
    VPUX_THROW_WHEN(parallelConsumersToMerge.empty(), "parallel consumers are empty");

    // check if total slot count required by parallelProducers <= available slots for producers
    if (canMergeBarriersForTasks(parallelProducers, totalSlots - availableSlotsForConsumer)) {
        log.nest().trace("Can merge '{0}' parallel producers for task '{1}'", parallelProducers.size(), taskInd);
        mergeLegalParallelProducers(taskInd, parallelProducers, parallelConsumersToMerge);
    } else {
        log.nest().trace("Have to linearize '{0}' parallel producers for task '{1}'", parallelProducers.size(),
                         taskInd);
        linearizeLegalParallelProducers(taskInd, parallelProducers, parallelConsumersToMerge, availableSlots,
                                        considerTaskExecutorType);
    }

    for (auto& id : parallelConsumersToMerge) {
        log.trace("Unlink parallel barriers from task '{0}'", id);
    }
    unlinkTaskFromParallelBarriers(parallelConsumersToMerge, waitBarriers);
    log.trace("Task '{0}' now has one wait barrier", taskInd);
    return true;
}

SmallVector<BarrierInfo::TaskSet> BarrierInfo::createProducerBatches(const BarrierInfo::TaskSet& waitBarriers,
                                                                     size_t availableSlots,
                                                                     bool considerTaskExecutorType) {
    // try to create batches of producers using existing barriers if producers satisfy order constraint
    SmallVector<BarrierInfo::TaskSet> legalBatches;

    auto prevLastUserInd = std::numeric_limits<size_t>::max();  // last user index in the previous batch
    for (auto& waitBarrierInd : waitBarriers) {
        const auto& barrierProducers = getBarrierProducers(waitBarrierInd);
        if (legalBatches.empty()) {
            prevLastUserInd = VPURT::getMaxEntry(barrierProducers);
            legalBatches =
                    canMergeBarriersForTasks(barrierProducers, availableSlots)
                            ? SmallVector<BarrierInfo::TaskSet>{barrierProducers}
                            : createLegalVariantBatches(barrierProducers, availableSlots, considerTaskExecutorType);
            continue;
        }

        const auto firstUserInd = VPURT::getMinEntry(barrierProducers);
        if (prevLastUserInd >= firstUserInd) {
            // producers do not satisfy order constraint
            return SmallVector<BarrierInfo::TaskSet>{};
        }

        prevLastUserInd = VPURT::getMaxEntry(barrierProducers);
        auto currentBatchPlusBarrierProducers = legalBatches.back();
        llvm::set_union(currentBatchPlusBarrierProducers, barrierProducers);

        if (canMergeBarriersForTasks(currentBatchPlusBarrierProducers, availableSlots)) {
            // can add to the same batch
            legalBatches.back() = std::move(currentBatchPlusBarrierProducers);
        } else {
            // possible improvement: if the entire batch doesn't fit, consider filling-in the last batch with a portion
            // of tasks from barrierProducers (E140023).
            // Need to create a new batch
            if (canMergeBarriersForTasks(barrierProducers, availableSlots)) {
                legalBatches.push_back(barrierProducers);
            } else {
                legalBatches.append(
                        createLegalVariantBatches(barrierProducers, availableSlots, considerTaskExecutorType));
            }
        }
    }

    return legalBatches;
}

/*
    Linearize wait barriers if task wait barrier count > 1 by linearizing legal batches of parallel wait barriers

        x..xn    y..ym             x..xn
          |        |              /     \
        Bar0     Bar1     -->   Bar0     Bar1
       /    \   /    \           |        |
    u..ui   o..oj   a..ak      u..ui    y..ym
                                       /     \
                                     Bar2   Bar3
                                      |      |
                                    o..oj   a..ak
*/
void vpux::BarrierInfo::linearizeLegalParallelProducers(size_t taskInd, const BarrierInfo::TaskSet& parallelProducers,
                                                        const BarrierInfo::TaskSet& parallelConsumers,
                                                        size_t availableSlots, bool considerTaskExecutorType) {
    // create legal batches of barrier producers
    auto legalBatches = createProducerBatches(getWaitBarriers(taskInd), availableSlots, considerTaskExecutorType);
    if (legalBatches.empty()) {
        // create new batches of producers that satisfy order constraint
        legalBatches = createLegalVariantBatches(parallelProducers, availableSlots, considerTaskExecutorType);
    }

    // last batch is the current task
    BarrierInfo::TaskSet nextBatch = parallelConsumers;
    auto taskOp = getTaskOpAtIndex(taskInd);
    mlir::OpBuilder builder(taskOp);

    // create a new barrier between batches
    for (const auto& batch : legalBatches | reversed) {
        auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(taskOp.getLoc());
        addNewBarrier(newBarrier);

        for (auto& barrierProducer : batch) {
            addProducer(newBarrier, barrierProducer);
        }

        for (auto& barrierConsumer : nextBatch) {
            addConsumer(newBarrier, barrierConsumer);
        }

        nextBatch = batch;
    }
}

// merge wait barriers if task wait barrier count > 1 by
// creating a new barrier replacing parallel wait barriers
/*
        x..xn    y..ym             x..xn    y..ym
          |        |              /    \   /    \
        Bar0     Bar1     -->   Bar0    Bar1    Bar2
       /    \   /    \           |       |       |
    u..ui   o..oj   a..ak      u..ui   o..oj   a..ak
*/
void vpux::BarrierInfo::mergeLegalParallelProducers(size_t taskInd, const BarrierInfo::TaskSet& parallelProducers,
                                                    const BarrierInfo::TaskSet& parallelConsumers) {
    // create a new barrier for all parallel producers
    auto taskOp = getTaskOpAtIndex(taskInd);
    mlir::OpBuilder builder(taskOp);
    auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(taskOp.getLoc());
    addNewBarrier(newBarrier);

    // add all legal producers to new barrier
    for (auto& barrierProducer : parallelProducers) {
        addProducer(newBarrier, barrierProducer);
    }

    // add all parallel consumers with same wait barriers
    for (auto& consumer : parallelConsumers) {
        addConsumer(newBarrier, consumer);
    }
}

void vpux::BarrierInfo::shareWaitAndUpdateBarriers(size_t availableSlots) {
    // This method requires task queue map built with buildTaskQueueTypeMap()
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "TaskQueueTypeMap has not been created");
    for (const auto& entry : _taskQueueTypeMap) {
        // operate per queue
        const auto& dmaQueue = entry.second;
        // handle update barriers
        BarrierInfo::TaskSet tasksNoUpdateBarrier;
        for (const auto& taskInd : dmaQueue.set_bits()) {
            const auto taskUpdateBarriers = getUpdateBarriers(taskInd);
            if (taskUpdateBarriers.empty()) {
                tasksNoUpdateBarrier.insert(taskInd);
            } else if (!tasksNoUpdateBarrier.empty()) {
                linkTasksToBarriers(tasksNoUpdateBarrier, taskUpdateBarriers, false, availableSlots);
                tasksNoUpdateBarrier.clear();
            }
        }

        // handle wait barriers
        BarrierInfo::TaskSet tasksNoWaitBarrier;
        for (auto taskInd = dmaQueue.find_last(); taskInd != -1;) {
            const auto taskWaitBarriers = getWaitBarriers(static_cast<size_t>(taskInd));
            if (taskWaitBarriers.empty()) {
                tasksNoWaitBarrier.insert(static_cast<size_t>(taskInd));
            } else if (!tasksNoWaitBarrier.empty()) {
                linkTasksToBarriers(tasksNoWaitBarrier, taskWaitBarriers, true, availableSlots);
                tasksNoWaitBarrier.clear();
            }
            taskInd = dmaQueue.find_prev(taskInd);
        }
    }
}

// link tasks to barriers and create legal variant batches if needed
void vpux::BarrierInfo::linkTasksToBarriers(const BarrierInfo::TaskSet& tasksToAdd,
                                            const BarrierInfo::TaskSet& newBarriers, bool waitBarriers,
                                            size_t availableSlots) {
    VPUX_THROW_UNLESS(_fifoDependencies.size() == 0,
                      "Cannot create new barriers when FIFO dependencies are created but not represented in IR");
    mlir::OpBuilder builder(getBarrierOpAtIndex(*newBarriers.begin()));
    BarrierInfo::TaskSet newBarrierBatch;

    for (const auto& barrierIdn : newBarriers) {
        BarrierInfo::TaskSet barrierTasks;
        if (waitBarriers) {
            barrierTasks = getBarrierConsumers(barrierIdn);
        } else {
            barrierTasks = getBarrierProducers(barrierIdn);
        }

        llvm::set_union(barrierTasks, tasksToAdd);
        auto batches = createLegalVariantBatches(barrierTasks, availableSlots);
        if (batches.size() > 1) {
            auto insertionPoint = getBarrierOpAtIndex(barrierIdn);
            BarrierInfo::TaskSet prevBatch;
            if (waitBarriers) {
                prevBatch = getBarrierProducers(barrierIdn);
            } else {
                batches.push_back(getBarrierConsumers(barrierIdn));
            }
            for (const auto& batch : batches) {
                if (!prevBatch.empty()) {
                    builder.setInsertionPoint(insertionPoint);
                    const auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(insertionPoint->getLoc());
                    addNewBarrier(newBarrier);
                    const auto newBarrierIdn = getIndex(newBarrier);
                    newBarrierBatch.insert(newBarrierIdn);
                    addProducers(newBarrierIdn, prevBatch);
                    addConsumers(newBarrierIdn, batch);
                }
                prevBatch = batch;
            }
            resetBarrier(barrierIdn);
        } else if (waitBarriers) {
            addConsumers(barrierIdn, tasksToAdd);
            newBarrierBatch.insert(barrierIdn);
        } else {
            addProducers(barrierIdn, tasksToAdd);
            newBarrierBatch.insert(barrierIdn);
        }
    }
}

size_t vpux::BarrierInfo::getBarrierLatestProducer(size_t barInd) {
    auto barProducers = getBarrierProducers(barInd);
    if (barProducers.empty()) {
        VPUX_THROW("Barrier {0} does not have any producers.", barInd);
    }
    return *std::max_element(barProducers.begin(), barProducers.end());
}

size_t vpux::BarrierInfo::getBarriersLatestProducer(const TaskSet& barriers) {
    size_t latestTaskInd = 0;
    llvm::for_each(barriers, [&](auto barInd) {
        latestTaskInd = std::max(latestTaskInd, getBarrierLatestProducer(barInd));
    });

    return latestTaskInd;
};

size_t vpux::BarrierInfo::getBarrierLatestConsumer(size_t barInd) {
    auto barConsumers = getBarrierConsumers(barInd);
    if (barConsumers.empty()) {
        VPUX_THROW("Barrier {0} does not have any consumers.", barInd);
    }
    return *std::max_element(barConsumers.begin(), barConsumers.end());
}

size_t vpux::BarrierInfo::getBarrierEarliestConsumer(size_t barInd) {
    auto barConsumers = getBarrierConsumers(barInd);
    if (barConsumers.empty()) {
        VPUX_THROW("Barrier {0} does not have any consumers.", barInd);
    }
    return *std::min_element(barConsumers.begin(), barConsumers.end());
}

size_t vpux::BarrierInfo::getBarriersEarliestConsumer(const TaskSet& barriers) {
    size_t earliestTaskInd = std::numeric_limits<size_t>::max();
    llvm::for_each(barriers, [&](auto barInd) {
        earliestTaskInd = std::min(earliestTaskInd, getBarrierEarliestConsumer(barInd));
    });

    return earliestTaskInd;
};

unsigned vpux::BarrierInfo::createBarrierDependenciesImpliedByFIFO(
        size_t blockIdx, std::optional<mlir::DenseSet<vpux::config::ExecutorKind>> executorKind) {
    // This method requires task queue map built with buildTaskQueueTypeMap()
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "TaskQueueTypeMap has not been created");
    unsigned newDepsCount = 0;
    auto addTemporaryBarrierDependence = [&](size_t prevIndex, size_t currentIndex) {
        // Add barrier dependence between task with index prevIndex and task with index currentIndex but do not create a
        // new barrier in the IR. This creates a temporary state of mismatch between number of barriers in the IR and
        // number of barriers tracked by the _barrierProducerMap and _barrierConsumerMap maps. This state is used to
        // perform barrier optimizations taking into account FIFO dependencies. These dependencies will be
        // removed after optimization is done for the sake of internal self-consistency.
        size_t tempBarrierInd = _barrierProducerMap.size();
        _barrierProducerMap.push_back({});
        _barrierConsumerMap.push_back({});

        addProducer(tempBarrierInd, prevIndex);
        addConsumer(tempBarrierInd, currentIndex);
        newDepsCount += 2;
        return tempBarrierInd;
    };

    auto [blockStartInd, blockEndInd] =
            getControlGraphBlockTaskRange(blockIdx, /* blockStartSyncPoint */ false, /* blockEndSyncPoint */ false);

    // Go over the FIFOs, if the two tasks are in the same blockRange and does not have barrier dependency
    // add a temporary dependency between them
    for (const auto& entry : _taskQueueTypeMap) {
        if (executorKind.has_value() && !executorKind.value().contains(entry.first.type)) {
            continue;
        }
        const auto& fifoTasks = entry.second;

        if (fifoTasks.count() < 2) {
            continue;
        }
        int prevIndex =
                fifoTasks.find_first_in(static_cast<unsigned>(blockStartInd), static_cast<unsigned>(blockEndInd));
        if (prevIndex == -1) {  // no tasks in this block
            continue;
        }
        int currIndex = fifoTasks.find_next(prevIndex);

        while (currIndex <= fifoTasks.find_last() && currIndex != -1 && static_cast<size_t>(currIndex) <= blockEndInd) {
            size_t commonBarrier = -1;
            // Check if the two tasks do not have barrier dependency already and that they're in the same block
            if (!hasBarrierDependency(static_cast<size_t>(prevIndex), static_cast<size_t>(currIndex), commonBarrier)) {
                // add temporary dependency from task prevIndex to currIndex
                auto newBarrierIdx = addTemporaryBarrierDependence(prevIndex, currIndex);
                _fifoDependencies.emplace_back(newBarrierIdx, prevIndex, currIndex);
            }
            prevIndex = currIndex;
            currIndex = fifoTasks.find_next(currIndex);
        }
    }

    return newDepsCount;
}

bool vpux::BarrierInfo::createBarrierDependenciesForDescriptorFetchInNonWlm(size_t blockIdx,
                                                                            ExecutionGroupListMap& executionGroups) {
    // This method requires task queue map built with buildTaskQueueTypeMap()
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "TaskQueueTypeMap has not been created");

    auto barrierConsumerDependsOnGrandChildGroup =
            [&](const size_t firstTaskInGrandChildGroup, const size_t reusableWaitBarrier,
                std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                std::optional<size_t>& blockIdxOfTaskControlMap) {
                for (const auto& consumer : getBarrierConsumers(reusableWaitBarrier)) {
                    if (isDepFromTaskAToTaskB(firstTaskInGrandChildGroup, consumer, taskControlMapAndOffset,
                                              blockIdxOfTaskControlMap)) {
                        return true;
                    }
                }
                return false;
            };

    auto canWaitBarrierBeUpdatedWithoutDeadlock =
            [&](const size_t producer, const size_t consumer,
                std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                std::optional<size_t>& blockIdxOfTaskControlMap, size_t& reusableWaitBarrier) {
                // Check if there exist a consumer wait barrier that can be updated by producer without creating a
                // deadlock. A wait barrier can be updated without deadlock if the barrier new producer does not depend
                // on any of the barrier consumers (except when the consumer is the producer itself).
                const auto availableWaitBarriers = getWaitBarriers(consumer);

                for (auto waitBarrier : availableWaitBarriers) {
                    bool reuseWaitBar = true;
                    const auto waitBarrierConsumers = getBarrierConsumers(waitBarrier);
                    for (const auto& barConsumer : waitBarrierConsumers) {
                        if (isDepFromTaskAToTaskB(barConsumer, producer, taskControlMapAndOffset,
                                                  blockIdxOfTaskControlMap)) {
                            reuseWaitBar = false;
                            break;
                        }
                    }
                    if (!reuseWaitBar) {
                        continue;
                    }
                    reusableWaitBarrier = waitBarrier;
                    return true;
                }
                return false;
            };

    auto compareBarriersByLatestConsumer = [](const std::pair<size_t, size_t>& bar1,
                                              const std::pair<size_t, size_t>& bar2) {
        return bar1.second > bar2.second;
    };
    std::set<std::pair<size_t, size_t>, decltype(compareBarriersByLatestConsumer)> legalizationBarriers(
            compareBarriersByLatestConsumer);
    auto updatingBarrierInvalidatesPreviousExecGroup = [&](const size_t barrierIdx,
                                                           std::pair<SmallVector<llvm::BitVector>, size_t>&
                                                                   taskControlMapAndOffset,
                                                           std::optional<size_t>& blockIdxOfTaskControlMap) {
        // if path exists from any consumer of the barrier barrierIdx to any consumer of the previous exec group
        // legalization barriers (clean_after barriers) then updating this barrier would invalidate the previous
        // execution group
        for (auto [legalizationBarrier, legalizationBarrierLatestConsumer] : legalizationBarriers) {
            if (getBarrierEarliestConsumer(barrierIdx) > legalizationBarrierLatestConsumer) {
                // no need to check further as legalization barriers are sorted by latest consumer in descending order
                return false;
            }
            for (const auto& consumer : getBarrierConsumers(barrierIdx)) {
                for (const auto& consumerOfCleanAfterBar : getBarrierConsumers(legalizationBarrier)) {
                    if (isDepFromTaskAToTaskB(consumer, consumerOfCleanAfterBar, taskControlMapAndOffset,
                                              blockIdxOfTaskControlMap)) {
                        _log.debug("Updating barrier {0} invalidates previous execution group because of path from "
                                   "consumer {1} to consumer {2} of legalization barrier {3}",
                                   barrierIdx, consumer, consumerOfCleanAfterBar, legalizationBarrier);
                        return true;
                    }
                }
            }
        }
        return false;
    };

    auto createNewBarrier = [&](const size_t producer, size_t consumer) {
        size_t barIdx = _barrierProducerMap.size();
        _barrierProducerMap.push_back({});
        _barrierConsumerMap.push_back({});
        auto taskOp = getTaskOpAtIndex(producer);
        mlir::OpBuilder builder(taskOp);
        auto newBarrier = builder.create<VPURT::DeclareVirtualBarrierOp>(taskOp->getLoc());
        newBarrier.setBarrierIndex(barIdx);
        _allBarrierOps.push_back(newBarrier);
        addProducer(barIdx, producer);
        addConsumer(barIdx, consumer);
    };

    // build task control map for current block and all executor kinds
    auto taskControlMapAndOffset = buildTaskControlMap(blockIdx, /* considerTaskFifoDependency */ true);
    std::optional<size_t> blockIdxOfTaskControlMap = blockIdx;
    // Insert barrier between execution groups when required
    bool newDependenciesCreated = false;
    for (auto [queue, fifoExecGroups] : executionGroups) {
        legalizationBarriers.clear();
        _log.trace("Processing queue: {0}{1} for block '{2}' with '{3}' execution groups",
                   vpux::config::stringifyExecutorKind(queue.type), queue.id, blockIdx, fifoExecGroups.size());
        if (fifoExecGroups.size() <= 1) {
            continue;  // No dependencies need to be created if there are no groups or all tasks fit into a single
                       // group.
        }

        int numberOfGroups = static_cast<int>(fifoExecGroups.size());
        for (int grpIdx = 0; grpIdx < numberOfGroups - 1; grpIdx++) {
            std::optional<size_t> legalizationBarrierIdx(std::nullopt);
            if (!hasBarrierDependencyRequiredForDescriptorFetch(grpIdx, blockIdx, fifoExecGroups,
                                                                taskControlMapAndOffset, blockIdxOfTaskControlMap,
                                                                false, legalizationBarrierIdx)) {
                _log.debug("{2}{3} group index: '{0}' in block '{1}' does not have required barrier dependency", grpIdx,
                           blockIdx, vpux::config::stringifyExecutorKind(queue.type), queue.id);
                // define last task in group as producer of the new barrier
                auto newBarProducer = *fifoExecGroups[grpIdx].rbegin();
                // define first task in child group as consumer of the new barrier
                auto newBarConsumer = *fifoExecGroups[grpIdx + 1].begin();

                if (getWaitBarriers(newBarConsumer).empty()) {
                    _log.debug("Consumer task {0} has no wait barriers, creating new barrier with producer {1}",
                               newBarConsumer, newBarProducer);
                    createNewBarrier(newBarProducer, newBarConsumer);
                    newDependenciesCreated = true;
                    continue;
                }

                // If new barrier consumer has at least one wait barrier, then reuse it by updating it from the
                // last task in current group, if it doesn't generate a deadlock and if the barrier does not depend on
                // grand child group
                VPUX_THROW_UNLESS(grpIdx + 2 < numberOfGroups,
                                  "Grand child group does not exist for group index '{0}' in block '{1}' and "
                                  "legalization barrier is requested",
                                  grpIdx, blockIdx);
                auto firstTaskInGrandChildGroup = *fifoExecGroups[grpIdx + 2].begin();
                size_t reusableWaitBarrier = -1;  // dummy value
                if (canWaitBarrierBeUpdatedWithoutDeadlock(newBarProducer, newBarConsumer, taskControlMapAndOffset,
                                                           blockIdxOfTaskControlMap, reusableWaitBarrier) &&
                    !barrierConsumerDependsOnGrandChildGroup(firstTaskInGrandChildGroup, reusableWaitBarrier,
                                                             taskControlMapAndOffset, blockIdxOfTaskControlMap) &&
                    !updatingBarrierInvalidatesPreviousExecGroup(reusableWaitBarrier, taskControlMapAndOffset,
                                                                 blockIdxOfTaskControlMap)) {
                    _log.debug("Consumer task {0} has reusable wait barrier {1} - adding new producer ({2}) to it.",
                               newBarConsumer, reusableWaitBarrier, newBarProducer);
                    addProducer(reusableWaitBarrier, newBarProducer);
                    legalizationBarriers.insert(
                            std::make_pair(reusableWaitBarrier, getBarrierLatestConsumer(reusableWaitBarrier)));
                    newDependenciesCreated = true;
                    continue;
                }

                // Create barrier between tasks on the same queue. This will generate parallel wait barriers
                _log.trace("Creating new barrier between producer task {0} and consumer task {1}", newBarConsumer,
                           newBarProducer);
                createNewBarrier(newBarProducer, newBarConsumer);
                newDependenciesCreated = true;
            } else {
                if (legalizationBarrierIdx.has_value()) {
                    legalizationBarriers.insert(std::make_pair(
                            legalizationBarrierIdx.value(), getBarrierLatestConsumer(legalizationBarrierIdx.value())));
                }
            }
        }
    }
    return newDependenciesCreated;
}

unsigned vpux::BarrierInfo::removeBarrierDependenciesImpliedByFIFO() {
    unsigned removedDepsCount = 0;

    // Remove barriers implicitly implied by FIFO dependency that were temporarily added between neighbour tasks
    // on the same FIFO
    for (const auto& dep : _fifoDependencies) {
        size_t commonBarrier = -1;
        const auto [barrierIdx, parentTaskIdx, childTaskIdx] = dep;

        if (hasBarrierDependency(parentTaskIdx, childTaskIdx, commonBarrier) && commonBarrier == barrierIdx) {
            removeProducer(barrierIdx, parentTaskIdx);
            removeConsumer(barrierIdx, childTaskIdx);
            removedDepsCount += 2;
        }
    }

    // restore original sizes of barrier producer and consumer maps after removing temporary barriers associated
    // with FIFO
    auto removeTemporaryBarriers = [&](SmallVector<TaskSet>& barrierMap, bool producerMap) {
        auto numUsersOfTemporaryFIFObarriers =
                std::accumulate(barrierMap.end() - _fifoDependencies.size(), barrierMap.end(), 0U,
                                [](const size_t& a, const TaskSet& b) {
                                    return a + b.size();
                                });
        VPUX_THROW_UNLESS(numUsersOfTemporaryFIFObarriers == 0, "Found {0} for temporary barrier.",
                          producerMap ? "producers" : "consumers");
        // To prevent unnecessary data copy and movement during temporary barrier creation, we directly truncate the
        // barrierMap while preserving its allocated capacity. This approach minimizes overhead by avoiding costly
        // reallocations.
        barrierMap.truncate(barrierMap.size() - _fifoDependencies.size());
    };

    removeTemporaryBarriers(_barrierProducerMap, /* producerMap */ true);
    removeTemporaryBarriers(_barrierConsumerMap, /* producerMap */ false);
    _fifoDependencies.clear();
    return removedDepsCount;
}

// iteratively merge parallel wait barriers
bool vpux::BarrierInfo::ensureTasksDrivenBySingleBarrier(size_t availableSlots, bool mergeWaitBarriersIteratively,
                                                         bool considerTaskExecutorType) {
    auto getTasksWithParallelWaitBarriers = [&]() {
        BarrierInfo::TaskSet tasksWithParallelWaitBarriers;
        for (size_t taskInd = 0; taskInd < getNumOfTasks(); taskInd++) {
            if (getWaitBarriers(taskInd).size() >= 2) {
                tasksWithParallelWaitBarriers.insert(taskInd);
            }
        }
        return tasksWithParallelWaitBarriers;
    };

    bool modifiedIR = false;
    auto tasksWithParallelWaitBarriers = getTasksWithParallelWaitBarriers();
    size_t iteration = 0;
    while (tasksWithParallelWaitBarriers.size() > 0) {
        _log.trace("Encountered {0} tasks with more than one wait barrier in iteration {1}",
                   tasksWithParallelWaitBarriers.size(), iteration);

        // Reverse "for" loop over tasks would reduce the required number of iterations, but
        // the schedule may not be exactly equivalent (E140023).
        _func->walk([&](VPURT::TaskOp taskOp) {
            const auto taskInd = getIndex(taskOp);
            if (getWaitBarriers(taskInd).size() < 2) {
                // valid configuration
                return;
            }

            modifiedIR |= eliminateParallelWaitBarriers(taskInd, availableSlots, considerTaskExecutorType);
        });

        auto tasksWithParallelWaitBarriersPrevIt = tasksWithParallelWaitBarriers;
        tasksWithParallelWaitBarriers = getTasksWithParallelWaitBarriers();
        VPUX_THROW_UNLESS(tasksWithParallelWaitBarriers != tasksWithParallelWaitBarriersPrevIt,
                          "Merging parallel wait barriers encountered an identical batch of tasks as in previous "
                          "iteration. Cannot enforce single wait barrier per task.");
        iteration++;

        if (!mergeWaitBarriersIteratively) {
            break;
        }
    }

    return modifiedIR;
}

bool vpux::BarrierInfo::inRange(const size_t low, const size_t high, const size_t val) const {
    return val >= low && val <= high;
}

void vpux::BarrierInfo::setBarrierMask(llvm::BitVector& mask, const BarrierInfo::TaskSet& barriers, size_t offset) {
    for (auto bar : barriers) {
        mask.set(static_cast<unsigned>(bar - offset));
    }
}

void vpux::BarrierInfo::initializeTaskQueueTypeMap(const mlir::DenseSet<vpux::config::ExecutorKind>& executorKind) {
    for (auto execKind : executorKind) {
        const auto module = _func->getParentOfType<mlir::ModuleOp>();
        VPURT::TaskQueueType taskQueueType;
        taskQueueType.type = execKind;

        switch (execKind) {
        case config::ExecutorKind::DMA_NN: {
            const auto dmaPortNum = config::getAvailableExecutor(module, config::ExecutorKind::DMA_NN).getCount();
            auto dmaChannels = getDMAChannelsWithIndependentLinkAgents(config::getArch(module));
            for (auto dmaPortIdx : irange(dmaPortNum)) {
                for (auto dmaChannel : dmaChannels) {
                    taskQueueType.id = getDMAQueueIdEncoding(dmaPortIdx, dmaChannel);
                    llvm::BitVector taskList(checked_cast<uint32_t>(getNumOfTasks()));
                    _taskQueueTypeMap.insert(std::make_pair(taskQueueType, std::move(taskList)));
                }
            }
            break;
        }

        case config::ExecutorKind::DPU: {
            auto numClusters = config::getTileExecutor(module).getCount();
            for (auto clusterId : irange(numClusters)) {
                taskQueueType.id = clusterId;
                llvm::BitVector taskList(checked_cast<uint32_t>(getNumOfTasks()));
                _taskQueueTypeMap.insert(std::make_pair(taskQueueType, std::move(taskList)));
            }
            break;
        }

        case config::ExecutorKind::SHAVE_ACT: {
            auto tileOp = config::getTileExecutor(module);
            VPUX_THROW_UNLESS(tileOp != nullptr, "Expected tileOp executor in order to query {0} executor.", execKind);

            auto numClusters = tileOp.getCount();
            auto numExecutorsPerTile = [&] {
                return config::isFifoPerShaveEngineEnabled(_func)
                               ? static_cast<size_t>(tileOp.getSubExecutor(execKind).getCount())
                               : static_cast<size_t>(1);
            }();

            for (auto clusterId : irange(numClusters)) {
                for (auto shaveId : irange(numExecutorsPerTile)) {
                    taskQueueType.id = getShaveQueueIdEncoding(numClusters, clusterId, shaveId);
                    llvm::BitVector taskList(checked_cast<uint32_t>(getNumOfTasks()));
                    _taskQueueTypeMap.insert(std::make_pair(taskQueueType, taskList));
                }
            }
            break;
        }

        default: {
            taskQueueType.id = 0;
            llvm::BitVector taskList(checked_cast<uint32_t>(getNumOfTasks()));
            _taskQueueTypeMap.insert(std::make_pair(taskQueueType, taskList));
        }
        }
    }
}

void vpux::BarrierInfo::clearTaskQueueTypeMap() {
    _taskQueueTypeMap.clear();
}

void vpux::BarrierInfo::buildTaskQueueTypeMap() {
    if (_taskQueueTypeMap.empty()) {
        mlir::DenseSet<vpux::config::ExecutorKind> executors = {config::ExecutorKind::DMA_NN,
                                                                config::ExecutorKind::DPU};
        if (config::isFifoPerShaveEngineEnabled(_func)) {
            executors.insert(config::ExecutorKind::SHAVE_ACT);
        }
        initializeTaskQueueTypeMap(executors);
    }

    for (const auto& taskOp : _allTaskOps | reversed) {
        auto taskInd = getIndex(taskOp);
        auto taskQueueType = VPURT::getTaskQueueType(taskOp, false);
        if (_taskQueueTypeMap.find(taskQueueType) != _taskQueueTypeMap.end()) {
            _taskQueueTypeMap[taskQueueType].set(taskInd);
        }
    }
}

//
// buildTaskControlMap
//
// Generate a map of task-to-task dependencies for a given block index.
// Dependencies between tasks are represented through barriers and task order on given HW FIFO
// Resulting map for a given task will contain all tasks that are reachable from the given task
// directly or through other tasks and barriers.
std::pair<SmallVector<llvm::BitVector>, size_t> vpux::BarrierInfo::buildTaskControlMap(
        size_t blockIdx, bool considerTaskFifoDependency, bool ignoreOutOfBlockDependencies) {
    SmallVector<llvm::BitVector> taskControlMap;
    VPUX_THROW_WHEN(blockIdx >= getControlGraphBlockCount(), "Invalid task block index ({0})", blockIdx);

    auto [blockStartInd, blockEndInd] =
            getControlGraphBlockTaskRange(blockIdx, /* blockStartSyncPoint */ true, /* blockEndSyncPoint */ true);

    VPUX_THROW_UNLESS(blockStartInd <= blockEndInd, "Invalid range of tasks for building task control map [{0}, {1}]",
                      blockStartInd, blockEndInd);

    _log.trace("Build task control map for task range [{0}, {1}]", blockStartInd, blockEndInd);
    auto newTaskControlMapSize = blockEndInd - blockStartInd + 1;
    if (taskControlMap.size() != newTaskControlMapSize) {
        resizeBitMap(taskControlMap, newTaskControlMapSize, checked_cast<uint32_t>(newTaskControlMapSize));
    }
    resetBitMap(taskControlMap);

    unsigned newDeps = 0;
    if (considerTaskFifoDependency) {
        // For every initialized queue in _taskQueueTypeMap, create temporary barriers between tasks on the same FIFO in
        // order to build taskControlMap with information about tasks-to-task dependencies stored exclusively on FIFO.
        // The these temporary barriers are removed afterwards.
        newDeps = createBarrierDependenciesImpliedByFIFO(blockIdx);
    }

    for (auto taskInd = blockEndInd + 1; taskInd-- > blockStartInd;) {
        for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
            for (auto consumerIdx : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                if (inRange(blockStartInd, blockEndInd, consumerIdx)) {
                    taskControlMap[taskInd - blockStartInd].set(static_cast<unsigned>(consumerIdx - blockStartInd));
                }
            }
        }
    }

    // In case task[i] has dep to task[j] and j < i then store index i to reiterate over this index later
    // Without that control map would not contain complete data on dependencies and to overcome this
    // IR would need to be reordered and BarrierInfo recreated
    // TODO: Analyze if this approach could be improved. For example by first generating order
    // of tasks and then iterating over this order (from last to first) instead of iterating through
    // task indexes directly in case task index no longer is aligned with order of execution.
    SmallVector<size_t> tasksWithDepsToSmallerIndexes;

    for (auto taskInd = blockEndInd + 1; taskInd-- > blockStartInd;) {
        if (taskInd == blockEndInd && isSyncPoint(taskInd)) {
            continue;
        }
        for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
            for (auto childTaskInd : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                if (inRange(blockStartInd, blockEndInd, childTaskInd)) {
                    if (childTaskInd < taskInd) {
                        tasksWithDepsToSmallerIndexes.push_back(taskInd);
                    }
                    taskControlMap[taskInd - blockStartInd] |=
                            taskControlMap[static_cast<size_t>(childTaskInd - blockStartInd)];
                } else {
                    VPUX_THROW_UNLESS(
                            ignoreOutOfBlockDependencies,
                            "Task {0} has update barrier {1} with a consumer ({2}) from outside of the current "
                            "range [{3}, {4}].",
                            taskInd, updateBarrierInd, childTaskInd, blockStartInd, blockEndInd);
                }
            }
        }
    }

    if (!tasksWithDepsToSmallerIndexes.empty()) {
        // Order tasks based on smaller index and update taskControlMap
        llvm::sort(tasksWithDepsToSmallerIndexes);
        auto highestTaskInd = tasksWithDepsToSmallerIndexes.back();
        // First process tasks with dependencies to smaller indexes to update their deps
        for (auto taskInd : tasksWithDepsToSmallerIndexes) {
            for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
                for (auto childTaskInd : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                    if (inRange(blockStartInd, blockEndInd, childTaskInd)) {
                        taskControlMap[taskInd - blockStartInd] |=
                                taskControlMap[static_cast<size_t>(childTaskInd - blockStartInd)];
                    }
                }
            }
        }
        // Reiterate all tasks again to update other tasks that would be affected by the
        // previous step
        // For example in first step it was update that task C has control path to task B
        // and in this step it could be update that task A that has path to task C should also have path to task B
        for (auto taskInd = highestTaskInd; taskInd-- > blockStartInd;) {
            for (auto updateBarrierInd : _taskUpdateBarriers[taskInd]) {
                for (auto childTaskInd : _barrierConsumerMap[static_cast<size_t>(updateBarrierInd)]) {
                    if (inRange(blockStartInd, blockEndInd, childTaskInd)) {
                        taskControlMap[taskInd - blockStartInd] |=
                                taskControlMap[static_cast<size_t>(childTaskInd - blockStartInd)];
                    }
                }
            }
        }
    }

    if (considerTaskFifoDependency) {
        auto removedDeps = removeBarrierDependenciesImpliedByFIFO();
        VPUX_THROW_WHEN(newDeps != removedDeps, "Temporary barriers have not been correctly removed");
    }

    return std::make_pair(std::move(taskControlMap), blockStartInd);
}

//
// controlPathExistsBetweenTasksInSameBlock
//

bool vpux::BarrierInfo::controlPathExistsBetweenTasksInSameBlock(const SmallVector<llvm::BitVector>& taskControlMap,
                                                                 size_t taskAInd, size_t taskBInd,
                                                                 bool biDirection) const {
    // ensure that taskControlMap is build at given time with buildTaskControlMap() for the correct task block
    VPUX_THROW_WHEN(taskControlMap.empty(), "Task control map not initialized");

    VPUX_THROW_WHEN(taskAInd >= taskControlMap.size(), "taskAInd out of range {0}", taskAInd);
    VPUX_THROW_WHEN(taskBInd >= taskControlMap.size(), "taskBInd out of range {0}", taskBInd);

    if (taskAInd == taskBInd) {
        return true;
    }

    if (biDirection) {
        return taskControlMap[taskAInd][taskBInd] || taskControlMap[taskBInd][taskAInd];
    }
    return taskControlMap[taskAInd][taskBInd];
}

// Check if there is a dependency from taskA to taskB taking into account barrier dependency
// and HW FIFOs.
// taskControlMapAndOffset and blockIdxOfTaskControlMap arguments are input and output parameters. If it turns
// out that provided control map for control graph block does not match block index of taskA and taskB function will
// build new control map and update those arguments.
bool vpux::BarrierInfo::isDepFromTaskAToTaskB(size_t taskA, size_t taskB,
                                              std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                                              std::optional<size_t>& blockIdxOfTaskControlMap) {
    auto taskABlock = getControlGraphBlockIndex(taskA);
    auto taskBBlock = getControlGraphBlockIndex(taskB);
    // If taskB is on later block than taskA then by definition it is dependent on taskA
    if (taskABlock < taskBBlock) {
        return true;
    }

    // If taskB is on earlier block than taskA then it cannot dependent on taskA
    if (taskABlock > taskBBlock) {
        return false;
    }

    // If taskA and taskB are on the same HW FIFO and taskA has smaller index than taskB
    // then dependency exists
    auto taskAQueueType = getTaskQueueType(taskA);
    auto taskBQueueType = getTaskQueueType(taskB);
    if (taskAQueueType.type != config::ExecutorKind::UNKNOWN && taskBQueueType.type != config::ExecutorKind::UNKNOWN &&
        taskAQueueType == taskBQueueType && taskA < taskB) {
        return true;
    }

    if (!blockIdxOfTaskControlMap.has_value() || blockIdxOfTaskControlMap.value() != taskABlock) {
        blockIdxOfTaskControlMap = taskABlock;
        taskControlMapAndOffset = buildTaskControlMap(taskABlock);
    }
    auto& [taskControlMap, taskControlMapOffset] = taskControlMapAndOffset;

    // Check if there is required dep
    return controlPathExistsBetweenTasksInSameBlock(taskControlMap, taskA - taskControlMapOffset,
                                                    taskB - taskControlMapOffset, false);
}

// Check if there is a dependency from task to barrier by checking if there is any dependency
// from this task to any producer of this barrier. If yes then there is a guarantee that task
// needs to execute before barrier is produced - there is topological dependency.
// taskControlMapAndOffset and blockIdxOfTaskControlMap arguments are input and output parameters. If it turns
// out that provided control map for control graph block does not match block index of taskInd and barInd function will
// build new control map and update those arguments.
bool vpux::BarrierInfo::isDepFromTaskToBarrier(size_t taskInd, size_t barInd,
                                               std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                                               std::optional<size_t>& blockIdxOfTaskControlMap) {
    auto barProdTasks = getBarrierProducers(barInd);
    return llvm::any_of(barProdTasks, [&](const auto barProdTask) {
        return isDepFromTaskAToTaskB(taskInd, barProdTask, taskControlMapAndOffset, blockIdxOfTaskControlMap);
    });
}

// Check if there is a dependency from barrier to task by checking if there is a dependency
// from any consumer of this barrier to this task. If yes then there is a guarantee that task
// can execute only after barrier is produced - there is topological dependency.
// taskControlMapAndOffset and blockIdxOfTaskControlMap arguments are input and output parameters. If it turns
// out that provided control map for control graph block does not match block index of taskInd and barInd function will
// build new control map and update those arguments.
bool vpux::BarrierInfo::isDepFromBarrierToTask(size_t barInd, size_t taskInd,
                                               std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                                               std::optional<size_t>& blockIdxOfTaskControlMap) {
    auto barConsTasks = getBarrierConsumers(barInd);
    return llvm::any_of(barConsTasks, [&](const auto barConsTask) {
        return isDepFromTaskAToTaskB(barConsTask, taskInd, taskControlMapAndOffset, blockIdxOfTaskControlMap);
    });
}

// Check if there is a dependency from barA to barB by checking if there is any dependency
// from barA consumer to barB producer.
// taskControlMapAndOffset and blockIdxOfTaskControlMap arguments are input and output parameters. If it turns
// out that provided control map for control graph block does not match block index of barA and barB function will
// build new control map and update those arguments.
bool vpux::BarrierInfo::isDepFromBarAToBarB(size_t barA, size_t barB,
                                            std::pair<SmallVector<llvm::BitVector>, size_t>& taskControlMapAndOffset,
                                            std::optional<size_t>& blockIdxOfTaskControlMap) {
    if (barA > barB) {
        return false;
    }

    auto barAConsumers = getBarrierConsumers(barA);
    auto barBProducers = getBarrierProducers(barB);

    for (auto barAConsumer : barAConsumers) {
        for (auto barBProducer : barBProducers) {
            if (isDepFromTaskAToTaskB(barAConsumer, barBProducer, taskControlMapAndOffset, blockIdxOfTaskControlMap)) {
                return true;
            }
        }
    }
    return false;
}

//
// updateIR
//

void vpux::BarrierInfo::updateIR() {
    // update IR using stored dependencies
    for (size_t taskInd = 0; taskInd < getNumOfTasks(); ++taskInd) {
        auto taskOp = getTaskOpAtIndex(taskInd);

        taskOp.getWaitBarriersMutable().clear();
        for (auto barrierInd : _taskWaitBarriers[taskInd]) {
            auto barrierOp = getBarrierOpAtIndex(static_cast<size_t>(barrierInd));
            taskOp.getWaitBarriersMutable().append(barrierOp.getBarrier());
        }

        taskOp.getUpdateBarriersMutable().clear();
        for (auto barrierInd : _taskUpdateBarriers[taskInd]) {
            auto barrierOp = getBarrierOpAtIndex(static_cast<size_t>(barrierInd));
            taskOp.getUpdateBarriersMutable().append(barrierOp.getBarrier());
        }
    }
}

//
// logBarrierInfo
//

void vpux::BarrierInfo::logBarrierInfo() {
    // useful for logging and debugging barrier dependencies
    _log.setName("barrier-info-log");
    _log.trace("Logging BarrierInfo");

    for (size_t taskInd = 0; taskInd < getNumOfTasks(); ++taskInd) {
        _log.nest().trace("Task '{0}'", taskInd);
        for (const auto& barrierOp : getWaitBarriers(taskInd)) {
            _log.nest(2).trace("Task '{0}' waits for '{1}'", taskInd, barrierOp);
        }
        for (const auto& barrierOp : getUpdateBarriers(taskInd)) {
            _log.nest(2).trace("Task '{0}' updates '{1}'", taskInd, barrierOp);
        }
    }
}

/**
 * hasBarrierDependency - Checks for common barriers between two tasks
 *
 *        0
 *        |
 *       bar0
 *        |
 *        1
 *        |
 *       bar1
 *        |
 *        2
 *        |
 *       bar2
 *
 * For the above example for Task 0 and Task 1 the function returns true & bar0
 * Similarly for Task 0 and Task 2 it returns false as there is no direct barrier dependency
 */

bool vpux::BarrierInfo::hasBarrierDependency(size_t taskOneIdx, size_t taskTwoIdx, size_t& commonBarrier) {
    const auto& taskOneUpdateBarriers = getUpdateBarriers(taskOneIdx);
    const auto& taskTwoWaitBarriers = getWaitBarriers(taskTwoIdx);

    auto hasCommonBarrier = [&commonBarrier](const TaskSet& setOne, const TaskSet& setTwo) {
        for (const auto& barrier : setOne) {
            if (setTwo.count(barrier)) {
                commonBarrier = barrier;
                return true;
            }
        }
        return false;
    };

    return hasCommonBarrier(taskOneUpdateBarriers, taskTwoWaitBarriers);
}

//
// createLegalVariantBatches
//

SmallVector<BarrierInfo::TaskSet> vpux::BarrierInfo::createLegalVariantBatches(const TaskSet& tasks,
                                                                               size_t availableSlots,
                                                                               bool considerTaskExecutorType) {
    // store batches of tasks
    SmallVector<TaskSet> legalBatches(1);

    // store total slot count used by batch
    size_t totalSlotCount = 0;

    const auto isLegalVariantCountWith = [&](size_t numSlotsUsedByTask) -> bool {
        return (totalSlotCount + numSlotsUsedByTask) <= availableSlots;
    };

    // create batches for new barriers
    auto orderedTasks = std::set<size_t>(tasks.begin(), tasks.end());
    auto curBatchExecutorKind = getTaskOpAtIndex(*orderedTasks.begin()).getExecutorKind();
    const auto isMatchingExecutorKind = [&](config::ExecutorKind taskExecutorKind) {
        if (!considerTaskExecutorType) {
            return true;
        }
        return taskExecutorKind == curBatchExecutorKind;
    };

    for (const auto& taskInd : orderedTasks) {
        // find number of slots consumed by this task
        auto numSlotsUsedByTask = getNumOfSlotsUsed(getTaskOpAtIndex(taskInd));
        auto curTaskExecKind = getTaskOpAtIndex(taskInd).getExecutorKind();

        // check if new batch needs to be created
        if (!isLegalVariantCountWith(numSlotsUsedByTask) || !isMatchingExecutorKind(curTaskExecKind)) {
            legalBatches.push_back({});
            totalSlotCount = 0;
            curBatchExecutorKind = curTaskExecKind;
            // TODO: possible improvement - merge batches of same executor type if the slot count permits (E140023)
        }

        legalBatches.rbegin()->insert(taskInd);
        totalSlotCount += numSlotsUsedByTask;
    }

    return legalBatches;
}

//
// haveSameImplicitDependencyTaskQueueType
//

std::optional<VPURT::TaskQueueType> vpux::BarrierInfo::haveSameImplicitDependencyTaskQueueType(
        const TaskSet& taskInds) {
    if (taskInds.empty()) {
        return std::nullopt;
    }
    auto firstTaskInd = *taskInds.begin();
    // ensure that _taskQueueTypeMap is build at given time with buildTaskQueueTypeMap()
    VPUX_THROW_WHEN(_taskQueueTypeMap.empty(), "Task queue map not initialized");
    for (const auto& item : _taskQueueTypeMap) {
        // get the task queue type for the first task, then check all the other tasks have same task queue type
        if (!item.second.test(firstTaskInd)) {
            continue;
        }

        for (const auto& taskInd : taskInds) {
            if (!item.second.test(static_cast<unsigned>(taskInd)) || isSyncPoint(taskInd)) {
                return std::nullopt;
            }
        }
        return item.first;
    }
    return std::nullopt;
}

std::optional<config::ExecutorKind> vpux::BarrierInfo::haveSameExecutorKind(const TaskSet& taskInds) {
    if (taskInds.empty()) {
        return std::nullopt;
    }

    if (_func == nullptr) {
        // for test purposes, use information from _taskQueueTypeMap
        std::set<config::ExecutorKind> executors;
        for (const auto& item : _taskQueueTypeMap) {
            for (const auto& taskInd : taskInds) {
                if (item.second.test(static_cast<unsigned>(taskInd))) {
                    executors.insert(item.first.type);
                    if (executors.size() > 1) {
                        return std::nullopt;
                    }
                }
            }
        }
        if (executors.empty()) {
            return std::nullopt;
        }
        return *executors.begin();
    }

    auto firstTaskExecKind = getTaskOpAtIndex(*taskInds.begin()).getExecutorKind();
    for (const auto& taskInd : taskInds) {
        if (getTaskOpAtIndex(taskInd).getExecutorKind() != firstTaskExecKind) {
            return std::nullopt;
        }
    }
    return firstTaskExecKind;
}

//
// canBarriersBeMerged
//

bool vpux::BarrierInfo::canBarriersBeMerged(const TaskSet& barrierProducersA, const TaskSet& barrierConsumersA,
                                            const TaskSet& barrierProducersB, const TaskSet& barrierConsumersB,
                                            ArrayRef<TaskSet> origWaitBarriersMap) {
    // Barriers A and B can be merged if
    // 1. producers of barrier A collectively control consumers of barrier B (it is enough to only consider consumers of
    // B not present in A).
    // The current implementation takes a simplified approach that does not thoroughly check whether some producers of A
    // can run in parallel with some of the consumers of B and it also relies on a heuristics that differentiate
    // between different executor types. Approach to be revisited in E#190087 in the context of FWLM.
    if (!canBarProducersControlBarConsumers(
                barrierProducersA, llvm::set_difference(barrierConsumersB, barrierConsumersA), origWaitBarriersMap)) {
        return false;
    }

    // 2. producers of barrier B collectively control consumers of barrier A (it is enough to only consider consumers of
    // A not present in B)
    if (!canBarProducersControlBarConsumers(
                barrierProducersB, llvm::set_difference(barrierConsumersA, barrierConsumersB), origWaitBarriersMap)) {
        return false;
    }

    return true;
}

// Remove redundant consumers and/or producers which are controlled
// by FIFO dependency. OP[<FIFO>][<opIndex>]
// After this optimization barriers should not have producers or consumers
// from same HW FIFO
/*
    OP[A][0] OP[A][1] OP[B][0]   OP[A][1]  OP[B][0]
        \      |     /                \    /
              Bar           =>          Bar
        /      |     \                /    \
    OP[A][2] OP[B][1] OP[B][2]   OP[A][2]  OP[B][1]
*/
void vpux::BarrierInfo::removeRedundantBarrierProducersAndConsumers(bool considerTaskFifoDependency) {
    const auto findRedundantDependencies = [&](const SmallVector<llvm::BitVector>& taskControlMap,
                                               size_t taskControlMapOffset, const BarrierInfo::TaskSet& dependencies,
                                               bool producer = true) {
        SmallVector<size_t> sortedDependencies = to_small_vector(dependencies);
        llvm::sort(sortedDependencies);

        // find dependencies to remove
        BarrierInfo::TaskSet dependenciesToRemove;
        for (auto taskIndexIter = sortedDependencies.begin(); taskIndexIter != sortedDependencies.end();
             ++taskIndexIter) {
            if (isSyncPoint(*taskIndexIter)) {
                continue;
            }
            for (auto nextIndex = std::next(taskIndexIter); nextIndex != sortedDependencies.end(); ++nextIndex) {
                if (isSyncPoint(*nextIndex) ||
                    !controlPathExistsBetweenTasksInSameBlock(taskControlMap, *taskIndexIter - taskControlMapOffset,
                                                              *nextIndex - taskControlMapOffset,
                                                              /* biDirection */ false)) {
                    continue;
                }

                if (producer) {
                    dependenciesToRemove.insert(*taskIndexIter);
                } else {
                    dependenciesToRemove.insert(*nextIndex);
                }
            }
        }
        return dependenciesToRemove;
    };

    // Perform optimization in tasks blocks matching the distribution of synchronization points.
    for (size_t taskBlockIndex = 0; taskBlockIndex < getControlGraphBlockCount(); ++taskBlockIndex) {
        // Build or update the control relationship between any two tasks. Note that the relationship includes the
        // dependency by the barriers as well as the implicit dependence by FIFO
        auto [taskControlMap, controlMapOffset] = buildTaskControlMap(taskBlockIndex, considerTaskFifoDependency);

        // get update barriers range for current block
        auto blockUpdateBarriers = getBarriersForTaskBlock(taskBlockIndex, /* blockStartSyncPoint */ true,
                                                           /* blockEndSyncPoint */ false, /* updateBarriers */ true);

        if (blockUpdateBarriers.empty()) {
            _log.trace("No update barriers found in tasks block {0}", taskBlockIndex);
        } else {
            _log.trace("Optimize producers for barrier range [{0}, {1}] in tasks block {2}",
                       blockUpdateBarriers.front(), blockUpdateBarriers.back(), taskBlockIndex);
        }

        for (auto barrierIdx : blockUpdateBarriers) {
            // find producers to remove
            const auto barrierProducers = getBarrierProducers(barrierIdx);
            const auto producersToRemove =
                    findRedundantDependencies(taskControlMap, controlMapOffset, barrierProducers);
            removeProducers(barrierIdx, producersToRemove);
        }

        // get wait barriers range for current block
        auto blockWaitBarriers = getBarriersForTaskBlock(taskBlockIndex, /* blockStartSyncPoint */ false,
                                                         /* blockEndSyncPoint */ true, /* updateBarriers */ false);
        if (blockWaitBarriers.empty()) {
            _log.trace("No wait barriers found in tasks block {0}", taskBlockIndex);
        } else {
            _log.trace("Optimize consumers for barrier range [{0}, {1}] in tasks block {2}", blockWaitBarriers.front(),
                       blockWaitBarriers.back(), taskBlockIndex);
        }

        for (auto barrierIdx : blockWaitBarriers) {
            // find consumers to remove
            const auto barrierConsumers = getBarrierConsumers(barrierIdx);
            const auto consumersToRemove =
                    findRedundantDependencies(taskControlMap, controlMapOffset, barrierConsumers, false);
            removeConsumers(barrierIdx, consumersToRemove);
        }
    }
}

// Remove barrier dependence between producer and consumer
// 1) if a barrier only has OPs using same queue as its producer,
//    remove all OPs using the same queue from its consumers.
/*
  OP[{FIFO}][<opIndex>]
  OP[A][0] OP[A][1]    OP[A][0] OP[A][1]
       \   /               \   /
        Bar         =>      Bar
       /   \                 |
  OP[A][2] OP[B][0]       OP[B][0]
*/
// 2) if a barrier only has OPs using same queue as its consumer,
//    remove all OPs using the same queue from its producers.
/*
 OP[A][0] OP[B][0]        OP[B][0]
       \   /                 |
        Bar         =>      Bar
       /   \               /   \
 OP[A][1] OP[A][2]    OP[A][1] OP[A][2]
*/
void vpux::BarrierInfo::removeExplicitDependencies() {
    const auto findExplicitDependencies = [&](const BarrierInfo::TaskSet& dependencies,
                                              const VPURT::TaskQueueType& type) {
        BarrierInfo::TaskSet dependenciesToRemove;
        for (auto& taskInd : dependencies) {
            if (isSyncPoint(taskInd)) {
                continue;
            }
            if (type == VPURT::getTaskQueueType(getTaskOpAtIndex(taskInd), false)) {
                dependenciesToRemove.insert(taskInd);
            }
        }
        return dependenciesToRemove;
    };

    // Perform optimization in tasks blocks matching the distribution of synchronization points.
    for (size_t taskBlockIndex = 0; taskBlockIndex < getControlGraphBlockCount(); ++taskBlockIndex) {
        // get update barriers range for current block
        auto blockUpdateBarriers = getBarriersForTaskBlock(taskBlockIndex, /* blockStartSyncPoint */ true,
                                                           /* blockEndSyncPoint */ false, /* updateBarriers */ true);

        for (auto barrierIdx : blockUpdateBarriers) {
            // try to optimize consumers (1)
            const auto& barrierProducers = getBarrierProducers(barrierIdx);
            auto producerTaskQueueType = haveSameImplicitDependencyTaskQueueType(barrierProducers);
            if (producerTaskQueueType.has_value()) {
                // barrier produced by tasks with same type
                auto consumersToRemove =
                        findExplicitDependencies(getBarrierConsumers(barrierIdx), producerTaskQueueType.value());
                // remove consumers
                removeConsumers(barrierIdx, consumersToRemove);
            }

            // try to optimize producers (2)
            const auto& barrierConsumers = getBarrierConsumers(barrierIdx);
            auto consumerTaskQueueType = haveSameImplicitDependencyTaskQueueType(barrierConsumers);
            if (consumerTaskQueueType.has_value() || barrierConsumers.empty()) {
                // barrier consumed by tasks with same type
                BarrierInfo::TaskSet producersToRemove;
                // find producers to remove
                if (barrierConsumers.empty()) {
                    for (const auto& barProd : barrierProducers) {
                        if (isSyncPoint(barProd)) {
                            continue;
                        }
                        producersToRemove.insert(barProd);
                    }
                } else {
                    producersToRemove = findExplicitDependencies(barrierProducers, consumerTaskQueueType.value());
                }

                // remove producers
                removeProducers(barrierIdx, producersToRemove);
            }

            bool invalidOptimization =
                    getBarrierConsumers(barrierIdx).empty() ^ getBarrierProducers(barrierIdx).empty();
            VPUX_THROW_WHEN(invalidOptimization,
                            "Invalid optimization : Only barrier {0} became empty for barrier '{1}'",
                            barrierProducers.empty() ? "producers" : "consumers", getBarrierOpAtIndex(barrierIdx));
        }
    }
}

//
// getWaitBarriersMap
//

SmallVector<BarrierInfo::TaskSet> vpux::BarrierInfo::getWaitBarriersMap() {
    return _taskWaitBarriers;
}

BarrierInfoTest::BarrierInfoTest(mlir::func::FuncOp func): BarrierInfo(func) {
}

BarrierInfoTest::BarrierInfoTest(BarrierInfoMaps& barrierMaps) {
    initializeBarrierMaps(barrierMaps);
}

void BarrierInfoTest::initializeBarrierMaps(BarrierInfoMaps& barrierMaps) {
    BarrierInfo::_barrierProducerMap = toTaskSet(barrierMaps.barrierProducerMap);
    BarrierInfo::_barrierConsumerMap = toTaskSet(barrierMaps.barrierConsumerMap);
    BarrierInfo::_taskUpdateBarriers = toTaskSet(barrierMaps.taskUpdateBarriers);
    BarrierInfo::_taskWaitBarriers = toTaskSet(barrierMaps.taskWaitBarriers);
    BarrierInfo::_allBarrierOps.resize(barrierMaps.barrierProducerMap.size());
    BarrierInfo::_allTaskOps.resize(barrierMaps.taskUpdateBarriers.size());
    BarrierInfo::_syncTasksIds = barrierMaps.syncTasksIds;
    if (!BarrierInfo::_syncTasksIds.empty()) {
        buildTaskBlockMap();
    }
    setTaskQueueTypeMap(barrierMaps.taskQueueTypeMap);
}

void vpux::BarrierInfoTest::setTaskQueueTypeMap(
        const std::map<VPURT::TaskQueueType, SmallVector<uint32_t>>& taskQueueMaps) {
    llvm::BitVector taskList(checked_cast<uint32_t>(getNumOfTasks()), false);
    for (const auto& [taskQueueType, taskIndVec] : taskQueueMaps) {
        _taskQueueTypeMap.insert(std::make_pair(taskQueueType, taskList));
        for (const auto taskInd : taskIndVec) {
            VPUX_THROW_UNLESS(taskInd < getNumOfTasks(), "Invalid task index '{0}', limit is '{1}'", taskInd,
                              getNumOfTasks());
            _taskQueueTypeMap[taskQueueType].set(taskInd);
        }
    }
}

void vpux::BarrierInfoTest::setMaxVariantCountPerBarrier(size_t variantCount) {
    _maxVariantCountPerBarrier = variantCount;
}

size_t vpux::BarrierInfoTest::getNumOfTasks() const {
    return std::max(_allTaskOps.size(),
                    _taskUpdateBarriers.size());  // BarrierInfoTest class may operate without input funcOp
}

size_t vpux::BarrierInfoTest::getBarrierMaxVariantSum() const {
    return _maxVariantCountPerBarrier;
}

size_t vpux::BarrierInfoTest::getNumOfSlotsUsedByTask([[maybe_unused]] VPURT::TaskOp op) const {
    return 1;  // for tests, assume single slot per task (as for eg. config::ExecutorKind::DMA_NN)
}

VPURT::TaskOp vpux::BarrierInfoTest::getTaskOpAtIndex([[maybe_unused]] size_t opIdx) const {
    if (_func != nullptr) {
        return BarrierInfo::getTaskOpAtIndex(opIdx);
    }
    return nullptr;
}

BarrierInfoMaps vpux::BarrierInfoTest::getBarrierMaps() {
    BarrierInfoMaps barrierMaps;
    barrierMaps.taskWaitBarriers = toTaskVec(BarrierInfo::_taskWaitBarriers);
    barrierMaps.taskUpdateBarriers = toTaskVec(BarrierInfo::_taskUpdateBarriers);
    barrierMaps.barrierConsumerMap = toTaskVec(BarrierInfo::_barrierConsumerMap);
    barrierMaps.barrierProducerMap = toTaskVec(BarrierInfo::_barrierProducerMap);
    return barrierMaps;
}

BarrierInfoMaps vpux::BarrierInfoTest::optimizeBarrierProducers(size_t blockIdx) {
    BarrierInfo::optimizeBarrierProducers(blockIdx);
    return getBarrierMaps();
}

BarrierInfoMaps vpux::BarrierInfoTest::optimizeBarrierConsumers(size_t blockIdx) {
    BarrierInfo::optimizeBarrierConsumers(blockIdx);
    return getBarrierMaps();
}

BarrierInfoMaps vpux::BarrierInfoTest::optimizeBarriersWithSameProducers(size_t blockIdx, bool checkValidSlotCount) {
    BarrierInfo::optimizeBarriersWithSameProducers(blockIdx, checkValidSlotCount);
    return getBarrierMaps();
}

BarrierInfoMaps vpux::BarrierInfoTest::optimizeBarriers(bool checkValidSlotCount, bool considerTaskFifoDependency) {
    BarrierInfo::optimizeBarriers(checkValidSlotCount, considerTaskFifoDependency);
    return getBarrierMaps();
}

BarrierInfoMaps vpux::getBarrierMaps(BarrierInfo& barrierInfo) {
    BarrierInfoMaps barrierMaps;
    barrierMaps.syncTasksIds = barrierInfo._syncTasksIds;
    barrierMaps.taskWaitBarriers = toTaskVec(barrierInfo._taskWaitBarriers);
    barrierMaps.taskUpdateBarriers = toTaskVec(barrierInfo._taskUpdateBarriers);
    barrierMaps.barrierConsumerMap = toTaskVec(barrierInfo._barrierConsumerMap);
    barrierMaps.barrierProducerMap = toTaskVec(barrierInfo._barrierProducerMap);
    return barrierMaps;
}

SmallVector<BarrierInfo::TaskSet> vpux::toTaskSet(const SmallVector<SmallVector<size_t>>& map) {
    SmallVector<BarrierInfo::TaskSet> convertedMap(map.size());
    for (auto deps : map | indexed) {
        for (auto dep : deps.value()) {
            convertedMap[deps.index()].insert(dep);
        }
    }
    return convertedMap;
}

BarrierMap vpux::toTaskVec(const SmallVector<BarrierInfo::TaskSet>& map) {
    BarrierMap convertedMap(map.size());
    for (auto deps : map | indexed) {
        if (deps.value().size() == 0) {
            continue;
        }
        convertedMap[deps.index()].reserve(deps.value().size());
        for (auto dep : deps.value()) {
            convertedMap[deps.index()].push_back(dep);
        }
        llvm::sort(convertedMap[deps.index()]);
    }
    return convertedMap;
}

/**
 * Calculate barrier producers/consumers map from task update/wait barriers map
 */
void vpux::fillProducersAndConsumers(BarrierInfoMaps& barrierMaps) {
    barrierMaps.barrierProducerMap.clear();
    barrierMaps.barrierConsumerMap.clear();

    // calculate barrier maximal ID
    llvm::DenseSet<size_t> barriers;
    for (auto taskBarriers : barrierMaps.taskUpdateBarriers) {
        barriers.insert(taskBarriers.begin(), taskBarriers.end());
    }
    for (auto taskBarriers : barrierMaps.taskWaitBarriers) {
        barriers.insert(taskBarriers.begin(), taskBarriers.end());
    }
    if (barriers.empty()) {
        return;
    }
    size_t maxBarrierID = *std::max_element(barriers.begin(), barriers.end());

    auto fillBarrierMap = [&](const BarrierMap& barriersPerTaskMap, BarrierMap& tasksPerBarrierMap) {
        // clear producers / consumers
        tasksPerBarrierMap.resize(maxBarrierID + 1);
        for (auto& tasks : tasksPerBarrierMap) {
            tasks.clear();
        }

        // fill-in tasks for each barrier
        for (auto taskBarriers : barriersPerTaskMap | indexed) {
            for (auto barrier : taskBarriers.value()) {
                tasksPerBarrierMap[barrier].push_back(taskBarriers.index());
            }
        }

        // sort maps
        for (auto& tasks : tasksPerBarrierMap) {
            llvm::sort(tasks);
        }
    };

    fillBarrierMap(barrierMaps.taskUpdateBarriers, barrierMaps.barrierProducerMap);
    fillBarrierMap(barrierMaps.taskWaitBarriers, barrierMaps.barrierConsumerMap);
}
