better throughput scheduling

update results_comparison.csv

AGENTS.md

@@ -1,5 +1,20 @@
 * Always read the full README.md before doing anything
 
+# Communication: ASD-STE100
+
+* Use ASD-STE100 Simplified Technical English (STE) for all user-facing English prose.
+* Apply this rule to plans, progress updates, questions, explanations, code reviews, and final responses.
+* Use approved STE words with their approved meanings and parts of speech, as far as the available standard permits.
+* Use consistent technical terms when approved general words cannot express an engineering concept accurately.
+* Use short, direct sentences. Use active voice in procedures and one instruction per sentence.
+* Keep procedural sentences to 20 words or fewer and descriptive sentences to 25 words or fewer.
+* Avoid idioms, slang, unnecessary words, and ambiguous expressions.
+* Before sending a response, check its prose against these rules and correct identifiable violations.
+* Preserve technical accuracy if an STE restriction would change the meaning.
+* Do not alter source code, identifiers, APIs, commands, file paths, logs, or quoted text to satisfy STE.
+* Do not rewrite existing documentation or code comments solely to satisfy STE. Apply STE to new prose when practical.
+* Do not claim verified ASD-STE100 compliance without independent validation against the official standard.
+
 # Required project invariants
 
 Before modifying the relevant subsystem, read:

docs/compiler/scheduling-implementation.md

@@ -6,6 +6,24 @@ there; the details below describe current implementation choices.
 
 ## Publication, order, and local materialization
 
+Deferred transfer planning allocates one host slot for repeated full-fragment
+requirements with the same physical producer, payload type, and source slot.
+One immutable producer result identifies one logical inference generation.
+Scalar source projections retain separate publications until their exact
+geometry has a common representation in both analyses.
+
+Boundary planning moves a shared publication's notifications to its earliest
+existing send boundary. The first notification stores the payload before READY;
+later notifications emit READY only. Each reader retains its load and READY
+count. FREE planning includes every physical reader before the first overwrite
+of the shared slot. Existing unique writer-reader releases, register allocation,
+capacity-driven FREE-group merging, and restart barriers remain in force.
+
+Host-slot allocation and shared-boundary planning take expected linear time
+and space in the number of physical consumer transfers. Structured send runs
+retain compact per-lane store flags. They do not duplicate payload shaping
+for notification-only readers.
+
 [`materializeScheduledCompute`](../../src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp) records
 graph IDs for publication identity, not execution order. Class lanes retain
 graph encounter order independently of physical core IDs, preserving regular
@@ -65,10 +83,14 @@ remain approximations, not extra inferred memory traffic.
 
 [`PipelineScheduling.cpp`](../../src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp) caches a per-template
 adjustment between physical execution cost and nominal volume priority, scaling
-it by chunk lanes. Partitioning and packing use the nominal priorities;
+it by chunk lanes. The original partitioning and packing candidate uses nominal
-transposes count as bulk moves there. PEFT, final schedule evaluation, and
+priorities; transposes count as bulk moves there. The memory-aware candidate
-reported per-core demand use physical estimates. These quantities must not be
+uses execution estimates for partitioning, retaining the existing core-budget
-substituted for one another when interpreting results.
+and packing heuristics. Final
+selection evaluates bounded alternatives after physical placement, retaining
+the original candidate when the combined resource estimates worsen. PEFT,
+final schedule evaluation, and reported per-core demand use physical estimates.
+These quantities must not be substituted for measured simulator periods.
 
 Placement compares maximum per-core serial compute plus incident same-stage
 transfer demand, then completion/span and compute load. Each participating core

docs/compiler/scheduling.md

@@ -57,8 +57,37 @@ paths exclude core-order constraints; completion includes them. Stage spans
 exclude cross-stage routing, restart/host synchronization, and contention.
 Simulator latency and measured throughput remain the performance gates.
 
-Pipeline partitioning uses ordered-packing/dependency-tail priorities and
+Pipeline scheduling retains the original ordered-packing/dependency-tail
-nominal movement volumes. PEFT and final evaluation use physical cost estimates.
+heuristic using nominal movement volumes. A second search uses execution
+estimates and aggregate serialized shared-memory service, minimizing their
+maximum. Its DP keeps at most eight compute/memory trade-offs per state;
+four final alternatives and the initial partition undergo the existing core
+allocation, topology placement, and capacity verification. The phase-1 winner
+remains a fallback. A second bounded search accounts for shared publications.
+Selection retains the fallback unless the modeled bottleneck improves. Selection also requires that
+shared-memory service plus maximum stage span does not increase. The same
+check applies to shared-memory service plus maximum per-core compute work.
+These guards limit compute imbalance; their sums do not predict simulator
+periods. Sharing-aware alternatives must also preserve maximum stage span,
+because placement-dependent baseline loads remain approximate. PEFT is unchanged.
+
+`ComputeGraph` retains physical publication and consumer identities separately
+from aggregated dependency edges. Cross-stage demands charge one store per
+shared physical publication and one load per consumer transfer. Publication
+identity includes the immutable result generation, exact payload type, and
+physical source fragment. Scalar projections retain separate stores
+conservatively. Same-stage local/direct transfers incur no host traffic.
+Function inputs, constants, and outputs supply a separate baseline estimate.
+The target adapter supplies access width and read/write service times; request
+rounding precedes aggregation. Serialized service is shared by all stages and
+is not added to individual stage spans. Baseline loads remain approximate
+before placement, bufferization, aliasing, and load folding.
+Scalar publication stores round each contiguous source run separately.
+Consumer load counts can remain conservative before load folding.
+The scheduler evaluates boundary demand again after physical placement.
+Full graph fragments resolve to unique physical producers, so their sharing
+does not depend on reader placement. Deferred planning uses the same identity
+rule and retains separate loads, including repeated loads on one reader core.
 Placement ranks maximum per-core compute/incident-transfer demand, then stage
 completion and compute load; within 1% demand, stage span breaks the tie.
 These are placement heuristics, not a steady-state timing model.
@@ -81,7 +110,8 @@ proved hardware infeasibility.
 ## Complexity and limits
 
 Let `T` be tasks, `E` dependencies, `P` cores, `G` groups, `S` stages,
-`R` weight references, `W` distinct weights, and `C` modeled work.
+`R` weight references, `W` distinct weights, `C` modeled work, `F` physical
+fragment/consumer demands, and bounded frontier size `K = 8`.
 
 | Analysis | Time / storage |
 |---|---|
@@ -89,13 +119,18 @@ Let `T` be tasks, `E` dependencies, `P` cores, `G` groups, `S` stages,
 | Weight indexing | Up to `O(R W)` equality comparisons |
 | Weight lower bounds | `O(R + W)` / `O(W)` bits |
 | Constructive best fit | `O(T log T + P R + P W)` / `O(P W)` bits plus `O(T + P)`; bounded retries only on failure |
-| Partition DP | Up to `O(S G² ((T + R) log C + W))` plus dependency traversal / `O(S G²)` cache |
+| Partition DP | Up to `O(S G² ((T + R) log C + W + K³))` plus dependency traversal / `O(S G² + S G K)`; at most five additional physical placements per memory objective |
+| Cut memory service | Expected `O(F + G)` time / `O(F + G)` scratch and retained demands |
 | Topology clustering / cached PEFT distances | `O(P²)` time |
 | Concat cost extraction | `O(A D)` for `A` inputs of rank `D`; constant auxiliary space |
 | Cached pipeline priorities | `O(T + D)` time / `O(T + G)` space for `D` template-body operations |
-| Exact projections | `O(Q + L)` time/temporary space for `Q` slot visits and `L` producer lanes; publication maps cached |
+| Exact projections | `O(Q + L)` time for `Q` slot visits and `L` producer lanes; repeated demands aggregate by projection and fragment |
 | Conservative producer tracing | Linear in the input projection DAG, including captures |
 
 Latency insertion-gap searches can remain quadratic. Issue overhead,
 eliminated assembly copies, contention, and pipeline synchronization remain
 cost approximations. Changes require matched execution measurements.
+Loop-invariant projections use trip-count multiplication. Loop-invariant
+resident weights are collected once. Loop-dependent projections and weights
+can still require iteration visits; retained projection demands scale with
+distinct physical fragments instead of repeated visits.

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -82,6 +82,10 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() {
   target.matrixInputReadMultiplier = 8;
   target.matrixColumns = crossbarSize.getValue();
   target.synchronizationRegisterCount = kPimEventRegisterCount;
+  target.sharedMemoryWidthBytes = 16;
+  target.sharedMemoryReadLatencyNs = 8;
+  target.sharedMemoryWriteLatencyNs = 10;
+  target.zeroSplatUsesSharedMemory = false;
 
   setDefaultPimInterProcessorLatencies(target);
   return target;
@@ -266,6 +270,18 @@ spatial::SchedulingTarget getPimSchedulingTarget() {
     requireObject(core, "local_memory_config", "chip_config.core_config");
   const llvm::json::Object& network =
     requireObject(chip, "network_config", "chip_config");
+  // Older target configurations omit global memory and use adapter defaults.
+  if (chip.get("global_memory_config")) {
+    const llvm::json::Object& sharedMemory =
+      requireObject(chip, "global_memory_config", "chip_config");
+    Cost memoryPeriod = getConfigCost(sharedMemory, "period", Cost{1});
+    target.sharedMemoryWidthBytes = getConfigCost(
+      sharedMemory, "data_width", target.sharedMemoryWidthBytes);
+    target.sharedMemoryReadLatencyNs = checkedMultiply(memoryPeriod,
+      getConfigCost(sharedMemory, "read_latency_cycle", Cost{8}, true));
+    target.sharedMemoryWriteLatencyNs = checkedMultiply(memoryPeriod,
+      getConfigCost(sharedMemory, "write_latency_cycle", Cost{10}, true));
+  }
 
   std::optional<int64_t> coreCount = chip.getInteger("core_cnt");
   if (!coreCount || *coreCount <= 0)

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp

@@ -514,6 +514,20 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
   DenseMap<BoundaryKey, unsigned> indices;
   DenseMap<DeferredExchangePlan *, unsigned> resultSteps;
   DenseMap<RequirementFamily *, LaneSet> coverage;
+  DenseMap<int64_t, unsigned> publicationSteps;
+  DenseMap<ExternalTransferFamily *, SmallVector<int64_t>> storeFlags;
+  DenseSet<int64_t> storedPublications;
+  for (const ScheduledTransferSlice &slice : schedule.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    if (!family.hostRouted)
+      continue;
+    storeFlags.try_emplace(&family, family.targetCores.size(), 0);
+    for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+      int64_t slot = family.hostOffsets.valueAt(slice.familyOffset + offset);
+      auto [entry, inserted] = publicationSteps.try_emplace(slot, slice.sourceInsertionStep);
+      entry->second = std::min(entry->second, slice.sourceInsertionStep);
+    }
+  }
 
   for (const std::unique_ptr<DeferredExchangePlan> &exchange :
        transfers.exchanges) {
@@ -537,10 +551,23 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
 
   for (const ScheduledTransferSlice &slice : schedule.slices) {
     ExternalTransferFamily &family = *slice.family;
-    BoundaryProgram &source = getBoundary(
+    if (!family.hostRouted) {
-      boundaries, indices,
+      BoundaryProgram &source = getBoundary(boundaries, indices,
-      {family.sourceScheduled, slice.sourceInsertionStep});
+        {family.sourceScheduled, slice.sourceInsertionStep});
-    appendSend(source, slice);
+      appendSend(source, slice);
+    } else {
+      // Notify every reader at the first existing publication boundary. This
+      // also puts all reader lifetimes before the shared slot's reuse wait.
+      for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+        size_t index = slice.familyOffset + offset;
+        int64_t slot = family.hostOffsets.valueAt(index);
+        BoundaryProgram &source = getBoundary(boundaries, indices,
+          {family.sourceScheduled, publicationSteps.lookup(slot)});
+        appendSend(source, {&family, index, 1,
+          publicationSteps.lookup(slot), slice.targetInsertionStep});
+        storeFlags[&family][index] = storedPublications.insert(slot).second;
+      }
+    }
     BoundaryProgram &target = getBoundary(
       boundaries, indices,
       {family.targetScheduled, slice.targetInsertionStep});
@@ -557,6 +584,8 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
                            coverage)))
       return failure();
   }
+  for (auto &[family, flags] : storeFlags)
+    family->hostStoreFlags = StaticIntSequence::fromValues(flags);
 
   for (const std::unique_ptr<DeferredExchangePlan> &exchange :
        transfers.exchanges) {

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp

@@ -22,6 +22,7 @@ struct LogicalTransferMetadataView {
   StaticIntSequenceChain sourceCores;
   StaticIntSequenceChain targetCores;
   StaticIntSequenceChain hostOffsets;
+  StaticIntSequenceChain hostStoreFlags;
   StaticIntSequenceChain eventRegisters;
   StaticIntSequenceChain waitValues;
   StaticIntSequenceChain targetLanes;
@@ -115,6 +116,7 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
   metadata.targetCores.append(family.targetCores, familyIndex, count);
   if (family.hostRouted) {
     metadata.hostOffsets.append(family.hostOffsets, familyIndex, count);
+    metadata.hostStoreFlags.append(family.hostStoreFlags, familyIndex, count);
     metadata.eventRegisters.append(
       family.eventRegisters, familyIndex, count);
     metadata.waitValues.append(family.waitValues, familyIndex, count);
@@ -225,6 +227,7 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     return failure();
   std::optional<StaticIntGrid> hostOffsets;
   std::optional<StaticIntGrid> eventRegisters;
+  std::optional<StaticIntGrid> storeFlags;
   if (firstFamily.hostRouted) {
     auto offsets = buildGrid(
       &LogicalTransferMetadataView::hostOffsets,
@@ -232,10 +235,12 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
     auto events = buildGrid(
       &LogicalTransferMetadataView::eventRegisters,
       logical.eventRegisters.valueAt(0));
-    if (failed(offsets) || failed(events))
+    auto flags = buildGrid(&LogicalTransferMetadataView::hostStoreFlags, 0);
+    if (failed(offsets) || failed(events) || failed(flags))
       return failure();
     hostOffsets = std::move(*offsets);
     eventRegisters = std::move(*events);
+    storeFlags = std::move(*flags);
   }
   GridGeometry projectionGrids;
   for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) {
@@ -265,29 +270,47 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
   auto emitOne = [&](Value action, Value runtimeLane) -> LogicalResult {
     Value localOffset = localOffsets->emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc);
     MixedSliceGeometry projection = lookupGeometry(projectionGrids, action, runtimeLane, anchor, context, loc);
-    auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc);
-    if (failed(payload))
-      return failure();
     Value sourceCore = transferGrids->sourceCores().emitLookup(
       action, runtimeLane, anchor, context.constants, context.rewriter, loc);
     Value targetCore = transferGrids->targetCores().emitLookup(
       action, runtimeLane, anchor, context.constants, context.rewriter, loc);
-    Operation *send;
+    if (firstFamily.hostRouted) {
-    if (firstFamily.hostRouted)
+      Value flag = storeFlags->emitLookup(action, runtimeLane, anchor,
-      send = SpatHostStoreSyncOp::create(
+        context.constants, context.rewriter, loc);
+      Value condition = arith::CmpIOp::create(context.rewriter, loc,
+        arith::CmpIPredicate::ne, flag, context.constants.getIndex(0));
+      auto conditional = scf::IfOp::create(context.rewriter, loc,
+        TypeRange{}, condition, true);
+      OpBuilder::InsertionGuard guard(context.rewriter);
+      context.rewriter.setInsertionPoint(conditional.thenYield());
+      auto payload = materializeSendPayload(requirement, localOffset,
+        projectionGrids[0].empty() ? nullptr : &projection, context, loc);
+      if (failed(payload))
+        return failure();
+      auto send = SpatHostStoreSyncOp::create(
         context.rewriter, loc, sourceCore, targetCore,
         hostOffsets->emitLookup(
           action, runtimeLane, anchor, context.constants, context.rewriter, loc),
         eventRegisters->emitLookup(
           action, runtimeLane, anchor, context.constants, context.rewriter, loc),
         *payload);
-    else
+      setLogicalTransferMetadata(send, logical);
-      send = SpatChannelSendOp::create(
+      context.rewriter.setInsertionPoint(conditional.elseYield());
+      SpatSyncOp::create(context.rewriter, loc, targetCore,
+        eventRegisters->emitLookup(action, runtimeLane, anchor,
+          context.constants, context.rewriter, loc));
+    } else {
+      auto payload = materializeSendPayload(requirement, localOffset,
+        projectionGrids[0].empty() ? nullptr : &projection, context, loc);
+      if (failed(payload))
+        return failure();
+      auto send = SpatChannelSendOp::create(
         context.rewriter, loc,
         transferGrids->channels().emitLookup(
           action, runtimeLane, anchor, context.constants, context.rewriter, loc),
         sourceCore, targetCore, *payload);
-    setLogicalTransferMetadata(send, logical);
+      setLogicalTransferMetadata(send, logical);
+    }
     return success();
   };
   Value runtimeLane = lane ? lane : context.constants.getIndex(0);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationModel.hpp

@@ -235,6 +235,7 @@ struct ExternalTransferFamily {
   StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
   StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
   StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence hostStoreFlags = StaticIntSequence::uniform(1, 1);
   StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
   StaticIntSequence waitValues = StaticIntSequence::uniform(1, 1);
   bool hostRouted = false;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp

@@ -4,6 +4,7 @@
 
 #include "DeferredProjectionAnalysis.hpp"
 #include "DeferredTransferPlanning.hpp"
+#include "HostPublicationIdentity.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
@@ -361,30 +362,6 @@ static LogicalResult buildAvailabilityFamilies(
           family.targetCores = StaticIntSequence::fromValues(targetCores);
           family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
           family.hostRouted = runAvailability == Availability::Host;
-          if (family.hostRouted) {
-            auto fragmentType = dyn_cast<ShapedType>(
-              requirement.publicationFragmentType);
-            auto fragmentBytes = fragmentType
-              ? pim::getCheckedShapedTypeSizeInBytes(
-                  fragmentType, exchange.deferred,
-                  "pipeline host transfer fragment")
-              : FailureOr<uint64_t>(failure());
-            if (failed(fragmentBytes))
-              return failure();
-            auto bytes = pim::checkedMul<size_t>(
-              count, static_cast<size_t>(*fragmentBytes), exchange.deferred,
-              "pipeline host transfer storage");
-            if (failed(bytes))
-              return failure();
-            family.hostOffsets = StaticIntSequence::affine(
-              plan.pipelineHostBufferBytes, *fragmentBytes, count);
-            auto endOffset = pim::checkedAdd<size_t>(
-              plan.pipelineHostBufferBytes, *bytes, exchange.deferred,
-              "pipeline host transfer storage");
-            if (failed(endOffset))
-              return failure();
-            plan.pipelineHostBufferBytes = *endOffset;
-          }
           nextChannel += count;
           exchange.externalTransferCount += count;
           exchange.external.push_back(std::move(family));
@@ -450,6 +427,48 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
   return success();
 }
 
+static LogicalResult allocateHostPublications(DeferredTransferPlan &plan) {
+  DenseMap<HostPublicationIdentity, int64_t> offsets;
+  for (auto &exchange : plan.exchanges)
+    for (ExternalTransferFamily &family : exchange->external) {
+      if (!family.hostRouted)
+        continue;
+      RequirementFamily &requirement = *family.requirement;
+      auto type = dyn_cast<ShapedType>(requirement.publicationFragmentType);
+      auto bytes = type ? pim::getCheckedShapedTypeSizeInBytes(
+        type, exchange->deferred, "pipeline host transfer fragment")
+        : FailureOr<uint64_t>(failure());
+      if (failed(bytes))
+        return failure();
+      SmallVector<int64_t> hostOffsets;
+      unsigned begin = family.targetLanes.intervals().front().begin;
+      unsigned requirementBegin = requirement.targetLanes.intervals().front().begin;
+      for (size_t index = 0; index < family.targetCores.size(); ++index) {
+        int64_t fragment = requirement.producerLocalOffsets
+          ? requirement.producerLocalOffsets->valueAt(begin + index - requirementBegin) : 0;
+        auto identity = getHostPublicationIdentity(requirement.producer,
+          fragment, requirement.publicationFragmentType,
+          !requirement.producerProjection);
+        auto found = identity ? offsets.find(*identity) : offsets.end();
+        if (found != offsets.end()) {
+          hostOffsets.push_back(found->second);
+          continue;
+        }
+        hostOffsets.push_back(plan.pipelineHostBufferBytes);
+        if (identity)
+          offsets.try_emplace(*identity, plan.pipelineHostBufferBytes);
+        auto end = pim::checkedAdd<size_t>(plan.pipelineHostBufferBytes, *bytes,
+          exchange->deferred, "pipeline host transfer storage");
+        if (failed(end))
+          return failure();
+        plan.pipelineHostBufferBytes = *end;
+      }
+      family.hostOffsets = StaticIntSequence::fromValues(hostOffsets);
+      family.hostStoreFlags = StaticIntSequence::uniform(1, hostOffsets.size());
+    }
+  return success();
+}
+
 static LogicalResult
 retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) {
   if (!isFragmentAssembly(blueprint.getMode()))
@@ -528,7 +547,8 @@ FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(
   if (failed(collectScheduledOperations(
         materialization, plan, pipelineStages, processorCount))
       || failed(collectProducedValues(materialization, plan))
-      || failed(buildExchanges(funcOp, plan)))
+      || failed(buildExchanges(funcOp, plan))
+      || failed(allocateHostPublications(plan)))
     return failure();
   return std::move(plan);
 }

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/HostPublicationIdentity.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "mlir/IR/Types.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include <optional>
+#include <tuple>
+
+namespace onnx_mlir::spatial {
+
+// The owner identifies one immutable result generation. Fragment indices are
+// physical source slots, never logical reconstruction coordinates.
+using HostPublicationIdentity = std::tuple<const void *, int64_t, mlir::Type>;
+
+inline std::optional<HostPublicationIdentity> getHostPublicationIdentity(
+    const void *owner, int64_t fragment, mlir::Type payloadType,
+    bool completeFragment) {
+  // shortcut: scalar projections stay separate; share them when both analyses
+  // can represent their exact source geometry.
+  if (!completeFragment)
+    return std::nullopt;
+  return HostPublicationIdentity{owner, fragment, payloadType};
+}
+
+} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

@@ -124,6 +124,15 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me
   for (auto [stage, demand] : llvm::enumerate(costs.stageSerialDemandNs))
     os << "  estimated stage " << stage << " maximum core compute/incident-transfer demand (ns): "
        << demand << "\n";
+  if (costs.baselineMemory.readRequests || costs.boundaryMemory.readRequests) {
+    auto printMemory = [&](StringRef name, const SharedMemoryDemand &demand) {
+      os << "  estimated " << name << " shared memory: read_bytes=" << demand.readBytes
+         << " write_bytes=" << demand.writeBytes << " read_requests=" << demand.readRequests
+         << " write_requests=" << demand.writeRequests << " service_ns=" << demand.serviceNs << "\n";
+    };
+    printMemory("baseline", costs.baselineMemory);
+    printMemory("boundary", costs.boundaryMemory);
+  }
   llvm::MapVector<size_t, SmallVector<size_t>> stageCores;
   for (auto [core, stage] : llvm::enumerate(schedule.processorStages))
     stageCores[stage].push_back(core);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -17,6 +17,7 @@
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
+#include <functional>
 #include <iterator>
 #include <limits>
 #include <optional>
@@ -297,18 +298,22 @@ Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) {
   return checkedAdd(SchedulerCostModel::kFallbackOperationCost, nestedCost);
 }
 
+static Cost getPositiveStepTripCount(int64_t lower, int64_t upper, int64_t step) {
+  assert(step > 0);
+  if (upper <= lower)
+    return 0;
+  Cost distance = static_cast<Cost>(upper) - static_cast<Cost>(lower);
+  Cost stride = static_cast<Cost>(step);
+  return distance / stride + (distance % stride != 0);
+}
+
 std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
   auto lb = getConstantIntValue(loop.getLowerBound());
   auto ub = getConstantIntValue(loop.getUpperBound());
   auto step = getConstantIntValue(loop.getStep());
   if (!lb || !ub || !step || *step <= 0)
     return std::nullopt;
-  if (*ub <= *lb)
+  return getPositiveStepTripCount(*lb, *ub, *step);
-    return 0;
-
-  uint64_t distance = static_cast<uint64_t>(*ub - *lb);
-  uint64_t stride = static_cast<uint64_t>(*step);
-  return (distance + stride - 1) / stride;
 }
 
 Cost getComputeBodyCost(Region& body, const SchedulerCostModel& costModel) {
@@ -410,20 +415,23 @@ static FailureOr<bool>
 collectProjectedEdges(SpatComputeBatch producer, Value input,
                       const ComputeInstance& consumerInstance,
                       size_t processorCount, size_t targetIndex,
-                      const ComputeGraph& graph,
+                      ComputeGraph& graph,
                       const SchedulerCostModel& costModel,
                       GraphBatchPublicationCache& publicationCache,
                       SmallVectorImpl<ComputeGraphEdge>& rawEdges) {
   auto consumer = dyn_cast<SpatComputeBatch>(consumerInstance.op);
-  if (!consumer)
+  auto inputs = consumer ? consumer.getInputs()
+    : cast<SpatCompute>(consumerInstance.op).getInputs();
+  auto inputIt = llvm::find(inputs, input);
+  if (inputIt == inputs.end())
     return false;
-  auto inputIt = llvm::find(consumer.getInputs(), input);
+  size_t inputIndex = std::distance(inputs.begin(), inputIt);
-  if (inputIt == consumer.getInputs().end())
+  auto inputArg = consumer ? consumer.getInputArgument(inputIndex)
+    : cast<SpatCompute>(consumerInstance.op).getInputArgument(inputIndex);
+  if (consumer && !consumer.getLaneArgument())
     return false;
-  size_t inputIndex = std::distance(consumer.getInputs().begin(), inputIt);
+  Value laneArg = consumer ? Value(*consumer.getLaneArgument()) : Value();
-  std::optional<BlockArgument> inputArg = consumer.getInputArgument(inputIndex);
+  if (!inputArg)
-  std::optional<BlockArgument> laneArg = consumer.getLaneArgument();
-  if (!inputArg || !laneArg)
     return false;
   auto publication = getGraphBatchPublicationMap(
     producer, cast<OpResult>(input).getResultNumber(), publicationCache);
@@ -431,61 +439,227 @@ collectProjectedEdges(SpatComputeBatch producer, Value input,
     return failure();
 
   SmallVector<ComputeGraphEdge, 4> edges;
+  SmallVector<PublicationTransfer> publications;
+  DenseMap<std::pair<Value, int64_t>, size_t> publicationIndices;
   llvm::SmallBitVector selectedSlots(producer.getLaneCount());
+  Cost fragmentBytes = costModel.getByteSize((*publication)->publicationFragmentType,
+    costModel.getComputeBitwidth((*publication)->publicationFragmentType));
   TransferCost fragmentCost = costModel.getTransferCostFromBytes(
     costModel.getByteSize((*publication)->publicationFragmentType,
       costModel.getComputeBitwidth((*publication)->publicationFragmentType)));
+  bool exactDependencies = static_cast<bool>(consumer);
   DenseMap<Value, int64_t> bindings;
   for (Operation* user : inputArg->getUsers()) {
     auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
     if (!extract || extract.getSource() != *inputArg || extract.getMixedOffsets().empty()
         || extract.getMixedSizes().empty() || extract.getMixedStrides().empty())
       return false;
+    SmallVector<scf::ForOp> loops;
+    for (Operation *parent = extract->getParentOp(); parent != consumerInstance.op;
+         parent = parent->getParentOp()) {
+      auto loop = dyn_cast<scf::ForOp>(parent);
+      if (!loop) return false;
+      loops.push_back(loop);
+    }
+    std::reverse(loops.begin(), loops.end());
     for (uint32_t lane = consumerInstance.laneStart;
-         lane < consumerInstance.laneStart + consumerInstance.laneCount;
+         lane < consumerInstance.laneStart + consumerInstance.laneCount; ++lane) {
-         ++lane) {
+      auto recordProjection = [&](bool dependency, Cost repetitions = 1) {
-      FailureOr<int64_t> offset =
+        auto offset = evaluateIndexLike(extract.getMixedOffsets().front(), bindings, lane, laneArg);
-        evaluateIndexLike(extract.getMixedOffsets().front(), bindings, lane, *laneArg);
+        auto size = evaluateIndexLike(extract.getMixedSizes().front(), bindings, lane, laneArg);
-      FailureOr<int64_t> size =
+        auto stride = evaluateIndexLike(extract.getMixedStrides().front(), bindings, lane, laneArg);
-        evaluateIndexLike(extract.getMixedSizes().front(), bindings, lane, *laneArg);
+        if (failed(offset) || failed(size) || failed(stride) || *offset < 0 || *size <= 0 || *stride <= 0)
-      FailureOr<int64_t> stride =
-        evaluateIndexLike(extract.getMixedStrides().front(), bindings, lane, *laneArg);
-      if (failed(offset) || failed(size) || failed(stride) || *offset < 0 || *size <= 0 || *stride <= 0)
-        return false;
-      for (int64_t index = 0; index < *size; ++index) {
-        int64_t slot = *offset + index * *stride;
-        if (slot < 0 || slot >= producer.getLaneCount())
           return false;
-        if (selectedSlots.test(slot))
+        if (*offset >= producer.getLaneCount()
-          continue;
+            || *size > 1 + (producer.getLaneCount() - 1 - *offset) / *stride)
-        selectedSlots.set(slot);
+          return false;
-        int64_t producerLane = (*publication)->physicalSlotToGraphLane[slot];
+        for (int64_t index = 0; index < *size; ++index) {
-        ComputeInstance instance = getBatchChunkForLane(
+          int64_t slot = *offset + index * *stride;
-          producer, static_cast<uint32_t>(producerLane), processorCount);
+          if (slot < 0 || slot >= producer.getLaneCount()) return false;
-        auto source = graph.instanceToIndex.find(instance);
+          int64_t producerLane = (*publication)->physicalSlotToGraphLane[slot];
-        if (source == graph.instanceToIndex.end())
+          ComputeInstance instance = getBatchChunkForLane(producer, producerLane, processorCount);
-          return producer.emitOpError("projected producer has no compute instance"), failure();
+          auto source = graph.instanceToIndex.find(instance);
-        edges.push_back({source->second, targetIndex, fragmentCost});
+          if (source == graph.instanceToIndex.end()) return false;
+          if (!dependency) {
+            auto [entry, inserted] = publicationIndices.try_emplace(
+              std::make_pair(extract.getResult(), producerLane), publications.size());
+            if (inserted)
+              publications.push_back({source->second, targetIndex, input,
+                producerLane, extract.getResult(), fragmentBytes, repetitions,
+                (*publication)->publicationFragmentType});
+            else
+              publications[entry->second].count = checkedAdd(
+                publications[entry->second].count, repetitions);
+          } else if (!selectedSlots.test(slot)) {
+            selectedSlots.set(slot);
+            edges.push_back({source->second, targetIndex, fragmentCost});
+          }
+        }
+        return true;
+      };
+      if (consumer && !recordProjection(true)) exactDependencies = false;
+      // Repeated invariant projections need one demand with a multiplicity.
+      Cost repetitions = 1;
+      bool invariantBounds = true;
+      for (scf::ForOp loop : loops) {
+        auto lower = evaluateIndexLike(loop.getLowerBound(), bindings, lane, laneArg);
+        auto upper = evaluateIndexLike(loop.getUpperBound(), bindings, lane, laneArg);
+        auto step = evaluateIndexLike(loop.getStep(), bindings, lane, laneArg);
+        if (failed(lower) || failed(upper) || failed(step) || *step <= 0) {
+          invariantBounds = false;
+          break;
+        }
+        Cost trips = getPositiveStepTripCount(*lower, *upper, *step);
+        repetitions = checkedMultiply(repetitions, trips);
       }
+      if (invariantBounds && (repetitions == 0 || recordProjection(false, repetitions)))
+        continue;
+      std::function<bool(size_t)> visit = [&](size_t depth) {
+        if (depth < loops.size()) {
+          auto loop = loops[depth];
+          auto lower = evaluateIndexLike(loop.getLowerBound(), bindings, lane, laneArg);
+          auto upper = evaluateIndexLike(loop.getUpperBound(), bindings, lane, laneArg);
+          auto step = evaluateIndexLike(loop.getStep(), bindings, lane, laneArg);
+          if (failed(lower) || failed(upper) || failed(step) || *step <= 0) return false;
+          int64_t iv = *lower;
+          for (Cost remaining = getPositiveStepTripCount(*lower, *upper, *step);
+               remaining; --remaining) {
+            bindings[loop.getInductionVar()] = iv;
+            if (!visit(depth + 1)) return false;
+            if (remaining > 1) iv += *step;
+          }
+          bindings.erase(loop.getInductionVar());
+          return true;
+        }
+        return recordProjection(false);
+      };
+      if (!visit(0)) return false;
     }
   }
-  llvm::append_range(rawEdges, edges);
+  // Memory accounting can resolve loop projections without changing the
+  // conservative dependency edges used by latency PEFT.
+  if (exactDependencies)
+    llvm::append_range(rawEdges, edges);
+  llvm::append_range(graph.publicationTransfers, publications);
+  return exactDependencies;
+}
+
+struct InputFragment {
+  Type type;
+  Cost count;
+  Value projection;
+};
+
+static bool isSharedMemoryInput(Value input, const SchedulingTarget &target) {
+  if (isa<BlockArgument>(input))
+    return true;
+  Attribute value;
+  if (!matchPattern(input, m_Constant(&value)) || !isa<ElementsAttr>(value))
+    return false;
+  auto dense = dyn_cast<DenseElementsAttr>(value);
+  if (!target.zeroSplatUsesSharedMemory && dense && dense.isSplat()) {
+    auto element = dense.getSplatValue<Attribute>();
+    if (auto floating = dyn_cast<FloatAttr>(element))
+      return !floating.getValue().isZero() || floating.getValue().isNegative();
+    if (auto integer = dyn_cast<IntegerAttr>(element))
+      return !integer.getValue().isZero();
+  }
   return true;
 }
 
-std::optional<TransferCost> getProjectedInputTransferCost(
+static std::pair<Cost, Cost> getProjectionMemoryRequests(Type type, Value projection,
-    Value input, const ComputeInstance& instance,
+    const SchedulerCostModel &costModel) {
-    const SchedulerCostModel& costModel) {
+  Cost bytes = costModel.getByteSize(type, costModel.getComputeBitwidth(type));
+  Cost requests = 1;
+  if (auto extract = projection.getDefiningOp<tensor::ExtractSliceOp>();
+      extract && extract.getSourceType().hasStaticShape()
+      && !ShapedType::isDynamicShape(extract.getStaticSizes())
+      && !ShapedType::isDynamicShape(extract.getStaticStrides())) {
+    auto strides = computeRowMajorStrides(extract.getSourceType().getShape());
+    Cost run = 1;
+    for (size_t dim = strides.size(); dim-- > 0;) {
+      Cost size = extract.getStaticSizes()[dim];
+      if (size == 1) continue;
+      if (checkedMultiply(static_cast<Cost>(strides[dim]),
+            static_cast<Cost>(extract.getStaticStrides()[dim])) != run)
+        break;
+      run = checkedMultiply(run, size);
+    }
+    Cost elements = cast<ShapedType>(type).getNumElements();
+    if (elements) {
+      requests = checkedMultiply(requests, elements / run);
+      bytes = SchedulerCostModel::ceilDiv(
+        checkedMultiply(run, costModel.getComputeBitwidth(type)), Cost{8});
+    }
+  }
+  return {bytes, requests};
+}
+
+static SharedMemoryDemand getInputMemoryDemand(const InputFragment &fragment,
+    const SchedulerCostModel &costModel) {
+  auto [bytes, runs] = getProjectionMemoryRequests(fragment.type, fragment.projection, costModel);
+  return costModel.target.getSharedMemoryDemand(bytes,
+    checkedMultiply(runs, fragment.count), false);
+}
+
+static SharedMemoryDemand getCapturedConstantMemoryDemand(
+    Operation *owner, const SchedulerCostModel &costModel) {
+  // shortcut: later aliasing/load folding can change baseline requests; exact
+  // accounting requires a placed, bufferized program.
+  SharedMemoryDemand demand;
+  owner->walk([&](Operation *user) {
+    if (user == owner) return;
+    for (Value value : user->getOperands()) {
+      if (isa<BlockArgument>(value) || !isSharedMemoryInput(value, costModel.target)) continue;
+      if (auto extract = dyn_cast<tensor::ExtractOp>(user)) {
+        DenseMap<Value, int64_t> bindings;
+        for (Operation *parent = user->getParentOp(); parent != owner;
+             parent = parent->getParentOp())
+          if (auto loop = dyn_cast<scf::ForOp>(parent))
+            if (auto lower = getConstantIntValue(loop.getLowerBound()))
+              bindings[loop.getInductionVar()] = *lower;
+        auto batch = dyn_cast<SpatComputeBatch>(owner);
+        Value lane = batch && batch.getLaneArgument() ? Value(*batch.getLaneArgument()) : Value();
+        if (succeeded(evaluateIndexLike(extract.getResult(), bindings, 0, lane)))
+          continue;
+      }
+      if (auto vmm = dyn_cast<SpatVMMOp>(user); vmm && vmm.getWeight() == value) continue;
+      if (auto extract = dyn_cast<tensor::ExtractSliceOp>(user)) {
+        if (llvm::all_of(extract.getResult().getUses(), [&](OpOperand &use) {
+              auto vmm = dyn_cast<SpatVMMOp>(use.getOwner());
+              return vmm && vmm.getWeight() == extract.getResult();
+            })) continue;
+      }
+      Cost count = 1;
+      for (Operation *parent = user->getParentOp(); parent != owner;
+           parent = parent->getParentOp()) {
+        if (auto loop = dyn_cast<scf::ForOp>(parent)) {
+          auto trips = getStaticTripCount(loop);
+          if (trips) count = checkedMultiply(count, static_cast<Cost>(*trips));
+        }
+      }
+      auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
+      Type type = extract ? extract.getResult().getType() : value.getType();
+      demand += getInputMemoryDemand({type, count, extract ? extract.getResult() : value}, costModel);
+    }
+  });
+  return demand;
+}
+
+std::optional<SmallVector<InputFragment>> getProjectedInputFragments(
+    Value input, const ComputeInstance &instance) {
   auto batch = dyn_cast<SpatComputeBatch>(instance.op);
-  if (!batch)
+  auto inputs = batch ? batch.getInputs()
+    : cast<SpatCompute>(instance.op).getInputs();
+  auto found = llvm::find(inputs, input);
+  if (found == inputs.end())
     return std::nullopt;
-  auto found = llvm::find(batch.getInputs(), input);
+  size_t index = found - inputs.begin();
-  if (found == batch.getInputs().end())
+  auto argument = batch ? batch.getInputArgument(index)
-    return std::nullopt;
+    : cast<SpatCompute>(instance.op).getInputArgument(index);
-  auto argument = batch.getInputArgument(found - batch.getInputs().begin());
   if (!argument)
     return std::nullopt;
-  TransferCost cost;
+  SmallVector<InputFragment> fragments;
   for (Operation* user : argument->getUsers()) {
     auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
     if (!extract || extract.getSource() != *argument)
@@ -499,12 +673,26 @@ std::optional<TransferCost> getProjectedInputTransferCost(
         return std::nullopt;
       count = checkedMultiply(count, static_cast<Cost>(*trips));
     }
-    auto type = extract.getResult().getType();
+    fragments.push_back({extract.getResult().getType(), count, extract.getResult()});
+  }
+  return fragments;
+}
+
+std::optional<TransferCost> getProjectedInputTransferCost(
+    Value input, const ComputeInstance& instance,
+    const SchedulerCostModel& costModel) {
+  if (!isa<SpatComputeBatch>(instance.op))
+    return std::nullopt;
+  auto fragments = getProjectedInputFragments(input, instance);
+  if (!fragments)
+    return std::nullopt;
+  TransferCost cost;
+  for (const InputFragment &inputFragment : *fragments) {
     TransferCost fragment = costModel.getTransferCostFromBytes(
-      costModel.getByteSize(type, costModel.getComputeBitwidth(type)));
+      costModel.getByteSize(inputFragment.type, costModel.getComputeBitwidth(inputFragment.type)));
     cost = addTransferCosts(cost,
-      {checkedMultiply(fragment.fixed, count),
+      {checkedMultiply(fragment.fixed, inputFragment.count),
-       checkedMultiply(fragment.networkFlits, count)});
+       checkedMultiply(fragment.networkFlits, inputFragment.count)});
   }
   return cost;
 }
@@ -578,8 +766,12 @@ TransferCost getProducerTransferCost(Value input,
     const ComputeInstance& consumerInstance, const ProducerValueRef& producerRef,
     const SchedulerCostModel& costModel);
 
+static LogicalResult collectPublicationTransfer(Value input, size_t targetIndex,
+    const ProducerValueRef &producer, size_t sourceIndex, ComputeGraph &graph,
+    const SchedulerCostModel &costModel, GraphBatchPublicationCache &cache);
+
 LogicalResult collectBlueprintEdges(Value input, size_t targetIndex,
-                           const ComputeGraph& graph, size_t processorCount,
+                           ComputeGraph& graph, size_t processorCount,
                            const SchedulerCostModel& costModel,
                            GraphBatchPublicationCache& publicationCache,
                            SmallVectorImpl<ComputeGraphEdge>& edges) {
@@ -607,6 +799,9 @@ LogicalResult collectBlueprintEdges(Value input, size_t targetIndex,
           return blueprint.emitOpError("fragment source has no compute instance");
         edges.push_back({indexed->second, targetIndex,
           getProducerTransferCost(source, graph.nodes[targetIndex].instance, producer, costModel)});
+        if (failed(collectPublicationTransfer(source, targetIndex, producer,
+              indexed->second, graph, costModel, publicationCache)))
+          return failure();
       }
       continue;
     }
@@ -629,6 +824,9 @@ LogicalResult collectBlueprintEdges(Value input, size_t targetIndex,
       elements = checkedMultiply(elements, static_cast<Cost>(size));
     auto fragmentType = (*publication)->publicationFragmentType;
     Cost bytes = costModel.getByteSize(fragmentType, costModel.getComputeBitwidth(fragmentType));
+    graph.publicationTransfers.push_back({indexed->second, targetIndex, source,
+      lane, input, bytes, graph.nodes[targetIndex].instance.laneCount,
+      fragmentType});
     TransferCost transfer = costModel.getTransferCostFromBytes(bytes);
     auto sourceOffsets = blueprint.getFragmentSourceOffsets();
     if (!sourceOffsets)
@@ -643,6 +841,41 @@ LogicalResult collectBlueprintEdges(Value input, size_t targetIndex,
   return success();
 }
 
+static LogicalResult collectPublicationTransfer(Value input, size_t targetIndex,
+    const ProducerValueRef &producer, size_t sourceIndex, ComputeGraph &graph,
+    const SchedulerCostModel &costModel, GraphBatchPublicationCache &cache) {
+  const ComputeInstance &consumer = graph.nodes[targetIndex].instance;
+  auto batch = dyn_cast<SpatComputeBatch>(producer.instance.op);
+  Value value = producer.instance.op->getResult(producer.resultIndex);
+  if (batch) {
+    auto publication = getGraphBatchPublicationMap(batch, producer.resultIndex, cache);
+    if (failed(publication))
+      return failure();
+    auto type = (*publication)->publicationFragmentType;
+    Cost bytes = costModel.getByteSize(type, costModel.getComputeBitwidth(type));
+    for (uint32_t lane = producer.instance.laneStart;
+         lane < producer.instance.laneStart + producer.instance.laneCount; ++lane)
+      graph.publicationTransfers.push_back({sourceIndex, targetIndex, value,
+        lane, input, bytes, consumer.laneCount, type});
+    return success();
+  }
+  if (auto fragments = getProjectedInputFragments(input, consumer)) {
+    for (const InputFragment &fragment : *fragments) {
+      Cost bytes = costModel.getByteSize(fragment.type, costModel.getComputeBitwidth(fragment.type));
+      auto [storeBytes, storeRuns] = getProjectionMemoryRequests(
+        fragment.type, fragment.projection, costModel);
+      graph.publicationTransfers.push_back({sourceIndex, targetIndex, value,
+        0, fragment.projection, bytes, fragment.count, {}, storeBytes, storeRuns});
+    }
+  } else {
+    Cost bytes = costModel.getByteSize(input.getType(), costModel.getComputeBitwidth(input.getType()));
+    graph.publicationTransfers.push_back({sourceIndex, targetIndex, value,
+      0, input, bytes, consumer.laneCount,
+      input == value ? value.getType() : Type()});
+  }
+  return success();
+}
+
 TransferCost getProducerTransferCost(
     Value input, const ComputeInstance& consumerInstance,
     const ProducerValueRef& producerRef,
@@ -846,10 +1079,35 @@ static void collectResidentWeightsFromOp(Operation* op,
     if (!lb || !ub || !step || *step <= 0)
       return;
 
-    for (int64_t iv = *lb; iv < *ub; iv += *step) {
+    ResidentWeightSet invariantWeights;
+    bool invariant = true, hasWeights = false;
+    loop.walk([&](Operation *nested) {
+      if (isa<scf::ForOp>(nested) && nested != loop.getOperation())
+        invariant = false;
+      auto vmm = dyn_cast<SpatVMMOp>(nested);
+      if (!vmm)
+        return;
+      hasWeights = true;
+      auto weight = getStaticResidentWeight(owner, vmm.getWeight(), bindings, lane, laneArg);
+      if (failed(weight))
+        invariant = false;
+      else
+        addResidentWeight(invariantWeights, *weight);
+    });
+    if (invariant || !hasWeights) {
+      if (*lb < *ub)
+        for (ResidentWeight &weight : invariantWeights)
+          addResidentWeight(usage, std::move(weight));
+      return;
+    }
+
+    int64_t iv = *lb;
+    for (Cost remaining = getPositiveStepTripCount(*lb, *ub, *step);
+         remaining; --remaining) {
       bindings[loop.getInductionVar()] = iv;
       for (Operation& nested : loop.getBody()->without_terminator())
         collectResidentWeightsFromOp(&nested, owner, bindings, usage, laneArg, lane);
+      if (remaining > 1) iv += *step;
     }
     bindings.erase(loop.getInductionVar());
     return;
@@ -1053,9 +1311,29 @@ FailureOr<ComputeGraph> buildComputeGraph(Operation* entryOp,
 
   llvm::SmallVector<ComputeGraphEdge, 16> rawEdges;
   GraphBatchPublicationCache publicationCache;
+  DenseMap<Operation *, SharedMemoryDemand> capturedMemory;
   for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) {
+    auto [captured, inserted] = capturedMemory.try_emplace(node.instance.op);
+    if (inserted)
+      captured->second = getCapturedConstantMemoryDemand(node.instance.op, costModel);
+    auto memory = captured->second;
+    memory.readBytes = checkedMultiply(memory.readBytes, static_cast<Cost>(node.instance.laneCount));
+    memory.readRequests = checkedMultiply(memory.readRequests, static_cast<Cost>(node.instance.laneCount));
+    memory.serviceNs = checkedMultiply(memory.serviceNs, static_cast<Cost>(node.instance.laneCount));
+    graph.baselineMemory += memory;
     llvm::SmallVector<Value, 4> inputs = getComputeInstanceInputs(node.instance);
     for (Value input : inputs) {
+      size_t publicationsBefore = graph.publicationTransfers.size();
+      if (isSharedMemoryInput(input, target)) {
+        if (auto fragments = getProjectedInputFragments(input, node.instance)) {
+          for (const InputFragment &fragment : *fragments)
+            graph.baselineMemory += getInputMemoryDemand(fragment, costModel);
+        } else {
+          graph.baselineMemory += target.getSharedMemoryDemand(
+            costModel.getByteSize(input.getType(), costModel.getComputeBitwidth(input.getType())),
+            node.instance.laneCount, false);
+        }
+      }
       if (auto producer = input.getDefiningOp<SpatComputeBatch>()) {
         auto projected = collectProjectedEdges(producer, input, node.instance,
           computePartitionCount, targetIndex, graph, costModel, publicationCache, rawEdges);
@@ -1074,6 +1352,7 @@ FailureOr<ComputeGraph> buildComputeGraph(Operation* entryOp,
       auto producers = collectProducerValueRefs(input, computePartitionCount, publicationCache);
       if (failed(producers))
         return failure();
+      bool memoryCollected = graph.publicationTransfers.size() != publicationsBefore;
       for (const ProducerValueRef& producerRef : *producers) {
         auto producerIt = graph.instanceToIndex.find(producerRef.instance);
         if (producerIt == graph.instanceToIndex.end())
@@ -1081,10 +1360,21 @@ FailureOr<ComputeGraph> buildComputeGraph(Operation* entryOp,
         rawEdges.push_back({producerIt->second,
                             targetIndex,
                             getProducerTransferCost(input, node.instance, producerRef, costModel)});
+        if (!memoryCollected && failed(collectPublicationTransfer(input, targetIndex, producerRef,
+              producerIt->second, graph, costModel, publicationCache)))
+          return failure();
       }
     }
   }
 
+  for (Region &region : entryOp->getRegions())
+    for (Block &block : region)
+      if (block.getTerminator()->hasTrait<OpTrait::ReturnLike>())
+        for (Value output : block.getTerminator()->getOperands())
+          graph.baselineMemory += target.getSharedMemoryDemand(
+            costModel.getByteSize(output.getType(), costModel.getComputeBitwidth(output.getType())),
+            1, true);
+
   std::vector<ComputeGraphEdge> aggregatedEdges = aggregateEdges(rawEdges);
   graph.edges.insert(graph.edges.end(), aggregatedEdges.begin(), aggregatedEdges.end());
   graph.successors.assign(graph.nodes.size(), {});
@@ -1097,6 +1387,24 @@ FailureOr<ComputeGraph> buildComputeGraph(Operation* entryOp,
   return graph;
 }
 
+SharedMemoryDemand getPipelineBoundaryMemoryDemand(const ComputeGraph &graph,
+    ArrayRef<size_t> taskStages, const SchedulingTarget &target,
+    bool sharePublications) {
+  assert(taskStages.size() == graph.nodes.size());
+  SharedMemoryDemand demand;
+  llvm::DenseSet<HostPublicationIdentity> stored;
+  for (const PublicationTransfer &transfer : graph.publicationTransfers)
+    if (taskStages[transfer.source] != taskStages[transfer.target]) {
+      demand += target.getSharedMemoryDemand(transfer.bytes, transfer.count, false);
+      auto identity = sharePublications ? transfer.sharedIdentity() : std::nullopt;
+      if (!identity || stored.insert(*identity).second)
+        demand += sharePublications
+          ? transfer.getStoreMemoryDemand(target, identity ? 1 : transfer.count)
+          : target.getSharedMemoryDemand(transfer.bytes, transfer.count, true);
+    }
+  return demand;
+}
+
 bool verifyAcyclic(const ComputeGraph& graph) {
   std::vector<size_t> remainingParents(graph.nodes.size(), 0);
   std::queue<size_t> readyNodes;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp

@@ -1,4 +1,5 @@
 #pragma once
+#include "../HostPublicationIdentity.hpp"
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Value.h"
@@ -63,14 +64,45 @@ struct ComputeGraphEdge {
   TransferCost transferCost;
 };
 
+// One physical fragment demand, before routing.
+struct PublicationTransfer {
+  size_t source = 0;
+  size_t target = 0;
+  mlir::Value publication;
+  int64_t graphLane = 0;
+  mlir::Value projection;
+  Cost bytes = 0;
+  Cost count = 1;
+  mlir::Type completeFragmentType = {};
+  Cost storeRunBytes = 0;
+  Cost storeRuns = 1;
+
+  std::optional<HostPublicationIdentity> sharedIdentity() const {
+    return getHostPublicationIdentity(publication.getAsOpaquePointer(),
+      graphLane, completeFragmentType, static_cast<bool>(completeFragmentType));
+  }
+
+  SharedMemoryDemand getStoreMemoryDemand(const SchedulingTarget &target,
+      Cost publications) const {
+    return target.getSharedMemoryDemand(storeRunBytes ? storeRunBytes : bytes,
+      checkedMultiply(publications, storeRuns), true);
+  }
+};
+
 struct ComputeGraph {
   std::vector<ComputeGraphNode> nodes;
   std::vector<ComputeGraphEdge> edges;
   std::vector<std::vector<std::pair<size_t, TransferCost>>> successors;
   std::vector<std::vector<std::pair<size_t, TransferCost>>> predecessors;
   llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
+  std::vector<PublicationTransfer> publicationTransfers;
+  SharedMemoryDemand baselineMemory;
 };
 
+SharedMemoryDemand getPipelineBoundaryMemoryDemand(const ComputeGraph &graph,
+    llvm::ArrayRef<size_t> taskStages, const SchedulingTarget &target,
+    bool sharePublications = true);
+
 mlir::FailureOr<ComputeGraph> buildComputeGraph(mlir::Operation* entryOp,
                                const SchedulingTarget& target,
                                size_t computePartitionCount = 0);

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp

@@ -8,6 +8,7 @@
 #include <vector>
 
 #include "ComputeInstance.hpp"
+#include "SchedulingTarget.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
@@ -20,6 +21,8 @@ struct ScheduleCostSummary {
   uint64_t completionNs = 0;
   std::vector<uint64_t> stageSpanNs;
   std::vector<uint64_t> stageSerialDemandNs;
+  SharedMemoryDemand baselineMemory;
+  SharedMemoryDemand boundaryMemory;
 };
 
 struct MergeScheduleResult {

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp

@@ -28,7 +28,9 @@ using TaskCosts = std::vector<Cost>;
 using namespace mlir;
 
 struct PipelineTaskModel {
-  PipelineTaskModel(const ComputeGraph &graph, const SchedulingTarget &target) : weights(graph) {
+  PipelineTaskModel(const ComputeGraph &graph, const SchedulingTarget &target,
+      bool sharePublications = true)
+      : weights(graph), graph(graph), target(target), sharePublications(sharePublications) {
     llvm::DenseMap<Operation *, Cost> transposeAdjustments;
     schedulingCosts.reserve(graph.nodes.size());
     executionCosts.reserve(graph.nodes.size());
@@ -52,6 +54,9 @@ struct PipelineTaskModel {
   TaskCosts schedulingCosts;
   TaskCosts executionCosts;
   IndexedResidentWeights weights;
+  const ComputeGraph &graph;
+  const SchedulingTarget &target;
+  bool sharePublications;
 };
 
 static Cost estimateOrderedPackCost(
@@ -189,6 +194,7 @@ struct PipelineGroupModel {
   std::vector<size_t> taskToGroup;
   TaskList groupOrder;
   TaskCosts groupCosts;
+  TaskCosts groupExecutionCosts;
   size_t minimumStageZeroEnd = 0;
   std::vector<size_t> furthestSuccessorBefore;
   bool hasCrossingEdge = false;
@@ -281,15 +287,22 @@ static FailureOr<PipelineGroupModel> buildPipelineGroups(
   }
 
   TaskCosts groupCosts(groups.size());
+  TaskCosts groupExecutionCosts(groups.size());
   TaskCosts internalEnds(graph.nodes.size());
+  TaskCosts executionEnds(graph.nodes.size());
   for (size_t group : groupOrder)
     for (size_t task : groups[group].tasks) {
       Cost start = 0;
+      Cost executionStart = 0;
       for (const auto &[predecessor, transferCost] : graph.predecessors[task])
-        if (taskToGroup[predecessor] == group)
+        if (taskToGroup[predecessor] == group) {
           start = std::max(start, internalEnds[predecessor]);
+          executionStart = std::max(executionStart, executionEnds[predecessor]);
+        }
       internalEnds[task] = checkedAdd(start, model.schedulingCosts[task]);
       groupCosts[group] = std::max(groupCosts[group], internalEnds[task]);
+      executionEnds[task] = checkedAdd(executionStart, model.executionCosts[task]);
+      groupExecutionCosts[group] = std::max(groupExecutionCosts[group], executionEnds[task]);
     }
 
   std::vector<size_t> position(groups.size());
@@ -316,136 +329,259 @@ static FailureOr<PipelineGroupModel> buildPipelineGroups(
   }
 
   return PipelineGroupModel{std::move(groups), std::move(taskToGroup),
-    std::move(groupOrder), std::move(groupCosts), minimumStageZeroEnd,
+    std::move(groupOrder), std::move(groupCosts), std::move(groupExecutionCosts), minimumStageZeroEnd,
     std::move(furthestSuccessorBefore), hasCrossingEdge};
 }
 
-static FailureOr<std::vector<size_t>> partitionPipelineGroups(
+struct PipelinePartitionCandidate {
+  Cost priority = 0;
+  Cost execution = 0;
+  Cost memory = 0;
+  size_t start = 0;
+  size_t parent = 0;
+};
+
+using PartitionFrontier = SmallVector<PipelinePartitionCandidate, 8>;
+
+static void retainPartitionCandidates(PartitionFrontier &frontier,
+    PipelinePartitionCandidate candidate, Cost baselineMemory, bool memoryAware) {
+  if (!memoryAware) {
+    if (frontier.empty() || candidate.priority < frontier.front().priority)
+      frontier.assign(1, candidate);
+    return;
+  }
+  frontier.push_back(candidate);
+  size_t legacy = 0;
+  for (size_t i = 1; i < frontier.size(); ++i)
+    if (frontier[i].priority < frontier[legacy].priority)
+      legacy = i;
+  PartitionFrontier kept;
+  for (size_t i = 0; i < frontier.size(); ++i) {
+    bool dominated = false;
+    for (size_t j = 0; j < frontier.size() && i != legacy; ++j)
+      if (i != j && frontier[j].execution <= frontier[i].execution
+          && frontier[j].memory <= frontier[i].memory
+          && (frontier[j].execution < frontier[i].execution
+              || frontier[j].memory < frontier[i].memory || j < i)) {
+        dominated = true;
+        break;
+      }
+    if (!dominated)
+      kept.push_back(frontier[i]);
+  }
+  // Bound the frontier, preserving the legacy path and both resource extremes.
+  constexpr size_t limit = 8;
+  if (kept.size() > limit) {
+    llvm::stable_sort(kept, [](const auto &a, const auto &b) {
+      return std::tie(a.execution, a.memory, a.priority)
+        < std::tie(b.execution, b.memory, b.priority);
+    });
+    auto score = [&](const auto &a) {
+      return std::make_tuple(std::max(a.execution, checkedAdd(baselineMemory, a.memory)),
+        a.execution, a.memory, a.priority);
+    };
+    size_t best = 0, lowestPriority = 0, lowestMemory = 0;
+    for (size_t i = 1; i < kept.size(); ++i) {
+      if (score(kept[i]) < score(kept[best])) best = i;
+      if (kept[i].priority < kept[lowestPriority].priority) lowestPriority = i;
+      if (kept[i].memory < kept[lowestMemory].memory) lowestMemory = i;
+    }
+    SmallVector<size_t> indices;
+    auto add = [&](size_t i) {
+      if (!llvm::is_contained(indices, i) && indices.size() < limit)
+        indices.push_back(i);
+    };
+    for (size_t i : {size_t{0}, lowestMemory, lowestPriority, best}) add(i);
+    for (size_t i = 0; i < limit; ++i) add(i * (kept.size() - 1) / (limit - 1));
+    for (size_t i = 0; i < kept.size(); ++i) add(i);
+    frontier.clear();
+    for (size_t i : indices) frontier.push_back(kept[i]);
+  } else {
+    frontier = std::move(kept);
+  }
+}
+
+static TaskCosts getPipelineCutMemoryCosts(const PipelineGroupModel &groups,
+    const PipelineTaskModel &model) {
+  size_t count = groups.groups.size();
+  std::vector<size_t> position(count);
+  for (auto [index, group] : llvm::enumerate(groups.groupOrder)) position[group] = index;
+  TaskCosts starts(count + 2), ends(count + 2), costs(count + 1);
+  struct PublicationRange { size_t source, target; Cost service; };
+  llvm::DenseMap<HostPublicationIdentity, PublicationRange> publications;
+  for (const PublicationTransfer &transfer : model.graph.publicationTransfers) {
+    size_t source = position[groups.taskToGroup[transfer.source]];
+    size_t target = position[groups.taskToGroup[transfer.target]];
+    if (source == target) continue;
+    assert(source < target);
+    Cost cost = model.target.getSharedMemoryDemand(
+      transfer.bytes, transfer.count, false).serviceNs;
+    if (auto identity = model.sharePublications ? transfer.sharedIdentity() : std::nullopt) {
+      auto [entry, inserted] = publications.try_emplace(*identity,
+        PublicationRange{source, target,
+          transfer.getStoreMemoryDemand(model.target, 1).serviceNs});
+      entry->second.target = std::max(entry->second.target, target);
+    } else {
+      cost = checkedAdd(cost, model.sharePublications
+        ? transfer.getStoreMemoryDemand(model.target, transfer.count).serviceNs
+        : model.target.getSharedMemoryDemand(transfer.bytes, transfer.count, true).serviceNs);
+    }
+    starts[source + 1] = checkedAdd(starts[source + 1], cost);
+    ends[target + 1] = checkedAdd(ends[target + 1], cost);
+  }
+  for (const auto &[identity, range] : publications) {
+    starts[range.source + 1] = checkedAdd(starts[range.source + 1], range.service);
+    ends[range.target + 1] = checkedAdd(ends[range.target + 1], range.service);
+  }
+  Cost active = 0;
+  for (size_t cut = 0; cut <= count; ++cut) {
+    active = checkedAdd(active, starts[cut]);
+    assert(active >= ends[cut]);
+    active -= ends[cut];
+    costs[cut] = active;
+  }
+  return costs;
+}
+
+static FailureOr<std::vector<std::vector<size_t>>> partitionPipelineGroups(
     const PipelineGroupModel &constraints, const PipelineTaskModel &model,
     ArrayRef<size_t> coreCounts, size_t residentWeightCapacity,
-    std::string &error) {
+    std::string &error, bool memoryAware) {
   const auto &groups = constraints.groups;
   const auto &groupOrder = constraints.groupOrder;
-  const auto &groupCosts = constraints.groupCosts;
   const auto &taskWeightIds = model.weights.taskWeightIds;
-  size_t minimumStageZeroEnd = constraints.minimumStageZeroEnd;
+  size_t count = groups.size();
-  const auto &furthestSuccessorBefore = constraints.furthestSuccessorBefore;
-  bool hasCrossingEdge = constraints.hasCrossingEdge;
   const Cost infinity = std::numeric_limits<Cost>::max();
-  const size_t noCut = std::numeric_limits<size_t>::max();
+  std::vector<std::vector<PartitionFrontier>> best(coreCounts.size() + 1,
-  std::vector<std::vector<Cost>> best(
+    std::vector<PartitionFrontier>(count + 1));
-    coreCounts.size() + 1,
+  using SegmentCost = std::pair<Cost, Cost>;
-    std::vector<Cost>(groups.size() + 1, infinity));
-  std::vector<std::vector<size_t>> parent(
-    coreCounts.size() + 1,
-    std::vector<size_t>(groups.size() + 1, noCut));
   std::vector<size_t> stageCostCache(coreCounts.size());
   std::vector<size_t> cachedCoreCounts;
-  std::vector<std::vector<std::vector<Cost>>> segmentCostCaches;
+  std::vector<std::vector<std::vector<SegmentCost>>> caches;
   for (auto [stage, coreCount] : llvm::enumerate(coreCounts)) {
     auto cached = llvm::find(cachedCoreCounts, coreCount);
     if (cached == cachedCoreCounts.end()) {
-      stageCostCache[stage] = segmentCostCaches.size();
+      stageCostCache[stage] = caches.size();
       cachedCoreCounts.push_back(coreCount);
-      segmentCostCaches.emplace_back(
+      caches.emplace_back(count + 1, std::vector<SegmentCost>(count + 1, {infinity, infinity}));
-        groups.size() + 1,
-        std::vector<Cost>(groups.size() + 1, infinity));
     } else {
       stageCostCache[stage] = cached - cachedCoreCounts.begin();
     }
   }
-  best[0][0] = 0;
+  TaskCosts cutMemory = getPipelineCutMemoryCosts(constraints, model);
-
+  Cost baselineMemory = model.graph.baselineMemory.serviceNs;
-  // ponytail: operation groups are small; replace this quadratic partition
+  best[0][0].push_back({});
-  // only if scheduling profiles show it matters.
   for (size_t stage = 0; stage < coreCounts.size(); ++stage) {
     size_t coreCount = coreCounts[stage];
-    size_t stageWeightCapacity = checkedMultiply(
+    size_t capacity = checkedMultiply(coreCount, residentWeightCapacity);
-      coreCount, residentWeightCapacity);
+    for (size_t start = 0; start < count; ++start) {
-    for (size_t start = 0; start < groups.size(); ++start) {
+      if (best[stage][start].empty()) continue;
-      if (best[stage][start] == infinity)
-        continue;
       TaskList segmentTasks;
-      TaskCosts segmentEnds(groups.size());
+      TaskCosts priorityEnds(count), executionEnds(count);
-      Cost segmentSpan = 0;
+      Cost prioritySpan = 0, executionSpan = 0, largestPriority = 0, largestExecution = 0;
-      Cost largestTask = 0;
+      llvm::SmallBitVector weights(model.weights.weightCount);
-      llvm::SmallBitVector segmentWeights(model.weights.weightCount);
+      size_t weightCount = 0;
-      size_t segmentWeightCount = 0;
+      for (size_t end = start + 1; end <= count; ++end) {
-      for (size_t end = start + 1; end <= groups.size(); ++end) {
+        size_t index = groupOrder[end - 1];
-        const PipelineGroup &group = groups[groupOrder[end - 1]];
+        const PipelineGroup &group = groups[index];
-        size_t groupIndex = groupOrder[end - 1];
+        priorityEnds[index] = checkedAdd(priorityEnds[index], constraints.groupCosts[index]);
-        segmentEnds[groupIndex] = checkedAdd(
+        executionEnds[index] = checkedAdd(executionEnds[index], constraints.groupExecutionCosts[index]);
-          segmentEnds[groupIndex], groupCosts[groupIndex]);
+        prioritySpan = std::max(prioritySpan, priorityEnds[index]);
-        segmentSpan = std::max(segmentSpan, segmentEnds[groupIndex]);
+        executionSpan = std::max(executionSpan, executionEnds[index]);
-        largestTask = std::max(largestTask, groupCosts[groupIndex]);
+        largestPriority = std::max(largestPriority, constraints.groupCosts[index]);
-        for (size_t successor : group.successors)
+        largestExecution = std::max(largestExecution, constraints.groupExecutionCosts[index]);
-          segmentEnds[successor] = std::max(
+        for (size_t successor : group.successors) {
-            segmentEnds[successor], segmentEnds[groupIndex]);
+          priorityEnds[successor] = std::max(priorityEnds[successor], priorityEnds[index]);
+          executionEnds[successor] = std::max(executionEnds[successor], executionEnds[index]);
+        }
         llvm::append_range(segmentTasks, group.tasks);
         for (size_t weight : group.weightIds)
-          if (!segmentWeights.test(weight)) {
+          if (!weights.test(weight)) { weights.set(weight); ++weightCount; }
-            segmentWeights.set(weight);
+        if (weightCount > capacity) break;
-            ++segmentWeightCount;
+        if (stage == 0 && end < constraints.minimumStageZeroEnd) continue;
-          }
+        if (stage != 0 && constraints.hasCrossingEdge
-        if (segmentWeightCount > stageWeightCapacity)
+            && constraints.furthestSuccessorBefore[start] >= end) continue;
-          break;
+        auto &[priority, execution] = caches[stageCostCache[stage]][start][end];
-        if (stage == 0 && end < minimumStageZeroEnd)
+        if (priority == infinity) {
-          continue;
+          priority = checkedAdd(estimateOrderedPackCost(model.schedulingCosts,
-        if (stage != 0 && hasCrossingEdge
+            taskWeightIds, model.weights.weightCount, segmentTasks,
-            && furthestSuccessorBefore[start] >= end)
+            residentWeightCapacity, coreCount), prioritySpan - largestPriority);
-          continue;
+          execution = memoryAware ? checkedAdd(estimateOrderedPackCost(model.executionCosts,
-        Cost &segmentCost =
+            taskWeightIds, model.weights.weightCount, segmentTasks,
-          segmentCostCaches[stageCostCache[stage]][start][end];
+            residentWeightCapacity, coreCount), executionSpan - largestExecution) : priority;
-        if (segmentCost == infinity) {
-          segmentCost = estimateOrderedPackCost(
-            model.schedulingCosts, taskWeightIds, model.weights.weightCount,
-            segmentTasks, residentWeightCapacity, coreCount);
-          // Partition priority includes a dependency tail to avoid packing
-          // dependent groups onto the same cores; this is not predicted time.
-          segmentCost = checkedAdd(segmentCost, segmentSpan - largestTask);
-        }
-        Cost maximumLoad = std::max(best[stage][start], segmentCost);
-        if (maximumLoad < best[stage + 1][end]) {
-          best[stage + 1][end] = maximumLoad;
-          parent[stage + 1][end] = start;
         }
+        for (auto [parent, previous] : llvm::enumerate(best[stage][start]))
+          retainPartitionCandidates(best[stage + 1][end], {
+            std::max(previous.priority, priority), std::max(previous.execution, execution),
+            checkedAdd(previous.memory, cutMemory[start]), start, parent},
+            baselineMemory, memoryAware);
       }
     }
   }
-
+  using Terminal = std::pair<size_t, size_t>;
-  size_t usedStages = 0;
+  std::vector<Terminal> terminals;
-  Cost bestLoad = infinity;
+  auto candidate = [&](Terminal terminal) -> const PipelinePartitionCandidate & {
+    return best[terminal.first][count][terminal.second];
+  };
+  auto score = [&](Terminal terminal) {
+    const auto &a = candidate(terminal);
+    return memoryAware ? std::make_tuple(std::max(a.execution, checkedAdd(baselineMemory, a.memory)),
+      a.execution, a.memory) : std::make_tuple(a.priority, Cost{0}, Cost{0});
+  };
   for (size_t stages = 1; stages <= coreCounts.size(); ++stages)
-    if (best[stages][groups.size()] != infinity
+    for (size_t index = 0; index < best[stages][count].size(); ++index)
-        && best[stages][groups.size()] <= bestLoad) {
+      terminals.emplace_back(stages, index);
-      bestLoad = best[stages][groups.size()];
+  if (terminals.empty()) {
-      usedStages = stages;
-    }
-  if (usedStages == 0) {
     error = "pipeline scheduling cannot split operations into dependency-adjacent "
             "stages within the physical crossbar limit";
     return failure();
   }
-
+  llvm::sort(terminals, [&](Terminal a, Terminal b) {
-  std::vector<size_t> groupStages(groups.size());
+    return score(a) != score(b) ? score(a) < score(b) : a > b;
-  size_t end = groups.size();
+  });
-  for (size_t stage = usedStages; stage > 0; --stage) {
+  SmallVector<Terminal, 4> selected{terminals.front()};
-    size_t start = parent[stage][end];
+  auto add = [&](Terminal terminal) {
-    assert(start != noCut && "selected pipeline partition has no parent");
+    if (!llvm::is_contained(selected, terminal) && selected.size() < 4)
-    for (size_t position = start; position < end; ++position)
+      selected.push_back(terminal);
-      groupStages[groupOrder[position]] = stage - 1;
+  };
-    end = start;
+  if (memoryAware) {
+    for (unsigned resource = 0; resource < 3; ++resource) {
+      auto resourceScore = [&](Terminal terminal) {
+        const auto &a = candidate(terminal);
+        return resource == 0 ? std::make_tuple(a.execution, a.memory, a.priority)
+          : resource == 1 ? std::make_tuple(a.memory, a.execution, a.priority)
+          : std::make_tuple(a.priority, a.execution, a.memory);
+      };
+      add(*llvm::min_element(terminals, [&](Terminal a, Terminal b) {
+        return resourceScore(a) < resourceScore(b);
+      }));
+    }
+    for (Terminal terminal : terminals) add(terminal);
   }
-  return groupStages;
+  std::vector<std::vector<size_t>> partitions;
+  for (auto [stages, index] : selected) {
+    std::vector<size_t> groupStages(count);
+    size_t end = count;
+    for (size_t stage = stages; stage > 0; --stage) {
+      auto entry = best[stage][end][index];
+      for (size_t position = entry.start; position < end; ++position)
+        groupStages[groupOrder[position]] = stage - 1;
+      end = entry.start;
+      index = entry.parent;
+    }
+    if (!llvm::is_contained(partitions, groupStages))
+      partitions.push_back(std::move(groupStages));
+  }
+  return partitions;
 }
 
-static FailureOr<PipelineStageAssignment> assignPipelineStages(
+static FailureOr<std::vector<PipelineStageAssignment>> assignPipelineStages(
     const ComputeGraph &graph, const PipelineTaskModel &model,
     const PipelineCoreLayout &layout, size_t residentWeightCapacity,
-    std::string &error) {
+    std::string &error, bool memoryAware) {
   if (graph.nodes.empty())
-    return PipelineStageAssignment {
+    return std::vector<PipelineStageAssignment>{{
-      {}, std::vector<size_t>(layout.getStageSizes())};
+      {}, std::vector<size_t>(layout.getStageSizes())}};
   auto constraints = buildPipelineGroups(graph, model, error);
   if (failed(constraints))
     return failure();
@@ -455,17 +591,17 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
   const auto &taskWeightIds = model.weights.taskWeightIds;
   auto partitionGroups = [&](ArrayRef<size_t> coreCounts) {
     return partitionPipelineGroups(*constraints, model, coreCounts,
-                                   residentWeightCapacity, error);
+                                   residentWeightCapacity, error, memoryAware);
   };
-  FailureOr<std::vector<size_t>> initialGroupStages =
+  auto initialPartitions = partitionGroups(layout.getStageSizes());
-    partitionGroups(layout.getStageSizes());
+  if (failed(initialPartitions))
-  if (failed(initialGroupStages))
     return failure();
+  const auto &initialGroupStages = initialPartitions->front();
 
   std::vector<TaskList> tasksByStage(layout.getStageCount());
   for (size_t group : groupOrder)
     llvm::append_range(
-      tasksByStage[(*initialGroupStages)[group]], groups[group].tasks);
+      tasksByStage[initialGroupStages[group]], groups[group].tasks);
   std::vector<size_t> stageSizes(layout.getStageCount(), 1);
   size_t assignedCores = stageSizes.size();
   for (size_t stage = 0; stage < stageSizes.size(); ++stage) {
@@ -496,7 +632,8 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
   }
 
   TaskCosts stageDependencyTails = getPipelineStageDependencyTails(
-    groups, groupOrder, *initialGroupStages, constraints->groupCosts, stageSizes.size());
+    groups, groupOrder, initialGroupStages,
+    constraints->groupCosts, stageSizes.size());
   auto getStageCost = [&](size_t stage, size_t coreCount) {
     Cost schedulingCost = estimateOrderedPackCost(
       model.schedulingCosts, taskWeightIds, model.weights.weightCount,
@@ -535,15 +672,19 @@ static FailureOr<PipelineStageAssignment> assignPipelineStages(
       getStageCost(bestStage, stageSizes[bestStage] + 1);
     ++assignedCores;
   }
-  FailureOr<std::vector<size_t>> refinedGroupStages =
+  auto refinedPartitions = partitionGroups(stageSizes);
-    partitionGroups(stageSizes);
+  if (failed(refinedPartitions))
-  if (failed(refinedGroupStages))
     return failure();
-  std::vector<size_t> taskStages(graph.nodes.size());
+  if (memoryAware && !llvm::is_contained(*refinedPartitions, initialGroupStages))
-  for (size_t task = 0; task < graph.nodes.size(); ++task)
+    refinedPartitions->push_back(initialGroupStages);
-    taskStages[task] = (*refinedGroupStages)[taskToGroup[task]];
+  std::vector<PipelineStageAssignment> assignments;
-  return PipelineStageAssignment {
+  for (const auto &groupStages : *refinedPartitions) {
-    std::move(taskStages), std::move(stageSizes)};
+    std::vector<size_t> taskStages(graph.nodes.size());
+    for (size_t task = 0; task < graph.nodes.size(); ++task)
+      taskStages[task] = groupStages[taskToGroup[task]];
+    assignments.push_back({std::move(taskStages), stageSizes});
+  }
+  return assignments;
 }
 
 static bool packPipelineStage(
@@ -841,34 +982,21 @@ static LogicalResult verifyPipelineStageAssignment(
   return success();
 }
 
-mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
+static LogicalResult packPipelineAssignment(const ComputeGraph &graph,
-                                        MergeScheduleResult& schedule,
+    const PipelineTaskModel &taskModel, PipelineStageAssignment assignment,
-                                        size_t pipelineStages,
+    MergeScheduleResult &schedule, size_t pipelineStages,
-                                        const SchedulingTarget& physicalTarget,
+    const SchedulingTarget &physicalTarget, std::string &error) {
-                                        std::string& error) {
-  PipelineCoreLayout balancedLayout(
-    physicalTarget.processorCount, pipelineStages);
-  if (!balancedLayout.isValid()) {
-    error = "pipeline scheduling received an incompatible physical core layout";
-    return mlir::failure();
-  }
-  PipelineTaskModel taskModel(graph, physicalTarget);
-  FailureOr<PipelineStageAssignment> assignment = assignPipelineStages(
-    graph, taskModel, balancedLayout,
-    physicalTarget.residentWeightCapacity, error);
-  if (failed(assignment))
-    return failure();
   std::vector<TaskList> tasksByPhysicalCpu(physicalTarget.processorCount);
   std::vector<size_t> searchCoreFloors(pipelineStages, 1);
   std::string packingError;
   bool packed = false;
   for (size_t attempt = 0; attempt < physicalTarget.processorCount; ++attempt) {
-    PipelineCoreLayout candidateLayout(assignment->stageSizes, physicalTarget);
+    PipelineCoreLayout candidateLayout(assignment.stageSizes, physicalTarget);
     for (TaskList &tasks : tasksByPhysicalCpu)
       tasks.clear();
     for (size_t task = 0; task < graph.nodes.size(); ++task) {
       PipelineStageRange range =
-        candidateLayout.getStageRange(assignment->taskStages[task]);
+        candidateLayout.getStageRange(assignment.taskStages[task]);
       tasksByPhysicalCpu[candidateLayout.getPhysicalCore(range.begin)].push_back(task);
     }
     size_t failedStage = 0;
@@ -881,26 +1009,26 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
     }
     searchCoreFloors[failedStage] = std::max(
       searchCoreFloors[failedStage],
-      assignment->stageSizes[failedStage] + 1);
+      assignment.stageSizes[failedStage] + 1);
     std::optional<size_t> donor;
     for (size_t stage = 0; stage < pipelineStages; ++stage)
       if (stage != failedStage
-          && assignment->stageSizes[stage]
+          && assignment.stageSizes[stage]
                > searchCoreFloors[stage]
           && (!donor
-              || assignment->stageSizes[stage]
+              || assignment.stageSizes[stage]
-                   > assignment->stageSizes[*donor]))
+                   > assignment.stageSizes[*donor]))
         donor = stage;
     if (!donor)
       break;
-    --assignment->stageSizes[*donor];
+    --assignment.stageSizes[*donor];
-    ++assignment->stageSizes[failedStage];
+    ++assignment.stageSizes[failedStage];
   }
   if (!packed) {
     error = packingError;
     return failure();
   }
-  PipelineCoreLayout pipelineLayout(assignment->stageSizes, physicalTarget);
+  PipelineCoreLayout pipelineLayout(assignment.stageSizes, physicalTarget);
   auto evaluated = evaluatePipelinePlacement(
     graph, taskModel, tasksByPhysicalCpu, pipelineLayout, physicalTarget, error);
   if (failed(evaluated))
@@ -943,10 +1071,49 @@ mlir::LogicalResult assignPipelineCores(const ComputeGraph& graph,
   recordScheduleCosts(evaluated->timing, schedule);
   for (const auto &stage : evaluated->stages)
     schedule.costs.stageSerialDemandNs.push_back(std::get<0>(stage));
+  schedule.costs.baselineMemory = graph.baselineMemory;
+  schedule.costs.boundaryMemory = getPipelineBoundaryMemoryDemand(
+    graph, assignment.taskStages, physicalTarget, taskModel.sharePublications);
   populateScheduleEquivalenceClasses(graph, tasksByPhysicalCpu, schedule);
   return mlir::success();
 }
 
+
+static LogicalResult assignPipelineCores(const ComputeGraph &graph,
+    MergeScheduleResult &schedule, size_t pipelineStages,
+    const SchedulingTarget &physicalTarget, std::string &error, bool memoryAware,
+    bool sharePublications = true) {
+  PipelineCoreLayout balancedLayout(physicalTarget.processorCount, pipelineStages);
+  if (!balancedLayout.isValid()) {
+    error = "pipeline scheduling received an incompatible physical core layout";
+    return failure();
+  }
+  PipelineTaskModel taskModel(graph, physicalTarget, sharePublications);
+  auto assignments = assignPipelineStages(graph, taskModel, balancedLayout,
+    physicalTarget.residentWeightCapacity, error, memoryAware);
+  if (failed(assignments)) return failure();
+  bool packed = false;
+  for (const auto &assignment : *assignments) {
+    MergeScheduleResult candidate;
+    if (failed(packPipelineAssignment(graph, taskModel, assignment, candidate,
+          pipelineStages, physicalTarget, error)))
+      continue;
+    auto maximumStageSpan = [](const MergeScheduleResult &result) {
+      return result.costs.stageSpanNs.empty() ? Cost{0}
+        : *llvm::max_element(result.costs.stageSpanNs);
+    };
+    // Store savings do not cover a longer stage when placement-dependent
+    // baseline loads remain approximate. Keep the calibrated phase-1 cut.
+    bool preservesStageSpan = !sharePublications
+      || maximumStageSpan(candidate) <= maximumStageSpan(schedule);
+    if (!memoryAware || (preservesStageSpan
+          && isBetterPipelineSchedule(candidate.costs, schedule.costs)))
+      schedule = std::move(candidate);
+    packed = true;
+  }
+  return success(packed);
+}
+
 } // namespace
 
 PipelineCoreLayout::PipelineCoreLayout(ArrayRef<size_t> sizes,
@@ -1002,7 +1169,21 @@ mlir::LogicalResult runPipelineScheduler(const ComputeGraph& graph,
     error = "pipeline scheduling requires a valid balanced physical core layout";
     return mlir::failure();
   }
-  return assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, error);
+  if (failed(assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, error, false, false)))
+    return failure();
+  if (!physicalTarget.sharedMemoryWidthBytes || graph.publicationTransfers.empty())
+    return success();
+  std::string candidateError;
+  // Retain the calibrated phase-1 winner before evaluating sharing-aware cuts.
+  (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, false);
+  std::vector<size_t> taskStages;
+  for (const auto &node : graph.nodes)
+    taskStages.push_back(schedule.processorStages[
+      schedule.computeToCpuMap.lookup(node.instance)]);
+  schedule.costs.boundaryMemory = getPipelineBoundaryMemoryDemand(
+    graph, taskStages, physicalTarget);
+  (void)assignPipelineCores(graph, schedule, pipelineStages, physicalTarget, candidateError, true, true);
+  return success();
 }
 
 } // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp

@@ -32,6 +32,26 @@ inline bool isBetterStagePlacement(const StagePlacementScore &candidate,
   return candidate < current;
 }
 
+inline bool isBetterPipelineSchedule(const ScheduleCostSummary &candidate,
+    const ScheduleCostSummary &current) {
+  auto score = [](const ScheduleCostSummary &result) {
+    Cost stage = 0;
+    for (Cost span : result.stageSpanNs) stage = std::max(stage, span);
+    Cost memory = checkedAdd(result.baselineMemory.serviceNs, result.boundaryMemory.serviceNs);
+    return std::make_tuple(std::max(stage, memory), stage, memory);
+  };
+  auto currentScore = score(current);
+  auto candidateScore = score(candidate);
+  // Stage spans omit contention: do not exchange a larger compute regression
+  // for a smaller memory saving, even when memory remains the largest resource.
+  Cost currentDemand = checkedAdd(std::get<1>(currentScore), std::get<2>(currentScore));
+  Cost candidateDemand = checkedAdd(std::get<1>(candidateScore), std::get<2>(candidateScore));
+  Cost currentCoreDemand = checkedAdd(static_cast<Cost>(current.maximumCoreWorkNs), std::get<2>(currentScore));
+  Cost candidateCoreDemand = checkedAdd(static_cast<Cost>(candidate.maximumCoreWorkNs), std::get<2>(candidateScore));
+  return candidateScore < currentScore && candidateDemand <= currentDemand
+    && candidateCoreDemand <= currentCoreDemand;
+}
+
 struct PipelineStageRange {
   size_t begin;
   size_t size;

src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp

@@ -9,6 +9,23 @@
 namespace onnx_mlir {
 namespace spatial {
 
+struct SharedMemoryDemand {
+  Cost readBytes = 0;
+  Cost writeBytes = 0;
+  Cost readRequests = 0;
+  Cost writeRequests = 0;
+  Cost serviceNs = 0;
+
+  SharedMemoryDemand &operator+=(const SharedMemoryDemand &other) {
+    readBytes = checkedAdd(readBytes, other.readBytes);
+    writeBytes = checkedAdd(writeBytes, other.writeBytes);
+    readRequests = checkedAdd(readRequests, other.readRequests);
+    writeRequests = checkedAdd(writeRequests, other.writeRequests);
+    serviceNs = checkedAdd(serviceNs, other.serviceNs);
+    return *this;
+  }
+};
+
 struct SchedulingTarget {
   size_t processorCount = 0;
   size_t residentWeightCapacity = 0;
@@ -25,6 +42,31 @@ struct SchedulingTarget {
   Cost vectorLatencyCycles = 4;
   size_t synchronizationRegisterCount = 0;
 
+  // Zero width means this target has no modeled serialized memory resource.
+  Cost sharedMemoryWidthBytes = 0;
+  Cost sharedMemoryReadLatencyNs = 0;
+  Cost sharedMemoryWriteLatencyNs = 0;
+  bool zeroSplatUsesSharedMemory = true;
+
+  SharedMemoryDemand getSharedMemoryDemand(
+      Cost bytesPerRequest, Cost requests, bool write) const {
+    SharedMemoryDemand demand;
+    if (write) {
+      demand.writeBytes = checkedMultiply(bytesPerRequest, requests);
+      demand.writeRequests = requests;
+    } else {
+      demand.readBytes = checkedMultiply(bytesPerRequest, requests);
+      demand.readRequests = requests;
+    }
+    if (sharedMemoryWidthBytes) {
+      Cost accesses = bytesPerRequest / sharedMemoryWidthBytes
+        + (bytesPerRequest % sharedMemoryWidthBytes != 0);
+      demand.serviceNs = checkedMultiply(requests, checkedMultiply(accesses,
+        write ? sharedMemoryWriteLatencyNs : sharedMemoryReadLatencyNs));
+    }
+    return demand;
+  }
+
   Cost matrixRows = 128;
   Cost matrixInputReadMultiplier = 1;
   Cost matrixColumns = 128;

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);
+    mlir::scf::ForOp loop;
-    SpatYieldOp::create(builder, loc, mlir::ValueRange{result.getResult()});
+    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()

validation/tools/pim/pimcomp/compare/results_comparison.csv

@@ -1,40 +1,40 @@
 arch,model,mode,raptor_pipeline,pimcomp_pipeline,raptor_functional_validation,pimcomp_functional_validation,raptor_latency_ms,pimcomp_latency_ms,raptor_throughput_samples_s,pimcomp_throughput_samples_s,raptor_power_mw,pimcomp_power_mw,raptor_energy_pj,pimcomp_energy_pj,better_compiler,speedup
 arch-a,vgg8,latency,1,element,PASS,PASS,1.52,7.99,NA,NA,321.05,200.11,487455139.04,1597904071.12,raptor,5.26
-arch-a,vgg8,throughput,2,batch,PASS,PASS,0.91,1.04,1100.00,965.00,423.80,482.60,431626597.00,541691691.00,raptor,1.14
+arch-a,vgg8,throughput,2,batch,PASS,PASS,0.91,1.04,1100.00,965.00,429.31,482.60,407972136.20,541691691.00,raptor,1.14
-arch-a,vgg8,throughput,4,batch,PASS,PASS,0.86,1.04,1160.00,965.00,432.99,482.60,416224904.20,541691691.00,raptor,1.20
+arch-a,vgg8,throughput,4,batch,PASS,PASS,0.78,1.04,1280.00,965.00,457.04,482.60,415761605.20,541691691.00,raptor,1.33
-arch-a,vgg8,throughput,8,batch,PASS,PASS,1.63,1.04,615.00,965.00,308.12,482.60,616026711.90,541691691.00,pimcomp,1.57
+arch-a,vgg8,throughput,8,batch,PASS,PASS,0.78,1.04,1290.00,965.00,464.07,482.60,396424007.70,541691691.00,raptor,1.34
 arch-a,resnet18,latency,1,element,PASS,FAIL,14.03,58.18,NA,NA,455.67,238.39,6394498463.12,13868866808.12,raptor,4.15
 arch-a,resnet18,throughput,2,batch,PASS,FAIL,11.73,13.46,85.20,74.30,503.11,481.60,6628073894.00,6563451004.00,raptor,1.15
-arch-a,resnet18,throughput,4,batch,PASS,FAIL,19.62,13.46,51.00,74.30,376.10,481.60,8024732847.00,6563451004.00,pimcomp,1.46
+arch-a,resnet18,throughput,4,batch,PASS,FAIL,19.21,13.46,52.10,74.30,380.10,481.60,7957033516.00,6563451004.00,pimcomp,1.43
-arch-a,resnet18,throughput,8,batch,PASS,FAIL,19.92,13.46,50.20,74.30,368.24,481.60,9023016592.00,6563451004.00,pimcomp,1.48
+arch-a,resnet18,throughput,8,batch,PASS,FAIL,19.93,13.46,50.20,74.30,361.80,481.60,7763790072.00,6563451004.00,pimcomp,1.48
 arch-a,resnet34,latency,1,element,PASS,FAIL,29.72,66.71,NA,NA,411.81,277.49,12238321706.68,18511425563.68,raptor,2.24
 arch-a,resnet34,throughput,2,batch,PASS,FAIL,22.26,23.58,44.90,42.40,473.84,515.68,16417303180.00,12283627410.00,raptor,1.06
 arch-a,resnet34,throughput,4,batch,PASS,FAIL,27.37,23.58,36.50,42.40,415.90,515.68,17019112320.00,12283627410.00,pimcomp,1.16
 arch-a,resnet34,throughput,8,batch,PASS,FAIL,41.04,23.58,24.40,42.40,346.15,515.68,17351854700.00,12283627410.00,pimcomp,1.74
 arch-a,googlenet,latency,1,element,PASS,PASS,9.78,62.92,NA,NA,563.88,231.19,5513234559.92,14547744283.24,raptor,6.44
-arch-a,googlenet,throughput,2,batch,PASS,FAIL,10.16,18.26,98.50,54.80,544.21,435.11,6029999329.00,8164787803.00,raptor,1.80
+arch-a,googlenet,throughput,2,batch,PASS,FAIL,8.42,18.26,119.00,54.80,610.74,435.11,5777550058.00,8164787803.00,raptor,2.17
-arch-a,googlenet,throughput,4,batch,PASS,FAIL,11.96,18.26,83.60,54.80,473.16,435.11,6741923526.00,8164787803.00,raptor,1.53
+arch-a,googlenet,throughput,4,batch,PASS,FAIL,10.73,18.26,93.20,54.80,528.02,435.11,6592912716.00,8164787803.00,raptor,1.70
-arch-a,googlenet,throughput,8,batch,PASS,FAIL,20.66,18.26,48.40,54.80,348.76,435.11,8421537146.00,8164787803.00,pimcomp,1.13
+arch-a,googlenet,throughput,8,batch,PASS,FAIL,11.03,18.26,90.70,54.80,505.94,435.11,6622163369.00,8164787803.00,raptor,1.66
 arch-a,yolo11n,latency,1,element,PASS,FAIL,147.49,NA,NA,NA,290.17,NA,42797722129.00,NA,NA,NA
 arch-a,yolo11n,throughput,2,batch,PASS,FAIL,120.93,NA,8.27,NA,313.20,NA,41830757130.00,NA,NA,NA
-arch-a,yolo11n,throughput,4,batch,PASS,FAIL,92.06,NA,10.90,NA,351.03,NA,39451980450.00,NA,NA,NA
+arch-a,yolo11n,throughput,4,batch,PASS,FAIL,87.91,NA,11.40,NA,358.20,NA,38729846260.00,NA,NA,NA
 arch-a,yolo11n,throughput,8,batch,PASS,FAIL,104.07,NA,9.61,NA,330.76,NA,44002712680.00,NA,NA,NA
 arch-b,vgg8,latency,1,element,PASS,PASS,1.53,7.15,NA,NA,292.56,173.77,448581679.04,1242814988.12,raptor,4.66
 arch-b,vgg8,throughput,2,batch,PASS,PASS,1.18,1.07,846.00,932.00,327.05,409.47,429183159.10,439288835.80,pimcomp,1.10
-arch-b,vgg8,throughput,4,batch,PASS,PASS,0.92,1.07,1090.00,932.00,383.82,409.47,464794094.50,439288835.80,raptor,1.17
+arch-b,vgg8,throughput,4,batch,PASS,PASS,0.67,1.07,1490.00,932.00,448.90,409.47,413526781.30,439288835.80,raptor,1.60
-arch-b,vgg8,throughput,8,batch,PASS,PASS,1.62,1.07,616.00,932.00,276.49,409.47,524048148.70,439288835.80,pimcomp,1.51
+arch-b,vgg8,throughput,8,batch,PASS,PASS,0.75,1.07,1330.00,932.00,425.74,409.47,394457860.80,439288835.80,raptor,1.43
 arch-b,resnet18,latency,1,element,PASS,FAIL,15.38,64.38,NA,NA,400.67,201.66,6160441550.12,12983432281.12,raptor,4.19
 arch-b,resnet18,throughput,2,batch,PASS,FAIL,13.08,13.43,76.40,74.40,440.79,484.80,6399942695.00,6569655551.00,raptor,1.03
-arch-b,resnet18,throughput,4,batch,PASS,FAIL,20.68,13.43,48.40,74.40,327.57,484.80,7777606323.00,6569655551.00,pimcomp,1.54
+arch-b,resnet18,throughput,4,batch,PASS,FAIL,12.68,13.43,78.90,74.40,451.19,484.80,7096979122.00,6569655551.00,raptor,1.06
-arch-b,resnet18,throughput,8,batch,PASS,FAIL,21.17,13.43,47.20,74.40,317.64,484.80,8897039159.00,6569655551.00,pimcomp,1.58
+arch-b,resnet18,throughput,8,batch,PASS,FAIL,20.59,13.43,48.60,74.40,321.62,484.80,7163273595.00,6569655551.00,pimcomp,1.53
 arch-b,resnet34,latency,1,element,PASS,FAIL,27.95,107.20,NA,NA,397.21,206.48,11103517952.68,22134174654.68,raptor,3.83
 arch-b,resnet34,throughput,2,batch,PASS,FAIL,22.28,22.85,44.90,43.80,438.24,493.33,14218683390.00,11392215660.00,raptor,1.03
 arch-b,resnet34,throughput,4,batch,PASS,FAIL,39.61,22.85,25.20,43.80,315.01,493.33,14629161660.00,11392215660.00,pimcomp,1.74
 arch-b,resnet34,throughput,8,batch,PASS,FAIL,49.34,22.85,20.30,43.80,285.76,493.33,16591020470.00,11392215660.00,pimcomp,2.16
 arch-b,googlenet,latency,1,element,PASS,PASS,11.22,37.94,NA,NA,483.27,242.04,5422799605.92,9182062336.24,raptor,3.38
-arch-b,googlenet,throughput,2,batch,PASS,FAIL,9.83,21.39,102.00,46.80,520.79,417.87,5669213312.00,9014726090.00,raptor,2.18
+arch-b,googlenet,throughput,2,batch,PASS,FAIL,9.21,21.39,109.00,46.80,545.82,417.87,5591022849.00,9014726090.00,raptor,2.33
-arch-b,googlenet,throughput,4,batch,PASS,FAIL,14.47,21.39,69.10,46.80,401.01,417.87,6835460477.00,9014726090.00,raptor,1.48
+arch-b,googlenet,throughput,4,batch,PASS,FAIL,14.23,21.39,70.30,46.80,411.90,417.87,6718342521.00,9014726090.00,raptor,1.50
-arch-b,googlenet,throughput,8,batch,PASS,FAIL,37.28,21.39,26.80,46.80,245.38,417.87,11453848400.00,9014726090.00,pimcomp,1.75
+arch-b,googlenet,throughput,8,batch,PASS,FAIL,16.19,21.39,61.80,46.80,355.70,417.87,6726459859.00,9014726090.00,raptor,1.32
 arch-b,yolo11n,latency,1,element,PASS,FAIL,168.14,NA,NA,NA,245.31,NA,41248047514.00,NA,NA,NA
 arch-b,yolo11n,throughput,2,batch,PASS,FAIL,118.15,NA,8.46,NA,287.59,NA,37062981240.00,NA,NA,NA
 arch-b,yolo11n,throughput,4,batch,PASS,FAIL,94.18,NA,10.60,NA,316.45,NA,37424491250.00,NA,NA,NA