better throughput scheduling

update results_comparison.csv

test/PIM/CMakeLists.txt

@@ -40,9 +40,12 @@ add_pim_unittest(SpatialSchedulingTargetTest
 add_test(
   NAME PimHostReuseSynchronizationTest
   COMMAND "${PYTHON_EXECUTABLE}"
-          "${CMAKE_SOURCE_DIR}/validation/tools/pim/pimcomp/compare/test_PIMCOMP_adversarial_memory_sync.py"
+          "${CMAKE_CURRENT_SOURCE_DIR}/test_HostPublication.py"
           --compiler "$<TARGET_FILE:onnx-mlir>"
-          --simple-model "${CMAKE_SOURCE_DIR}/validation/operations/relu/after_conv/relu_after_conv.onnx"
-          --grouped-model "${CMAKE_SOURCE_DIR}/validation/operations/conv/relu_conv_store/conv_relu_conv_store.onnx"
 )
 set_tests_properties(PimHostReuseSynchronizationTest PROPERTIES LABELS pim-unittest)
+
+add_test(NAME PimTargetConfigTest
+  COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PimTargetConfig.py"
+          --compiler "$<TARGET_FILE:onnx-mlir>")
+set_tests_properties(PimTargetConfigTest PROPERTIES LABELS pim-unittest)

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();

test/PIM/test_HostPublication.py

@@ -0,0 +1,177 @@
+import argparse
+from collections import defaultdict, deque
+import json
+import os
+from pathlib import Path
+import struct
+import subprocess
+import sys
+import tempfile
+from types import SimpleNamespace
+
+
+REPO = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(REPO / "validation/tools/pim/pimcomp/correctness/synchronization"))
+import global_memory as sync
+
+
+def read_events(artifact):
+    streams = {}
+    for path in artifact.glob("core_*.pim"):
+        core = int(path.stem.split("_")[1])
+        registers, events = {}, []
+        data = path.read_bytes()
+        assert data[:4] == b"PIMB"
+        for pc, (op, rd, rs, flags, arg, selection, offset, size) in enumerate(
+                struct.iter_unpack("<BBBBiiii", data[12:])):
+            if op == 1:
+                registers[rd] = arg
+            elif op in (25, 26):
+                register, bit = (rs, 2) if op == 25 else (rd, 1)
+                address = registers[register] + (offset if selection & bit else 0)
+                events.append((pc, op, address, size))
+            elif op in (29, 30, 31, 32):
+                events.append((pc, op, arg, selection, offset))
+        if events:
+            streams[core] = events
+    return streams
+
+
+def check_lifetimes(streams):
+    graph, signals, waits = defaultdict(list), defaultdict(list), {}
+    sends, receives, stores, loads = defaultdict(list), defaultdict(list), {}, defaultdict(list)
+    for core, events in streams.items():
+        for iteration in range(3):
+            nodes = [(core, iteration, event[0]) for event in events]
+            for left, right in zip(nodes, nodes[1:]):
+                graph[left].append(right)
+            if iteration:
+                graph[(core, iteration - 1, events[-1][0])].append(nodes[0])
+            for node, event in zip(nodes, events):
+                pc, op, *args = event
+                if op == 32:
+                    target, register, _ = args
+                    assert 0 <= register < 32, (core, event)
+                    signals[target, iteration, register].append(node)
+                elif op == 31:
+                    _, register, expected = args
+                    key = core, iteration, register
+                    assert key not in waits and 0 <= register < 32
+                    waits[key] = node, expected
+                elif op in (29, 30):
+                    peer = args[0]
+                    key = (core, peer, iteration) if op == 29 else (peer, core, iteration)
+                    (sends if op == 29 else receives)[key].append(node)
+                elif op == 26:
+                    key = iteration, *args
+                    assert key not in stores, ("duplicate slot store", key)
+                    stores[key] = node
+                elif op == 25:
+                    loads[iteration, *args].append(node)
+    assert signals.keys() == waits.keys()
+    for key, (wait, expected) in waits.items():
+        assert len(signals[key]) == expected, (key, expected, signals[key])
+        for signal in signals[key]:
+            graph[signal].append(wait)
+    assert sends.keys() == receives.keys()
+    for key, send_nodes in sends.items():
+        assert len(send_nodes) == len(receives[key])
+        for send, receive in zip(send_nodes, receives[key]):
+            graph[send].append(receive)
+    indegree = defaultdict(int)
+    for source, successors in list(graph.items()):
+        indegree[source] += 0
+        for successor in successors:
+            indegree[successor] += 1
+    pending = deque(node for node, degree in indegree.items() if degree == 0)
+    visited = 0
+    while pending:
+        visited += 1
+        for successor in graph[pending.popleft()]:
+            indegree[successor] -= 1
+            if indegree[successor] == 0:
+                pending.append(successor)
+    assert visited == len(indegree), "communication happens-before cycle"
+
+    def reaches(source, target):
+        pending, seen = [source], {source}
+        while pending:
+            for successor in graph[pending.pop()]:
+                if successor == target:
+                    return True
+                if successor not in seen:
+                    seen.add(successor)
+                    pending.append(successor)
+        return False
+
+    shared, stalled_reader = 0, None
+    for (iteration, address, size), writer in stores.items():
+        readers = loads.get((iteration, address, size), [])
+        if not readers:
+            continue
+        shared += iteration == 0 and len({reader[0] for reader in readers}) > 1
+        if len({reader[0] for reader in readers}) > 1:
+            stalled_reader = readers[-1]
+        for reader in readers:
+            assert reaches(writer, reader), ("missing READY", writer, reader)
+            if iteration < 2:
+                assert reaches(reader, stores[iteration + 1, address, size]), (
+                    "unsafe overwrite", reader, address)
+    assert shared, "fixture must exercise a publication with multiple physical readers"
+    return shared, stalled_reader
+
+
+def run_functional(artifact, inputs, references, outputs, out, policy, seed, reader=None):
+    out.mkdir()
+    command = [str(sync.RUST_BINARY), "--folder", str(artifact),
+        "--output", str(out / "output.bin"), "--dump",
+        sync.compare.build_dump_ranges(artifact / "config.json", outputs),
+        "--mode", "throughput", "--batch-size", "8", "--input-dir", str(inputs[0].parent),
+        "--batch-output-dir", str(out / "iterations"),
+        "--diagnostic-schedule-policy", policy, "--diagnostic-schedule-seed", str(seed)]
+    if reader:
+        command += ["--diagnostic-target-stall", f"{reader[0] + 1}:{reader[2]}:1:10000"]
+    result = subprocess.run(command, capture_output=True, text=True, timeout=120)
+    assert result.returncode == 0, result.stdout + result.stderr
+    for index, reference in enumerate(references[:8]):
+        result = sync.compare.compare_simulator_outputs(
+            out / "iterations" / f"output_{index:06d}.bin", outputs, reference,
+            threshold=1e-3, rtol=1e-5)
+        assert result.passed, result
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--compiler", required=True)
+    args = parser.parse_args()
+    with tempfile.TemporaryDirectory(prefix="raptor-publication-") as temporary:
+        root = Path(temporary)
+        model = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
+        options = SimpleNamespace(timeout=120, threshold=1e-3, rtol=1e-5)
+        samples, inputs, _, _ = sync.make_inputs(model, 8, 0, root)
+        references = sync.make_references(model, samples, root, options)
+        _, outputs = sync.compare.onnx_io(model)
+        for arch, cores, weights in (("arch-a", 168, 96), ("arch-b", 138, 128)):
+            out = root / arch
+            out.mkdir()
+            command = [args.compiler, str(model), "-o", str(out / "model"),
+                "--maccel=PIM", "--EmitPimCodegen", "--crossbar-size=128",
+                f"--crossbar-count={weights}", f"--core-count={cores}", "--pipeline=4",
+                f"--pim-target-config={sync.CONFIG_ROOT / arch / 'throughput_config_1000ms.json'}",
+                "--mlir-disable-threading", "--pim-emit-json", "--pim-detect-communication-deadlock"]
+            compiled = subprocess.run(command, capture_output=True, text=True,
+                timeout=120, env=dict(os.environ, OMP_NUM_THREADS="1"))
+            assert compiled.returncode == 0, compiled.stdout + compiled.stderr
+            artifact = out / "pim"
+            streams = read_events(artifact)
+            shared, reader = check_lifetimes(streams)
+            for policy, seed in (("greedy", 0), ("randomized", 1), ("randomized", 7)):
+                run_functional(artifact, inputs, references, outputs,
+                    out / f"{policy}-{seed}", policy, seed)
+            run_functional(artifact, inputs, references, outputs,
+                out / "stalled", "randomized", 11, reader)
+            print(f"{arch}: {shared} shared publications; three-iteration lifetime proof; eight inputs; randomized and stalled execution passed")
+
+
+if __name__ == "__main__":
+    main()

test/PIM/test_PimTargetConfig.py

@@ -0,0 +1,50 @@
+import argparse
+import copy
+import json
+from pathlib import Path
+import subprocess
+import tempfile
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--compiler", required=True)
+    args = parser.parse_args()
+    repo = Path(__file__).resolve().parents[2]
+    config_path = repo / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
+    config = json.loads(config_path.read_text())
+    network = config["chip_config"]["network_config"]
+    network["net_config_file_path"] = str(config_path.parent / network["net_config_file_path"])
+    with tempfile.TemporaryDirectory(prefix="raptor-target-config-") as temporary:
+        root = Path(temporary)
+        for name, memory, valid in (
+            ("absent", None, True),
+            ("defaults", {}, True),
+            ("explicit", config["chip_config"]["global_memory_config"], True),
+            ("zero-service", {"read_latency_cycle": 0, "write_latency_cycle": 0}, True),
+            ("bad-object", [], False),
+            ("bad-width", {"data_width": 0}, False),
+            ("bad-period", {"period": -1}, False),
+        ):
+            candidate = copy.deepcopy(config)
+            if memory is None:
+                del candidate["chip_config"]["global_memory_config"]
+            else:
+                candidate["chip_config"]["global_memory_config"] = memory
+            path = root / f"{name}.json"
+            path.write_text(json.dumps(candidate))
+            command = [args.compiler,
+                str(repo / "validation/operations/relu/after_conv/relu_after_conv.onnx"),
+                "--maccel=PIM", "--EmitSpatial", "--crossbar-size=128",
+                "--crossbar-count=96", "--core-count=168", "--pipeline=1",
+                "--mlir-disable-threading", f"--pim-target-config={path}",
+                "-o", str(root / name)]
+            result = subprocess.run(command, stdout=subprocess.PIPE,
+                                    stderr=subprocess.STDOUT, text=True)
+            assert (result.returncode == 0) == valid, (name, result.stdout)
+            if not valid:
+                assert "Pim target config" in result.stdout, result.stdout
+
+
+if __name__ == "__main__":
+    main()