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

#include "vpux/compiler/NPU40XX/conversion.hpp"
#include "vpux/compiler/NPU50XX/conversion.hpp"
#include "vpux/compiler/conversion.hpp"
#include "vpux/compiler/dialect/ELF/transforms/passes.hpp"
#include "vpux/compiler/dialect/VPUIPDPU/passes.hpp"
#include "vpux/compiler/dialect/VPUMI40XX/passes.hpp"
#include "vpux/compiler/dialect/VPURT/transforms/passes.hpp"
#include "vpux/compiler/dialect/VPURegMapped/passes.hpp"

#include <npu_40xx_nnrt.hpp>

#include <mlir/Transforms/Passes.h>

using namespace vpux;

//
// buildLowerVPUIP2ELFPipeline
//

void vpux::arch50xx::buildLowerVPUIP2ELFPipeline(mlir::OpPassManager& pm,
                                                 const BackendCompilationOptions50XX& backendCompilationOptions,
                                                 Logger log) {
    log.info("BackendCompilationOptions:\n"
             "  workloadManagementEnable = {0}\n"
             "  workloadManagementMode = {1}\n"
             "  workloadManagementBarrierCountThreshold = {2}\n"
             "  workloadManagementBarrierProgrammingMode = {3}\n"
             "  workloadManagementDmaFifoType = {4}\n"
             "  enableMemorySideCache = {5}\n"
             "  enableDMAProfiling = {6}\n"
             "  enableShaveDDRAccessOptimization = {7}\n",
             backendCompilationOptions.workloadManagementEnable,
             stringifyEnum(backendCompilationOptions.workloadManagementMode),
             backendCompilationOptions.workloadManagementBarrierCountThreshold,
             stringifyEnum(backendCompilationOptions.workloadManagementBarrierProgrammingMode),
             stringifyEnum(backendCompilationOptions.workloadManagementDmaFifoType),
             backendCompilationOptions.enableMemorySideCache, backendCompilationOptions.enableDMAProfiling,
             backendCompilationOptions.enableShaveDDRAccessOptimization);

    pm.addPass(VPUMI40XX::createAddPlatformInfoPass(log));

    // Below VPURT WLM passes are placed in LowerVPUIP2ELF pipeline to leverage BarrierInfo
    // working on VPURT dialect and order barriers in a way suitable for WLM. These passes can be moved to the end of
    // VPURT.
    if (backendCompilationOptions.workloadManagementEnable &&
        backendCompilationOptions.workloadManagementMode != WorkloadManagementMode::FWLM_V1_PAGES &&
        backendCompilationOptions.workloadManagementMode != WorkloadManagementMode::PWLM_V0_1_PAGES) {
        if (backendCompilationOptions.workloadManagementMode > WorkloadManagementMode::PWLM_V0_LCA) {
            pm.addPass(VPURT::createFindWlmEnqueueBarrierPass(
                    backendCompilationOptions.workloadManagementMode,
                    backendCompilationOptions.workloadManagementDmaFifoType == DMAFifoType::HW, log));
        } else {
            pm.addPass(VPURT::createOrderBarriersForWlmPass(log));
        }
    }

    pm.addPass(createConvertVPUIP2VPUMI40XXPass(log, backendCompilationOptions.enableMemorySideCache,
                                                backendCompilationOptions.allocateShaveStackFrames));
    pm.addPass(VPUMI40XX::createSetupProfilingVPUMI40XXPass(backendCompilationOptions.enableDMAProfiling, log));
    pm.addPass(mlir::createCanonicalizerPass());
    pm.addPass(ELF::createAddABIVersionPass(log));
    arch40xx::elfSubsetPipelineVPUMI(pm, backendCompilationOptions.workloadManagementEnable,
                                     backendCompilationOptions.workloadManagementMode,
                                     backendCompilationOptions.enableDumpStatisticsOfWlmOps,
                                     backendCompilationOptions.workloadManagementBarrierProgrammingMode, log);

    pm.addPass(VPUMI40XX::createAddMappedInferenceVersionOpPass(
            log, VPU_NNRT_40XX_API_VER_MAJOR, VPU_NNRT_40XX_API_VER_MINOR, VPU_NNRT_40XX_API_VER_PATCH));

    arch40xx::elfSubsetPipelineVPUASM(pm, backendCompilationOptions.workloadManagementEnable,
                                      backendCompilationOptions.workloadManagementDmaFifoType == DMAFifoType::HW, log);

    pm.addPass(VPUIPDPU::createExpandDPUConfigPass(log, backendCompilationOptions.npu5PPEBackwardsCompatibilityMode));
    pm.addPass(ELF::createUpdateELFSectionFlagsPass(log, backendCompilationOptions.enableShaveDDRAccessOptimization));
    pm.addPass(createConvertVPUASM2NPUReg50XXPass(log, backendCompilationOptions.modelIdentifier));
    pm.addPass(createConvertVPUIPDPU2NPUReg50XXPass(log, backendCompilationOptions.npu5PPEBackwardsCompatibilityMode));

    pm.addPass(VPURegMapped::createDeduceDynamicMappedInferenceVersionPass(log));

    pm.addPass(ELF::createHandleAlignmentRequirementsPass(log));
    pm.addPass(ELF::createSetOpOffsetsPass(log));
    pm.addPass(ELF::createSetCMXSymbolValuePass(
            log, npu40xx::nn_public::VPU_WORKSPACE_ADDR, npu40xx::nn_public::VPU_WORKSPACE_SIZE,
            npu40xx::VPU_METADATA_STORAGE_START, npu40xx::nn_public::VPU_METADATA_SIZE));
    pm.addPass(ELF::createAddRelocationsForDynamicStridesDMAsPass(log));
    pm.addPass(ELF::createAddELFRelocationsPass(log));
    pm.addPass(ELF::createRemoveEmptyELFSectionsPass(log));
}

void vpux::arch50xx::registerConversionPipeline() {
    mlir::PassPipelineRegistration<BackendCompilationOptions50XX>(
            "lower-VPUIP-to-ELF", "Performs full lowering from the VPUIP Dialect to ELF for NPU50XX arch IR",
            [](mlir::OpPassManager& pm, const BackendCompilationOptions50XX& options) {
                vpux::arch50xx::buildLowerVPUIP2ELFPipeline(pm, options);
            });
}
