finally fast googlenet with correct latency artifacts for fair comparison

src/PIM/Dialect/Spatial/Transforms/TrivialGraphComputeMergePass.cpp

@@ -36,10 +36,14 @@ static bool hasOnlyStructuralAttrs(ComputeOp op) {
   });
 }
 
-static bool hasCapacityFor(Operation* producer, Operation* consumer) {
-  CrossbarUsage producerWeights = collectDistinctCrossbarWeights(producer);
-  CrossbarUsage consumerWeights = collectDistinctCrossbarWeights(consumer);
-  return getCrossbarUnionSize(producerWeights, consumerWeights) <= static_cast<size_t>(crossbarCountInCore.getValue());
+static bool hasCapacityFor(Operation* producer, Operation* consumer,
+                           size_t residentWeightCapacity) {
+  ResidentWeightSet producerWeights =
+    collectDistinctResidentWeights(producer);
+  ResidentWeightSet consumerWeights =
+    collectDistinctResidentWeights(consumer);
+  return getResidentWeightUnionSize(producerWeights, consumerWeights)
+      <= residentWeightCapacity;
 }
 
 template <typename ConsumerOp>
@@ -156,8 +160,11 @@ static void mapExternalArguments(OldOp oldOp, NewOp newOp, IRMapping& mapper, bo
 }
 
 struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
-  MergeTrivialScalarComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialScalarComputes(MLIRContext *context,
+                             TrivialGraphMergeStats *stats,
+                             size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
     SpatGraphCompute producer;
@@ -166,7 +173,8 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
       if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasNoNestedArgumentCaptures(candidate)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasNoNestedArgumentCaptures(candidate)
           && hasNoNestedArgumentCaptures(consumer)) {
         producer = candidate;
         break;
@@ -201,6 +209,7 @@ struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 static bool isLaneIndex(Value value, Value lane, int64_t laneCount) {
@@ -397,8 +406,11 @@ static bool hasDirectLaneConsumers(SpatGraphComputeBatch producer, SpatGraphComp
 }
 
 struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
-  MergeTrivialBatchComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
-      : OpRewritePattern(context), stats(stats) {}
+  MergeTrivialBatchComputes(MLIRContext *context,
+                            TrivialGraphMergeStats *stats,
+                            size_t residentWeightCapacity)
+      : OpRewritePattern(context), stats(stats),
+        residentWeightCapacity(residentWeightCapacity) {}
 
   LogicalResult matchAndRewrite(SpatGraphComputeBatch consumer, PatternRewriter& rewriter) const override {
     SpatGraphComputeBatch producer;
@@ -409,7 +421,8 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
           && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
           && hasOnlyStructuralAttrs(consumer) && isUniqueGraphComputePredecessor(candidate, consumer)
           && isExclusivelyConsumedBy(candidate, consumer)
-          && hasCapacityFor(candidate, consumer) && hasDirectLaneConsumers(candidate, consumer)
+          && hasCapacityFor(candidate, consumer, residentWeightCapacity)
+          && hasDirectLaneConsumers(candidate, consumer)
           && succeeded(fragments = collectPublishedFragments(candidate))) {
         producer = candidate;
         break;
@@ -469,11 +482,16 @@ struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
 
 private:
   TrivialGraphMergeStats *stats;
+  size_t residentWeightCapacity;
 };
 
 struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMergePass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrivialGraphComputeMergePass)
 
+  TrivialGraphComputeMergePass() = default;
+  explicit TrivialGraphComputeMergePass(size_t residentWeightCapacity)
+      : residentWeightCapacity(residentWeightCapacity) {}
+
   StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
   StringRef getDescription() const override {
     return "Inline linear exclusive graph compute chains while preserving fan-in boundaries.";
@@ -481,10 +499,17 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
+    if (residentWeightCapacity == 0) {
+      module.emitError(
+        "TrivialGraphComputeMerge requires an explicit valid resident-weight capacity");
+      signalPassFailure();
+      return;
+    }
     TrivialGraphMergeStats stats;
     std::tie(stats.scalarBefore, stats.batchBefore) = countGraphComputes(module);
     RewritePatternSet patterns(&getContext());
-    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(&getContext(), &stats);
+    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(
+      &getContext(), &stats, residentWeightCapacity);
     if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
       signalPassFailure();
       return;
@@ -510,6 +535,9 @@ struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMerge
         signalPassFailure();
     }
   }
+
+private:
+  size_t residentWeightCapacity = 0;
 };
 
 } // namespace
@@ -519,4 +547,10 @@ std::unique_ptr<Pass> createTrivialGraphComputeMergePass() {
   return std::make_unique<spatial::TrivialGraphComputeMergePass>();
 }
 
+std::unique_ptr<Pass> createTrivialGraphComputeMergePass(
+    size_t residentWeightCapacity) {
+  return std::make_unique<spatial::TrivialGraphComputeMergePass>(
+    residentWeightCapacity);
+}
+
 } // namespace onnx_mlir