Almost all model compile and they are faster then pimcomp

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

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "DeferredBoundaryRealization.hpp"
+#include "DeferredProjectionAnalysis.hpp"
 #include "DeferredResultRealization.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
@@ -428,9 +429,9 @@ static FailureOr<Value> insertProjectionFragment(Value fragment, Value specializ
                                                  DeferredExchangePlan &exchange, bool grouped, DeferredEmissionContext &context) {
   Value shaped = fragment;
   if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
-    SmallVector<int64_t> shape(leaf.leadingRankReduced ? leaf.reconstructedType.getShape() : leaf.reconstructedType.getShape().drop_front());
+    RankedTensorType projectedType = getDeferredProjectedFragmentType(leaf);
     shaped = extractMixedSliceOrIdentity(context.rewriter, exchange.deferred.getLoc(), shaped,
-                                         RankedTensorType::get(shape, leaf.reconstructedType.getElementType()),
+                                         projectedType,
                                          lookupGeometry(geometry, geometryRow, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc()));
     if (!shaped) return failure();
   }
@@ -494,8 +495,26 @@ static LogicalResult emitLeafCollectionUpdate(const EmitReceiveAssemblyRun &run,
   return emitCollectionUpdate(run.lanes, lane, laneCount, key, current, exchange.deferred, context, emit);
 }
 
-static FailureOr<Value> transformAssemblySource(Value fragment, const DeferredInsertAssemblyEntryTemplate &entry,
-                                                DeferredExchangePlan &exchange, DeferredEmissionContext &context) {
+static FailureOr<Value> transformAssemblySource(
+    Value fragment, const DeferredInsertAssemblyEntryTemplate &entry,
+    Value runtimeLane, const DeferredResultPlan &resultPlan,
+    DeferredExchangePlan &exchange, DeferredEmissionContext &context) {
+  if (entry.coordinate.leafIndex >= exchange.program.leaves.size()
+      || entry.coordinate.leafIndex >= resultPlan.innerGeometry.size())
+    return failure();
+  const DeferredProjectionLeafTemplate &leaf =
+    exchange.program.leaves[entry.coordinate.leafIndex];
+  if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+    fragment = extractMixedSliceOrIdentity(
+      context.rewriter, exchange.deferred.getLoc(), fragment,
+      getDeferredProjectedFragmentType(leaf),
+      lookupGeometry(resultPlan.innerGeometry[entry.coordinate.leafIndex],
+                     context.constants.getIndex(0), runtimeLane,
+                     exchange.deferred, context,
+                     exchange.deferred.getLoc()));
+    if (!fragment)
+      return failure();
+  }
   switch (entry.sourceTransform) {
   case DeferredAssemblySourceTransform::Identity:
     return fragment.getType() == entry.sourceType ? FailureOr<Value>(fragment) : FailureOr<Value>(failure());
@@ -510,12 +529,15 @@ static FailureOr<Value> transformAssemblySource(Value fragment, const DeferredIn
 static FailureOr<Value> materializeLoopedLocalAssemblySource(
     RequirementFamily &requirement,
     const DeferredInsertAssemblyEntryTemplate &entry,
-    Value localOffset, DeferredExchangePlan &exchange,
+    Value localOffset, Value runtimeLane,
+    const DeferredResultPlan &resultPlan, DeferredExchangePlan &exchange,
     DeferredEmissionContext &context) {
   Value payload = requirement.producer->payload;
   auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
   RankedTensorType sourceType = entry.sourceType;
-  if (entry.sourceTransform
+  if (exchange.program.leaves[entry.coordinate.leafIndex].form
+        != DeferredLeafForm::GraphBatchProjection
+      && entry.sourceTransform
         == DeferredAssemblySourceTransform::RemoveLeadingUnitDimension
       && payloadType && sourceType
       && payloadType.getRank() > sourceType.getRank()
@@ -542,7 +564,8 @@ static FailureOr<Value> materializeLoopedLocalAssemblySource(
     exchange.deferred.getLoc());
   if (failed(fragment))
     return failure();
-  return transformAssemblySource(*fragment, entry, exchange, context);
+  return transformAssemblySource(
+    *fragment, entry, runtimeLane, resultPlan, exchange, context);
 }
 
 static LogicalResult emitLoopedLocalCollectionUpdate(
@@ -590,7 +613,8 @@ static LogicalResult emitLoopedLocalCollectionUpdate(
       RequirementFamily &requirement = *run.families.front()->requirement;
       FailureOr<Value> source = entry
         ? materializeLoopedLocalAssemblySource(
-            requirement, *entry, localOffset, exchange, context)
+            requirement, *entry, localOffset, runtimeLane, resultPlan,
+            exchange, context)
         : materializeSendPayload(
             requirement, localOffset, nullptr, context, loc);
       if (failed(source))
@@ -650,7 +674,8 @@ static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run,
   auto emit = [&](Value initial) {
     return emitReceiveAssembly(run, lane, laneCount, initial, context,
       [&](Value fragment, Value position, Value, Value runtimeLane, Value assembled) -> FailureOr<Value> {
-        auto shaped = transformAssemblySource(fragment, sourceEntry, exchange, context);
+        auto shaped = transformAssemblySource(
+          fragment, sourceEntry, runtimeLane, resultPlan, exchange, context);
         if (failed(shaped) || shaped->getType() != sourceEntry.sourceType) return failure();
         return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled,
           lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc()));
@@ -775,7 +800,9 @@ static LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &upd
     auto emit = [&](Value assembled) -> FailureOr<Value> {
       auto fragment = materialize();
       if (failed(fragment)) return failure();
-      auto shaped = transformAssemblySource(*fragment, entry, exchange, context);
+      auto shaped = transformAssemblySource(
+        *fragment, entry, lane ? lane : context.constants.getIndex(0),
+        resultPlan, exchange, context);
       if (failed(shaped) || shaped->getType() != entry.sourceType) return failure();
       return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled,
         lookupGeometry(resultPlan.assemblyGeometry, context.constants.getIndex(update.collectionPosition), lane ? lane : context.constants.getIndex(0),

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

@@ -354,6 +354,16 @@ FailureOr<int64_t> evaluateDeferredIndex(
   return evaluate(value, environment, visiting);
 }
 
+RankedTensorType getDeferredProjectedFragmentType(
+    const DeferredProjectionLeafTemplate &leaf) {
+  if (leaf.form == DeferredLeafForm::GraphBatchProjection
+      && !leaf.leadingRankReduced)
+    return RankedTensorType::get(
+      leaf.reconstructedType.getShape().drop_front(),
+      leaf.reconstructedType.getElementType());
+  return leaf.reconstructedType;
+}
+
 FailureOr<int64_t> evaluateDeferredIndex(
     OpFoldResult value, const StaticIndexEnvironment &environment) {
   if (auto attr = dyn_cast<Attribute>(value))

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

@@ -18,6 +18,9 @@ mlir::FailureOr<int64_t> evaluateDeferredIndex(
 mlir::FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
   SpatDeferredCommunicationOp deferred);
 
+mlir::RankedTensorType getDeferredProjectedFragmentType(
+  const DeferredProjectionLeafTemplate &leaf);
+
 class DeferredLaneValueEvaluator {
 public:
   DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,

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

@@ -72,11 +72,7 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
              << ": publication fragment types differ or are unranked";
     RankedTensorType normalized = fragmentType;
     if (leaf.form == DeferredLeafForm::GraphBatchProjection)
-      normalized = leaf.leadingRankReduced
-        ? leaf.reconstructedType
-        : RankedTensorType::get(
-            leaf.reconstructedType.getShape().drop_front(),
-            leaf.reconstructedType.getElementType());
+      normalized = getDeferredProjectedFragmentType(leaf);
     bool direct = positionCount == 1 && normalized == leaf.reconstructedType;
     bool leading = leaf.reconstructedType.getRank() == normalized.getRank() + 1
                 && leaf.reconstructedType.getDimSize(0) == positionCount
@@ -113,11 +109,20 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
 }
 
 static bool supportsAssemblyTransform(
-    Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry) {
+    Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry,
+    const DeferredProjectionLeafTemplate &leaf) {
   auto publication = dyn_cast<RankedTensorType>(publicationType);
   RankedTensorType source = entry.sourceType;
   if (!publication || !source)
     return false;
+  if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+    auto physical = dyn_cast<RankedTensorType>(leaf.sourceRoot.getType());
+    if (!physical || physical.getRank() != publication.getRank() + 1
+        || physical.getElementType() != publication.getElementType()
+        || physical.getShape().drop_front() != publication.getShape())
+      return false;
+    publication = getDeferredProjectedFragmentType(leaf);
+  }
   switch (entry.sourceTransform) {
   case DeferredAssemblySourceTransform::Identity:
     return publication == source;
@@ -150,11 +155,16 @@ static LogicalResult buildInsertAssemblyCollection(
     for (RequirementFamily &requirement : exchange.requirements) {
       if (!(requirement.coordinate == entry.coordinate))
         continue;
+      const DeferredProjectionLeafTemplate &leaf =
+        exchange.program.leaves[entry.coordinate.leafIndex];
       if (!supportsAssemblyTransform(
-            requirement.publicationFragmentType, entry))
+            requirement.publicationFragmentType, entry, leaf))
         return exchange.deferred.emitOpError(
                  "insert assembly source transform does not match publication type at entry ")
-               << position;
+               << position << ": publication "
+               << requirement.publicationFragmentType << ", assembly source "
+               << entry.sourceType << ", transform "
+               << static_cast<unsigned>(entry.sourceTransform);
       if (!collected.insert(&requirement).second)
         return exchange.deferred.emitOpError(
                  "insert assembly requirement is owned by multiple entries at entry ")

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

@@ -25,6 +25,7 @@
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
 namespace onnx_mlir {
@@ -534,6 +535,27 @@ evaluateIndexLike(Value value, const DenseMap<Value, int64_t>& bindings, std::op
     return evaluateAffineApply(affineApply,
                                [&](Value operand) { return evaluateIndexLike(operand, bindings, lane, laneArg); });
 
+  Operation* op = value.getDefiningOp();
+  if (!op || !isPureIndexComputationOp(op))
+    return failure();
+  SmallVector<Attribute> operands;
+  Builder builder(op->getContext());
+  for (Value operand : op->getOperands()) {
+    FailureOr<int64_t> folded =
+      evaluateIndexLike(operand, bindings, lane, laneArg);
+    if (failed(folded))
+      return failure();
+    operands.push_back(builder.getIntegerAttr(operand.getType(), *folded));
+  }
+  SmallVector<OpFoldResult> results;
+  if (failed(op->fold(operands, results)) || results.size() != 1)
+    return failure();
+  if (auto attribute = dyn_cast<Attribute>(results.front()))
+    if (auto integer = dyn_cast<IntegerAttr>(attribute))
+      return integer.getInt();
+  if (auto folded = dyn_cast<Value>(results.front()))
+    return evaluateIndexLike(folded, bindings, lane, laneArg);
+
   return failure();
 }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp

@@ -12,10 +12,13 @@
 #include <vector>
 
 #include "PeftScheduler.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
+using namespace mlir;
+
 namespace {
 
 // Pressure means distinct weights exceed half the fleet's one-copy capacity.
@@ -209,6 +212,79 @@ std::vector<CrossbarUsage> planCrossbarReservations(const ComputeGraph& graph,
   return reservations;
 }
 
+using LanePublicationSignatures = llvm::SmallVector<llvm::SmallVector<int64_t, 8>, 8>;
+
+FailureOr<LanePublicationSignatures>
+buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCache& publicationCache) {
+  LanePublicationSignatures signatures(batch.getLaneCount());
+  for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) {
+    auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache);
+    if (failed(publicationMap))
+      return failure();
+
+    for (auto [useIndex, use] : llvm::enumerate(result.getUses())) {
+      auto blueprint = dyn_cast<SpatBlueprintOp>(use.getOwner());
+      if (!blueprint || blueprint.getMode() != "fragment_assembly")
+        continue;
+      auto operandIndices = blueprint.getFragmentOperandIndices();
+      auto sourceSlots = blueprint.getFragmentSourceSlots();
+      auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+      auto fragmentStrides = blueprint.getFragmentStrides();
+      auto outputType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+      if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides
+          || !outputType || !outputType.hasStaticShape())
+        return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"),
+               failure();
+
+      llvm::ArrayRef<int64_t> fragmentOffsets = blueprint.getFragmentOffsets();
+      llvm::ArrayRef<int64_t> fragmentSizes = blueprint.getFragmentSizes();
+      int64_t rank = outputType.getRank();
+      if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size()
+          || fragmentOffsets.size() != fragmentStrides->size()
+          || fragmentOffsets.size() != operandIndices->size() * rank
+          || sourceSlots->size() != operandIndices->size()
+          || sourceOffsets->size() != operandIndices->size())
+        return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"),
+               failure();
+
+      llvm::SmallVector<llvm::SmallVector<size_t, 2>, 8> fragmentsByLane(batch.getLaneCount());
+      for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+        if (operandIndex != static_cast<int64_t>(use.getOperandNumber()))
+          continue;
+        int64_t slot = (*sourceSlots)[fragmentIndex];
+        if (slot < 0 || slot >= static_cast<int64_t>((*publicationMap)->physicalSlotToGraphLane.size()))
+          return blueprint.emitOpError("PEFT publication fragment source slot is out of range"), failure();
+        int64_t graphLane = (*publicationMap)->physicalSlotToGraphLane[slot];
+        if (graphLane < 0 || graphLane >= batch.getLaneCount())
+          return blueprint.emitOpError("PEFT publication fragment has no graph lane owner"), failure();
+        fragmentsByLane[graphLane].push_back(fragmentIndex);
+      }
+
+      for (auto [lane, fragments] : llvm::enumerate(fragmentsByLane)) {
+        if (fragments.empty())
+          continue;
+        llvm::SmallVector<int64_t, 8>& signature = signatures[lane];
+        signature.push_back(resultIndex);
+        signature.push_back(useIndex);
+        signature.push_back(fragments.size());
+        signature.push_back(rank);
+        size_t firstFragment = fragments.front();
+        for (size_t fragmentIndex : fragments) {
+          signature.push_back((*sourceOffsets)[fragmentIndex] - (*sourceOffsets)[firstFragment]);
+          for (int64_t dim = 0; dim < rank; ++dim) {
+            size_t index = fragmentIndex * rank + dim;
+            size_t firstIndex = firstFragment * rank + dim;
+            signature.push_back(fragmentOffsets[index] - fragmentOffsets[firstIndex]);
+            signature.push_back(fragmentSizes[index]);
+            signature.push_back((*fragmentStrides)[index]);
+          }
+        }
+      }
+    }
+  }
+  return signatures;
+}
+
 } // namespace
 
 Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) {
@@ -450,6 +526,28 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
 
   // 5. Check if equal schedule in two level
   llvm::DenseMap<size_t, mlir::SmallVector<size_t, 5>> equivalentClass;
+  GraphBatchPublicationCache publicationCache;
+  llvm::DenseMap<Operation *, LanePublicationSignatures> publicationSignatures;
+  llvm::DenseSet<Operation *> invalidPublicationSignatures;
+  auto haveCompatiblePublications = [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
+    auto batch = dyn_cast<SpatComputeBatch>(lhs.op);
+    if (!batch || batch.getNumResults() == 0)
+      return true;
+    auto signatures = publicationSignatures.find(lhs.op);
+    if (signatures == publicationSignatures.end() && !invalidPublicationSignatures.contains(lhs.op)) {
+      auto built = buildLanePublicationSignatures(batch, publicationCache);
+      if (failed(built))
+        invalidPublicationSignatures.insert(lhs.op);
+      else
+        signatures = publicationSignatures.try_emplace(lhs.op, std::move(*built)).first;
+    }
+    if (invalidPublicationSignatures.contains(lhs.op))
+      return false;
+    for (uint32_t lane = 0; lane < lhs.laneCount; ++lane)
+      if (signatures->second[lhs.laneStart + lane] != signatures->second[rhs.laneStart + lane])
+        return false;
+    return true;
+  };
   for (size_t currentProcessor = 0; currentProcessor < processorCount - 1; ++currentProcessor) {
     for (size_t controlProcessor = currentProcessor; controlProcessor < processorCount; ++controlProcessor) {
       if (tasksByProcessor[currentProcessor].size() != tasksByProcessor[controlProcessor].size())
@@ -462,7 +560,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu
         const ComputeInstance currentComputeInstance = graph.nodes[currentTask].instance;
         const ComputeInstance controlComputeInstance = graph.nodes[controlTask].instance;
         if (currentComputeInstance.op != controlComputeInstance.op
-            || currentComputeInstance.laneCount != controlComputeInstance.laneCount) {
+            || currentComputeInstance.laneCount != controlComputeInstance.laneCount
+            || !haveCompatiblePublications(currentComputeInstance, controlComputeInstance)) {
           equalSchedule = false;
           break;
         }