#include <cassert>
#include <cstdlib>
#include <string>
#include <vector>

#include "mlir/IR/BuiltinOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"

#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"

using namespace onnx_mlir::spatial;

static void testResidentWeightViewCost() {
  mlir::MLIRContext context;
  context.loadDialect<SpatialDialect, mlir::tensor::TensorDialect>();
  mlir::Block operands;
  mlir::OwningOpRef<mlir::ModuleOp> owner =
    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(&context);
  auto loc = owner->getLoc();
  auto weightType = mlir::RankedTensorType::get({128, 128}, builder.getF32Type());
  auto inputType = mlir::RankedTensorType::get({1, 128}, builder.getF32Type());
  auto weight = operands.addArgument(weightType, loc);
  auto input = operands.addArgument(inputType, loc);
  auto wideType = mlir::RankedTensorType::get({1, 256}, builder.getF32Type());
  auto wideInput = operands.addArgument(wideType, loc);
  auto makeCompute = [&](bool weightView, unsigned inputView) {
    builder.setInsertionPointToEnd(owner->getBody());
    auto compute = SpatGraphCompute::create(builder, loc,
      mlir::TypeRange{inputType}, mlir::ValueRange{weight},
      mlir::ValueRange{inputView == 2 ? wideInput : input});
    auto *block = new mlir::Block();
    compute.getBody().push_back(block);
    mlir::Value w = block->addArgument(weightType, loc);
    mlir::Value x = block->addArgument(inputView == 2 ? wideType : inputType, loc);
    builder.setInsertionPointToStart(block);
    auto slice = [&](mlir::Value value, mlir::RankedTensorType type, bool strided = false) {
      llvm::SmallVector<mlir::OpFoldResult> offsets(2, builder.getIndexAttr(0));
      llvm::SmallVector<mlir::OpFoldResult> strides(2, builder.getIndexAttr(1));
      if (strided)
        strides.back() = builder.getIndexAttr(2);
      llvm::SmallVector<mlir::OpFoldResult> sizes;
      for (int64_t size : type.getShape())
        sizes.push_back(builder.getIndexAttr(size));
      return mlir::tensor::ExtractSliceOp::create(
        builder, loc, type, value, offsets, sizes, strides).getResult();
    };
    if (weightView)
      w = slice(w, weightType);
    if (inputView)
      x = slice(x, inputType, inputView == 2);
    auto result = SpatVMMOp::create(builder, loc, inputType, w, x);
    SpatYieldOp::create(builder, loc, mlir::ValueRange{result.getResult()});
    return ComputeInstance{compute, 0, 1};
  };
  auto direct = makeCompute(false, false);
  auto weightView = makeCompute(true, false);
  auto inputView = makeCompute(false, true);
  auto stridedInput = makeCompute(false, 2);
  for (Cost readLatency : {1, 20}) {
    SchedulingTarget target;
    target.localMemoryReadLatencyCycles = readLatency;
    target.localMemoryWriteLatencyCycles = 3;
    target.localMemoryWidthBytes = 8;
    target.processorPeriodNs = 1;
    target.computeBitwidth = 8;
    assert(getComputeInstanceCost(weightView, target) == getComputeInstanceCost(direct, target));
    assert(getComputeInstanceCost(inputView, target) == getComputeInstanceCost(direct, target));
    assert(getComputeInstanceCost(stridedInput, target)
      == getComputeInstanceCost(direct, target) + 64 * readLatency);
  }
}

static void testPipelineStageSpan() {
  using namespace mlir;
  MLIRContext context;
  context.loadDialect<SpatialDialect>();
  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(owner->getBodyRegion());
  ComputeGraph graph;
  graph.successors.resize(12);
  graph.predecessors.resize(12);
  for (size_t group = 0; group < 6; ++group) {
    auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
      TypeRange{}, builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{});
    for (uint32_t lane = 0; lane < 2; ++lane) {
      size_t task = graph.nodes.size();
      ComputeInstance instance{batch, lane, 1};
      graph.nodes.push_back({instance, group == 0 ? 40ULL : 10ULL, {}, task});
      graph.instanceToIndex[instance] = task;
      if (group != 0)
        for (size_t predecessor = task - lane - 2; predecessor < task - lane; ++predecessor) {
          graph.predecessors[task].push_back({predecessor, {}});
          graph.successors[predecessor].push_back({task, {}});
        }
    }
  }
  for (size_t coreCount : {6, 12, 20}) {
    SchedulingTarget target;
    target.processorCount = coreCount;
    target.residentWeightCapacity = 1;
    target.interProcessorLatencyNs.assign(coreCount * coreCount, 0);
    auto schedule = runPeftScheduler(graph, {target, &context});
    std::string error;
    assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
    Cost stageCosts[2] = {};
    Cost stageWork[2] = {};
    for (size_t task = 0; task < graph.nodes.size(); ++task) {
      size_t core = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
      size_t stage = schedule.processorStages[core];
      stageWork[stage] += graph.nodes[task].cost;
      if (task % 2 == 0)
        stageCosts[stage] += graph.nodes[task].cost;
    }
    // Both 40/50 and 50/40 are optimal cuts for this 90-cycle chain.
    assert(std::max(stageCosts[0], stageCosts[1]) == 50);
    if (coreCount == 20)
      for (size_t stage = 0; stage < 2; ++stage)
        assert(stageWork[stage] <= 10 * static_cast<size_t>(
          llvm::count(schedule.processorStages, stage)));
  }
}

static void testPipelinePackingCapacity() {
  using namespace mlir;
  MLIRContext context;
  context.loadDialect<SpatialDialect>();
  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(owner->getBodyRegion());
  auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
    TypeRange{}, builder.getI32IntegerAttr(4), ValueRange{}, ValueRange{});
  ComputeGraph graph;
  graph.successors.resize(4);
  graph.predecessors.resize(4);
  for (uint32_t task = 0; task < 4; ++task) {
    ResidentWeightSet weights;
    for (uint32_t index = 0; index < (task == 3 ? 3u : 1u); ++index) {
      ResidentWeight weight;
      weight.opaqueLane = task == 3 ? index + 1 : 0;
      weights.push_back(weight);
    }
    ComputeInstance instance{batch, task, 1};
    graph.nodes.push_back({instance, 10, weights, task});
    graph.instanceToIndex[instance] = task;
  }
  SchedulingTarget target;
  target.processorCount = 4;
  target.residentWeightCapacity = 3;
  target.interProcessorLatencyNs.assign(16, 0);
  auto schedule = runPeftScheduler(graph, {target, &context});
  std::string error;
  assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
  std::vector<Cost> loads(4);
  for (const auto &node : graph.nodes) {
    size_t core = schedule.computeToCpuMap.lookup(node.instance);
    assert(schedule.processorStages[core] == 0);
    loads[core] += node.cost;
  }
  assert(*std::max_element(loads.begin(), loads.end()) == 20);
}

static void testPhysicalPeft() {
  TransferCost transfer {.fixed = 50, .networkFlits = 4};
  SchedulingTarget line;
  line.processorCount = 3;
  line.interProcessorLatencyNs = {
    0,
    1,
    10,
    1,
    0,
    1,
    10,
    1,
    0,
  };
  line.averageInterProcessorLatencyNs = 4;
  assert(getPeftTransferTime(transfer, 0, 1, line) == 54);
  assert(getPeftTransferTime(transfer, 0, 2, line) == 90);
  SchedulingTarget reversed = line;
  reversed.interProcessorLatencyNs = {0, 10, 1, 10, 0, 1, 1, 1, 0};
  assert(getPeftTransferTime(transfer, 0, 1, reversed) == 90);
  assert(getPeftTransferTime(transfer, 0, 2, reversed) == 54);

  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> owner =
    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  ComputeGraph chain;
  chain.successors.resize(3);
  chain.predecessors.resize(3);
  for (uint32_t task = 0; task < 3; ++task) {
    ComputeInstance instance {owner->getOperation(), task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    chain.nodes.push_back({instance, 10, {weight}, task});
    if (task != 0) {
      chain.successors[task - 1].push_back({task, transfer});
      chain.predecessors[task].push_back({task - 1, transfer});
    }
  }
  line.residentWeightCapacity = reversed.residentWeightCapacity = 1;
  for (const SchedulingTarget &target : {line, reversed}) {
    auto scheduled = runPeftScheduler(chain, {target, &context});
    size_t firstCore = scheduled.computeToCpuMap.lookup(chain.nodes[0].instance);
    assert(firstCore == (target.interProcessorLatencyNs == line.interProcessorLatencyNs ? 1 : 2));
    for (size_t task = 1; task < chain.nodes.size(); ++task) {
      auto previous = chain.nodes[task - 1].instance;
      auto current = chain.nodes[task].instance;
      size_t source = scheduled.computeToCpuMap.lookup(previous);
      size_t destination = scheduled.computeToCpuMap.lookup(current);
      assert(source != destination);
      assert(scheduled.computeToAestMap.lookup(current)
             == scheduled.computeToAestMap.lookup(previous) + 10
                  + getPeftTransferTime(transfer, source, destination, target));
    }
  }
}

static void testPeftReservationCapacity() {
  mlir::MLIRContext context;
  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  ComputeGraph graph;
  graph.successors.resize(4);
  graph.predecessors.resize(4);
  for (uint32_t task = 0; task < 4; ++task) {
    ResidentWeightSet weights;
    for (uint32_t index = 0; index < (task == 3 ? 4u : 1u); ++index) {
      ResidentWeight weight;
      weight.opaqueLane = task == 3 ? index + 1 : 0;
      weights.push_back(weight);
    }
    graph.nodes.push_back({{owner, task, 1}, 10, weights, task});
  }
  SchedulingTarget target;
  target.processorCount = 3;
  target.residentWeightCapacity = 4;
  target.interProcessorLatencyNs.assign(9, 0);
  auto schedule = runPeftScheduler(graph, {target, &context});
  size_t largeCore = schedule.computeToCpuMap.lookup(graph.nodes[3].instance);
  for (size_t task : {0, 1, 2})
    assert(schedule.computeToCpuMap.lookup(graph.nodes[task].instance) != largeCore);
  owner.erase();
}

static void testPeftGapOrder() {
  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> owner =
    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  ComputeGraph graph;
  graph.successors.resize(3);
  graph.predecessors.resize(3);
  for (uint32_t task = 0; task < 3; ++task) {
    ResidentWeight weight;
    weight.opaqueLane = task == 0 ? 0 : 1;
    graph.nodes.push_back({{owner->getOperation(), task, 1},
      task == 2 ? 20u : 10u, {weight}, task});
  }
  TransferCost transfer {.fixed = 100};
  graph.successors[0].push_back({1, transfer});
  graph.predecessors[1].push_back({0, transfer});
  SchedulingTarget target;
  target.processorCount = 2;
  target.residentWeightCapacity = 1;
  target.interProcessorLatencyNs.assign(4, 0);
  auto schedule = runPeftScheduler(graph, {target, &context});
  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[2].instance) == 0);
  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[1].instance) == 1);
  assert(schedule.computeToAestMap.lookup(graph.nodes[2].instance) == 0);
}

static void testPipelineInputArrival() {
  ComputeGraph graph;
  graph.successors.resize(4);
  graph.predecessors.resize(4);
  MergeScheduleResult schedule;
  schedule.processorCount = 4;
  for (uint32_t task = 0; task < 4; ++task) {
    ComputeInstance instance {nullptr, task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    graph.nodes.push_back({instance, 1, {weight}, task});
    schedule.dominanceOrderCompute.push_back(instance);
    schedule.computeToCpuMap[instance] = task;
    schedule.computeToCpuSlotMap[instance] = 0;
    if (task != 0) {
      TransferCost transfer {.fixed = 0, .networkFlits = 100};
      graph.predecessors[task].push_back({task - 1, transfer});
      graph.successors[task - 1].push_back({task, transfer});
    }
  }
  SchedulingTarget target;
  target.processorCount = 4;
  target.residentWeightCapacity = 4;
  target.interProcessorLatencyNs.assign(16, 1);
  for (size_t core = 0; core < 4; ++core)
    target.interProcessorLatencyNs[core * 4 + core] = 0;
  std::string error;
  assert(mlir::succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
  for (size_t first : {0, 2}) {
    size_t source = schedule.computeToCpuMap.lookup(graph.nodes[first].instance);
    size_t destination = schedule.computeToCpuMap.lookup(graph.nodes[first + 1].instance);
    assert(schedule.processorStages[source] == first / 2);
    // An idle remote core would finish later than waiting for the local core.
    assert(source == destination);
  }
}

static void testReadyCommunication() {
  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> owner =
    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  auto type = mlir::RankedTensorType::get({1}, mlir::Float32Type::get(&context));
  mlir::Block payloads;
  auto payload = payloads.addArgument(type, owner->getLoc());
  DeferredTransferPlan plan;
  plan.stepCounts = {3, 1, 1};
  plan.scheduled.resize(3);
  for (auto [index, info] : llvm::enumerate(plan.scheduled)) {
    info.op = owner->getOperation();
    info.streamIds.push_back(index);
  }
  ProducedValue producer;
  producer.scheduled = &plan.scheduled[0];
  producer.payload = payload;
  // Two consumers exercise removal of already-scheduled ready-queue entries.
  for (unsigned target = 1; target <= 2; ++target) {
    auto exchange = std::make_unique<DeferredExchangePlan>();
    exchange->target = &plan.scheduled[target];
    exchange->exchangeId = target;
    RequirementFamily requirement;
    requirement.producer = &producer;
    requirement.publicationFragmentType = type;
    exchange->requirements.push_back(requirement);
    ExternalTransferFamily transfer;
    transfer.requirement = &exchange->requirements.front();
    transfer.sourceScheduled = producer.scheduled;
    transfer.targetScheduled = exchange->target;
    transfer.targetLanes = LaneSet::all(1);
    transfer.targetStreams = onnx_mlir::StaticIntSequence::uniform(target, 1);
    exchange->external.push_back(transfer);
    plan.exchanges.push_back(std::move(exchange));
  }
  auto scheduled = scheduleDeferredCommunication({}, plan);
  assert(mlir::succeeded(scheduled));
  assert(scheduled->slices.size() == 2);
  for (const auto &slice : scheduled->slices)
    assert(slice.sourceInsertionStep == 1);

  plan.stepCounts = {3, 3, 3};
  for (auto &exchange : plan.exchanges)
    exchange->consumerStep = 2;
  scheduled = scheduleDeferredCommunication({}, plan);
  assert(mlir::succeeded(scheduled));
  for (const auto &slice : scheduled->slices)
    assert(slice.targetInsertionStep == 2);
}

static void testPipelineRestartOrder() {
  mlir::MLIRContext context;
  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  DeferredTransferPlan transfers;
  for (int64_t core : {0, 7, 3, 5}) {
    ScheduledInfo info;
    info.op = owner;
    info.cores.push_back(core);
    info.pipelineStages.push_back(core == 0 ? 0 : 1);
    transfers.scheduled.push_back(info);
  }
  auto boundaries = buildDeferredBoundaryPlan(transfers, {}, 16);
  assert(mlir::succeeded(boundaries));
  // Restarting the producer must not wait for a later consumer on core 3.
  assert((transfers.downstreamCores == llvm::SmallVector<int64_t>{7, 3, 5}));
  owner.erase();
}

int main() {
  testResidentWeightViewCost();
  testPipelineStageSpan();
  testPipelinePackingCapacity();
  testPipelineRestartOrder();
  testPipelineInputArrival();
  testReadyCommunication();
  PipelineCoreLayout unevenLayout(138, 4);
  assert(unevenLayout.isValid());
  assert(unevenLayout.getStageRange(0).begin == 0);
  assert(unevenLayout.getStageRange(0).size == 35);
  assert(unevenLayout.getStageRange(1).begin == 35);
  assert(unevenLayout.getStageRange(1).size == 35);
  assert(unevenLayout.getStageRange(2).begin == 70);
  assert(unevenLayout.getStageRange(2).size == 34);
  assert(unevenLayout.getStageRange(3).begin == 104);
  assert(unevenLayout.getStageRange(3).size == 34);
  assert(unevenLayout.getStageForCore(34) == 0);
  assert(unevenLayout.getStageForCore(35) == 1);
  assert(unevenLayout.getStageForCore(69) == 1);
  assert(unevenLayout.getStageForCore(70) == 2);
  assert(unevenLayout.getStageForCore(137) == 3);
  assert(!unevenLayout.getStageForCore(138));

  PipelineCoreLayout dynamicLayout(std::vector<size_t> {2, 4, 1, 3});
  assert(dynamicLayout.isValid());
  assert(dynamicLayout.getProcessorCount() == 10);
  assert(dynamicLayout.getStageRange(0).begin == 0);
  assert(dynamicLayout.getStageRange(1).begin == 2);
  assert(dynamicLayout.getStageRange(2).begin == 6);
  assert(dynamicLayout.getStageRange(3).begin == 7);
  assert(dynamicLayout.getStageForCore(1) == 0);
  assert(dynamicLayout.getStageForCore(2) == 1);
  assert(dynamicLayout.getStageForCore(6) == 2);
  assert(dynamicLayout.getStageForCore(9) == 3);

  TransferCost transfer {.fixed = 50, .networkFlits = 4};

  SchedulingTarget fast;
  fast.processorCount = 2;
  fast.interProcessorLatencyNs = {0, 3, 3, 0};
  fast.averageInterProcessorLatencyNs = 3;

  SchedulingTarget slow = fast;
  slow.interProcessorLatencyNs = {0, 10, 10, 0};
  slow.averageInterProcessorLatencyNs = 10;

  assert(getPeftTransferTime(transfer, 0, 0, fast) == 0);
  assert(getPeftTransferTime(transfer, 0, 1, fast) == 62);
  assert(getPeftTransferTime(transfer, 0, 1, slow) == 90);
  assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
  assert(slow.getInterProcessorLatencyNs(0, 1) == 10);

  testPhysicalPeft();
  testPeftGapOrder();
  testPeftReservationCapacity();

  ComputeGraph graph;
  graph.successors.resize(6);
  graph.predecessors.resize(6);
  graph.successors[1].push_back({2, TransferCost {.fixed = 1, .networkFlits = 1}});
  graph.predecessors[2].push_back({1, TransferCost {.fixed = 1, .networkFlits = 1}});
  const Cost costs[] = {6, 4, 6, 4, 8, 8};
  for (uint32_t task = 0; task < 6; ++task) {
    ComputeInstance instance {nullptr, task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    graph.nodes.push_back({instance, costs[task], {weight}, task});
    graph.instanceToIndex[instance] = task;
  }

  MergeScheduleResult physicalSchedule;
  physicalSchedule.processorCount = 4;
  physicalSchedule.dominanceOrderCompute.reserve(graph.nodes.size());
  for (size_t task = 0; task < graph.nodes.size(); ++task) {
    const ComputeInstance& instance = graph.nodes[task].instance;
    physicalSchedule.dominanceOrderCompute.push_back(instance);
    size_t cpu = task < 4 ? 0 : 1;
    physicalSchedule.computeToCpuMap[instance] = cpu;
    physicalSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : task - 4;
    physicalSchedule.computeToAestMap[instance] = task;
  }

  SchedulingTarget physical = fast;
  physical.processorCount = 4;
  physical.residentWeightCapacity = 2;
  physical.interProcessorLatencyNs = {
    0, 3, 3, 3,
    3, 0, 3, 3,
    3, 3, 0, 3,
    3, 3, 3, 0,
  };
  SchedulingTarget clustered = physical;
  clustered.interProcessorLatencyNs = {
    0, 20, 1, 20,
    20, 0, 20, 1,
    1, 20, 0, 20,
    20, 1, 20, 0,
  };
  PipelineCoreLayout clusteredLayout(std::vector<size_t> {2, 2}, clustered);
  assert(clusteredLayout.getPhysicalCore(0) == 0);
  assert(clusteredLayout.getPhysicalCore(1) == 2);
  assert(clusteredLayout.getStageForCore(0) == 0);
  assert(clusteredLayout.getStageForCore(2) == 0);
  assert(clusteredLayout.getStageForCore(1) == 1);
  assert(clusteredLayout.getStageForCore(3) == 1);
  PipelineCoreLayout uniformLayout(std::vector<size_t> {2, 2}, physical);
  assert(uniformLayout.getPhysicalCore(1) == 1);
  std::string pipelineError;
  ComputeGraph preparationGraph = graph;
  ResidentWeight extraWeight;
  extraWeight.opaqueLane = graph.nodes.size();
  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
  auto preparation = preparePipelineWorkload(
    preparationGraph, 2, physical, pipelineError);
  assert(mlir::succeeded(preparation));
  assert(*preparation == PipelineWorkloadPreparation::Ready);
  extraWeight.opaqueLane++;
  preparationGraph.nodes[2].residentWeights.push_back(extraWeight);
  assert(mlir::failed(preparePipelineWorkload(
    preparationGraph, 2, physical, pipelineError)));
  assert(pipelineError.find("physical core's crossbars") != std::string::npos);
  pipelineError.clear();
  ComputeGraph emptyGraph;
  MergeScheduleResult emptySchedule;
  emptySchedule.processorCount = 4;
  assert(mlir::succeeded(applyPipelineScheduling(
    emptyGraph, emptySchedule, 2, physical, pipelineError)));
  assert(emptySchedule.processorCount == physical.processorCount);
  assert(emptySchedule.processorStages == std::vector<size_t>({0, 0, 1, 1}));

  MergeScheduleResult pipelineSchedule = physicalSchedule;
  assert(mlir::succeeded(applyPipelineScheduling(
    graph, pipelineSchedule, 2, physical, pipelineError)));
  assert(pipelineSchedule.processorCount == 4);
  size_t predecessorCore = pipelineSchedule.computeToCpuMap.lookup(
    graph.nodes[1].instance);
  size_t successorCore = pipelineSchedule.computeToCpuMap.lookup(
    graph.nodes[2].instance);
  assert(pipelineSchedule.computeToAestMap.lookup(graph.nodes[2].instance)
         >= pipelineSchedule.computeToAestMap.lookup(graph.nodes[1].instance)
              + graph.nodes[1].cost
              + getPeftTransferTime(
                  TransferCost {.fixed = 1, .networkFlits = 1},
                  predecessorCore, successorCore, physical));
  assert(pipelineSchedule.processorStages[predecessorCore]
         <= pipelineSchedule.processorStages[successorCore]);
  assert(pipelineSchedule.processorStages[successorCore]
         <= pipelineSchedule.processorStages[predecessorCore] + 1);
  assert(pipelineSchedule.equivalentClass.empty());

  graph.successors[0].push_back(
    {5, TransferCost {.fixed = 1, .networkFlits = 1}});
  graph.predecessors[5].push_back(
    {0, TransferCost {.fixed = 1, .networkFlits = 1}});
  SchedulingTarget fourStagePhysical = physical;
  fourStagePhysical.processorCount = 8;
  fourStagePhysical.interProcessorLatencyNs.assign(64, 3);
  for (size_t core = 0; core < 8; ++core)
    fourStagePhysical.interProcessorLatencyNs[core * 8 + core] = 0;
  MergeScheduleResult fourStageSchedule = physicalSchedule;
  fourStageSchedule.processorCount = 8;
  assert(mlir::succeeded(applyPipelineScheduling(
    graph, fourStageSchedule, 4, fourStagePhysical, pipelineError)));
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    for (const auto &[predecessor, cost] : graph.predecessors[task]) {
      (void)cost;
      size_t sourceStage = fourStageSchedule.processorStages[
        fourStageSchedule.computeToCpuMap.lookup(
          graph.nodes[predecessor].instance)];
      size_t targetStage = fourStageSchedule.processorStages[
        fourStageSchedule.computeToCpuMap.lookup(graph.nodes[task].instance)];
      assert(sourceStage <= targetStage);
      assert(targetStage <= sourceStage + 1);
    }

  ComputeGraph communicationGraph;
  communicationGraph.successors.resize(5);
  communicationGraph.predecessors.resize(5);
  communicationGraph.successors[4].push_back(
    {3, TransferCost {.fixed = 0, .networkFlits = 1}});
  communicationGraph.predecessors[3].push_back(
    {4, TransferCost {.fixed = 0, .networkFlits = 1}});
  const Cost communicationCosts[] = {6, 4, 6, 4, 1};
  MergeScheduleResult communicationSchedule;
  communicationSchedule.processorCount = 4;
  for (uint32_t task = 0; task < 5; ++task) {
    ComputeInstance instance {nullptr, task, 1};
    ResidentWeight weight;
    weight.opaqueLane = task;
    communicationGraph.nodes.push_back(
      {instance, communicationCosts[task], {weight}, task});
    communicationGraph.instanceToIndex[instance] = task;
    communicationSchedule.dominanceOrderCompute.push_back(instance);
    size_t cpu = task < 4 ? 0 : 1;
    communicationSchedule.computeToCpuMap[instance] = cpu;
    communicationSchedule.computeToCpuSlotMap[instance] = task < 4 ? task : 0;
    communicationSchedule.computeToAestMap[instance] = task;
  }

  SchedulingTarget fastPipeline = physical;
  fastPipeline.residentWeightCapacity = 4;
  MergeScheduleResult fastCommunicationSchedule = communicationSchedule;
  assert(mlir::succeeded(applyPipelineScheduling(
    communicationGraph, fastCommunicationSchedule, 2, fastPipeline, pipelineError)));

  SchedulingTarget slowPipeline = fastPipeline;
  slowPipeline.averageInterProcessorLatencyNs = 10;
  slowPipeline.interProcessorLatencyNs.assign(16, 10);
  for (size_t core = 0; core < 4; ++core)
    slowPipeline.interProcessorLatencyNs[core * 4 + core] = 0;
  MergeScheduleResult slowCommunicationSchedule = communicationSchedule;
  assert(mlir::succeeded(applyPipelineScheduling(
    communicationGraph, slowCommunicationSchedule, 2, slowPipeline, pipelineError)));
  size_t sourceCore = slowCommunicationSchedule.computeToCpuMap.lookup(
    communicationGraph.nodes[4].instance);
  size_t targetCore = slowCommunicationSchedule.computeToCpuMap.lookup(
    communicationGraph.nodes[3].instance);
  assert(slowCommunicationSchedule.processorStages[sourceCore]
         <= slowCommunicationSchedule.processorStages[targetCore]);
  assert(slowCommunicationSchedule.processorStages[targetCore]
         <= slowCommunicationSchedule.processorStages[sourceCore] + 1);
  return EXIT_SUCCESS;
}