better throughput scheduling

update results_comparison.csv

test/PIM/SpatialSchedulingTargetTest.cpp

@@ -149,7 +149,8 @@ static void testBatchedReceiveForwarding() {
 
 static void testResidentWeightViewCost() {
   mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
-  context.loadDialect<SpatialDialect, mlir::tensor::TensorDialect>();
+  context.loadDialect<SpatialDialect, mlir::tensor::TensorDialect,
+                      mlir::arith::ArithDialect, mlir::scf::SCFDialect>();
   mlir::Block operands;
   mlir::OwningOpRef<mlir::ModuleOp> owner =
     mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
@@ -161,7 +162,7 @@ static void testResidentWeightViewCost() {
   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) {
+  auto makeCompute = [&](bool weightView, unsigned inputView, int64_t iterations = 0) {
     builder.setInsertionPointToEnd(owner->getBody());
     auto compute = SpatGraphCompute::create(builder, loc,
       mlir::TypeRange{inputType}, mlir::ValueRange{weight},
@@ -186,14 +187,42 @@ static void testResidentWeightViewCost() {
       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()});
+    mlir::scf::ForOp loop;
+    if (iterations) {
+      auto lower = mlir::arith::ConstantIndexOp::create(builder, loc, 0);
+      auto upper = mlir::arith::ConstantIndexOp::create(builder, loc, iterations);
+      auto step = mlir::arith::ConstantIndexOp::create(builder, loc, 1);
+      loop = mlir::scf::ForOp::create(builder, loc, lower, upper, step, mlir::ValueRange{x});
+      builder.setInsertionPointToStart(loop.getBody());
+      x = loop.getRegionIterArg(0);
+    }
+    mlir::Value result = SpatVMMOp::create(builder, loc, inputType, w, x);
+    if (loop) {
+      mlir::scf::YieldOp::create(builder, loc, mlir::ValueRange{result});
+      builder.setInsertionPointAfter(loop);
+      result = loop.getResult(0);
+    }
+    SpatYieldOp::create(builder, loc, mlir::ValueRange{result});
     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);
+  auto invariantLoop = makeCompute(false, false, 1000000000);
+  assert(collectDistinctResidentWeights(invariantLoop.op, 0).size() == 1);
+  Cost loopBodyCost = getComputeInstanceCost(invariantLoop, SchedulingTarget{}) / 1000000000;
+  auto loop = *mlir::cast<SpatGraphCompute>(invariantLoop.op)
+                 .getBody().front().getOps<mlir::scf::ForOp>().begin();
+  builder.setInsertionPoint(loop);
+  auto lower = mlir::arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::min());
+  auto upper = mlir::arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::max());
+  auto step = mlir::arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::max());
+  loop.getLowerBoundMutable().assign(lower.getResult());
+  loop.getUpperBoundMutable().assign(upper.getResult());
+  loop.getStepMutable().assign(step.getResult());
+  assert(getComputeInstanceCost(invariantLoop, SchedulingTarget{})
+    == 3 * loopBodyCost);
   for (Cost readLatency : {1, 20}) {
     SchedulingTarget target;
     target.localMemoryReadLatencyCycles = readLatency;
@@ -373,6 +402,53 @@ static void testScalarProjectedTransferCost() {
   }
 }
 
+static void testStridedPublicationMemory() {
+  using namespace mlir;
+  MLIRContext context(MLIRContext::Threading::DISABLED);
+  context.loadDialect<SpatialDialect, func::FuncDialect, tensor::TensorDialect>();
+  auto module = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(module.getBodyRegion());
+  auto loc = module.getLoc();
+  auto image = RankedTensorType::get({2, 16}, builder.getF32Type());
+  auto fragment = RankedTensorType::get({2, 8}, builder.getF32Type());
+  auto function = func::FuncOp::create(builder, loc, "strided_publication",
+    builder.getFunctionType({image}, {}));
+  builder.setInsertionPointToStart(function.addEntryBlock());
+  auto producer = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
+    TypeRange{image}, ValueRange{}, ValueRange{function.getArgument(0)});
+  builder.setInsertionPointToEnd(&producer.getBody().front());
+  Value value = SpatReluOp::create(builder, loc, image, *producer.getInputArgument(0));
+  SpatYieldOp::create(builder, loc, ValueRange{value});
+  builder.setInsertionPointAfter(producer);
+  for (unsigned reader = 0; reader < 2; ++reader) {
+    auto consumer = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
+      TypeRange{fragment}, ValueRange{}, ValueRange{producer.getResult(0)});
+    builder.setInsertionPointToEnd(&consumer.getBody().front());
+    value = tensor::ExtractSliceOp::create(builder, loc, fragment,
+      *consumer.getInputArgument(0),
+      SmallVector<OpFoldResult>(2, builder.getIndexAttr(0)),
+      SmallVector<OpFoldResult>{builder.getIndexAttr(2), builder.getIndexAttr(8)},
+      SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
+    value = SpatReluOp::create(builder, loc, fragment, value);
+    SpatYieldOp::create(builder, loc, ValueRange{value});
+    builder.setInsertionPointAfter(consumer);
+  }
+  func::ReturnOp::create(builder, loc);
+  for (Cost width : {4, 32}) {
+    SchedulingTarget target;
+    target.processorCount = 3;
+    target.sharedMemoryWidthBytes = width;
+    target.sharedMemoryReadLatencyNs = 7;
+    target.sharedMemoryWriteLatencyNs = 11;
+    auto graph = buildCheckedComputeGraph(function, target);
+    auto demand = getPipelineBoundaryMemoryDemand(graph, {0, 1, 1}, target);
+    assert(demand.readBytes == 32 && demand.writeBytes == 32);
+    assert(demand.readRequests == 2 && demand.writeRequests == 4);
+    assert(demand.serviceNs == (width == 4 ? 144 : 58));
+  }
+  module.erase();
+}
+
 static void testProducerChunkTransferCost() {
   using namespace mlir;
   MLIRContext context(MLIRContext::Threading::DISABLED);
@@ -908,6 +984,100 @@ static void testMaterializedPeftGapOrder() {
     }
 }
 
+static void checkPublicationMemoryTraffic(mlir::func::FuncOp function,
+    const ComputeGraph &graph, const SchedulingTarget &target,
+    mlir::PatternRewriter &builder) {
+  MergeScheduleResult schedule;
+  schedule.processorCount = target.processorCount;
+  schedule.processorStages = {0, 0, 1, 1};
+  std::vector<std::vector<size_t>> tasks(4);
+  std::vector<size_t> stages;
+  for (auto [task, node] : llvm::enumerate(graph.nodes)) {
+    bool producer = node.instance.op == graph.nodes.front().instance.op;
+    size_t core = producer ? node.instance.laneStart % 2 : 2 + task % 2;
+    stages.push_back(producer ? 0 : 1);
+    schedule.dominanceOrderCompute.push_back(node.instance);
+    schedule.computeToCpuMap[node.instance] = core;
+    schedule.computeToCpuSlotMap[node.instance] = tasks[core].size();
+    tasks[core].push_back(task);
+  }
+  auto expected = getPipelineBoundaryMemoryDemand(graph, stages, target);
+  populateScheduleEquivalenceClasses(graph, tasks, schedule);
+  auto materialized = materializeScheduledCompute(function, schedule, builder);
+  assert(mlir::succeeded(materialized));
+  auto plan = buildDeferredTransferPlan(function, *materialized, 2, 4);
+  assert(mlir::succeeded(plan));
+  SharedMemoryDemand actual;
+  llvm::DenseSet<int64_t> stored;
+  for (const auto &exchange : plan->exchanges)
+    for (const auto &transfer : exchange->external)
+      if (transfer.hostRouted) {
+        auto type = mlir::cast<mlir::ShapedType>(transfer.requirement->publicationFragmentType);
+        Cost bytes = (type.getNumElements() * target.computeBitwidth + 7) / 8;
+        Cost requests = transfer.targetLanes.size();
+        actual += target.getSharedMemoryDemand(bytes, requests, false);
+        for (size_t index = 0; index < requests; ++index)
+          if (stored.insert(transfer.hostOffsets.valueAt(index)).second)
+            actual += target.getSharedMemoryDemand(bytes, 1, true);
+      }
+  assert(actual.readBytes == expected.readBytes && actual.writeBytes == expected.writeBytes);
+  assert(actual.readRequests == expected.readRequests && actual.writeRequests == expected.writeRequests);
+  assert(actual.serviceNs == expected.serviceNs);
+  for (unsigned registers : {3u, 16u}) {
+    auto shared = buildDeferredTransferPlan(function, *materialized, 2, 4);
+    assert(mlir::succeeded(shared));
+    auto communication = scheduleDeferredCommunication(function, *shared);
+    assert(mlir::succeeded(communication));
+    auto boundaries = buildDeferredBoundaryPlan(*shared, *communication, registers);
+    assert(mlir::succeeded(boundaries));
+    size_t stores = 0, signals = 0, readyWaits = 0, freeWaits = 0, releases = 0;
+    llvm::DenseSet<int64_t> published;
+    for (const auto &boundary : boundaries->boundaries)
+      for (const auto &instruction : boundary.instructions) {
+        if (auto send = std::get_if<EmitSendRun>(&instruction))
+          for (const auto &slice : send->slices)
+            if (slice.family->hostRouted)
+              for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+                size_t index = slice.familyOffset + offset;
+                bool first = published.insert(slice.family->hostOffsets.valueAt(index)).second;
+                assert(bool(slice.family->hostStoreFlags.valueAt(index)) == first);
+                stores += first;
+                ++signals;
+              }
+        if (auto wait = std::get_if<EmitHostReuseWait>(&instruction))
+          for (auto interval : wait->lanes.intervals())
+            for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
+              assert(wait->eventRegisters.valueAt(lane) < registers);
+              freeWaits += wait->waitValues.valueAt(lane);
+            }
+      }
+    for (const auto &exchange : shared->exchanges)
+      for (const auto &family : exchange->external)
+        if (family.hostRouted)
+          for (size_t index = 0; index < family.targetCores.size(); ++index) {
+            readyWaits += family.waitValues.valueAt(index);
+            assert(family.eventRegisters.valueAt(index) < registers);
+          }
+    for (const auto &[reader, releaseSignals] : shared->hostReleaseSignals) {
+      llvm::DenseSet<int64_t> writers;
+      for (const auto &release : releaseSignals) {
+        assert(writers.insert(release.writerCore).second);
+        assert(release.eventRegister < registers);
+        ++releases;
+      }
+    }
+    assert(stores == expected.writeRequests && signals == expected.readRequests);
+    assert(readyWaits == signals && freeWaits == releases && releases == 4);
+  }
+  auto insufficient = buildDeferredTransferPlan(function, *materialized, 2, 4);
+  assert(mlir::succeeded(insufficient));
+  auto communication = scheduleDeferredCommunication(function, *insufficient);
+  assert(mlir::succeeded(communication));
+  mlir::ScopedDiagnosticHandler suppress(function.getContext(),
+    [](mlir::Diagnostic &) { return mlir::success(); });
+  assert(mlir::failed(buildDeferredBoundaryPlan(*insufficient, *communication, 2)));
+}
+
 static void testPhysicalPublicationDependencies() {
   using namespace mlir;
   MLIRContext context(MLIRContext::Threading::DISABLED);
@@ -923,7 +1093,7 @@ static void testPhysicalPublicationDependencies() {
   Block *body = function.addEntryBlock();
   builder.setInsertionPointToStart(body);
   auto makeBatch = [&](Value input, uint32_t lanes, int64_t projectionSize,
-                       bool reverse) {
+                       bool reverse, int64_t iterations = 0, bool invariantProjection = false) {
     auto inputType = cast<RankedTensorType>(input.getType());
     auto fragment = projectionSize
       ? RankedTensorType::get({projectionSize, 1}, builder.getF32Type())
@@ -934,6 +1104,16 @@ static void testPhysicalPublicationDependencies() {
     assert(succeeded(batch));
     Value lane = *batch->getLaneArgument();
     Value value = *batch->getInputArgument(0);
+    scf::ForOp loop;
+    if (iterations) {
+      Value lower = arith::ConstantIndexOp::create(builder, loc, 0);
+      Value upper = arith::ConstantIndexOp::create(builder, loc, iterations);
+      Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+      Value empty = tensor::EmptyOp::create(builder, loc, fragment.getShape(), builder.getF32Type());
+      loop = scf::ForOp::create(builder, loc, lower, upper, step, ValueRange{empty});
+      builder.setInsertionPointToStart(loop.getBody());
+      lane = invariantProjection ? lower : loop.getInductionVar();
+    }
     if (projectionSize) {
       Value size = arith::ConstantIndexOp::create(builder, loc, projectionSize);
       Value offset = arith::MulIOp::create(builder, loc, lane, size);
@@ -944,6 +1124,12 @@ static void testPhysicalPublicationDependencies() {
         SmallVector<OpFoldResult>(2, builder.getIndexAttr(1)));
     }
     value = SpatReluOp::create(builder, loc, fragment, value);
+    if (loop) {
+      scf::YieldOp::create(builder, loc, ValueRange{value});
+      builder.setInsertionPointAfter(loop);
+      value = loop.getResult(0);
+      lane = *batch->getLaneArgument();
+    }
     Value slot = lane;
     if (reverse) {
       Value last = arith::ConstantIndexOp::create(builder, loc, lanes - 1);
@@ -965,6 +1151,8 @@ static void testPhysicalPublicationDependencies() {
   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);
+  auto looped = makeBatch(producer.getResult(0), 1, 2, false, 2);
+  auto repeated = makeBatch(producer.getResult(0), 1, 2, false, 2, true);
   SmallVector<onnx_mlir::FragmentAssemblyEntry> entries{
     {0, 0, {0}, {1}}, {2, 0, {1}, {1}}};
   auto logicalType = RankedTensorType::get({2}, builder.getF32Type());
@@ -1004,7 +1192,19 @@ static void testPhysicalPublicationDependencies() {
     SchedulingTarget target;
     target.processorCount = 4;
     target.processorPeriodNs = period;
+    target.sharedMemoryWidthBytes = 32;
+    target.sharedMemoryReadLatencyNs = 7;
+    target.sharedMemoryWriteLatencyNs = 11;
     ComputeGraph graph = buildCheckedComputeGraph(function, target);
+    std::vector<size_t> memoryStages(graph.nodes.size(), 1);
+    for (size_t task = 0; task < graph.nodes.size(); ++task)
+      if (graph.nodes[task].instance.op == producer)
+        memoryStages[task] = 0;
+    auto memory = getPipelineBoundaryMemoryDemand(graph, memoryStages, target);
+    assert(memory.readBytes == 26 && memory.writeBytes == 4);
+    assert(memory.readRequests == 26 && memory.writeRequests == 4);
+    assert(memory.serviceNs == 26 * 7 + 4 * 11);
+    assert(graph.baselineMemory.readBytes == 4 && graph.baselineMemory.readRequests == 4);
     auto predecessorLanes = [&](ComputeInstance instance) {
       std::vector<uint32_t> lanes;
       for (const auto &[task, cost] : graph.predecessors[graph.instanceToIndex.lookup(instance)])
@@ -1016,6 +1216,8 @@ static void testPhysicalPublicationDependencies() {
     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({looped, 0, 1}) == std::vector<uint32_t>{0, 1, 2, 3}));
+    assert((predecessorLanes({repeated, 0, 1}) == std::vector<uint32_t>{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);
@@ -1031,8 +1233,24 @@ static void testPhysicalPublicationDependencies() {
     assert(dependencies.size() == 1);
     assert(dependencies.front().second.fixed == 8 * period);
     assert(dependencies.front().second.networkFlits == 12);
+    if (period == 3)
+      checkPublicationMemoryTraffic(function, graph, target, builder);
   }
 
+  builder.setInsertionPointToEnd(owner->getBody());
+  auto largeLoop = func::FuncOp::create(builder, loc, "invariant_loop",
+    builder.getFunctionType({fragmentType}, {}));
+  builder.setInsertionPointToStart(largeLoop.addEntryBlock());
+  auto loopProducer = makeBatch(largeLoop.getArgument(0), 4, 0, false);
+  makeBatch(loopProducer.getResult(0), 1, 2, false, 1000000000, true);
+  func::ReturnOp::create(builder, loc);
+  SchedulingTarget loopTarget;
+  loopTarget.processorCount = 4;
+  auto compactGraph = buildCheckedComputeGraph(largeLoop, loopTarget);
+  assert(compactGraph.publicationTransfers.size() == 2);
+  for (const auto &transfer : compactGraph.publicationTransfers)
+    assert(transfer.count == 1000000000);
+
   builder.setInsertionPointToEnd(owner->getBody());
   auto independent = func::FuncOp::create(builder, loc, "batch_order",
     builder.getFunctionType({fragmentType}, {}));
@@ -1303,6 +1521,152 @@ static void testPipelineInputArrival() {
   }
 }
 
+static void testBaselineConstantMemory() {
+  using namespace mlir;
+  MLIRContext context(MLIRContext::Threading::DISABLED);
+  context.loadDialect<SpatialDialect, func::FuncDialect, arith::ArithDialect,
+    tensor::TensorDialect, scf::SCFDialect>();
+  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(module->getBodyRegion());
+  auto loc = module->getLoc();
+  auto type = RankedTensorType::get({16}, builder.getF32Type());
+  auto function = func::FuncOp::create(builder, loc, "constant_memory",
+    builder.getFunctionType({type}, {type}));
+  builder.setInsertionPointToStart(function.addEntryBlock());
+  Value external = arith::ConstantOp::create(builder, loc, type,
+    DenseElementsAttr::get(type, builder.getF32FloatAttr(2.5)));
+  auto compute = onnx_mlir::createEmptySpatGraphCompute(builder, loc,
+    TypeRange{type}, ValueRange{}, ValueRange{function.getArgument(0), external});
+  Value start = arith::ConstantIndexOp::create(builder, loc, 0);
+  Value end = arith::ConstantIndexOp::create(builder, loc, 3);
+  Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+  Value captured = arith::ConstantOp::create(builder, loc, type,
+    DenseElementsAttr::get(type, builder.getF32FloatAttr(1.5)));
+  auto loop = scf::ForOp::create(builder, loc, start, end, step,
+    ValueRange{*compute.getInputArgument(0)});
+  builder.setInsertionPointToStart(loop.getBody());
+  Value result = SpatVAddOp::create(builder, loc, type,
+    loop.getRegionIterArg(0), captured);
+  scf::YieldOp::create(builder, loc, ValueRange{result});
+  builder.setInsertionPointAfter(loop);
+  result = SpatVAddOp::create(builder, loc, type, loop.getResult(0), *compute.getInputArgument(1));
+  SpatYieldOp::create(builder, loc, ValueRange{result});
+  builder.setInsertionPointAfter(compute);
+  func::ReturnOp::create(builder, loc, compute.getResults());
+  assert(succeeded(verify(function)));
+  SchedulingTarget target;
+  target.processorCount = 1;
+  target.sharedMemoryWidthBytes = 32;
+  target.sharedMemoryReadLatencyNs = 7;
+  target.sharedMemoryWriteLatencyNs = 11;
+  auto graph = buildCheckedComputeGraph(function, target);
+  assert(graph.baselineMemory.readBytes == 80);
+  assert(graph.baselineMemory.readRequests == 5);
+  assert(graph.baselineMemory.writeBytes == 16 && graph.baselineMemory.writeRequests == 1);
+  assert(graph.baselineMemory.serviceNs == 46);
+}
+
+static void testPipelineSharedMemory() {
+  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
+  auto fragment = mlir::RankedTensorType::get({33}, mlir::IntegerType::get(&context, 8));
+  mlir::Block generations;
+  auto firstGeneration = generations.addArgument(fragment, mlir::UnknownLoc::get(&context));
+  auto nextGeneration = generations.addArgument(fragment, mlir::UnknownLoc::get(&context));
+  auto shared = getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 0, fragment, true);
+  assert(shared == getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 0, fragment, true));
+  assert(shared != getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 1, fragment, true));
+  assert(shared != getHostPublicationIdentity(nextGeneration.getAsOpaquePointer(), 0, fragment, true));
+  assert(!getHostPublicationIdentity(firstGeneration.getAsOpaquePointer(), 0, fragment, false));
+  ScheduleCostSummary previous, candidate;
+  previous.stageSpanNs = {25000};
+  previous.maximumCoreWorkNs = 15000;
+  previous.baselineMemory.serviceNs = 50000;
+  candidate.stageSpanNs = {32000};
+  candidate.maximumCoreWorkNs = 27000;
+  candidate.baselineMemory.serviceNs = 40000;
+  // Reduced memory and stage demand must not hide concentrated core work.
+  assert(!isBetterPipelineSchedule(candidate, previous));
+  candidate.maximumCoreWorkNs = 24000;
+  assert(isBetterPipelineSchedule(candidate, previous));
+  candidate.stageSpanNs = {36000};
+  assert(!isBetterPipelineSchedule(candidate, previous));
+  assert(!isBetterPipelineSchedule(previous, previous));
+  ComputeGraph graph;
+  graph.predecessors.resize(4);
+  graph.successors.resize(4);
+  for (uint32_t task = 0; task < 4; ++task) {
+    graph.nodes.push_back({{nullptr, task, 1}, 1000, {}, task});
+    if (task) {
+      graph.predecessors[task].push_back({task - 1, {}});
+      graph.successors[task - 1].push_back({task, {}});
+      graph.publicationTransfers.push_back({task - 1, task, {}, 0, {},
+        task == 2 ? Cost{16384} : Cost{32}, 1});
+    }
+  }
+  SchedulingTarget target;
+  target.processorCount = 4;
+  target.residentWeightCapacity = 4;
+  target.interProcessorLatencyNs.assign(16, 0);
+  target.sharedMemoryWidthBytes = 32;
+  auto stages = [&](const MergeScheduleResult &schedule) {
+    std::vector<size_t> result;
+    for (const auto &node : graph.nodes)
+      result.push_back(schedule.processorStages[schedule.computeToCpuMap.lookup(node.instance)]);
+    return result;
+  };
+  for (Cost latency : {Cost{1}, Cost{10}}) {
+    target.sharedMemoryReadLatencyNs = latency;
+    target.sharedMemoryWriteLatencyNs = latency;
+    graph.baselineMemory = target.getSharedMemoryDemand(32, 16, false);
+    MergeScheduleResult schedule;
+    std::string error;
+    assert(mlir::succeeded(runPipelineScheduler(graph, schedule, 2, target, error)));
+    auto assignment = stages(schedule);
+    assert(assignment.front() == 0 && assignment.back() == 1);
+    for (size_t task = 1; task < 4; ++task)
+      assert(assignment[task] >= assignment[task - 1]
+        && assignment[task] <= assignment[task - 1] + 1);
+    auto memory = getPipelineBoundaryMemoryDemand(graph, assignment, target);
+    assert(memory.readRequests == 1 && memory.writeRequests == 1);
+    assert(memory.readBytes == (latency == 1 ? 16384 : 32));
+    assert(memory.writeBytes == memory.readBytes);
+    assert(schedule.costs.boundaryMemory.serviceNs == memory.serviceNs);
+  }
+  // Per-fragment rounding and fan-out must precede aggregation.
+  auto rounded = target.getSharedMemoryDemand(33, 3, true);
+  assert(rounded.writeBytes == 99 && rounded.writeRequests == 3);
+  assert(rounded.serviceNs == 60);
+  graph.publicationTransfers = {{0, 1, {}, 3, {}, 33, 2}, {0, 2, {}, 3, {}, 33, 1}};
+  auto fanout = getPipelineBoundaryMemoryDemand(graph, {0, 1, 1, 1}, target);
+  assert(fanout.readBytes == 99 && fanout.writeBytes == 99);
+  assert(fanout.readRequests == 3 && fanout.writeRequests == 3);
+  assert(fanout.serviceNs == 120);
+  graph.publicationTransfers = {
+    {0, 1, firstGeneration, 0, {}, 33, 2, fragment},
+    {0, 2, firstGeneration, 0, {}, 33, 1, fragment}};
+  fanout = getPipelineBoundaryMemoryDemand(graph, {0, 1, 1, 1}, target);
+  assert(fanout.readBytes == 99 && fanout.writeBytes == 33);
+  assert(fanout.readRequests == 3 && fanout.writeRequests == 1);
+  assert(fanout.serviceNs == 80);
+  graph.publicationTransfers.push_back({0, 3, nextGeneration, 0, {}, 33, 1, fragment});
+  graph.publicationTransfers.push_back({0, 3, firstGeneration, 1, {}, 33, 1, fragment});
+  fanout = getPipelineBoundaryMemoryDemand(graph, {0, 1, 1, 1}, target);
+  assert(fanout.readRequests == 5 && fanout.writeRequests == 3);
+  assert(getPipelineBoundaryMemoryDemand(graph, {0, 0, 0, 0}, target).serviceNs == 0);
+  graph.predecessors[3].push_back({0, {}});
+  graph.successors[0].push_back({3, {}});
+  graph.publicationTransfers.push_back({0, 3, {}, 0, {}, 4096, 1});
+  MergeScheduleResult constrained;
+  std::string error;
+  assert(mlir::succeeded(runPipelineScheduler(graph, constrained, 4, target, error)));
+  auto assignment = stages(constrained);
+  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    for (const auto &[predecessor, transfer] : graph.predecessors[task])
+      assert(assignment[task] >= assignment[predecessor]
+        && assignment[task] <= assignment[predecessor] + 1);
+  assert(mlir::failed(runPipelineScheduler(graph, constrained, 5, target, error)));
+}
+
 static void testReadyCommunication() {
   mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
   mlir::OwningOpRef<mlir::ModuleOp> owner =
@@ -1372,6 +1736,9 @@ static void testPipelineRestartOrder() {
 }
 
 int main() {
+  testStridedPublicationMemory();
+  testPipelineSharedMemory();
+  testBaselineConstantMemory();
   testResidentWeightPackingRepair();
   testProducerChunkTransferCost();
   testPeftReservationPassFailure();