Merge branch 'TestRottoConDeadLock' of chef.heaplab.deib.polimi.it:nnicolosi/Raptor into TestRottoConDeadLock

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

@@ -142,6 +142,71 @@ static std::optional<EmitLocalCollectionRun> buildLocalConcat(
        ? std::optional<EmitLocalCollectionRun>(std::move(run)) : std::nullopt;
 }
 
+static bool canLoopLocalCollection(
+    const EmitLocalCollectionRun &update, unsigned targetLaneCount) {
+  if (!update.collection || update.concatenatePayloads
+      || update.families.size() != 1
+      || !(update.lanes == LaneSet::all(targetLaneCount)))
+    return false;
+  RequirementFamily &requirement = *update.families.front()->requirement;
+  return requirement.targetLanes == update.lanes
+      && !requirement.producerProjection
+      && (!requirement.producerLocalOffsets
+          || requirement.producerLocalOffsets->size() == targetLaneCount);
+}
+
+static bool haveSameLocalCollectionContract(
+    const EmitLocalCollectionRun &lhs,
+    const EmitLocalCollectionRun &rhs) {
+  if (lhs.collection != rhs.collection)
+    return false;
+  RequirementFamily &left = *lhs.families.front()->requirement;
+  RequirementFamily &right = *rhs.families.front()->requirement;
+  if (left.producer->payload != right.producer->payload
+      || left.publicationFragmentType != right.publicationFragmentType)
+    return false;
+  if (lhs.collection->key.kind != FragmentCollectionKind::InsertAssembly)
+    return true;
+  const auto &entries =
+    lhs.collection->key.exchange->program.insertAssembly->entries;
+  const auto &leftEntry = entries[lhs.collectionPosition];
+  const auto &rightEntry = entries[rhs.collectionPosition];
+  return leftEntry.sourceTransform == rightEntry.sourceTransform
+      && leftEntry.sourceType == rightEntry.sourceType;
+}
+
+static void appendLocalUpdates(
+    BoundaryProgram &boundary,
+    SmallVectorImpl<EmitLocalCollectionRun> &updates,
+    unsigned targetLaneCount) {
+  for (size_t index = 0; index < updates.size();) {
+    EmitLocalCollectionRun &first = updates[index];
+    if (!canLoopLocalCollection(first, targetLaneCount)) {
+      boundary.instructions.push_back(std::move(first));
+      ++index;
+      continue;
+    }
+    size_t end = index + 1;
+    while (end < updates.size()
+           && canLoopLocalCollection(updates[end], targetLaneCount)
+           && haveSameLocalCollectionContract(first, updates[end]))
+      ++end;
+    if (end - index == 1) {
+      boundary.instructions.push_back(std::move(first));
+      ++index;
+      continue;
+    }
+    EmitLocalCollectionLoopRun run;
+    run.collection = first.collection;
+    run.lanes = first.lanes;
+    for (; index < end; ++index) {
+      run.positions.push_back(updates[index].collectionPosition);
+      run.families.push_back(updates[index].families.front());
+    }
+    boundary.instructions.push_back(std::move(run));
+  }
+}
+
 static void appendReceive(BoundaryProgram &boundary,
                           const ScheduledTransferSlice &slice,
                           CollectionTarget target) {
@@ -288,8 +353,7 @@ FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
       if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
         return failure();
     }
-    for (EmitLocalCollectionRun &update : localUpdates)
-      boundary.instructions.push_back(std::move(update));
+    appendLocalUpdates(boundary, localUpdates, exchange->targetLaneCount);
     for (RequirementFamily &requirement : exchange->requirements)
       if (!(coverage.lookup(&requirement) == requirement.targetLanes))
         return exchange->deferred.emitOpError(

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

@@ -20,6 +20,12 @@ struct EmitLocalCollectionRun {
   LaneSet lanes;
   bool concatenatePayloads = false;
 };
+struct EmitLocalCollectionLoopRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  llvm::SmallVector<unsigned> positions;
+  llvm::SmallVector<LocalAvailabilityFamily*> families;
+  LaneSet lanes;
+};
 struct EmitReceiveAssemblyRun {
   const FragmentCollectionPlan* collection = nullptr;
   llvm::SmallVector<ScheduledTransferSlice> slices;
@@ -33,7 +39,8 @@ struct ProduceDeferredResult {
 };
 
 using BoundaryInstruction =
-  std::variant<EmitSendRun, EmitLocalCollectionRun, EmitReceiveAssemblyRun,
+  std::variant<EmitSendRun, EmitLocalCollectionRun,
+               EmitLocalCollectionLoopRun, EmitReceiveAssemblyRun,
                ProduceDeferredResult>;
 struct BoundaryProgram {
   BoundaryKey key;

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

@@ -507,6 +507,137 @@ static FailureOr<Value> transformAssemblySource(Value fragment, const DeferredIn
   llvm_unreachable("unknown deferred assembly source transform");
 }
 
+static FailureOr<Value> materializeLoopedLocalAssemblySource(
+    RequirementFamily &requirement,
+    const DeferredInsertAssemblyEntryTemplate &entry,
+    Value localOffset, DeferredExchangePlan &exchange,
+    DeferredEmissionContext &context) {
+  Value payload = requirement.producer->payload;
+  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
+  RankedTensorType sourceType = entry.sourceType;
+  if (entry.sourceTransform
+        == DeferredAssemblySourceTransform::RemoveLeadingUnitDimension
+      && payloadType && sourceType
+      && payloadType.getRank() > sourceType.getRank()
+      && payloadType.getElementType() == sourceType.getElementType()
+      && payloadType.getShape().take_back(sourceType.getRank())
+           == sourceType.getShape()) {
+    MixedSliceGeometry geometry;
+    int64_t rankDifference = payloadType.getRank() - sourceType.getRank();
+    geometry.offsets.assign(payloadType.getRank(),
+                            context.rewriter.getIndexAttr(0));
+    if (payload.getType() != requirement.publicationFragmentType)
+      geometry.offsets.front() = localOffset;
+    geometry.sizes.assign(rankDifference, context.rewriter.getIndexAttr(1));
+    for (int64_t dimension : sourceType.getShape())
+      geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+    geometry.strides.assign(payloadType.getRank(),
+                            context.rewriter.getIndexAttr(1));
+    return extractMixedSliceOrIdentity(
+      context.rewriter, exchange.deferred.getLoc(), payload, sourceType,
+      geometry);
+  }
+  auto fragment = materializeSendPayload(
+    requirement, localOffset, nullptr, context,
+    exchange.deferred.getLoc());
+  if (failed(fragment))
+    return failure();
+  return transformAssemblySource(*fragment, entry, exchange, context);
+}
+
+static LogicalResult emitLoopedLocalCollectionUpdate(
+    const EmitLocalCollectionLoopRun &run, Value lane, unsigned laneCount,
+    const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
+  if (!run.collection || run.positions.size() < 2
+      || run.positions.size() != run.families.size())
+    return failure();
+  const FragmentCollectionPlan &collection = *run.collection;
+  DeferredExchangePlan &exchange = *collection.key.exchange;
+  const DeferredInsertAssemblyEntryTemplate *entry = nullptr;
+  if (collection.key.kind == FragmentCollectionKind::InsertAssembly)
+    entry = &exchange.program.insertAssembly
+               ->entries[run.positions.front()];
+  SmallVector<StaticIntSequence> offsetRows;
+  SmallVector<int64_t> positions;
+  offsetRows.reserve(run.families.size());
+  positions.reserve(run.positions.size());
+  for (auto [position, family] : llvm::zip_equal(run.positions,
+                                                 run.families)) {
+    RequirementFamily &requirement = *family->requirement;
+    offsetRows.push_back(requirement.producerLocalOffsets
+      ? *requirement.producerLocalOffsets
+      : StaticIntSequence::uniform(0, laneCount));
+    positions.push_back(position);
+  }
+  auto localOffsets = StaticIntGrid::fromRows(offsetRows);
+  if (failed(localOffsets))
+    return failure();
+  Operation *anchor = exchange.deferred;
+  Location loc = anchor->getLoc();
+  Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+  Value current = context.fragmentCollections.lookup(collection.key);
+  if (!current)
+    current = createCollectionInitial(collection, context);
+  auto loop = buildNormalizedScfFor(
+    context.rewriter, loc, context.constants.getIndex(0),
+    context.constants.getIndex(run.positions.size()),
+    context.constants.getIndex(1), ValueRange {current},
+    [&](OpBuilder &, Location, Value action, ValueRange iterArgs,
+        SmallVectorImpl<Value> &yielded) -> LogicalResult {
+      Value position = lookup(positions, action, anchor, context, loc);
+      Value localOffset = localOffsets->emitLookup(
+        action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+      RequirementFamily &requirement = *run.families.front()->requirement;
+      FailureOr<Value> source = entry
+        ? materializeLoopedLocalAssemblySource(
+            requirement, *entry, localOffset, exchange, context)
+        : materializeSendPayload(
+            requirement, localOffset, nullptr, context, loc);
+      if (failed(source))
+        return failure();
+      Value next;
+      if (entry) {
+        if (source->getType() != entry->sourceType)
+          return failure();
+        next = insertMixedSlice(
+          context.rewriter, loc, *source, iterArgs.front(),
+          lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane,
+                         anchor, context, loc));
+      } else {
+        bool grouped = collection.key.kind
+                    == FragmentCollectionKind::GroupedLeaf;
+        Value specialization = context.constants.getIndex(0);
+        Value leafPosition = position;
+        if (grouped) {
+          Value divisor = context.constants.getIndex(collection.positionCount);
+          specialization = arith::DivUIOp::create(
+            context.rewriter, loc, position, divisor);
+          leafPosition = arith::RemUIOp::create(
+            context.rewriter, loc, position, divisor);
+        }
+        unsigned leafIndex = collection.key.leafIndex;
+        const DeferredProjectionLeafTemplate &leaf =
+          exchange.program.leaves[leafIndex];
+        auto inserted = insertProjectionFragment(
+          *source, specialization, leafPosition,
+          grouped ? specialization : context.constants.getIndex(0),
+          runtimeLane, iterArgs.front(), leaf,
+          resultPlan.innerGeometry[leafIndex], exchange, grouped, context);
+        if (failed(inserted))
+          return failure();
+        next = *inserted;
+      }
+      if (!next)
+        return failure();
+      yielded.push_back(next);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  context.fragmentCollections[collection.key] = loop->results.front();
+  return success();
+}
+
 static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount,
                                               const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
   const FragmentCollectionPlan &collection = *run.collection;
@@ -683,6 +814,18 @@ static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstructi
       if (failed(emitted))
         return exchange->deferred.emitOpError(
           "failed to update fragment collection from local availability"), failure();
+    } else if (auto update =
+                 std::get_if<EmitLocalCollectionLoopRun>(&instruction)) {
+      DeferredExchangePlan *exchange = update->collection->key.exchange;
+      const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);
+      LogicalResult emitted = resultPlan
+        ? emitLoopedLocalCollectionUpdate(
+            *update, lane, laneCount, *resultPlan, context)
+        : failure();
+      if (failed(emitted))
+        return exchange->deferred.emitOpError(
+          "failed to update fragment collection from local assembly run"),
+          failure();
     } else if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
       DeferredExchangePlan *exchange = assembly->collection->key.exchange;
       const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);

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

@@ -168,12 +168,16 @@ static LogicalResult materializeResultfulBatchRun(
 
       IRMapping mapper;
       mapper.map(*batch.getLaneArgument(), originalLane);
-      Value localLane = runLaneCount == 1
-        ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0)
-        : arith::SubIOp::create(
-            builder, bodyLoc, originalLane,
-            getOrCreateIndexConstant(
-              rewriter, batch.getOperation(), first.laneStart));
+      Value localLane;
+      if (runLaneCount == 1)
+        localLane = getOrCreateIndexConstant(rewriter, batch.getOperation(), 0);
+      else if (first.laneStart == 0)
+        localLane = originalLane;
+      else
+        localLane = arith::SubIOp::create(
+          builder, bodyLoc, originalLane,
+          getOrCreateIndexConstant(
+            rewriter, batch.getOperation(), first.laneStart));
       for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
         mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
       SmallVector<DeferredInputPlan> inputPlans;