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

#include "vpux/utils/logger/logger.hpp"

#include "vpux/compiler/core/aliases_info.hpp"
#include "vpux/compiler/core/async_deps_info.hpp"
#include "vpux/compiler/core/feasible_memory_scheduler_spilling.hpp"
#include "vpux/compiler/core/linear_scan_handler.hpp"
#include "vpux/compiler/dialect/VPU/utils/cost_model/cost_model.hpp"
#include "vpux/compiler/dialect/VPUIP/IR/ops.hpp"
#include "vpux/compiler/dialect/config/IR/resources.hpp"
#include "vpux/compiler/utils/hw_settings.hpp"

#include "common/utils.hpp"

#include <mlir/Dialect/Async/IR/Async.h>
#include <mlir/Dialect/MemRef/IR/MemRef.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/MLIRContext.h>
#include <mlir/IR/Visitors.h>
#include <mlir/Parser/Parser.h>

#include <gtest/gtest.h>

// Run cmd: npuUnitTests --gtest_filter="MLIR_FeasibleMemorySchedulerSpilling.*"

using namespace vpux;

using MLIR_FeasibleMemorySchedulerSpilling = MLIR_UnitBase;

TEST_F(MLIR_FeasibleMemorySchedulerSpilling, RemoveComputeOpRelocationSpillsForInplace) {
    mlir::MLIRContext ctx(registry);

    // Create a simple IR with 2 CopyOps and 1 NCEOp which is an inplace
    // eltwise that uses input1 also as an output buffer
    constexpr llvm::StringLiteral inputIR = R"(
        #NCHW = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>

        !Type_DDR = memref<1x32x32x32xf16, #NCHW, @DDR>
        !Type_CMX = memref<1x32x32x32xf16, #NCHW, [@CMX_NN, 0]>

        module @test {
            func.func @main(%arg0: !Type_DDR, %arg1: !Type_DDR) -> !Type_CMX {

                %buf_cmx_1 = memref.alloc() : !Type_CMX
                %buf_cmx_2 = memref.alloc() : !Type_CMX

                %t0, %r0 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg0 : !Type_DDR) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg1 : !Type_DDR) outputs(%buf_cmx_2 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t2, %r2 = async.execute [%t0, %t1] (%r0 as %arg2: !async.value<!Type_CMX>, %r1 as %arg3: !async.value<!Type_CMX>) -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64} <{is_inplace = true, task_type = #VPUIP.nce_task_type<ELTWISE>}> input(%arg2 : !Type_CMX) weights(%arg3 : !Type_CMX) parent_input(%arg2 : !Type_CMX) parent_output(%buf_cmx_1 : !Type_CMX) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [31, 31, 31], outStart = [0, 0, 0], pad = #VPU.Padding<bottom = 0 : i64, left = 0 : i64, right = 0 : i64, top = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %1: !Type_CMX
                }

                %r = async.await %r2 : !async.value<!Type_CMX>
                return %r : !Type_CMX
            }
        }
    )";

    auto module = mlir::parseSourceString<mlir::ModuleOp>(inputIR, &ctx);
    ASSERT_TRUE(module.get() != nullptr);

    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    auto log = vpux::Logger::global();

    // Run helper utilities used by scheduler which allow it to gather data
    // represented in the IR
    AliasesInfo aliasesInfo{func};
    AsyncDepsInfo depsInfo{func};

    // Assume following buffers confgiuration of an eltwise inplace
    // input1: [    0 -  65536]
    // input2: [65536 - 131072]
    // output: [    0 -  65536] <- same buffer as input1, this buffer is spilled
    const size_t bufSize = 65536;
    const size_t buf1SpilledOffset = 0;
    const size_t buf2Offset = 65536;
    const int64_t newOutputOffset = 131072;

    VPUIP::NCEClusterTaskOp nceOp = nullptr;
    func.walk([&](mlir::Operation* op) {
        if (mlir::isa<VPUIP::NCEClusterTaskOp>(op)) {
            nceOp = mlir::cast<VPUIP::NCEClusterTaskOp>(op);
        }
    });
    ASSERT_TRUE(nceOp != nullptr);

    auto buf1Spilled = aliasesInfo.getRoot(nceOp.getInput());
    auto buf2 = aliasesInfo.getRoot(nceOp.getWeights());

    // Required components needed to construct FeasibleMemorySchedulerSpilling object
    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    LinearScan<mlir::Value, LinearScanHandler> scan(1024 * 1024, {}, alignment);

    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);

    // Initialize example scheduled ops structure that could have been produced during scheduling
    // This is a simple IR but by manually preparing it it is possible to create a scenario where there
    // is a case for performing relocation spilling optimization for inplace op
    // Below scheduledOps structure has following operations
    //  opId  info
    //  0     copyOp
    //  1     copyOp
    //  2     eltwise inplace
    //  3     spill write
    //  4     spill read
    // spill write/read pair can be optimized

    FeasibleMemoryScheduler::ScheduledOpInfoVec scheduledOps = {
            {0,                                                                 // opId
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,                     // opType
             0,                                                                 // cycleBegin
             1,                                                                 // cycleEnd
             {},                                                                // inputResources
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}}},  // outputResources
            {1, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 1, 2, {}, {{buf2Offset, buf2Offset + bufSize, buf2}}},
            {2,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             2,
             3,
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}, {buf2Offset, buf2Offset + bufSize, buf2}},
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}}},
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             3,
             4,
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}},
             {}},
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP,
             4,
             5,
             {},
             {{newOutputOffset, newOutputOffset + bufSize, buf1Spilled}}}};

    // Perform spilling relocation optimization. After it scheduledOps size should be 3 as spill
    // operations should no longer be present
    spilling.removeComputeOpRelocationSpills(scheduledOps);

    EXPECT_EQ(scheduledOps.size(), 3);
    // nceOp has index 2
    EXPECT_EQ(scheduledOps[2].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);

    // Check if output of nceOp is no longer the same as first input
    EXPECT_NE(scheduledOps[2].beginOutputResource(0), scheduledOps[2].beginInputResource(0));
    EXPECT_NE(scheduledOps[2].endOutputResource(0), scheduledOps[2].endInputResource(0));

    // Verify output of nceOp now points to location previosuly assigned for spill-read op
    EXPECT_EQ(scheduledOps[2].beginOutputResource(0), newOutputOffset);
    EXPECT_EQ(scheduledOps[2].endOutputResource(0), newOutputOffset + bufSize);

    // Verify inPlace attribute is no longer present in nceOp
    EXPECT_FALSE(nceOp.getIsInplace().has_value());

    // Check both inputs and output buffers are allocated by different allocation operations
    auto input1AllocOp = scheduledOps[2].getInputBuffer(0).getDefiningOp<mlir::memref::AllocOp>();
    auto input2AllocOp = scheduledOps[2].getInputBuffer(1).getDefiningOp<mlir::memref::AllocOp>();
    auto outputAllocOp = scheduledOps[2].getOutputBuffer(0).getDefiningOp<mlir::memref::AllocOp>();

    ASSERT_TRUE(input1AllocOp != nullptr);
    ASSERT_TRUE(input2AllocOp != nullptr);
    ASSERT_TRUE(outputAllocOp != nullptr);
    EXPECT_NE(input1AllocOp, input2AllocOp);
    EXPECT_NE(input1AllocOp, outputAllocOp);
    EXPECT_NE(input2AllocOp, outputAllocOp);
}

TEST_F(MLIR_FeasibleMemorySchedulerSpilling, RemoveComputeOpRelocationSpillsForRemovedInplaceAttr) {
    constexpr llvm::StringLiteral inputIR = R"(
        #NHWC = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>
        #NCHW = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>

        !Conv_In_DDR = memref<1x128x28x28xf16, #NHWC, @DDR>
        !Conv_In_CMX = memref<1x128x28x28xf16, #NHWC, [@CMX_NN, 0]>
        !Conv_WT_CMX = memref<512x1x1x4xsi32, #NCHW, [@CMX_NN, 0]>
        !Conv_Weights_CMX = memref<512x128x1x1xf16, #NHWC, [@CMX_NN, 0]>
        !Conv_Out_CMX = memref<1x512x28x28xf16, #NHWC, [@CMX_NN, 0]>

        !EltWise_In_DDR =  memref<1x512x28x28xf16, #NHWC, @DDR>
        !EltWise_CMX =  memref<1x512x28x28xf16, #NHWC, [@CMX_NN, 0]>

        !DWConv_In_CMX = memref<1x512x28x28xf16, #NHWC, [@CMX_NN, 0]>
        !DWConv_OUT_CMX = memref<1x512x28x28xf16, #NHWC, [@CMX_NN, 0]>
        !DWConv_WT_CMX = memref<512x1x1x4xsi32, #NCHW, [@CMX_NN, 0]>
        !DWConv_Weights_CMX = memref<512x16x1x1xf16, #NHWC, [@CMX_NN, 0]>

        module @test {
            func.func @main(%arg0: !Conv_In_DDR, %arg1: !Conv_Weights_CMX, %arg2: !Conv_WT_CMX, %arg3: !EltWise_In_DDR, %arg4: !DWConv_Weights_CMX, %arg5: !DWConv_WT_CMX) -> !DWConv_OUT_CMX {
                %conv1_in_cmx = memref.alloc() : !Conv_In_CMX
                %conv1_out_cmx = memref.alloc() : !Conv_Out_CMX

                %eltwise_in_cmx = memref.alloc() : !EltWise_CMX
                %dwconv_out_cmx = memref.alloc() : !DWConv_OUT_CMX

                %t0, %r0 = async.execute -> !async.value<!Conv_In_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.Copy inputs(%arg0 : !Conv_In_DDR) outputs(%conv1_in_cmx : !Conv_In_CMX) -> !Conv_In_CMX
                    async.yield %270 : !Conv_In_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!EltWise_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.Copy inputs(%arg3 : !EltWise_In_DDR) outputs(%eltwise_in_cmx : !EltWise_CMX) -> !EltWise_CMX
                    async.yield %270 : !EltWise_CMX
                }

                %t2, %r2 = async.execute [%t0] (%r0 as %arg8: !async.value<!Conv_In_CMX>) -> !async.value<!Conv_Out_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 11735 : i64, constantsFused = true} <{kernel_padding = #VPU.Padding<left = 0 : i64, right = 0 : i64, top = 0 : i64, bottom = 0 : i64>, kernel_size = [1, 1], kernel_strides = [1, 1], mpe_engine = #VPU.MPEEngine37XX<mode = <SCL>>, task_type = #VPUIP.nce_task_type<CONV>}>
                        input(%arg8 : !Conv_In_CMX) weights(%arg1 : !Conv_Weights_CMX)  weight_table(%arg2 : !Conv_WT_CMX) parent_input(%arg8 : !Conv_In_CMX) parent_output(%conv1_out_cmx : !Conv_Out_CMX) outputs(%conv1_out_cmx : !Conv_Out_CMX) -> !Conv_Out_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, inEnd = [27, 27, 127], inStart = [0, 0, 0], mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [27, 27, 511], outStart = [0, 0, 0], pad = #VPU.Padding<left = 0 : i64, right = 0 : i64, top = 0 : i64, bottom = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %270: !Conv_Out_CMX
                }

                %t4, %r4 = async.execute [%t1, %t2] (%r1 as %arg8: !async.value<!EltWise_CMX>, %r2 as %arg9: !async.value<!Conv_Out_CMX>) -> !async.value<!EltWise_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64} <{is_inplace = true, task_type = #VPUIP.nce_task_type<ELTWISE>}>
                        input(%arg9 : !Conv_Out_CMX) weights(%arg8 : !EltWise_CMX) parent_input(%arg9 : !Conv_Out_CMX) parent_output(%conv1_out_cmx : !Conv_Out_CMX) outputs(%conv1_out_cmx : !Conv_Out_CMX) -> !EltWise_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, inEnd = [27, 27, 511], inStart = [0, 0, 0], mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [27, 27, 511], outStart = [0, 0, 0], pad = #VPU.Padding<left = 0 : i64, right = 0 : i64, top = 0 : i64, bottom = 0 : i64>}
                            } PPE : {
                                PPETask {ppe = #VPU.PPEStub<>}
                            }

                    async.yield %270: !EltWise_CMX
                }

                %t5, %r5 = async.execute [%t4] (%r4 as %arg8: !async.value<!EltWise_CMX>) -> !async.value<!DWConv_OUT_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                        %270 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 8240 : i64, constantsFused = true} <{kernel_padding = #VPU.Padding<left = 0 : i64, right = 0 : i64, top = 0 : i64, bottom = 0 : i64>, kernel_size = [1, 1], kernel_strides = [1, 1], mpe_engine = #VPU.MPEEngine37XX<mode = <SCL>>, task_type = #VPUIP.nce_task_type<DWCONV>}>
                        input(%arg8 : !EltWise_CMX) weights(%arg4 : !DWConv_Weights_CMX)  weight_table(%arg5 : !DWConv_WT_CMX) parent_input(%arg8 : !EltWise_CMX) parent_output(%dwconv_out_cmx : !DWConv_OUT_CMX) outputs(%dwconv_out_cmx : !DWConv_OUT_CMX) -> !DWConv_OUT_CMX  variants : {
                        DPUTask {cluster_id = 0 : i64, inEnd = [27, 9, 63], inStart = [0, 0, 0], mpe_mode = #VPU.mpe_mode<CUBOID_16x16>, outEnd = [27, 9, 63], outStart = [0, 0, 0], pad = #VPU.Padding<left = 0 : i64, right = 0 : i64, top = 0 : i64, bottom = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }

                        async.yield %270 : !DWConv_OUT_CMX
                }

                %r = async.await %r5 : !async.value<!DWConv_OUT_CMX>
                return %r : !DWConv_OUT_CMX
            }
        }
    )";

    mlir::MLIRContext ctx(registry);
    auto module = mlir::parseSourceString<mlir::ModuleOp>(inputIR, &ctx);
    ASSERT_TRUE(module.get() != nullptr);

    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    AliasesInfo aliasesInfo{func};
    AsyncDepsInfo depsInfo{func};

    auto log = vpux::Logger::global();

    VPUIP::NCEClusterTaskOp convOp = nullptr;
    VPUIP::NCEClusterTaskOp eltwiseOp = nullptr;
    VPUIP::NCEClusterTaskOp dwConvOp = nullptr;

    func.walk([&](mlir::Operation* op) {
        if (auto nceOp = mlir::dyn_cast<VPUIP::NCEClusterTaskOp>(op)) {
            if (nceOp.getTaskType() == VPUIP::NCETaskType::ELTWISE) {
                eltwiseOp = nceOp;
            } else if (nceOp.getTaskType() == VPUIP::NCETaskType::CONV) {
                convOp = nceOp;
            } else if (nceOp.getTaskType() == VPUIP::NCETaskType::DWCONV) {
                dwConvOp = nceOp;
            }
        }
    });

    ASSERT_TRUE(convOp != nullptr);
    ASSERT_TRUE(eltwiseOp != nullptr);
    ASSERT_TRUE(dwConvOp != nullptr);

    auto convInputBuff = aliasesInfo.getRoot(convOp.getInput());
    auto convOutputBuff = convOp.getOutputBuff();
    auto eltwiseInputBuff = aliasesInfo.getRoot(eltwiseOp.getInput());
    auto eltwiseWeightsBuff = aliasesInfo.getRoot(eltwiseOp.getWeights());
    auto eltwiseOutputBuff = convOp.getOutputBuff();
    auto dwConvInputBuff = aliasesInfo.getRoot(dwConvOp.getInput());
    auto dwConvOpOutputBuff = dwConvOp.getOutputBuff();

    ASSERT_TRUE(convOutputBuff == eltwiseInputBuff);
    ASSERT_TRUE(eltwiseInputBuff == eltwiseOutputBuff);
    ASSERT_TRUE(eltwiseOutputBuff == dwConvInputBuff);

    // Required components needed to construct FeasibleMemorySchedulerSpilling object
    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    LinearScan<mlir::Value, LinearScanHandler> scan(1024 * 1024, {}, alignment);

    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);

    //
    //  opId  info
    //  0     copyOp
    //  1     copyOp
    //  2     Conv
    //  3     Eltwise
    //  4     Eltwise spill write
    //  5     Eltwise spill read
    //  6     DWConv
    //
    FeasibleMemoryScheduler::ScheduledOpInfoVec scheduledOps = {
            {0, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 1, 2, {}, {{338112, 538816, convInputBuff}}},
            {1, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 2, 3, {}, {{1182016, 1468736, eltwiseWeightsBuff}}},
            {2,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             3,
             4,
             {{338112, 538816, convInputBuff}},
             {{538816, 825536, convOutputBuff}}},  // Conv
            {3,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             4,
             5,
             {{538816, 825536, eltwiseInputBuff}, {1182016, 1468736, eltwiseWeightsBuff}},
             {{538816, 825536, eltwiseOutputBuff}}},  // Eltwise

            {3,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             5,
             6,
             {{538816, 825536, eltwiseOutputBuff}},
             {}},

            {3, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP, 6, 7, {}, {{0, 286720, dwConvInputBuff}}},
            {4,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             7,
             8,
             {{0, 286720, dwConvInputBuff}},
             {{286720, 1089536, dwConvOpOutputBuff}}},  // DWConv
    };

    // The spill buffer is allocated-deallocated-allocated due to splling and the offset is determined by the last
    // allocation.
    scan.handler().setAddress(dwConvInputBuff, 0);
    spilling.removeComputeOpRelocationSpills(scheduledOps);

    EXPECT_EQ(scan.handler().getAddress(convOutputBuff), 538816);
    EXPECT_EQ(scan.handler().getAddress(eltwiseInputBuff), 538816);
    EXPECT_EQ(scan.handler().getAddress(eltwiseOutputBuff), 538816);

    // The scheduled ops after handled by `removeComputeOpRelocationSpills`
    //
    // op = '0'  type = 'ORIGINAL'  inputs = '<none>' outputs = '[338112 538816]'
    // op = '1'  type = 'ORIGINAL'  inputs = '<none>' outputs = '[1182016 1468736]'
    // op = '2'  type = 'ORIGINAL'  inputs = '[338112 538816]' outputs = '[538816 825536]'
    // op = '3'  type = 'ORIGINAL'  inputs = '[538816 825536], [1182016 1468736]' outputs = '[0 286720]'
    // op = '4'  type = 'ORIGINAL'  inputs = '[0 286720]' outputs = '[286720 1089536]
    //
    EXPECT_EQ(scheduledOps.size(), 5);

    // Conv
    EXPECT_EQ(scheduledOps[2].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);

    // EltWise
    EXPECT_EQ(scheduledOps[3].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);

    // DWConv
    EXPECT_EQ(scheduledOps[4].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);

    EXPECT_EQ(scheduledOps[2].op_, 2);
    EXPECT_EQ(scheduledOps[3].op_, 3);
    EXPECT_EQ(scheduledOps[4].op_, 4);

    EXPECT_EQ(scheduledOps[0].beginOutputResource(0), scheduledOps[2].beginInputResource(0));
    EXPECT_EQ(scheduledOps[1].beginOutputResource(0), scheduledOps[3].beginInputResource(1));
    EXPECT_EQ(scheduledOps[2].beginOutputResource(0), scheduledOps[3].beginInputResource(0));
    EXPECT_EQ(scheduledOps[3].beginOutputResource(0), scheduledOps[4].beginInputResource(0));
}

TEST_F(MLIR_FeasibleMemorySchedulerSpilling, RemoveComputeOpRelocationSpillsForSharedOutputBuffer) {
    mlir::MLIRContext ctx(registry);

    //
    //   Add        ZeroCopy
    //    \           /
    //     UpsamplingDMA
    //          |
    //        SubView
    //          |
    //         Copy
    //
    constexpr llvm::StringLiteral inputIR = R"(
        #NHWC = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>

        !Type_Tensor = tensor<1x16x144x256xf16, {order = #NHWC}>

        !Type_DDR = memref<1x16x72x128xf16, #NHWC, @DDR>
        !Type_CMX = memref<1x16x72x128xf16, #NHWC, [@CMX_NN, 0]>

        !Type_DDR_upsampled = memref<1x16x144x256xf16, #NHWC, @DDR>
        !Type_CMX_upsampled = memref<1x16x144x256xf16, #NHWC, [@CMX_NN, 0]>

        !Type_CMX_subview = memref<1x16x72x256xf16, {order=#NHWC, strides = [589824, 1, 4096, 16]}, [@CMX_NN, 0]>
        !Type_CMX_out = memref<1x16x72x256xf16, #NHWC, [@CMX_NN, 0]>

        module @test {
            func.func @main(%arg0: !Type_DDR, %arg1: !Type_DDR) -> !Type_CMX_out {
                %zero = const.Declare !Type_DDR_upsampled = dense<0.000000e+00> : !Type_Tensor

                %buf_cmx_1 = memref.alloc() : !Type_CMX
                %buf_cmx_2 = memref.alloc() : !Type_CMX
                %buf_cmx_3 = memref.alloc() : !Type_CMX

                %buf_cmx_4 = memref.alloc() : !Type_CMX_upsampled
                %buf_cmx_out = memref.alloc() : !Type_CMX_out

                %t0, %r0 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.Copy inputs(%arg0 : !Type_DDR) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX
                    async.yield %270 : !Type_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.Copy inputs(%arg1 : !Type_DDR) outputs(%buf_cmx_2 : !Type_CMX) -> !Type_CMX
                    async.yield %270 : !Type_CMX
                }

                %t2, %r2 = async.execute [%t0, %t1] (%r0 as %arg2: !async.value<!Type_CMX>, %r1 as %arg3: !async.value<!Type_CMX>) -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64} <{is_inplace = true, task_type = #VPUIP.nce_task_type<ELTWISE>}> input(%arg2 : !Type_CMX) weights(%arg3 : !Type_CMX) parent_input(%arg2 : !Type_CMX) parent_output(%buf_cmx_3 : !Type_CMX) outputs(%buf_cmx_3 : !Type_CMX) -> !Type_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [15, 71, 127], outStart = [0, 0, 0], pad = #VPU.Padding<bottom = 0 : i64, left = 0 : i64, right = 0 : i64, top = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %270: !Type_CMX
                }

                %t3, %r3 = async.execute -> !async.value<!Type_CMX_upsampled> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %270 = VPUIP.Copy inputs(%zero : !Type_DDR_upsampled) outputs(%buf_cmx_4 : !Type_CMX_upsampled) -> !Type_CMX_upsampled
                    async.yield %270 : !Type_CMX_upsampled
                }

                %t4, %r4 = async.execute [%t2, %t3] (%r2 as %arg2: !async.value<!Type_CMX>)
                        -> !async.value<!Type_CMX_upsampled> attributes {VPUIP.executor = @DMA_NN} {
                        %270 = VPUIP.UpsamplingDMAOp <{port = 0 : i64, upsampling_factor = [1, 1, 2, 2]}> inputs(%arg2 : !Type_CMX) outputs(%buf_cmx_4 : !Type_CMX_upsampled) -> !Type_CMX_upsampled
                        async.yield %270 : !Type_CMX_upsampled
                }

                %t5, %r5 = async.execute [%t4] (%r4 as %arg3: !async.value<!Type_CMX_upsampled>) -> !async.value<!Type_CMX_out> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                   %269 = VPUIP.SubView %arg3 [0, 0, 0, 0] [1, 16, 72, 256] : !Type_CMX_upsampled to !Type_CMX_subview
                   %270 = VPUIP.Copy inputs(%269 : !Type_CMX_subview) outputs(%buf_cmx_out : !Type_CMX_out) -> !Type_CMX_out
                   async.yield %270 : !Type_CMX_out
                }
                %r = async.await %r5 : !async.value<!Type_CMX_out>
                return %r : !Type_CMX_out
            }
        }
    )";

    auto module = mlir::parseSourceString<mlir::ModuleOp>(inputIR, &ctx);
    ASSERT_TRUE(module.get() != nullptr);

    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    // Run helper utilities used by scheduler which allow it to gather data
    // represented in the IR
    AliasesInfo aliasesInfo{func};
    AsyncDepsInfo depsInfo{func};

    VPUIP::UpsamplingDMAOp upsamplingDMAOp = nullptr;
    VPUIP::CopyOp zeroCopyOp = nullptr;
    mlir::async::ExecuteOp execOpUpsampling = nullptr;
    mlir::async::ExecuteOp execOpZeroCopy = nullptr;
    func.walk([&](mlir::Operation* op) {
        if (mlir::isa<VPUIP::UpsamplingDMAOp>(op)) {
            upsamplingDMAOp = mlir::cast<VPUIP::UpsamplingDMAOp>(op);
            execOpUpsampling = upsamplingDMAOp->getParentOfType<mlir::async::ExecuteOp>();
        } else if (mlir::isa<VPUIP::CopyOp>(op)) {
            execOpZeroCopy = op->getParentOfType<mlir::async::ExecuteOp>();
            zeroCopyOp = mlir::cast<VPUIP::CopyOp>(op);
        }

        if (execOpUpsampling && execOpZeroCopy && execOpZeroCopy->isBeforeInBlock(execOpUpsampling)) {
            return mlir::WalkResult::interrupt();
        }
        return mlir::WalkResult::advance();
    });

    ASSERT_TRUE(upsamplingDMAOp != nullptr);
    ASSERT_TRUE(zeroCopyOp != nullptr);

    auto upsamplingOutputBuff = upsamplingDMAOp.getOutputBuff();
    auto rootsUpsamplingOutputBuff = aliasesInfo.getRoots(upsamplingDMAOp.getOutputBuff());
    ASSERT_TRUE(rootsUpsamplingOutputBuff.size() == 1);

    auto zeroCopyOutputBuff = zeroCopyOp.getOutputBuff();
    auto rootsZeroCopyOutputBuff = aliasesInfo.getRoots(zeroCopyOutputBuff);
    ASSERT_TRUE(rootsZeroCopyOutputBuff.size() == 1);

    auto rootBufferOfUpsampling = *rootsUpsamplingOutputBuff.begin();
    auto rootBufferOfZeroCopy = *rootsZeroCopyOutputBuff.begin();
    ASSERT_TRUE(rootBufferOfUpsampling == rootBufferOfZeroCopy);

    // Required components needed to construct FeasibleMemorySchedulerSpilling object
    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    LinearScan<mlir::Value, LinearScanHandler> scan(1982464, {}, alignment);

    auto log = vpux::Logger::global();
    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);

    // Initialize example scheduled ops structure that could have been produced during scheduling
    // This is a simple IR but by manually preparing it it is possible to create a scenario where there
    // is a case for performing relocation spilling optimization for ops which have shared output buffer.
    // Below scheduledOps structure has following operations
    //
    //  opId  info
    //  0     copyOp
    //  1     copyOp
    //  2     zero-copyOp
    //  3     zero-copy spill write
    //  4     eltwise
    //  5     zero-copy spill read
    //  6     upsamplingDMA
    //  7     upsamplingDMA spill write
    //  8     upsamplingDMA spill read
    //  9     copyOp
    //
    // upsamplingDMA spill ops should not be optimized
    //
    FeasibleMemoryScheduler::ScheduledOpInfoVec scheduledOps = {
            {0, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 1, 2, {}, {{1179648, 1474560}}},
            {1, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 2, 3, {}, {{1474560, 1769472}}},

            {2, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 3, 4, {}, {{0, 1179648, zeroCopyOutputBuff}}},
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             4,
             49178,
             {{0, 1179648, zeroCopyOutputBuff}},
             {}},

            {3,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             49178,
             49179,
             {{1179648, 1474560}, {1474560, 1769472}},
             {{0, 294912}}},
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP,
             49179,
             98353,
             {},
             {{294912, 1474560, zeroCopyOutputBuff}}},

            {4,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             98353,
             98354,
             {{0, 294912}},
             {{294912, 1474560, upsamplingOutputBuff}}},
            {4,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             98354,
             147528,
             {{294912, 1474560, upsamplingOutputBuff}},
             {}},
            {4,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP,
             147528,
             196702,
             {},
             {{589824, 1769472, upsamplingOutputBuff}}},
            {5, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 196702, 196703, {{589824, 1769472}}, {{0, 589824}}},
    };

    spilling.removeComputeOpRelocationSpills(scheduledOps);

    // spill operations of op-4 should still be present
    EXPECT_EQ(scheduledOps.size(), 10);

    // zero-copy op
    EXPECT_EQ(scheduledOps[5].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP);

    // upsamplingDMA op
    EXPECT_EQ(scheduledOps[6].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);
    EXPECT_EQ(scheduledOps[7].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP);
    EXPECT_EQ(scheduledOps[8].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP);

    EXPECT_EQ(scheduledOps[6].op_, 4);
    EXPECT_EQ(scheduledOps[7].op_, 4);
    EXPECT_EQ(scheduledOps[8].op_, 4);

    // zero-copy and upsamplingDMA have the same CMX range of output buffer
    EXPECT_EQ(scheduledOps[5].beginOutputResource(0), scheduledOps[6].beginOutputResource(0));
    EXPECT_EQ(scheduledOps[5].endOutputResource(0), scheduledOps[6].endOutputResource(0));
}

TEST_F(MLIR_FeasibleMemorySchedulerSpilling, OptimizeDataOpsSpillsForReRead) {
    constexpr llvm::StringLiteral inputIR = R"(
        #NCHW = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>

        !Type_DDR = memref<1x32x32x32xf16, #NCHW, @DDR>
        !Type_CMX = memref<1x32x32x32xf16, #NCHW, [@CMX_NN, 0]>

        module @test attributes {config.arch = #config.arch_kind<NPU50XX>, config.compilationMode = #config.compilation_mode<DefaultHW>} {
            config.ExecutorResource 2 of @DMA_NN
            config.Resources 3 of @NCE at 2.100000e+03 MHz {
                config.MemoryResource 1326182 bytes of @CMX_NN_FragmentationAware
                config.MemoryResource 1473536 bytes of @CMX_NN {config.bandwidth = 64 : i64, config.derateFactor = 1.000000e+00 : f64}
                config.ExecutorResource 2 of @SHAVE_ACT
                config.ExecutorResource 1 of @DPU
            }
            config.Resources 1 of @global {
                config.ExecutorResource 1 of @M2I
                config.ExecutorResource 2 of @DMA_NN
                config.MemoryResource 67108864000 bytes of @DDR {config.bandwidth = 64 : i64, config.derateFactor = 6.000000e-01 : f64}
            }
            func.func @main(%arg0: !Type_DDR, %arg1: !Type_DDR) -> !Type_CMX {

                %buf_cmx_1 = memref.alloc() : !Type_CMX
                %buf_cmx_2 = memref.alloc() : !Type_CMX

                %t0, %r0 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg0 : !Type_DDR) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg1 : !Type_DDR) outputs(%buf_cmx_2 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t2, %r2 = async.execute [%t0, %t1] (%r0 as %arg2: !async.value<!Type_CMX>, %r1 as %arg3: !async.value<!Type_CMX>)
                        -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64} <{is_inplace = true, task_type = #VPUIP.nce_task_type<ELTWISE>}>
                    input(%arg2 : !Type_CMX) weights(%arg3 : !Type_CMX) parent_input(%arg2 : !Type_CMX) parent_output(%buf_cmx_1 : !Type_CMX) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [31, 31, 31], outStart = [0, 0, 0], pad = #VPU.Padding<bottom = 0 : i64, left = 0 : i64, right = 0 : i64, top = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %1: !Type_CMX
                }

                %r = async.await %r2 : !async.value<!Type_CMX>
                return %r : !Type_CMX
            }
        }
    )";

    const auto arch = config::ArchKind::NPU50XX;
    VPU::initializeSingletons(registry, VPU::DeviceVersion{std::nullopt, arch});

    mlir::MLIRContext ctx(registry);
    ctx.loadDialect<vpux::VPU::VPUDialect>();

    auto module = mlir::parseSourceString<mlir::ModuleOp>((inputIR).str(), &ctx);
    ASSERT_TRUE(module.get() != nullptr);
    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    // Prepare variables
    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    auto aliasesInfo = AliasesInfoMemType<VPU::MemoryKind::CMX_NN>{func};
    AsyncDepsInfo depsInfo{func};
    auto liveRange = MemLiveRangeInfoMemType<VPU::MemoryKind::CMX_NN>(func, aliasesInfo);
    auto log = vpux::Logger::global().nest("feasible-memory-scheduler-spilling");
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    auto memKindAttr = mlir::SymbolRefAttr::get(&ctx, stringifyEnum(memKind));
    auto available = config::getAvailableMemory(*module, memKindAttr);
    const auto maxSize = available.size();
    auto reservedMemVec = config::getReservedMemOffsetAndSizeVec(*module, memKindAttr);
    LinearScan<mlir::Value, LinearScanHandler> scan(maxSize.count(), reservedMemVec, alignment);
    auto vpuDev = VPU::getVPUDeviceType(module->getOperation());
    auto costModel = VPU::CostModelConfig::createCostModel(&ctx);
    auto tileOp = config::getTileExecutor(*module);
    ASSERT_TRUE(tileOp != nullptr);
    auto tileCount = tileOp.getCount();
    auto dmaPorts = config::getAvailableExecutor(*module, config::ExecutorKind::DMA_NN);
    ASSERT_TRUE(dmaPorts != nullptr);
    auto dmaCount = dmaPorts.getCount();

    // init schedule
    FeasibleMemoryScheduler initSchedule(memKind, secondLvlMemKind, liveRange, depsInfo, log, scan, arch, vpuDev,
                                         costModel, tileCount, dmaCount, /*enableScheduleStatistics*/ false,
                                         /*optimizeFragmentation*/ true,
                                         /*activelySpillForPrefetching*/ false);
    auto scheduledOps = initSchedule.generateSchedule();
    // The initial scheduled ops
    // op = 0   cycles = 1 -> 2         inputs = <none>         outputs = [0 65536] size = 65536
    // op = 1   cycles = 1 -> 2         inputs = <none>         outputs = [65536 131072] size = 65536
    // op = 2   cycles = 2 -> 3         inputs = [0 65536] size = 65536, [65536 131072] size = 65536

    // Insert dummy spill ops to simulate re-read scenario
    // op = 0   cycles = 1 -> 2         inputs = <none>         outputs = [0 65536] size = 65536
    // op = 1   cycles = 1 -> 2         inputs = <none>         outputs = [65536 131072] size = 65536
    // op = 2   cycles = 2 -> 3         inputs = [0 65536] size = 65536, [65536 131072] size = 65536
    // op = 1   cycles = 3 -> 4         inputs = [0 65536] size = 65536, outputs = <none>
    // op = 1   cycles = 4 -> 5         inputs = <none>         outputs = [0 65536] size = 65536
    scheduledOps.emplace_back(1, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP, 3, 1342464, true);
    scheduledOps.emplace_back(1, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP, 4, 1342464, true);
    EXPECT_EQ(scheduledOps.size(), 5);
    scheduledOps[4].outputResourceInfo_ = scheduledOps[0].outputResourceInfo_;

    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);
    spilling.optimizeDataOpsSpills(scheduledOps);

    // spill_write should be removed
    EXPECT_EQ(scheduledOps.size(), 4);
    EXPECT_EQ(scheduledOps[0].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);             // copy op 1
    EXPECT_EQ(scheduledOps[1].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);             // copy op 2
    EXPECT_EQ(scheduledOps[2].opType_, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP);             // nce task
    EXPECT_EQ(scheduledOps[3].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP);  // reread copy

    spilling.insertSpillDmaOps(scheduledOps);
    // Count the number of VPUIP.Copy ops after inserting spill DMAs
    // Should have: 2 original copies + 1 reread copy = 3 total
    size_t numOfDMA = 0;
    func->walk([&](VPUIP::CopyOp /*copyOp*/) {
        numOfDMA++;
    });
    EXPECT_EQ(numOfDMA, 3);
}

TEST_F(MLIR_FeasibleMemorySchedulerSpilling, RemoveRedundantSpillWrite) {
    mlir::MLIRContext ctx(registry);

    // Create a simple IR with 2 CopyOps and 1 NCEOp
    constexpr llvm::StringLiteral inputIR = R"(
        #NCHW = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>

        !Type_DDR = memref<1x32x32x32xf16, #NCHW, @DDR>
        !Type_CMX = memref<1x32x32x32xf16, #NCHW, [@CMX_NN, 0]>

        module @test {
            func.func @main(%arg0: !Type_DDR, %arg1: !Type_DDR) -> !Type_CMX {

                %buf_cmx_1 = memref.alloc() : !Type_CMX
                %buf_cmx_2 = memref.alloc() : !Type_CMX
                %buf_cmx_3 = memref.alloc() : !Type_CMX

                %t0, %r0 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg0 : !Type_DDR) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.Copy inputs(%arg1 : !Type_DDR) outputs(%buf_cmx_2 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t2, %r2 = async.execute [%t0, %t1] (%r0 as %arg2: !async.value<!Type_CMX>, %r1 as %arg3: !async.value<!Type_CMX>) -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DPU, VPUIP.num_units = 1 : i64} {
                    %1 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64} <{task_type = #VPUIP.nce_task_type<ELTWISE>}> input(%arg2 : !Type_CMX) weights(%arg3 : !Type_CMX) parent_input(%arg2 : !Type_CMX) parent_output(%buf_cmx_3 : !Type_CMX) outputs(%buf_cmx_3 : !Type_CMX) -> !Type_CMX  variants : {
                            DPUTask {cluster_id = 0 : i64, mpe_mode = #VPU.mpe_mode<CUBOID_8x16>, outEnd = [31, 31, 31], outStart = [0, 0, 0], pad = #VPU.Padding<bottom = 0 : i64, left = 0 : i64, right = 0 : i64, top = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %1: !Type_CMX
                }

                %r = async.await %r2 : !async.value<!Type_CMX>
                return %r : !Type_CMX
            }
        }
    )";

    auto module = mlir::parseSourceString<mlir::ModuleOp>(inputIR, &ctx);
    ASSERT_TRUE(module.get() != nullptr);

    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    auto log = vpux::Logger::global();

    // Run helper utilities used by scheduler which allow it to gather data
    // represented in the IR
    AliasesInfo aliasesInfo{func};
    AsyncDepsInfo depsInfo{func};

    // Assume the following buffer configuration of an eltwise
    // input1: [     0 -  65536]
    // input2: [ 65536 - 131072]
    // output: [131072 - 196608]
    const size_t bufSize = 65536;
    const size_t buf1Offset = 0;
    const size_t buf2Offset = 65536;
    const size_t buf3Offset = 131072;

    VPUIP::NCEClusterTaskOp nceOp = nullptr;
    func.walk([&](mlir::Operation* op) {
        if (mlir::isa<VPUIP::NCEClusterTaskOp>(op)) {
            nceOp = mlir::cast<VPUIP::NCEClusterTaskOp>(op);
        }
    });
    ASSERT_TRUE(nceOp != nullptr);

    auto buf1 = aliasesInfo.getRoot(nceOp.getInput());
    auto buf2 = aliasesInfo.getRoot(nceOp.getWeights());
    auto buf3 = aliasesInfo.getRoot(nceOp.getOutputBuff());

    // Required components needed to construct FeasibleMemorySchedulerSpilling object
    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    LinearScan<mlir::Value, LinearScanHandler> scan(1024 * 1024, {}, alignment);

    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);

    // Initialize example scheduled ops structure that could have been produced during scheduling
    // This is a simple IR but by manually preparing it it is possible to create a scenario where there
    // is a case for performing redundant spill-write removal optimization
    // Below scheduledOps structure has following operations
    //  opId  info
    //  0     copyOp
    //  1     copyOp
    //  2     spill write
    //  3     eltwise
    // spill write can be optimized

    FeasibleMemoryScheduler::ScheduledOpInfoVec scheduledOps = {
            {0,                                              // opId
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,  // opType
             0,                                              // cycleBegin
             1,                                              // cycleEnd
             {},                                             // inputResources
             {{buf1Offset, buf1Offset + bufSize, buf1}}},    // outputResources
            {1, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 1, 2, {}, {{buf2Offset, buf2Offset + bufSize, buf2}}},
            {1,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             3,
             4,
             {{buf2Offset, buf2Offset + bufSize, buf2}},
             {}},
            {2,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             2,
             3,
             {{buf1Offset, buf1Offset + bufSize, buf1}, {buf2Offset, buf2Offset + bufSize, buf2}},
             {{buf3Offset, buf3Offset + bufSize, buf3}}}};

    // Perform redundant spill-write optimization. After it scheduledOps size should be 3 as
    // spill-write should no longer be present
    spilling.removeRedundantSpillWrites(scheduledOps);

    EXPECT_EQ(scheduledOps.size(), 3);

    // Verify offset for the previously redundant spill-write buffer
    EXPECT_EQ(scan.handler().getAddress(buf2), buf2Offset);
}

// Verifies that spill write removal optimization is skipped when there are multiple SPILL_READ operations
TEST_F(MLIR_FeasibleMemorySchedulerSpilling, SkipOptimizationForMultipleSpillReads) {
    mlir::MLIRContext ctx(registry);

    constexpr llvm::StringLiteral inputIR = R"(
        #NCHW = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>

        !Type_DDR = memref<1x32x32x32xf16, #NCHW, @DDR>
        !Type_CMX = memref<1x32x32x32xf16, #NCHW, [@CMX_NN, 0]>

        module @test {
            func.func @main(%arg0: !Type_DDR, %arg1: !Type_DDR) -> !Type_CMX {
                %buf_cmx_1 = memref.alloc() : !Type_CMX
                %buf_cmx_2 = memref.alloc() : !Type_CMX

                %t0, %r0 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN} {
                    %1 = VPUIP.Copy inputs(%arg0 : !Type_DDR) outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t1, %r1 = async.execute -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DMA_NN} {
                    %1 = VPUIP.Copy inputs(%arg1 : !Type_DDR) outputs(%buf_cmx_2 : !Type_CMX) -> !Type_CMX
                    async.yield %1 : !Type_CMX
                }

                %t2, %r2 = async.execute [%t0, %t1] (%r0 as %arg2: !async.value<!Type_CMX>, %r1 as %arg3: !async.value<!Type_CMX>)
                        -> !async.value<!Type_CMX> attributes {VPUIP.executor = @DPU} {
                    %1 = VPUIP.NCEClusterTask {minimumHardwareExecutionCost = 21125 : i64}
                        <{is_inplace = true, task_type = #VPUIP.nce_task_type<ELTWISE>}>
                        input(%arg2 : !Type_CMX) weights(%arg3 : !Type_CMX)
                        parent_input(%arg2 : !Type_CMX) parent_output(%buf_cmx_1 : !Type_CMX)
                        outputs(%buf_cmx_1 : !Type_CMX) -> !Type_CMX variants : {
                            DPUTask {cluster_id = 0 : i64, mpe_mode = #VPU.mpe_mode<CUBOID_8x16>,
                                    outEnd = [31, 31, 31], outStart = [0, 0, 0],
                                    pad = #VPU.Padding<bottom = 0 : i64, left = 0 : i64, right = 0 : i64, top = 0 : i64>}
                        } PPE : {
                            PPETask {ppe = #VPU.PPEStub<>}
                        }
                    async.yield %1: !Type_CMX
                }

                %r = async.await %r2 : !async.value<!Type_CMX>
                return %r : !Type_CMX
            }
        }
    )";

    auto module = mlir::parseSourceString<mlir::ModuleOp>(inputIR, &ctx);
    ASSERT_TRUE(module.get() != nullptr);

    auto func = module.get().lookupSymbol<mlir::func::FuncOp>("main");
    ASSERT_TRUE(func != nullptr);

    auto log = vpux::Logger::global();
    AliasesInfo aliasesInfo{func};
    AsyncDepsInfo depsInfo{func};

    const size_t bufSize = 65536;
    const size_t buf1SpilledOffset = 0;
    const size_t buf2Offset = 65536;
    const int64_t newOutputOffset = 131072;

    VPUIP::NCEClusterTaskOp nceOp = nullptr;
    func.walk([&](mlir::Operation* op) {
        if (mlir::isa<VPUIP::NCEClusterTaskOp>(op)) {
            nceOp = mlir::cast<VPUIP::NCEClusterTaskOp>(op);
        }
    });
    ASSERT_TRUE(nceOp != nullptr);

    auto buf1Spilled = aliasesInfo.getRoot(nceOp.getInput());
    auto buf2 = aliasesInfo.getRoot(nceOp.getWeights());

    const auto memKind = VPU::MemoryKind::CMX_NN;
    const auto secondLvlMemKind = VPU::MemoryKind::DDR;
    uint64_t alignment = vpux::DEFAULT_CMX_ALIGNMENT;
    LinearScan<mlir::Value, LinearScanHandler> scan(1024 * 1024, {}, alignment);

    FeasibleMemorySchedulerSpilling spilling(memKind, secondLvlMemKind, depsInfo, aliasesInfo, log, scan);

    // Scenario: One SPILL_WRITE followed by MULTIPLE SPILL_READs
    FeasibleMemoryScheduler::ScheduledOpInfoVec scheduledOps = {
            {0,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             0,
             1,
             {},
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}}},
            {1, FeasibleMemoryScheduler::EOpType::ORIGINAL_OP, 1, 2, {}, {{buf2Offset, buf2Offset + bufSize, buf2}}},
            {2,
             FeasibleMemoryScheduler::EOpType::ORIGINAL_OP,
             2,
             3,
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}, {buf2Offset, buf2Offset + bufSize, buf2}},
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}}},
            // One SPILL_WRITE
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP,
             3,
             4,
             {{buf1SpilledOffset, buf1SpilledOffset + bufSize, buf1Spilled}},
             {}},
            // First SPILL_READ
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP,
             4,
             5,
             {},
             {{newOutputOffset, newOutputOffset + bufSize, buf1Spilled}}},
            // Second SPILL_READ - this causes the optimization to be skipped
            {2,
             FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP,
             5,
             6,
             {},
             {{newOutputOffset + bufSize, newOutputOffset + 2 * bufSize, buf1Spilled}}}};

    size_t initialSize = scheduledOps.size();

    // Perform spilling relocation optimization
    spilling.removeComputeOpRelocationSpills(scheduledOps);

    // With multiple SPILL_READs, the optimization should be SKIPPED
    // Therefore, all operations should still be present
    EXPECT_EQ(scheduledOps.size(), initialSize);

    // Verify operations are still present
    EXPECT_EQ(scheduledOps[3].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_WRITE_OP);
    EXPECT_EQ(scheduledOps[4].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP);
    EXPECT_EQ(scheduledOps[5].opType_, FeasibleMemoryScheduler::EOpType::IMPLICIT_SPILL_READ_OP);

    // Verify inPlace attribute is still present (optimization was skipped)
    EXPECT_TRUE(nceOp.getIsInplace().has_value());
}
