Raptor/test/PIM/SpatialSchedulingTargetTest.cpp

#include <cassert>
#ifdef NDEBUG
#error "Spatial scheduling tests require assertions"
#endif
#include <cstdlib>
#include <string>
#include <vector>

#include "mlir/IR/BuiltinOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Pass/PassManager.h"
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"

#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/ScheduledComputeMaterialization.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationRealization.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ResidentWeightPacking.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"

using namespace onnx_mlir::spatial;

static ComputeGraph buildCheckedComputeGraph(mlir::Operation *operation,
    const SchedulingTarget &target, size_t partitions = 0) {
  auto graph = buildComputeGraph(operation, target, partitions);
  assert(mlir::succeeded(graph));
  return std::move(*graph);
}

static MergeScheduleResult runCheckedPeftScheduler(const ComputeGraph &graph,
    const PeftScheduleOptions &options) {
  std::string error;
  auto schedule = runPeftScheduler(graph, options, error);
  assert(mlir::succeeded(schedule));
  return std::move(*schedule);
}

static void testLocalFragmentInsertion() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<tensor::TensorDialect, arith::ArithDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto check = [&](ArrayRef<int64_t> sourceShape, ArrayRef<int64_t> destShape,
                   ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes,
                   ArrayRef<int64_t> strides, bool view) {
    Value source = tensor::EmptyOp::create(builder, loc, sourceShape, builder.getF32Type());
    Value dest = tensor::EmptyOp::create(builder, loc, destShape, builder.getF32Type());
    auto mixed = [&](ArrayRef<int64_t> values) {
      SmallVector<OpFoldResult> result;
      for (int64_t value : values)
        result.push_back(builder.getIndexAttr(value));
      return result;
    };
    Value result = insertLocalFragment(builder, loc, source, dest,
      mixed(offsets), mixed(sizes), mixed(strides));
    assert(isa<tensor::ExpandShapeOp>(result.getDefiningOp()) == view);
    assert(succeeded(verify(result.getDefiningOp())));
  };
  check({4}, {1, 4}, {0, 0}, {1, 4}, {1, 1}, true);
  check({2}, {1, 4}, {0, 0}, {1, 2}, {1, 1}, false);
  check({2}, {1, 4}, {0, 2}, {1, 2}, {1, 1}, false);
  check({2}, {1, 4}, {0, 0}, {1, 2}, {1, 2}, false);
  check({4}, {2, 4}, {1, 0}, {1, 4}, {1, 1}, false);
  check({}, {1}, {0}, {1}, {1}, true);
  check({}, {}, {}, {}, {}, false);
  check({1, 4}, {1, 4}, {0, 0}, {1, 4}, {1, 1}, false);
}

static void testBatchedReceiveForwarding() {
  using namespace mlir;
  using namespace onnx_mlir;
  for (bool genericCopy : {false, true})
    for (unsigned scenario : {0, 1, 2}) {
      bool barrier = scenario == 1;
      bool nonzeroReceive = scenario == 2;
      if (genericCopy && nonzeroReceive)
        continue;
      MLIRContext context(MLIRContext::Threading::DISABLED);
      context.loadDialect<func::FuncDialect, memref::MemRefDialect,
                          arith::ArithDialect, pim::PimDialect>();
      OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
      OpBuilder builder(module->getBodyRegion());
      auto loc = module->getLoc();
      auto function = func::FuncOp::create(builder, loc, "main_graph",
                                           builder.getFunctionType({}, {}));
      builder.setInsertionPointToStart(function.addEntryBlock());
      auto core = pim::PimCoreOp::create(builder, loc, ValueRange{}, builder.getI32IntegerAttr(0));
      auto *body = new Block;
      core.getBody().push_back(body);
      builder.setInsertionPointToStart(body);
      auto type = MemRefType::get({4}, builder.getF32Type());
      Value target = memref::AllocOp::create(builder, loc,
        MemRefType::get({8}, builder.getF32Type()));
      Value temporary = memref::AllocOp::create(builder, loc,
        nonzeroReceive ? MemRefType::get({8}, builder.getF32Type()) : type);
      Value table = memref::AllocOp::create(builder, loc,
        MemRefType::get({1}, builder.getIndexType()));
      Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
      Value receiveOffset = arith::ConstantIndexOp::create(
        builder, loc, nonzeroReceive ? 4 : 0);
      auto received = pim::PimReceiveOp::create(builder, loc, temporary.getType(), temporary,
        receiveOffset, builder.getI32IntegerAttr(16), zero);
      Value offset = memref::LoadOp::create(builder, loc, table, ValueRange{zero});
      if (barrier)
        pim::PimSyncOp::create(builder, loc, zero, zero);
      if (genericCopy) {
        auto destination = memref::SubViewOp::create(builder, loc, target,
          SmallVector<OpFoldResult>{offset},
          SmallVector<OpFoldResult>{builder.getIndexAttr(4)},
          SmallVector<OpFoldResult>{builder.getIndexAttr(1)});
        memref::CopyOp::create(builder, loc, received.getOutput(), destination);
      }
      else
        pim::PimMemCopyOp::create(builder, loc, target.getType(), offset, zero,
                                  target, received.getOutput(), builder.getI32IntegerAttr(16));
      pim::PimSendOp::create(builder, loc, target, builder.getI32IntegerAttr(16), zero);
      pim::PimHaltOp::create(builder, loc);
      builder.setInsertionPointAfter(core);
      func::ReturnOp::create(builder, loc);
      assert(succeeded(verify(*module)));
      PassManager manager(&context);
      manager.addPass(createPimMemoryNormalizationPass());
      assert(succeeded(manager.run(*module)));
      unsigned copies = 0, receives = 0;
      core.walk([&](pim::PimMemCopyOp) { ++copies; });
      core.walk([&](pim::PimReceiveOp receive) {
        ++receives;
        assert(receive.getOutputBuffer() == (barrier || nonzeroReceive ? temporary : target));
        if (!barrier && !nonzeroReceive && !genericCopy)
          assert(receive.getOutputOffset() == offset);
      });
      assert(receives == 1 && copies == static_cast<unsigned>(barrier || nonzeroReceive));
    }
}

static void testResidentWeightViewCost() {
  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  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) + 16 * (readLatency + 3));
    Cost singleRead = getComputeInstanceCost(direct, target);
    target.matrixInputReadMultiplier = 8;
    assert(getComputeInstanceCost(direct, target) == singleRead + 112 * readLatency);
  }
}

static void testTransposeCost() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, tensor::TensorDialect,
                      linalg::LinalgDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  for (int64_t rows : {1, 4, 8}) {
    builder.setInsertionPointToEnd(module->getBody());
    auto inputType = RankedTensorType::get({rows, 16}, builder.getF32Type());
    auto outputType = RankedTensorType::get({16, rows}, builder.getF32Type());
    auto function = func::FuncOp::create(builder, loc, "transpose_" + std::to_string(rows),
      builder.getFunctionType({inputType}, {outputType}));
    builder.setInsertionPointToStart(function.addEntryBlock());
    auto compute = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
      TypeRange{outputType}, ValueRange{}, ValueRange{function.getArgument(0)});
    Value init = tensor::EmptyOp::create(builder, loc, outputType.getShape(), outputType.getElementType());
    Value result = linalg::TransposeOp::create(builder, loc,
      *compute.getInputArgument(0), init, SmallVector<int64_t>{1, 0}).getResult()[0];
    SpatYieldOp::create(builder, loc, ValueRange{result});
    builder.setInsertionPointAfter(compute);
    func::ReturnOp::create(builder, loc, compute.getResults());
    assert(succeeded(verify(function)));
    for (Cost period : {1, 5}) {
      SchedulingTarget target;
      target.computeBitwidth = period == 1 ? 8 : 16;
      target.localMemoryWidthBytes = 8;
      target.localMemoryReadLatencyCycles = 2;
      target.localMemoryWriteLatencyCycles = 3;
      target.processorPeriodNs = period;
      target.vectorWidth = 16;
      target.vectorLatencyCycles = 4;
      Cost runBytes = rows * target.computeBitwidth / 8;
      Cost expected = rows == 1
        ? 16 * target.computeBitwidth / 8 / 8 * (2 + 3) * period
        : 16 * ((runBytes + 7) / 8 * (2 + 3) + 4) * period;
      assert(getComputeInstanceCost({compute, 0, 1}, target)
             == expected);
      Cost bytes = rows * 16 * target.computeBitwidth / 8;
      assert(getComputeInstanceCost({compute, 0, 1}, target, ComputeCostKind::VolumePriority)
             == bytes / 8 * (2 + 3) * period);
      target.processorCount = 4;
      target.residentWeightCapacity = 1;
      target.interProcessorLatencyNs.assign(16, 0);
      auto graph = buildCheckedComputeGraph(function, target);
      MergeScheduleResult schedule;
      std::string error;
      assert(succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
      assert(schedule.costs.totalWorkNs == expected);
      assert(schedule.costs.maximumCoreWorkNs == expected);
    }
  }
}

static void testConcatCost() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  for (int64_t axis : {0, 1, 2}) {
    SmallVector<int64_t> shape{2, 3, 5};
    auto inputType = RankedTensorType::get(shape, builder.getF32Type());
    shape[axis] *= 2;
    auto outputType = RankedTensorType::get(shape, builder.getF32Type());
    builder.setInsertionPointToEnd(module->getBody());
    auto function = func::FuncOp::create(builder, loc, "concat_" + std::to_string(axis),
      builder.getFunctionType({inputType, inputType}, {outputType}));
    builder.setInsertionPointToStart(function.addEntryBlock());
    auto compute = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
      TypeRange{outputType}, ValueRange{}, function.getArguments());
    Value result = SpatConcatOp::create(builder, loc, outputType, axis,
      ValueRange{*compute.getInputArgument(0), *compute.getInputArgument(1)});
    SpatYieldOp::create(builder, loc, ValueRange{result});
    builder.setInsertionPointAfter(compute);
    func::ReturnOp::create(builder, loc, compute.getResults());
    assert(succeeded(verify(function)));
    for (Cost period : {1, 5}) {
      SchedulingTarget target;
      target.computeBitwidth = period == 1 ? 8 : 16;
      target.localMemoryWidthBytes = 8;
      target.localMemoryReadLatencyCycles = 2;
      target.localMemoryWriteLatencyCycles = 3;
      target.processorPeriodNs = period;
      Cost blocks = axis == 0 ? 1 : axis == 1 ? 2 : 6;
      Cost bytes = 30 / blocks * target.computeBitwidth / 8;
      Cost expected = 2 * blocks * ((bytes + 7) / 8) * (2 + 3) * period;
      assert(getComputeInstanceCost({compute, 0, 1}, target) == expected);
    }
  }
}

static void testScalarProjectedTransferCost() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, tensor::TensorDialect,
                      scf::SCFDialect, arith::ArithDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto imageType = RankedTensorType::get({1024}, builder.getF32Type());
  auto fragment = RankedTensorType::get({1}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "scalar_projection",
    builder.getFunctionType({imageType}, {}));
  auto *body = function.addEntryBlock();
  builder.setInsertionPointToStart(body);
  auto scalar = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
    TypeRange{imageType}, ValueRange{}, ValueRange{body->getArgument(0)});
  builder.setInsertionPointToEnd(&scalar.getBody().front());
  Value image = SpatReluOp::create(builder, loc, imageType, *scalar.getInputArgument(0));
  SpatYieldOp::create(builder, loc, ValueRange{image});
  builder.setInsertionPointAfter(scalar);
  for (int64_t trips : {1, 3}) {
    auto batch = onnx_mlir::createEmptySpatGraphComputeBatch(builder, loc,
      TypeRange{RankedTensorType::get({2, 1}, builder.getF32Type())}, 2,
      ValueRange{}, ValueRange{scalar.getResult(0)});
    assert(succeeded(batch));
    builder.setInsertionPointToEnd(&batch->getBody().front());
    Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
    Value one = arith::ConstantIndexOp::create(builder, loc, 1);
    Value limit = arith::ConstantIndexOp::create(builder, loc, trips);
    Value initial = tensor::EmptyOp::create(builder, loc, fragment.getShape(), fragment.getElementType());
    auto loop = scf::ForOp::create(builder, loc, zero, limit, one, ValueRange{initial},
      [&](OpBuilder &nested, Location nestedLoc, Value iv, ValueRange) {
        Value offset = arith::AddIOp::create(nested, nestedLoc, iv, *batch->getLaneArgument());
        Value selected = tensor::ExtractSliceOp::create(nested, nestedLoc, fragment,
          *batch->getInputArgument(0), SmallVector<OpFoldResult>{offset},
          SmallVector<OpFoldResult>{nested.getIndexAttr(1)},
          SmallVector<OpFoldResult>{nested.getIndexAttr(1)});
        Value result = SpatReluOp::create(nested, nestedLoc, fragment, selected);
        scf::YieldOp::create(nested, nestedLoc, ValueRange{result});
      });
    auto parallel = SpatInParallelOp::create(builder, loc);
    builder.setInsertionPointToEnd(&parallel.getRegion().front());
    tensor::ParallelInsertSliceOp::create(builder, loc, loop.getResult(0), *batch->getOutputArgument(0),
      SmallVector<OpFoldResult>{*batch->getLaneArgument(), builder.getIndexAttr(0)},
      SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)),
      SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
    builder.setInsertionPointAfter(*batch);
  }
  func::ReturnOp::create(builder, loc);
  assert(succeeded(verify(*module)));
  SchedulingTarget target;
  target.processorCount = 2;
  for (size_t partitions : {1, 2}) {
    ComputeGraph graph = buildCheckedComputeGraph(function, target, partitions);
    for (const auto &edge : graph.edges) {
      auto instance = graph.nodes[edge.target].instance;
      if (instance.op == scalar)
        continue;
      auto batch = cast<SpatGraphComputeBatch>(instance.op);
      auto loop = *batch.getBody().front().getOps<scf::ForOp>().begin();
      Cost trips = static_cast<Cost>(
        loop.getUpperBound().getDefiningOp<arith::ConstantIndexOp>().value());
      assert(edge.transferCost.fixed == 2 * trips * instance.laneCount);
      assert(edge.transferCost.networkFlits == 3 * trips * instance.laneCount);
    }
  }
}

static void testProducerChunkTransferCost() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, tensor::TensorDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  for (unsigned bits : {1, 3, 4, 8, 16})
    for (int64_t elements : {0, 1, 3, 8}) {
      builder.setInsertionPointToEnd(module->getBody());
      auto fragment = RankedTensorType::get({elements}, builder.getIntegerType(bits));
      auto physical = getGraphBatchPhysicalResultType(3, fragment);
      auto function = func::FuncOp::create(builder, loc,
        "transfer_" + std::to_string(bits) + "_" + std::to_string(elements),
        builder.getFunctionType({fragment}, {}));
      builder.setInsertionPointToStart(function.addEntryBlock());
      auto producer = onnx_mlir::createEmptySpatGraphComputeBatch(builder, loc,
        TypeRange{physical}, 3, ValueRange{}, ValueRange{function.getArgument(0)});
      assert(succeeded(producer));
      auto parallel = SpatInParallelOp::create(builder, loc);
      builder.setInsertionPointToEnd(&parallel.getRegion().front());
      tensor::ParallelInsertSliceOp::create(builder, loc, *producer->getInputArgument(0),
        *producer->getOutputArgument(0),
        SmallVector<OpFoldResult>{*producer->getLaneArgument(), builder.getIndexAttr(0)},
        SmallVector<OpFoldResult>{builder.getIndexAttr(1), builder.getIndexAttr(elements)},
        SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
      builder.setInsertionPointAfter(*producer);
      auto consumer = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
        TypeRange{physical}, ValueRange{}, ValueRange{producer->getResult(0)});
      SpatYieldOp::create(builder, loc, ValueRange{*consumer.getInputArgument(0)});
      builder.setInsertionPointAfter(consumer);
      func::ReturnOp::create(builder, loc);
      SchedulingTarget target;
      target.processorCount = 3;
      target.computeBitwidth = 16;
      target.localMemoryWidthBytes = target.transferWidthBytes = 1;
      if (elements == 0) {
        // MLIR rejects physical insertions into a zero-sized fragment dimension.
        ScopedDiagnosticHandler handler(&context, [](Diagnostic &) { return success(); });
        assert(failed(verify(function)));
        auto empty = getTransferCostFromBytes(0, target);
        assert(empty.fixed == 0 && empty.networkFlits == 2);
        function.erase();
        continue;
      }
      assert(succeeded(verify(function)));
      for (size_t partitions : {1, 2, 3}) {
        auto graph = buildCheckedComputeGraph(function, target, partitions);
        const auto &edges = graph.predecessors[graph.instanceToIndex.lookup({consumer, 0, 1})];
        assert(edges.size() == partitions);
        for (const auto &[task, transfer] : edges) {
          Cost selectedBits = elements * bits * graph.nodes[task].instance.laneCount;
          Cost bytes = (selectedBits + 7) / 8;
          assert(transfer.fixed == 2 * bytes);
          assert(transfer.networkFlits == 2 + bytes);
        }
      }
    }
}

static void testPipelineStageSpan() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  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);
    MergeScheduleResult schedule;
    std::string error;
    assert(succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
    MergeScheduleResult repeated;
    assert(succeeded(runPipelineScheduler(graph, repeated, 2, target, error)));
    assert(schedule.processorStages == repeated.processorStages);
    assert(schedule.computeToCpuMap == repeated.computeToCpuMap);
    assert(schedule.computeToCpuSlotMap == repeated.computeToCpuSlotMap);
    for (size_t task = 0; task < graph.nodes.size(); ++task)
      for (const auto &[parent, transfer] : graph.predecessors[task]) {
        const auto &source = graph.nodes[parent].instance;
        const auto &destination = graph.nodes[task].instance;
        size_t from = schedule.computeToCpuMap.lookup(source);
        size_t to = schedule.computeToCpuMap.lookup(destination);
        assert(schedule.processorStages[from] <= schedule.processorStages[to]);
        assert(schedule.processorStages[to] <= schedule.processorStages[from] + 1);
        if (from == to)
          assert(schedule.computeToCpuSlotMap.lookup(source)
                 < schedule.computeToCpuSlotMap.lookup(destination));
      }
    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) {
      assert(!schedule.equivalentClass.empty());
      for (const auto &[core, peers] : schedule.equivalentClass)
        for (size_t peer : peers) {
          assert(peer != core);
          assert(schedule.processorStages[core] == schedule.processorStages[peer]);
        }
      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(MLIRContext::Threading::DISABLED);
  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);
  MergeScheduleResult schedule;
  std::string error;
  assert(succeeded(runPipelineScheduler(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 testRepeatedResidentWeightPacking() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  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(12), ValueRange{}, ValueRange{});
  ComputeGraph graph;
  graph.predecessors.resize(12);
  graph.successors.resize(12);
  std::vector<size_t> tasks;
  std::vector<Cost> costs(12, 10);
  for (uint32_t task = 0; task < 12; ++task) {
    ResidentWeight weight;
    weight.opaqueLane = task % 2;
    graph.nodes.push_back({{batch, task, 1}, 10, {weight}, task});
    tasks.push_back(task);
  }
  IndexedResidentWeights weights(graph);
  assert(weights.getPackCountLowerBound(tasks, 1) == 2);
  auto packing = weights.pack(tasks, costs, 1, 2);
  assert(packing);
  for (size_t task : tasks)
    assert((*packing)[task] == (*packing)[task % 2]);
  assert(!weights.pack(tasks, costs, 1, 1));
  assert(weights.pack({}, costs, 1, 0)->empty());

  for (size_t stages : {2, 4, 8})
    for (Cost linkLatency : {0, 10}) {
      SchedulingTarget target;
      target.processorCount = stages + 1;
      target.residentWeightCapacity = 1;
      target.interProcessorLatencyNs.assign(
        target.processorCount * target.processorCount, linkLatency);
      for (size_t core = 0; core < target.processorCount; ++core)
        target.interProcessorLatencyNs[core * target.processorCount + core] = 0;
      MergeScheduleResult schedule;
      std::string error;
      assert(succeeded(runPipelineScheduler(graph, schedule, stages, target, error)));
      assert(schedule.dominanceOrderCompute.size() == graph.nodes.size());
      std::vector<ResidentWeightSet> residents(target.processorCount);
      for (const auto &node : graph.nodes) {
        size_t cpu = schedule.computeToCpuMap.lookup(node.instance);
        insertResidentWeights(residents[cpu], node.residentWeights);
        assert(residents[cpu].size() <= target.residentWeightCapacity);
        assert(schedule.processorStages[cpu] == 0);
      }
    }
  graph.nodes[0].residentWeights.push_back(graph.nodes[1].residentWeights.front());
  IndexedResidentWeights oversized(graph);
  assert(oversized.getPackCountLowerBound(tasks, 1)
         == std::numeric_limits<size_t>::max());
  assert(!oversized.pack(tasks, costs, 1, 12));
}

static void testResidentWeightPackingRecovery() {
  ComputeGraph graph;
  for (const auto &ids : std::vector<std::vector<uint32_t>>{
         {0, 3}, {0, 1}, {0, 2}, {3, 4}}) {
    ResidentWeightSet residents;
    for (uint32_t id : ids) {
      ResidentWeight weight;
      weight.opaqueLane = id;
      residents.push_back(weight);
    }
    graph.nodes.push_back({{}, 10, residents, graph.nodes.size()});
  }
  IndexedResidentWeights weights(graph);
  std::vector<size_t> tasks{0, 1, 2, 3};
  std::vector<Cost> costs(4, 10);
  assert(weights.getPackCountLowerBound(tasks, 3) == 2);
  auto assignment = weights.pack(tasks, costs, 3, 2);
  assert(assignment && assignment == weights.pack(tasks, costs, 3, 2));
  assert(*assignment == std::vector<size_t>({0, 1, 1, 0}));
  assert((*assignment)[0] == (*assignment)[3]);
  assert((*assignment)[1] == (*assignment)[2]);
  assert((*assignment)[0] != (*assignment)[1]);
  std::vector<ResidentWeightSet> residents(2);
  for (size_t task : tasks) {
    insertResidentWeights(residents[(*assignment)[task]], graph.nodes[task].residentWeights);
    assert(residents[(*assignment)[task]].size() <= 3);
  }
  assert(!weights.pack(tasks, costs, 3, 1));
}

static void testResidentWeightPackingRepair() {
  ComputeGraph graph;
  for (const auto &ids : std::vector<std::vector<uint32_t>>{
         {1, 3}, {3, 4}, {4, 5}, {2, 3}}) {
    ResidentWeightSet residents;
    for (uint32_t id : ids) {
      ResidentWeight weight;
      weight.opaqueLane = id;
      residents.push_back(weight);
    }
    graph.nodes.push_back({{}, 10, residents, graph.nodes.size()});
  }
  IndexedResidentWeights weights(graph);
  std::vector<size_t> tasks{0, 1, 2, 3};
  std::vector<Cost> costs(4, 10);
  auto assignment = weights.pack(tasks, costs, 3, 2);
  assert(assignment);
  assert((*assignment)[0] == (*assignment)[3]);
  assert((*assignment)[1] == (*assignment)[2]);
  assert((*assignment)[0] != (*assignment)[1]);
  assert(assignment == weights.pack(tasks, costs, 3, 2));
  assert(!weights.pack(tasks, costs, 2, 2));

  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  mlir::OwningOpRef<mlir::ModuleOp> owner =
    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  graph.predecessors.resize(graph.nodes.size());
  graph.successors.resize(graph.nodes.size());
  for (auto [task, node] : llvm::enumerate(graph.nodes))
    node.instance = {owner->getOperation(), static_cast<uint32_t>(task), 1};
  SchedulingTarget target;
  target.processorCount = 2;
  target.residentWeightCapacity = 3;
  target.interProcessorLatencyNs.assign(4, 0);
  auto schedule = runCheckedPeftScheduler(graph, {target, &context});
  assert(schedule.computeToCpuMap.lookup(graph.nodes[0].instance)
         == schedule.computeToCpuMap.lookup(graph.nodes[3].instance));
  assert(schedule.computeToCpuMap.lookup(graph.nodes[1].instance)
         == schedule.computeToCpuMap.lookup(graph.nodes[2].instance));
  assert(schedule.costs.maximumCoreWorkNs == 20);
  std::string error;
  target.residentWeightCapacity = 2;
  assert(mlir::failed(runPeftScheduler(graph, {target, &context}, error)));
  assert(error.find("lower bound exceeds physical capacity") != std::string::npos);

  // The union lower bound fits, but two cores cannot hold five distinct two-weight spokes.
  graph.nodes.clear();
  for (uint32_t task = 0; task < 5; ++task) {
    ResidentWeight common, spoke;
    common.opaqueLane = 0;
    spoke.opaqueLane = task + 1;
    graph.nodes.push_back({{owner->getOperation(), task, 1}, 10, {common, spoke}, task});
  }
  graph.predecessors.resize(5);
  graph.successors.resize(5);
  target.residentWeightCapacity = 3;
  IndexedResidentWeights impossible(graph);
  std::vector<size_t> spokes{0, 1, 2, 3, 4};
  assert(impossible.getPackCountLowerBound(spokes, 3) == 2);
  assert(!impossible.pack(spokes, std::vector<Cost>(5, 10), 3, 2));
  MergeSchedulingAnalysis analysis(graph, target, &context, error);
  assert(mlir::failed(analysis.getResult()));
  assert(error.find("heuristic search exhausted its bounded repair") != std::string::npos);
  target.residentWeightCapacity = 0;
  assert(mlir::failed(runPeftScheduler(graph, {target, &context}, error)));
  assert(error.find("must be positive") != std::string::npos);
}

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(getTransferTime(transfer, 0, 1, line) == 54);
  assert(getTransferTime(transfer, 0, 2, line) == 90);
  SchedulingTarget reversed = line;
  reversed.interProcessorLatencyNs = {0, 10, 1, 10, 0, 1, 1, 1, 0};
  assert(getTransferTime(transfer, 0, 1, reversed) == 90);
  assert(getTransferTime(transfer, 0, 2, reversed) == 54);

  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  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 = runCheckedPeftScheduler(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
                  + getTransferTime(transfer, source, destination, target));
    }
  }
}

static void testPeftReservationPassFailure() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, arith::ArithDialect>();
  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(module->getBodyRegion());
  auto loc = module->getLoc();
  auto weight = RankedTensorType::get({128, 128}, builder.getF32Type());
  auto vector = RankedTensorType::get({1, 128}, builder.getF32Type());
  auto function = func::FuncOp::create(builder, loc, "main_graph",
    builder.getFunctionType({vector}, {vector}));
  builder.setInsertionPointToStart(function.addEntryBlock());
  SmallVector<Value> weights;
  for (float value : {1.f, 2.f})
    weights.push_back(arith::ConstantOp::create(builder, loc,
      DenseElementsAttr::get(weight, ArrayRef<float>{value})));
  auto compute = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
    TypeRange{vector}, weights, ValueRange{function.getArgument(0)});
  Value result = *compute.getInputArgument(0);
  for (unsigned index : {0, 1})
    result = SpatVMMOp::create(builder, loc, vector, *compute.getWeightArgument(index), result);
  SpatYieldOp::create(builder, loc, ValueRange{result});
  builder.setInsertionPointAfter(compute);
  func::ReturnOp::create(builder, loc, compute.getResults());
  assert(succeeded(verify(*module)));
  SchedulingTarget target;
  target.processorCount = target.residentWeightCapacity = 1;
  target.interProcessorLatencyNs = {0};
  std::string diagnostics;
  ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
    diagnostics += diagnostic.str();
    return success();
  });
  PassManager manager(&context);
  manager.addPass(createScheduleAndRealizeSpatialPass(target, SpatialDataflowExportStage::None));
  assert(failed(manager.run(*module)));
  assert(diagnostics.find("resident-weight lower bound exceeds physical capacity") != std::string::npos);
}

static void testPeftReservationCapacity() {
  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  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 = runCheckedPeftScheduler(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);
  assert(schedule.costs.maximumCoreWorkNs == 20);
  assert(schedule.costs.completionNs == 20);
  owner.erase();
}

static void testPeftGapOrder() {
  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  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 = runCheckedPeftScheduler(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 testMaterializedPeftGapOrder() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, scf::SCFDialect,
                      arith::ArithDialect, tensor::TensorDialect>();
  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
  OpBuilder builder(owner->getBodyRegion());
  auto loc = owner->getLoc();
  auto type = RankedTensorType::get({1}, builder.getF32Type());
  auto function = func::FuncOp::create(
    builder, loc, "gap_order", builder.getFunctionType({type}, {}));
  Block *body = function.addEntryBlock();
  builder.setInsertionPointToStart(body);
  ComputeGraph graph;
  graph.successors.resize(3);
  graph.predecessors.resize(3);
  for (uint32_t task = 0; task < 3; ++task) {
    Value input = body->getArgument(0);
    if (task == 1)
      input = graph.nodes[0].instance.op->getResult(0);
    auto compute = SpatGraphCompute::create(
      builder, loc, TypeRange{type}, ValueRange{}, ValueRange{input});
    auto *block = new Block();
    compute.getBody().push_back(block);
    Value argument = block->addArgument(type, loc);
    OpBuilder nested(block, block->begin());
    Value result = SpatReluOp::create(nested, loc, type, argument);
    SpatYieldOp::create(nested, loc, ValueRange{result});
    ResidentWeight weight;
    weight.opaqueLane = task == 0 ? 0 : 1;
    graph.nodes.push_back({{compute, 0, 1}, task == 2 ? 20u : 10u,
                           {weight}, task});
  }
  func::ReturnOp::create(builder, loc);
  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 = runCheckedPeftScheduler(graph, {target, &context});
  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[2].instance) == 0);
  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[1].instance) == 1);
  PatternRewriter rewriter(&context);
  auto materialized = materializeScheduledCompute(function, schedule, rewriter);
  assert(succeeded(materialized));
  size_t core = schedule.computeToCpuMap.lookup(graph.nodes[2].instance);
  const auto &instances = materialized->peftClassPlans[core].instancesByCpu[core];
  assert(instances.size() == 2);
  assert(instances[0] == graph.nodes[2].instance);
  assert(instances[1] == graph.nodes[1].instance);
  for (const auto &record : materialized->materializedSchedules)
    if (record.cpus.front() == core) {
      assert(record.runs[0].instances.front() == graph.nodes[2].instance);
      assert(record.stepAnchors[0]->isBeforeInBlock(record.stepAnchors[1]));
    }
}

static void testPhysicalPublicationDependencies() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, arith::ArithDialect,
                      tensor::TensorDialect, scf::SCFDialect>();
  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
  PatternRewriter builder(&context);
  builder.setInsertionPointToEnd(owner->getBody());
  auto loc = owner->getLoc();
  auto fragmentType = RankedTensorType::get({1}, builder.getF32Type());
  auto function = func::FuncOp::create(
    builder, loc, "publications", builder.getFunctionType({fragmentType}, {}));
  Block *body = function.addEntryBlock();
  builder.setInsertionPointToStart(body);
  auto makeBatch = [&](Value input, uint32_t lanes, int64_t projectionSize,
                       bool reverse) {
    auto inputType = cast<RankedTensorType>(input.getType());
    auto fragment = projectionSize
      ? RankedTensorType::get({projectionSize, 1}, builder.getF32Type())
      : inputType;
    auto physical = getGraphBatchPhysicalResultType(lanes, fragment);
    auto batch = onnx_mlir::createEmptySpatGraphComputeBatch(
      builder, loc, TypeRange{physical}, lanes, ValueRange{}, ValueRange{input});
    assert(succeeded(batch));
    Value lane = *batch->getLaneArgument();
    Value value = *batch->getInputArgument(0);
    if (projectionSize) {
      Value size = arith::ConstantIndexOp::create(builder, loc, projectionSize);
      Value offset = arith::MulIOp::create(builder, loc, lane, size);
      value = tensor::ExtractSliceOp::create(
        builder, loc, fragment, value,
        SmallVector<OpFoldResult>{offset, builder.getIndexAttr(0)},
        SmallVector<OpFoldResult>{builder.getIndexAttr(projectionSize), builder.getIndexAttr(1)},
        SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
    }
    value = SpatReluOp::create(builder, loc, fragment, value);
    Value slot = lane;
    if (reverse) {
      Value last = arith::ConstantIndexOp::create(builder, loc, lanes - 1);
      slot = arith::SubIOp::create(builder, loc, last, lane);
    }
    auto parallel = SpatInParallelOp::create(builder, loc);
    builder.setInsertionPointToEnd(&parallel.getRegion().front());
    SmallVector<OpFoldResult> offsets(fragment.getRank() + 1, builder.getIndexAttr(0));
    offsets.front() = slot;
    SmallVector<OpFoldResult> sizes{builder.getIndexAttr(1)};
    for (int64_t size : fragment.getShape())
      sizes.push_back(builder.getIndexAttr(size));
    tensor::ParallelInsertSliceOp::create(builder, loc, value,
      *batch->getOutputArgument(0), offsets, sizes,
      SmallVector<OpFoldResult>(offsets.size(), builder.getIndexAttr(1)));
    builder.setInsertionPointAfter(*batch);
    return *batch;
  };
  auto producer = makeBatch(body->getArgument(0), 4, 0, true);
  auto projected = makeBatch(producer.getResult(0), 2, 2, false);
  auto conservative = makeBatch(producer.getResult(0), 2, 0, false);
  SmallVector<onnx_mlir::FragmentAssemblyEntry> entries{
    {0, 0, {0}, {1}}, {2, 0, {1}, {1}}};
  auto logicalType = RankedTensorType::get({2}, builder.getF32Type());
  auto blueprint = onnx_mlir::createFragmentAssemblyBlueprint(
    producer.getResult(0), logicalType, entries, "dense_nchw",
    "fragment_assembly", builder, loc);
  assert(succeeded(blueprint));
  auto scalar = SpatGraphCompute::create(
    builder, loc, TypeRange{logicalType}, ValueRange{}, ValueRange{*blueprint});
  auto *block = new Block();
  scalar.getBody().push_back(block);
  Value argument = block->addArgument(logicalType, loc);
  OpBuilder nested(block, block->begin());
  Value value = SpatReluOp::create(nested, loc, logicalType, argument);
  SpatYieldOp::create(nested, loc, ValueRange{value});
  auto projectedBlueprint = onnx_mlir::createEmptySpatGraphComputeBatch(
    builder, loc, TypeRange{getGraphBatchPhysicalResultType(2, fragmentType)},
    2, ValueRange{}, ValueRange{*blueprint});
  assert(succeeded(projectedBlueprint));
  Value projectedLane = *projectedBlueprint->getLaneArgument();
  Value selectedBlueprint = tensor::ExtractSliceOp::create(builder, loc, fragmentType,
    *projectedBlueprint->getInputArgument(0),
    SmallVector<OpFoldResult>{projectedLane},
    SmallVector<OpFoldResult>{builder.getIndexAttr(1)},
    SmallVector<OpFoldResult>{builder.getIndexAttr(1)});
  selectedBlueprint = SpatReluOp::create(builder, loc, fragmentType, selectedBlueprint);
  auto projectedParallel = SpatInParallelOp::create(builder, loc);
  builder.setInsertionPointToEnd(&projectedParallel.getRegion().front());
  tensor::ParallelInsertSliceOp::create(builder, loc, selectedBlueprint,
    *projectedBlueprint->getOutputArgument(0),
    SmallVector<OpFoldResult>{projectedLane, builder.getIndexAttr(0)},
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)),
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
  builder.setInsertionPointAfter(*projectedBlueprint);
  func::ReturnOp::create(builder, loc);
  for (Cost period : {1, 3}) {
    SchedulingTarget target;
    target.processorCount = 4;
    target.processorPeriodNs = period;
    ComputeGraph graph = buildCheckedComputeGraph(function, target);
    auto predecessorLanes = [&](ComputeInstance instance) {
      std::vector<uint32_t> lanes;
      for (const auto &[task, cost] : graph.predecessors[graph.instanceToIndex.lookup(instance)])
        lanes.push_back(graph.nodes[task].instance.laneStart);
      llvm::sort(lanes);
      return lanes;
    };
    assert((predecessorLanes({projected, 0, 1}) == std::vector<uint32_t>{2, 3}));
    assert((predecessorLanes({projected, 1, 1}) == std::vector<uint32_t>{0, 1}));
    for (uint32_t lane : {0, 1})
      assert((predecessorLanes({conservative, lane, 1}) == std::vector<uint32_t>{0, 1, 2, 3}));
    assert((predecessorLanes({scalar, 0, 1}) == std::vector<uint32_t>{1, 3}));
    // Two local fragment copies and the vector operation use the same ns unit.
    assert(graph.nodes[graph.instanceToIndex.lookup({scalar, 0, 1})].cost == 10 * period);
    for (uint32_t lane : {0, 1})
      assert(graph.nodes[graph.instanceToIndex.lookup(
        {projectedBlueprint->getOperation(), lane, 1})].cost == 10 * period);
    for (const auto &[task, transfer] : graph.predecessors[graph.instanceToIndex.lookup({scalar, 0, 1})])
      assert(transfer.fixed == 0);
    assert(verifyAcyclic(graph));
    ComputeGraph chunked = buildCheckedComputeGraph(function, target, 1);
    const auto &dependencies = chunked.predecessors[
      chunked.instanceToIndex.lookup({projected, 0, 2})];
    assert(dependencies.size() == 1);
    assert(dependencies.front().second.fixed == 8 * period);
    assert(dependencies.front().second.networkFlits == 12);
  }

  builder.setInsertionPointToEnd(owner->getBody());
  auto independent = func::FuncOp::create(builder, loc, "batch_order",
    builder.getFunctionType({fragmentType}, {}));
  body = independent.addEntryBlock();
  builder.setInsertionPointToStart(body);
  auto first = makeBatch(body->getArgument(0), 2, 0, false);
  auto second = makeBatch(body->getArgument(0), 2, 0, false);
  func::ReturnOp::create(builder, loc);
  SchedulingTarget target;
  target.processorCount = 2;
  ComputeGraph graph = buildCheckedComputeGraph(independent, target);
  MergeScheduleResult schedule;
  schedule.processorCount = 2;
  std::vector<std::vector<size_t>> tasksByCore(2);
  for (const auto &node : graph.nodes)
    schedule.dominanceOrderCompute.push_back(node.instance);
  for (uint32_t core : {0, 1})
    for (auto [slot, batch] : llvm::enumerate(std::vector<SpatGraphComputeBatch>{second, first})) {
      ComputeInstance instance{batch, 1 - core, 1};
      schedule.computeToCpuMap[instance] = core;
      schedule.computeToCpuSlotMap[instance] = slot;
      tasksByCore[core].push_back(graph.instanceToIndex.lookup(instance));
    }
  schedule.processorStages = {0, 1};
  populateScheduleEquivalenceClasses(graph, tasksByCore, schedule);
  assert(schedule.equivalentClass.empty());
  schedule.processorStages.clear();
  populateScheduleEquivalenceClasses(graph, tasksByCore, schedule);
  assert((schedule.equivalentClass[0] == SmallVector<size_t, 5>{1}));
  auto materialized = materializeScheduledCompute(independent, schedule, builder);
  assert(succeeded(materialized));
  assert(materialized->materializedSchedules.size() == 1);
  const auto &record = materialized->materializedSchedules.front();
  assert(isa<SpatScheduledComputeBatch>(record.scheduledOp));
  assert(record.stepPlans.size() == 2);
  assert(record.stepPlans[0].stepTuple.instances.front().op == second);
  assert(record.stepPlans[1].stepTuple.instances.front().op == first);
  assert(record.cpus.size() == 2);
  assert(record.cpus[0] == 1 && record.cpus[1] == 0);
  unsigned fragmentViews = 0;
  record.scheduledOp->walk([&](tensor::ExpandShapeOp) { ++fragmentViews; });
  record.scheduledOp->walk([&](tensor::InsertSliceOp insert) {
    assert(insert.getDestType().getRank() != insert.getSourceType().getRank() + 1
           || insert.getDestType().getDimSize(0) != 1);
  });
  assert(fragmentViews == 2);

  builder.setInsertionPointToEnd(owner->getBody());
  auto forwarded = func::FuncOp::create(builder, loc, "physical_forwarding",
    builder.getFunctionType({fragmentType}, {fragmentType}));
  body = forwarded.addEntryBlock();
  builder.setInsertionPointToStart(body);
  auto permuted = makeBatch(body->getArgument(0), 4, 0, true);
  auto consumer = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
    TypeRange{fragmentType}, ValueRange{}, ValueRange{permuted.getResult(0)});
  Value selected = tensor::ExtractSliceOp::create(builder, loc, fragmentType,
    *consumer.getInputArgument(0),
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(0)),
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)),
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
  auto relu = SpatReluOp::create(builder, loc, fragmentType, selected);
  relu->setAttr("test.physical_consumer", builder.getUnitAttr());
  SpatYieldOp::create(builder, loc, ValueRange{relu.getResult()});
  builder.setInsertionPointAfter(consumer);
  func::ReturnOp::create(builder, loc, ValueRange{consumer.getResult(0)});
  target.processorCount = 4;
  graph = buildCheckedComputeGraph(forwarded, target);
  for (const auto &[task, transfer] : graph.predecessors[
         graph.instanceToIndex.lookup({consumer, 0, 1})]) {
    assert(transfer.fixed == 2);
    assert(transfer.networkFlits == 3);
  }
  schedule = MergeScheduleResult();
  schedule.processorCount = 4;
  for (auto [slot, node] : llvm::enumerate(graph.nodes)) {
    schedule.dominanceOrderCompute.push_back(node.instance);
    schedule.computeToCpuMap[node.instance] = 0;
    schedule.computeToCpuSlotMap[node.instance] = slot;
  }
  materialized = materializeScheduledCompute(forwarded, schedule, builder);
  assert(succeeded(materialized));
  assert(succeeded(realizeDeferredCommunication(forwarded, *materialized, target, 1)));
  bool checked = false;
  forwarded.walk([&](SpatReluOp op) {
    if (!op->hasAttr("test.physical_consumer"))
      return;
    auto slice = op.getInput().getDefiningOp<tensor::ExtractSliceOp>();
    assert(slice && getConstantIntValue(slice.getMixedOffsets().front()) == 3);
    checked = true;
  });
  assert(checked && succeeded(verify(forwarded)));
}

static void testBlueprintSourceProjections() {
  using namespace mlir;
  MLIRContext context(MLIRContext::Threading::DISABLED);
  context.loadDialect<SpatialDialect, func::FuncDialect, arith::ArithDialect,
                      tensor::TensorDialect>();
  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
  PatternRewriter builder(&context);
  builder.setInsertionPointToEnd(owner->getBody());
  auto loc = owner->getLoc();
  auto fragment = RankedTensorType::get({1}, builder.getF32Type());
  auto physical = getGraphBatchPhysicalResultType(4, fragment);
  auto function = func::FuncOp::create(builder, loc, "blueprint_views",
    builder.getFunctionType({fragment}, {}));
  Block *body = function.addEntryBlock();
  builder.setInsertionPointToStart(body);
  auto producer = onnx_mlir::createEmptySpatGraphComputeBatch(
    builder, loc, TypeRange{physical}, 4, ValueRange{}, ValueRange{body->getArgument(0)});
  assert(succeeded(producer));
  Value lane = *producer->getLaneArgument();
  Value value = SpatReluOp::create(builder, loc, fragment, *producer->getInputArgument(0));
  auto parallel = SpatInParallelOp::create(builder, loc);
  builder.setInsertionPointToEnd(&parallel.getRegion().front());
  auto publication = tensor::ParallelInsertSliceOp::create(builder, loc, value,
    *producer->getOutputArgument(0),
    SmallVector<OpFoldResult>{lane, builder.getIndexAttr(0)},
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)),
    SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
  builder.setInsertionPointAfter(*producer);
  Value casted = tensor::CastOp::create(builder, loc, physical, producer->getResult(0));
  auto flat = RankedTensorType::get({4}, builder.getF32Type());
  Value collapsed = tensor::CollapseShapeOp::create(builder, loc, flat, casted,
    SmallVector<ReassociationIndices>{{0, 1}});
  Value expanded = tensor::ExpandShapeOp::create(builder, loc, physical, collapsed,
    SmallVector<ReassociationIndices>{{0, 1}});
  auto makeSlice = [&](int64_t size) -> Value {
    return tensor::ExtractSliceOp::create(builder, loc,
      RankedTensorType::get({size, 1}, builder.getF32Type()), producer->getResult(0),
      SmallVector<OpFoldResult>{builder.getIndexAttr(2), builder.getIndexAttr(0)},
      SmallVector<OpFoldResult>{builder.getIndexAttr(size), builder.getIndexAttr(1)},
      SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
  };
  Value sliced = makeSlice(2), single = makeSlice(1);
  auto makeConsumer = [&](Value source) {
    int64_t count = cast<RankedTensorType>(source.getType()).getDimSize(0);
    auto logical = RankedTensorType::get({count}, builder.getF32Type());
    SmallVector<onnx_mlir::FragmentAssemblyEntry> entries;
    for (int64_t slot = 0; slot < count; ++slot)
      entries.push_back({slot, 0, {slot}, {1}});
    auto blueprint = onnx_mlir::createFragmentAssemblyBlueprint(source, logical,
      entries, "dense_nchw", "fragment_assembly", builder, loc);
    assert(succeeded(blueprint));
    auto consumer = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
      TypeRange{logical}, ValueRange{}, ValueRange{*blueprint});
    Value relu = SpatReluOp::create(builder, loc, logical, *consumer.getInputArgument(0));
    SpatYieldOp::create(builder, loc, ValueRange{relu});
    builder.setInsertionPointAfter(consumer);
    return consumer;
  };
  auto direct = makeConsumer(producer->getResult(0));
  auto viaCast = makeConsumer(casted);
  auto viaShape = makeConsumer(expanded);
  auto viaSlice = makeConsumer(sliced);
  auto one = makeConsumer(single);
  func::ReturnOp::create(builder, loc);
  assert(succeeded(verify(function)));
  SchedulingTarget target;
  target.processorCount = 4;
  auto graph = buildCheckedComputeGraph(function, target);
  for (auto consumer : {direct, viaCast, viaShape, viaSlice})
    assert(graph.predecessors[graph.instanceToIndex.lookup({consumer, 0, 1})].size() == 4);
  const auto &parents = graph.predecessors[graph.instanceToIndex.lookup({one, 0, 1})];
  assert(parents.size() == 1 && graph.nodes[parents[0].first].instance.laneStart == 2);
  // A malformed physical publication must diagnose and fail, not abort.
  builder.setInsertionPointToStart(&producer->getBody().front());
  Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
  publication.getOffsetsMutable().assign(ValueRange{zero});
  std::string diagnostics;
  ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
    diagnostics += diagnostic.str();
    return success();
  });
  assert(failed(buildComputeGraph(function, target)));
  assert(diagnostics.find("graph publication leading geometry is invalid") != std::string::npos);
}

static void testScheduleEvaluation() {
  ComputeGraph graph;
  graph.nodes = {{{nullptr, 0, 1}, 10, {}, 0},
                 {{nullptr, 1, 1}, 20, {}, 1},
                 {{nullptr, 2, 1}, 5, {}, 2}};
  graph.predecessors.resize(3);
  graph.predecessors[2] = {{0, {}}, {1, {.fixed = 5, .networkFlits = 1}}};
  for (Cost link : {2, 20}) {
    SchedulingTarget target;
    target.processorCount = 2;
    target.interProcessorLatencyNs = {0, link, link, 0};
    std::string error;
    auto costs = evaluateSchedule(graph, {{0, 2}, {1}}, target, {}, error);
    assert(mlir::succeeded(costs));
    assert(costs->totalWork == 35);
    assert(costs->coreWork == std::vector<Time>({15, 20}));
    assert(costs->communication == 5 + link);
    assert(costs->criticalPath == 30 + link);
    assert(costs->completion == costs->criticalPath);
    assert(costs->stageSpans.front() == costs->completion);
    costs = evaluateSchedule(graph, {{0}, {1, 2}}, target, {0, 1}, error);
    assert(mlir::succeeded(costs));
    assert(costs->stageSpans == std::vector<Time>({10, 25}));
    assert(costs->completion == 25);
    assert(mlir::failed(evaluateSchedule(graph, {{2, 0}, {1}}, target, {}, error)));
    assert(error.find("cyclic") != std::string::npos);
    assert(mlir::failed(evaluateSchedule(graph, {{0, 2}, {}}, target, {}, error)));
    assert(mlir::failed(evaluateSchedule(graph, {{0, 2}, {0, 1}}, target, {}, error)));
  }
}

static void testPipelinePlacementNearTie() {
  ComputeGraph graph;
  for (Cost cost : {1000, 1000, 1000, 2000, 999, 1})
    graph.nodes.push_back({{}, cost, {}, graph.nodes.size()});
  graph.predecessors.resize(6);
  graph.predecessors[1] = {{0, {}}};
  graph.predecessors[2] = {{1, {}}};
  SchedulingTarget target;
  target.processorCount = 2;
  target.interProcessorLatencyNs.assign(4, 0);
  std::string error;
  auto slow = evaluateSchedule(graph, {{3, 0}, {1, 2, 4, 5}}, target, {}, error);
  auto fast = evaluateSchedule(graph, {{0, 1, 2, 5}, {3, 4}}, target, {}, error);
  assert(mlir::succeeded(slow) && mlir::succeeded(fast));
  assert(slow->completion == 6000 && fast->completion == 3001);
  StagePlacementScore oldScore{3000, slow->completion, 3000};
  StagePlacementScore newScore{3001, fast->completion, 3001};
  assert(newScore > oldScore);
  assert(isBetterStagePlacement(newScore, oldScore));
  assert(!isBetterStagePlacement(oldScore, newScore));
  assert(isBetterStagePlacement({2900, 6000, 2900}, newScore));
  assert(!isBetterStagePlacement(newScore, newScore));
}

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(runPipelineScheduler(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::MLIRContext::Threading::DISABLED);
  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(mlir::MLIRContext::Threading::DISABLED);
  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() {
  testResidentWeightPackingRepair();
  testProducerChunkTransferCost();
  testPeftReservationPassFailure();
  testLocalFragmentInsertion();
  testResidentWeightPackingRecovery();
  testBlueprintSourceProjections();
  testPipelinePlacementNearTie();
  testTransposeCost();
  testConcatCost();
  testScalarProjectedTransferCost();
  testBatchedReceiveForwarding();
  testResidentWeightViewCost();
  testPipelineStageSpan();
  testPipelinePackingCapacity();
  testRepeatedResidentWeightPacking();
  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(getTransferTime(transfer, 0, 0, fast) == 0);
  assert(getTransferTime(transfer, 0, 1, fast) == 62);
  assert(getTransferTime(transfer, 0, 1, slow) == 90);
  assert(fast.getInterProcessorLatencyNs(0, 1) == 3);
  assert(slow.getInterProcessorLatencyNs(0, 1) == 10);

  testPhysicalPeft();
  testPeftGapOrder();
  testMaterializedPeftGapOrder();
  testPhysicalPublicationDependencies();
  testScheduleEvaluation();
  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(runPipelineScheduler(
    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(runPipelineScheduler(
    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
              + getTransferTime(
                  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(runPipelineScheduler(
    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(runPipelineScheduler(
    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(runPipelineScheduler(
    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;
}