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

#include "vpux/compiler/core/barrier_info.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/types.hpp"
#include "vpux/compiler/dialect/VPURT/IR/ops.hpp"
#include "vpux/compiler/dialect/VPURT/transforms/passes.hpp"
#include "vpux/compiler/dialect/VPURT/utils/barrier_legalization_utils.hpp"
#include "vpux/compiler/dialect/config/utils/config_option_utils.hpp"

namespace vpux::VPURT {
#define GEN_PASS_DECL_OPTIMIZEBARRIERSSLOTSUSAGE
#define GEN_PASS_DEF_OPTIMIZEBARRIERSSLOTSUSAGE
#include "vpux/compiler/dialect/VPURT/passes.hpp.inc"
}  // namespace vpux::VPURT

using namespace vpux;

namespace {

class OptimizeBarriersSlotsUsagePass final :
        public VPURT::impl::OptimizeBarriersSlotsUsageBase<OptimizeBarriersSlotsUsagePass> {
public:
    explicit OptimizeBarriersSlotsUsagePass(Logger log) {
        Base::initLogger(log, Base::getArgumentName());
    }

private:
    void safeRunOnFunc() final;
};

void OptimizeBarriersSlotsUsagePass::safeRunOnFunc() {
    auto func = getOperation();

    auto& barrierInfo = getAnalysis<BarrierInfo>();

    barrierInfo.buildTaskQueueTypeMap();

    barrierInfo.removeRedundantBarrierProducersAndConsumers(/* considerTaskFifoDependency */ true);

    barrierInfo.updateIR();
    barrierInfo.clearAttributes();

    VPUX_THROW_UNLESS(VPURT::verifyBarrierSlots(func, _log), "Barrier slot count check failed");
    VPUX_THROW_UNLESS(barrierInfo.verifyBarriersUsersCount(VPURT::countIndependentTaskExecutors(func)),
                      "Encountered unexpected number of barrier users.");
}
}  // namespace

//
// createOptimizeBarriersSlotsUsagePass
//

std::unique_ptr<mlir::Pass> vpux::VPURT::createOptimizeBarriersSlotsUsagePass(Logger log) {
    return std::make_unique<OptimizeBarriersSlotsUsagePass>(log);
}
