diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp index 012e8d9..4a35aee 100644 --- a/src/PIM/Compiler/PimCodeGen.cpp +++ b/src/PIM/Compiler/PimCodeGen.cpp @@ -270,7 +270,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) { } void PimMemory::allocateCore(Operation* op, std::optional lane) { - auto intervals = buildLocalAllocIntervals(op, lane); + auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull); SmallVector plannedSlots = planPhysicalSlots(intervals); SmallVector slotOrder(plannedSlots.size()); diff --git a/src/PIM/Compiler/PimMemoryLiveness.cpp b/src/PIM/Compiler/PimMemoryLiveness.cpp index 9e91ac9..09657df 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.cpp +++ b/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 visitedValues; SmallPtrSet 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]); - appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index)); + recordAlias(forOp.getRegionIterArgs()[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, - indexSwitch.getResult(index)); + recordAlias(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 onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp, - std::optional lane) { + std::optional lane, + bool includeAliasDescriptions) { SmallVector intervals; OperationOrdering ordering = buildOperationOrdering(coreLikeOp); if (ordering.position.empty()) @@ -428,7 +435,8 @@ SmallVector 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; diff --git a/src/PIM/Compiler/PimMemoryLiveness.hpp b/src/PIM/Compiler/PimMemoryLiveness.hpp index 736b59d..99c8da7 100644 --- a/src/PIM/Compiler/PimMemoryLiveness.hpp +++ b/src/PIM/Compiler/PimMemoryLiveness.hpp @@ -49,7 +49,8 @@ struct PlannedPhysicalSlot { }; llvm::SmallVector buildLocalAllocIntervals(mlir::Operation* coreLikeOp, - std::optional lane); + std::optional lane, + bool includeAliasDescriptions = true); llvm::SmallVector planPhysicalSlots(llvm::MutableArrayRef intervals); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp index f0a2fac..701c898 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp @@ -92,15 +92,6 @@ static unsigned getBoundaryIndex(SmallVectorImpl& boundaries, return index; } -static unsigned getResultBoundary(DeferredExchangePlan& exchange, - const ScheduledCommunicationPlan& schedule) { - std::optional 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(&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 } 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 assemblyEntries; @@ -552,8 +546,13 @@ FailureOr buildDeferredBoundaryPlan(DeferredTransferPlan& DeferredBoundaryPlan result; SmallVector boundaries; DenseMap boundaryIndices; + DenseMap 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 buildDeferredBoundaryPlan(DeferredTransferPlan& DenseMap> locals; DenseMap> exchanges; for (const std::unique_ptr& 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 buildDeferredBoundaryPlan(DeferredTransferPlan& localsBySemantic[semantic].push_back(local); } SmallVector eventSemantics; - DenseMap> receivesBySemantic; - for (auto [eventId, event] : llvm::enumerate(work.events)) { + for (BoundaryEvent& event : 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 emittedReceives; + llvm::SmallDenseSet emittedAvailabilities; for (auto [eventId, event] : llvm::enumerate(work.events)) { unsigned semantic = eventSemantics[eventId]; if (event.kind == BoundaryEventKind::Send) { @@ -625,15 +622,15 @@ FailureOr buildDeferredBoundaryPlan(DeferredTransferPlan& semantic, static_cast(eventId)}); continue; } - if (!emittedReceives.insert(semantic).second) - continue; - for (unsigned receive : receivesBySemantic[semantic]) - tokens.push_back({work.events[receive].activeLanes, semantic, receive}); - for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) { - tokens.push_back({local->targetLanes, semantic, local}); - emittedLocals.insert(local); + tokens.push_back( + {event.activeLanes, semantic, static_cast(eventId)}); + if (emittedAvailabilities.insert(semantic).second) { + for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) { + tokens.push_back({local->targetLanes, semantic, local}); + emittedLocals.insert(local); + } + localsBySemantic.erase(semantic); } - localsBySemantic.erase(semantic); } for (LocalAvailabilityFamily* local : locals[boundary.key]) if (!emittedLocals.contains(local)) { diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp index e0129a6..318a63c 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp @@ -198,6 +198,7 @@ static FailureOr parseRealizedOperation(Operation *op) { struct CoreTransferSequences { DenseMap sends; DenseMap receives; + DenseMap events; }; struct ExpectedFamily { @@ -221,6 +222,25 @@ static void appendByCore(DenseMap &result, }); } +static void appendEventsByCore( + DenseMap &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 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 &expected, @@ -254,6 +274,47 @@ static LogicalResult compareSequences( return success(); } +static LogicalResult compareEventSequences( + func::FuncOp funcOp, + const DenseMap &expected, + const DenseMap &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("none", -1); + int64_t event = cursor.value(); + return std::pair( + 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(); }