Faster codegen and merge

src/PIM/Compiler/PimCodeGen.cpp

@@ -270,7 +270,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }
 
 void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  auto intervals = buildLocalAllocIntervals(op, lane);
+  auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
   SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
 
   SmallVector<size_t> slotOrder(plannedSlots.size());

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -237,12 +237,19 @@ getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const Operati
 }
 
 static MemoryTouchInterval
-computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ordering, uint64_t fallbackEnd) {
+computeMemoryTouchInterval(memref::AllocOp allocOp,
+                           const OperationOrdering& ordering,
+                           uint64_t fallbackEnd,
+                           bool includeAliasDescriptions) {
   MemoryTouchInterval interval;
   interval.start = ordering.position.lookup(allocOp);
   interval.end = interval.start;
   interval.startOp = allocOp;
   interval.endOp = allocOp;
+  auto recordAlias = [&](mlir::Value value) {
+    if (includeAliasDescriptions)
+      appendAliasDescription(interval.aliasesFollowed, value);
+  };
 
   SmallPtrSet<mlir::Value, 16> visitedValues;
   SmallPtrSet<Operation*, 32> visitedUsers;
@@ -262,7 +269,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
       if (isSupportedAliasOp(user)) {
         for (mlir::Value result : user->getResults()) {
           pendingValues.push_back(result);
-          appendAliasDescription(interval.aliasesFollowed, result);
+          recordAlias(result);
         }
       }
 
@@ -272,7 +279,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
           if (!tiedOperand || tiedOperand->get() != value)
             continue;
           pendingValues.push_back(result);
-          appendAliasDescription(interval.aliasesFollowed, result);
+          recordAlias(result);
         }
       }
 
@@ -282,8 +289,8 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             continue;
           pendingValues.push_back(forOp.getRegionIterArgs()[index]);
           pendingValues.push_back(forOp.getResult(index));
-          appendAliasDescription(interval.aliasesFollowed, forOp.getRegionIterArgs()[index]);
+          recordAlias(forOp.getRegionIterArgs()[index]);
-          appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+          recordAlias(forOp.getResult(index));
           if (parentLoop && forOp != parentLoop)
             interval.escapesLoop = true;
         }
@@ -298,7 +305,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             if (operand != value)
               continue;
             pendingValues.push_back(ifOp.getResult(index));
-            appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index));
+            recordAlias(ifOp.getResult(index));
           }
         }
         else if (indexSwitch) {
@@ -306,8 +313,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             if (operand != value)
               continue;
             pendingValues.push_back(indexSwitch.getResult(index));
-            appendAliasDescription(interval.aliasesFollowed,
+            recordAlias(indexSwitch.getResult(index));
-                                   indexSwitch.getResult(index));
           }
         }
         else if (!forOp) {
@@ -318,7 +324,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             if (operand != value)
               continue;
             pendingValues.push_back(forOp.getResult(index));
-            appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+            recordAlias(forOp.getResult(index));
             if (parentLoop && forOp == parentLoop)
               interval.escapesLoop = true;
           }
@@ -411,7 +417,8 @@ static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<Lo
 } // namespace
 
 SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp,
-                                                                       std::optional<unsigned> lane) {
+                                                                       std::optional<unsigned> lane,
+                                                                       bool includeAliasDescriptions) {
   SmallVector<LocalAllocInterval, 0> intervals;
   OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
   if (ordering.position.empty())
@@ -428,7 +435,8 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
       llvm_unreachable("Failed to compute local allocation size");
     }
 
-    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
+    MemoryTouchInterval touchInterval =
+      computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
     LocalAllocInterval interval;
     interval.id = nextIntervalId++;
     interval.alloc = allocOp;

src/PIM/Compiler/PimMemoryLiveness.hpp

@@ -49,7 +49,8 @@ struct PlannedPhysicalSlot {
 };
 
 llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp,
-                                                                  std::optional<unsigned> lane);
+                                                                  std::optional<unsigned> lane,
+                                                                  bool includeAliasDescriptions = true);
 
 llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
 

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

@@ -92,15 +92,6 @@ static unsigned getBoundaryIndex(SmallVectorImpl<BoundaryWork>& boundaries,
   return index;
 }
 
-static unsigned getResultBoundary(DeferredExchangePlan& exchange,
-                                  const ScheduledCommunicationPlan& schedule) {
-  std::optional<unsigned> step;
-  for (const ScheduledTransferSlice& slice : schedule.slices)
-    if (slice.family->requirement->exchange == &exchange)
-      step = std::max(step.value_or(0), slice.targetInsertionStep);
-  return step.value_or(exchange.consumerStep);
-}
-
 static size_t hashSemanticKey(const SemanticKey& key) {
   if (key.kind == SemanticKind::Send)
     return llvm::hash_combine(key.kind,
@@ -191,9 +182,11 @@ static LogicalResult appendReplayToken(const PendingToken& token,
   if (auto eventId = std::get_if<unsigned>(&token.value)) {
     const BoundaryEvent& event = events[*eventId];
     if (event.kind == BoundaryEventKind::Send) {
-      action.slices.push_back(event.slice);
+      LaneSet active = event.activeLanes.intersect(lanes);
+      if (!active.empty())
+        action.slices.push_back(event.slice);
       action.instructionLanes =
-        action.instructionLanes.unite(event.activeLanes);
+        action.instructionLanes.unite(active);
     }
     else {
       llvm::append_range(action.slices, intersectReceiveSlice(event.slice, lanes));
@@ -366,7 +359,8 @@ static BoundaryInstructionList materializeInstructions(ArrayRef<CanonicalAction>
       }
       while (index < actions.size() && actions[index].key.kind == SemanticKind::Send
              && actions[index].key.emission == action.key.emission);
-      instructions.push_back(std::move(run));
+      if (!run.slices.empty())
+        instructions.push_back(std::move(run));
       continue;
     }
     SmallVector<unsigned> assemblyEntries;
@@ -552,8 +546,13 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan&
   DeferredBoundaryPlan result;
   SmallVector<BoundaryWork> boundaries;
   DenseMap<BoundaryKey, unsigned> boundaryIndices;
+  DenseMap<DeferredExchangePlan*, unsigned> resultBoundarySteps;
   for (const ScheduledTransferSlice& slice : schedule.slices) {
     ExternalTransferFamily& family = *slice.family;
+    auto [resultStep, inserted] =
+      resultBoundarySteps.try_emplace(family.requirement->exchange, slice.targetInsertionStep);
+    if (!inserted)
+      resultStep->second = std::max(resultStep->second, slice.targetInsertionStep);
     unsigned sourceIndex =
       getBoundaryIndex(boundaries, boundaryIndices,
                        {family.sourceScheduled, slice.sourceInsertionStep});
@@ -576,7 +575,8 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan&
   DenseMap<BoundaryKey, SmallVector<LocalAvailabilityFamily*>> locals;
   DenseMap<BoundaryKey, SmallVector<DeferredExchangePlan*>> exchanges;
   for (const std::unique_ptr<DeferredExchangePlan>& exchange : transfers.exchanges) {
-    unsigned step = getResultBoundary(*exchange, schedule);
+    auto resultStep = resultBoundarySteps.find(exchange.get());
+    unsigned step = resultStep == resultBoundarySteps.end() ? exchange->consumerStep : resultStep->second;
     BoundaryKey key {exchange->target, step};
     getBoundaryIndex(boundaries, boundaryIndices, key);
     for (LocalAvailabilityFamily& local : exchange->local)
@@ -608,15 +608,12 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan&
       localsBySemantic[semantic].push_back(local);
     }
     SmallVector<unsigned> eventSemantics;
-    DenseMap<unsigned, SmallVector<unsigned>> receivesBySemantic;
+    for (BoundaryEvent& event : work.events) {
-    for (auto [eventId, event] : llvm::enumerate(work.events)) {
       unsigned semantic =
         internSemantic(getEventKey(event), semantics, semanticsByHash);
       eventSemantics.push_back(semantic);
-      if (event.kind == BoundaryEventKind::Receive)
-        receivesBySemantic[semantic].push_back(eventId);
     }
-    llvm::SmallDenseSet<unsigned, 8> emittedReceives;
+    llvm::SmallDenseSet<unsigned, 8> emittedAvailabilities;
     for (auto [eventId, event] : llvm::enumerate(work.events)) {
       unsigned semantic = eventSemantics[eventId];
       if (event.kind == BoundaryEventKind::Send) {
@@ -625,15 +622,15 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan&
                           semantic, static_cast<unsigned>(eventId)});
         continue;
       }
-      if (!emittedReceives.insert(semantic).second)
+      tokens.push_back(
-        continue;
+        {event.activeLanes, semantic, static_cast<unsigned>(eventId)});
-      for (unsigned receive : receivesBySemantic[semantic])
+      if (emittedAvailabilities.insert(semantic).second) {
-        tokens.push_back({work.events[receive].activeLanes, semantic, receive});
+        for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) {
-      for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) {
+          tokens.push_back({local->targetLanes, semantic, local});
-        tokens.push_back({local->targetLanes, semantic, local});
+          emittedLocals.insert(local);
-        emittedLocals.insert(local);
+        }
+        localsBySemantic.erase(semantic);
       }
-      localsBySemantic.erase(semantic);
     }
     for (LocalAvailabilityFamily* local : locals[boundary.key])
       if (!emittedLocals.contains(local)) {

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp

@@ -198,6 +198,7 @@ static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
 struct CoreTransferSequences {
   DenseMap<int64_t, StaticIntSequenceChain> sends;
   DenseMap<int64_t, StaticIntSequenceChain> receives;
+  DenseMap<int64_t, StaticIntSequenceChain> events;
 };
 
 struct ExpectedFamily {
@@ -221,6 +222,25 @@ static void appendByCore(DenseMap<int64_t, StaticIntSequenceChain> &result,
     });
 }
 
+static void appendEventsByCore(
+    DenseMap<int64_t, StaticIntSequenceChain> &result,
+    const StaticIntSequence &channels, const StaticIntSequence &cores,
+    size_t begin, size_t count, bool send) {
+  size_t end = begin + count;
+  cores.forEachEqualRun(
+    [&](int64_t core, size_t runBegin, size_t runCount) {
+      size_t selectedBegin = std::max(begin, runBegin);
+      size_t selectedEnd = std::min(end, runBegin + runCount);
+      if (selectedBegin >= selectedEnd)
+        return;
+      SmallVector<int64_t> events;
+      events.reserve(selectedEnd - selectedBegin);
+      for (size_t index = selectedBegin; index < selectedEnd; ++index)
+        events.push_back(2 * channels.valueAt(index) + (send ? 0 : 1));
+      result[core].append(StaticIntSequence::fromValues(events));
+    });
+}
+
 static LogicalResult compareSequences(
     func::FuncOp funcOp,
     const DenseMap<int64_t, StaticIntSequenceChain> &expected,
@@ -254,6 +274,47 @@ static LogicalResult compareSequences(
   return success();
 }
 
+static LogicalResult compareEventSequences(
+    func::FuncOp funcOp,
+    const DenseMap<int64_t, StaticIntSequenceChain> &expected,
+    const DenseMap<int64_t, StaticIntSequenceChain> &actual) {
+  if (expected.size() != actual.size())
+    return funcOp.emitOpError(
+      "realized communication stream set differs from plan");
+  for (const auto &[core, sequence] : expected) {
+    auto found = actual.find(core);
+    if (found == actual.end())
+      return funcOp.emitOpError()
+             << "realized communication stream is missing on core " << core;
+    StaticIntSequenceChainCursor expectedCursor(sequence);
+    StaticIntSequenceChainCursor actualCursor(found->second);
+    uint64_t ordinal = 0;
+    while (!expectedCursor.done() && !actualCursor.done()
+           && expectedCursor.value() == actualCursor.value()) {
+      expectedCursor.advance();
+      actualCursor.advance();
+      ++ordinal;
+    }
+    if (expectedCursor.done() && actualCursor.done())
+      continue;
+    auto describe = [](StaticIntSequenceChainCursor &cursor) {
+      if (cursor.done())
+        return std::pair<StringRef, int64_t>("none", -1);
+      int64_t event = cursor.value();
+      return std::pair<StringRef, int64_t>(
+        event % 2 == 0 ? "send" : "receive", event / 2);
+    };
+    auto [expectedKind, expectedChannel] = describe(expectedCursor);
+    auto [actualKind, actualChannel] = describe(actualCursor);
+    return funcOp.emitOpError()
+           << "realized communication order differs on core " << core
+           << " at ordinal " << ordinal << ": expected " << expectedKind
+           << " channel " << expectedChannel << ", actual " << actualKind
+           << " channel " << actualChannel;
+  }
+  return success();
+}
+
 } // namespace
 
 LogicalResult verifyPlannedCommunicationDeadlockFree(
@@ -322,6 +383,10 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
                  slice.familyOffset, slice.transferCount);
     appendByCore(expected.receives, family.channelIds, family.targetCores,
                  slice.familyOffset, slice.transferCount);
+    appendEventsByCore(expected.events, family.channelIds, family.sourceCores,
+                       slice.familyOffset, slice.transferCount, true);
+    appendEventsByCore(expected.events, family.channelIds, family.targetCores,
+                       slice.familyOffset, slice.transferCount, false);
   }
 
   CoreTransferSequences actual;
@@ -378,13 +443,18 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
                  *actualChannels.back(),
                  realized->send ? realized->sources : realized->targets,
                  0, actualChannels.back()->size());
+    appendEventsByCore(actual.events, *actualChannels.back(),
+                       realized->send ? realized->sources : realized->targets,
+                       0, actualChannels.back()->size(), realized->send);
   });
   if (invalid)
     return failure();
   if (failed(compareSequences(
         funcOp, expected.sends, actual.sends, "send"))
       || failed(compareSequences(
-        funcOp, expected.receives, actual.receives, "receive")))
+        funcOp, expected.receives, actual.receives, "receive"))
+      || failed(compareEventSequences(
+        funcOp, expected.events, actual.events)))
     return failure();
   return success();
 }