ReduceMean + resnet

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -38,6 +38,12 @@ llvm::cl::opt<bool>
                  llvm::cl::init(false),
                  llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool>
+  pimDisableMemoryCoalescing("pim-disable-memory-coalescing",
+                             llvm::cl::desc("Skip the PIM memory coalescing pass (developer diagnostic option)"),
+                             llvm::cl::init(false),
+                             llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
                                             llvm::cl::desc("Use experimental implementation for convolution"),
                                             llvm::cl::init(false),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -36,6 +36,7 @@ extern llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler;
 extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
+extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
 

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -46,7 +46,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   if (pimEmissionTarget >= EmitPimCodegen) {
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
-    pm.addPass(createPimMemoryCoalescingPass());
+    if (!pimDisableMemoryCoalescing)
+      pm.addPass(createPimMemoryCoalescingPass());
     pm.addPass(createPimVerificationPass());
     pm.addPass(createMessagePass("Pim verified"));
     pm.addPass(createEmitPimCodePass());

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -22,6 +22,7 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Gather.cpp
   Patterns/Tensor/Resize.cpp
   Patterns/Tensor/Reshape.cpp
+  Patterns/Tensor/Slice.cpp
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
   ONNXToSpatialPass.cpp

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -138,7 +138,9 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXGatherOp>();
   target.addIllegalOp<ONNXReshapeOp>();
   target.addIllegalOp<ONNXResizeOp>();
+  target.addIllegalOp<ONNXSliceOp>();
   target.addIllegalOp<ONNXLRNOp>();
+  target.addIllegalOp<ONNXReduceMeanOp>();
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -22,6 +22,7 @@ void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);
+  populateSlicePatterns(patterns, ctx);
   populateSplitPatterns(patterns, ctx);
   populateTransposePatterns(patterns, ctx);
 }

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -29,6 +29,7 @@ void populateConcatPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext
 void populateGatherPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateResizePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReshapePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateSlicePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSplitPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateTransposePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

@@ -6,6 +6,8 @@
 
 #include <algorithm>
 #include <numeric>
+#include <optional>
+#include <type_traits>
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -19,6 +21,85 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+struct ReduceMeanSemantics {
+  SmallVector<int64_t> axes;
+  int64_t keepdims = 1;
+  bool isIdentity = false;
+};
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static FailureOr<SmallVector<int64_t>> normalizeAxesChecked(ArrayRef<int64_t> axes, int64_t rank) {
+  SmallVector<int64_t> normalizedAxes;
+  normalizedAxes.reserve(axes.size());
+  for (int64_t axis : axes) {
+    auto normalizedAxis = normalizeAxisChecked(axis, rank);
+    if (failed(normalizedAxis))
+      return failure();
+    normalizedAxes.push_back(*normalizedAxis);
+  }
+  llvm::sort(normalizedAxes);
+  normalizedAxes.erase(std::unique(normalizedAxes.begin(), normalizedAxes.end()), normalizedAxes.end());
+  return normalizedAxes;
+}
+
+template <typename ReduceMeanOp, typename ReduceMeanOpAdaptor>
+static FailureOr<ReduceMeanSemantics>
+getReduceMeanSemantics(ReduceMeanOp reduceMeanOp, ReduceMeanOpAdaptor adaptor, int64_t inputRank) {
+  ReduceMeanSemantics semantics;
+  semantics.keepdims = reduceMeanOp.getKeepdims();
+
+  if constexpr (std::is_same_v<ReduceMeanOp, ONNXReduceMeanV13Op>) {
+    auto axes = onnx_mlir::normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputRank);
+    if (failed(axes))
+      return failure();
+    semantics.axes = std::move(*axes);
+    return semantics;
+  }
+  else {
+    if (isNoneValueLike(adaptor.getAxes())) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto axes = getConstantIntValues(adaptor.getAxes());
+    if (failed(axes))
+      return failure();
+
+    if (axes->empty()) {
+      if (reduceMeanOp.getNoopWithEmptyAxes() != 0) {
+        semantics.isIdentity = true;
+        return semantics;
+      }
+
+      semantics.axes.reserve(inputRank);
+      for (int64_t axis = 0; axis < inputRank; ++axis)
+        semantics.axes.push_back(axis);
+      return semantics;
+    }
+
+    auto normalizedAxes = normalizeAxesChecked(*axes, inputRank);
+    if (failed(normalizedAxes))
+      return failure();
+    semantics.axes = std::move(*normalizedAxes);
+    return semantics;
+  }
+}
+
 static SmallVector<bool> buildReducedAxesMask(ArrayRef<int64_t> axes, int64_t rank) {
   SmallVector<bool> reducedAxes(rank, false);
   for (int64_t axis : axes) {
@@ -251,11 +332,13 @@ static Value squeezeReducedAxes(Value keepdimsValue,
   return squeezeCompute.getResult(0);
 }
 
-struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
-  using OpConversionPattern::OpConversionPattern;
+template <typename ReduceMeanOp>
+struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
+  using OpConversionPattern<ReduceMeanOp>::OpConversionPattern;
+  using Adaptor = typename ReduceMeanOp::Adaptor;
 
-  LogicalResult matchAndRewrite(ONNXReduceMeanV13Op reduceMeanOp,
-                                ONNXReduceMeanV13OpAdaptor adaptor,
+  LogicalResult matchAndRewrite(ReduceMeanOp reduceMeanOp,
+                                Adaptor adaptor,
                                 ConversionPatternRewriter& rewriter) const override {
     auto inputType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
     auto resultType = dyn_cast<RankedTensorType>(reduceMeanOp.getReduced().getType());
@@ -266,10 +349,18 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
       return success();
     }
 
-    auto axes = normalizeAxesChecked(std::optional<ArrayAttr>(reduceMeanOp.getAxesAttr()), inputType.getRank());
-    if (failed(axes))
-      return failure();
-    SmallVector<bool> reducedAxes = buildReducedAxesMask(*axes, inputType.getRank());
+    auto semantics = getReduceMeanSemantics(reduceMeanOp, adaptor, inputType.getRank());
+    if (failed(semantics))
+      return rewriter.notifyMatchFailure(reduceMeanOp, "requires compile-time constant, in-range ReduceMean axes");
+    if (semantics->isIdentity) {
+      if (inputType != resultType)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "noop_with_empty_axes identity requires the result type to match the input type");
+      rewriter.replaceOp(reduceMeanOp, adaptor.getData());
+      return success();
+    }
+
+    SmallVector<bool> reducedAxes = buildReducedAxesMask(semantics->axes, inputType.getRank());
     if (reducedAxes.empty() && inputType.getRank() != 0)
       return failure();
 
@@ -289,7 +380,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
     Value reducedKeepdims =
       buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
 
-    if (reduceMeanOp.getKeepdims() != 0) {
+    if (semantics->keepdims != 0) {
       rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
       return success();
     }
@@ -303,7 +394,7 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {
 } // namespace
 
 void populateReduceMeanPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<ReduceMeanToSpatialCompute>(ctx);
+  patterns.add<ReduceMeanToSpatialCompute<ONNXReduceMeanV13Op>, ReduceMeanToSpatialCompute<ONNXReduceMeanOp>>(ctx);
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Slice.cpp

@@ -0,0 +1,189 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <algorithm>
+#include <optional>
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<SmallVector<int64_t>> getConstantIntValues(Value value) {
+  auto denseAttr = dyn_cast_or_null<DenseIntElementsAttr>(getHostConstDenseElementsAttr(value));
+  if (!denseAttr)
+    return failure();
+  return SmallVector<int64_t>(denseAttr.getValues<int64_t>().begin(), denseAttr.getValues<int64_t>().end());
+}
+
+static bool isNoneValueLike(Value value) { return isa_and_nonnull<ONNXNoneOp>(value.getDefiningOp()); }
+
+static FailureOr<Value> buildSlice(Value data,
+                                   RankedTensorType dataType,
+                                   RankedTensorType resultType,
+                                   ArrayRef<int64_t> starts,
+                                   ArrayRef<int64_t> ends,
+                                   std::optional<ArrayRef<int64_t>> axes,
+                                   std::optional<ArrayRef<int64_t>> steps,
+                                   ConversionPatternRewriter& rewriter,
+                                   Location loc) {
+  int64_t rank = dataType.getRank();
+  if (!dataType.hasStaticShape() || !resultType.hasStaticShape() || resultType.getRank() != rank)
+    return failure();
+
+  if (starts.size() != ends.size())
+    return failure();
+  if (axes && axes->size() != starts.size())
+    return failure();
+  if (steps && steps->size() != starts.size())
+    return failure();
+
+  SmallVector<int64_t> normalizedAxes;
+  if (axes) {
+    SmallVector<bool> seenAxes(rank, false);
+    normalizedAxes.reserve(axes->size());
+    for (int64_t axis : *axes) {
+      auto normalizedAxis = normalizeAxisChecked(axis, rank);
+      if (failed(normalizedAxis))
+        return failure();
+      if (seenAxes[*normalizedAxis])
+        return failure();
+      seenAxes[*normalizedAxis] = true;
+      normalizedAxes.push_back(*normalizedAxis);
+    }
+  }
+  else {
+    if (starts.size() > static_cast<size_t>(rank))
+      return failure();
+    normalizedAxes.reserve(starts.size());
+    for (size_t i = 0; i < starts.size(); ++i)
+      normalizedAxes.push_back(static_cast<int64_t>(i));
+  }
+
+  SmallVector<int64_t> normalizedSteps;
+  if (steps)
+    normalizedSteps.assign(steps->begin(), steps->end());
+  else
+    normalizedSteps.assign(starts.size(), 1);
+
+  SmallVector<int64_t> computedShape(dataType.getShape().begin(), dataType.getShape().end());
+  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, rank);
+  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, dataType.getShape());
+  SmallVector<OpFoldResult> strides = getUnitStrides(rewriter, rank);
+
+  for (auto [sliceIndex, axis] : llvm::enumerate(normalizedAxes)) {
+    int64_t step = normalizedSteps[sliceIndex];
+    if (step <= 0)
+      return failure();
+
+    int64_t dimSize = dataType.getShape()[axis];
+    int64_t start = starts[sliceIndex];
+    int64_t end = ends[sliceIndex];
+
+    start = normalizeIndex(start, dimSize);
+    end = normalizeIndex(end, dimSize);
+
+    start = std::clamp(start, int64_t {0}, dimSize);
+    end = std::clamp(end, int64_t {0}, dimSize);
+
+    int64_t extent = std::max(end - start, int64_t {0});
+    int64_t size = (extent + step - 1) / step;
+
+    offsets[axis] = rewriter.getIndexAttr(start);
+    sizes[axis] = rewriter.getIndexAttr(size);
+    strides[axis] = rewriter.getIndexAttr(step);
+    computedShape[axis] = size;
+  }
+
+  if (llvm::ArrayRef(computedShape) != resultType.getShape())
+    return failure();
+
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, data, offsets, sizes, strides).getResult();
+}
+
+struct Slice final : OpConversionPattern<ONNXSliceOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXSliceOp sliceOp,
+                                ONNXSliceOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto dataType = dyn_cast<RankedTensorType>(adaptor.getData().getType());
+    auto resultType = dyn_cast<RankedTensorType>(sliceOp.getResult().getType());
+    if (!dataType || !resultType || !dataType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+
+    auto starts = getConstantIntValues(adaptor.getStarts());
+    auto ends = getConstantIntValues(adaptor.getEnds());
+    if (failed(starts))
+      return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant starts");
+    if (failed(ends))
+      return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant ends");
+
+    std::optional<SmallVector<int64_t>> axes;
+    if (!isNoneValueLike(adaptor.getAxes())) {
+      auto parsedAxes = getConstantIntValues(adaptor.getAxes());
+      if (failed(parsedAxes))
+        return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant axes when present");
+      axes = std::move(*parsedAxes);
+    }
+
+    std::optional<SmallVector<int64_t>> steps;
+    if (!isNoneValueLike(adaptor.getSteps())) {
+      auto parsedSteps = getConstantIntValues(adaptor.getSteps());
+      if (failed(parsedSteps))
+        return rewriter.notifyMatchFailure(sliceOp, "requires compile-time constant steps when present");
+      steps = std::move(*parsedSteps);
+      if (llvm::any_of(*steps, [](int64_t step) { return step <= 0; }))
+        return rewriter.notifyMatchFailure(sliceOp, "supports only positive constant steps");
+    }
+
+    ArrayRef<int64_t> startsRef = *starts;
+    ArrayRef<int64_t> endsRef = *ends;
+    std::optional<ArrayRef<int64_t>> axesRef = axes ? std::optional<ArrayRef<int64_t>>(ArrayRef<int64_t>(*axes))
+                                                    : std::nullopt;
+    std::optional<ArrayRef<int64_t>> stepsRef = steps ? std::optional<ArrayRef<int64_t>>(ArrayRef<int64_t>(*steps))
+                                                       : std::nullopt;
+
+    Location loc = sliceOp.getLoc();
+    auto tryBuildSlice = [&](Value data) {
+      return buildSlice(data, dataType, resultType, startsRef, endsRef, axesRef, stepsRef, rewriter, loc);
+    };
+
+    if (isCompileTimeComputable(adaptor.getData())) {
+      auto sliced = tryBuildSlice(adaptor.getData());
+      if (failed(sliced))
+        return rewriter.notifyMatchFailure(sliceOp, "failed to normalize static slice parameters");
+      rewriter.replaceOp(sliceOp, *sliced);
+      return success();
+    }
+
+    auto computeOp =
+      createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, adaptor.getData(), [&](Value data) {
+        auto sliced = tryBuildSlice(data);
+        if (failed(sliced))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, *sliced);
+        return success();
+      });
+    if (failed(computeOp))
+      return rewriter.notifyMatchFailure(sliceOp, "failed to build runtime tensor.extract_slice lowering");
+
+    rewriter.replaceOp(sliceOp, computeOp->getResults());
+    return success();
+  }
+};
+
+} // namespace
+
+void populateSlicePatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Slice>(ctx); }
+
+} // namespace onnx_mlir

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

@@ -11,6 +11,7 @@
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -117,6 +118,66 @@ struct ProducerKeyInfo {
   static bool isEqual(const ProducerKey& lhs, const ProducerKey& rhs) { return lhs == rhs; }
 };
 
+struct SameClassConsumerLookupKey {
+  Operation* sourceOp = nullptr;
+  size_t resultIndex = 0;
+  ClassId classId = 0;
+
+  bool operator==(const SameClassConsumerLookupKey& other) const {
+    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
+  }
+};
+
+struct SameClassConsumerLookupKeyInfo {
+  static SameClassConsumerLookupKey getEmptyKey() {
+    return {llvm::DenseMapInfo<Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
+            std::numeric_limits<ClassId>::max()};
+  }
+
+  static SameClassConsumerLookupKey getTombstoneKey() {
+    return {llvm::DenseMapInfo<Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
+            std::numeric_limits<ClassId>::max()};
+  }
+
+  static unsigned getHashValue(const SameClassConsumerLookupKey& key) {
+    return llvm::hash_combine(llvm::DenseMapInfo<Operation*>::getHashValue(key.sourceOp), key.resultIndex, key.classId);
+  }
+
+  static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) {
+    return lhs == rhs;
+  }
+};
+
+struct WholeBatchAssemblyLookupKey {
+  Operation* sourceOp = nullptr;
+  size_t resultIndex = 0;
+  ClassId classId = 0;
+
+  bool operator==(const WholeBatchAssemblyLookupKey& other) const {
+    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
+  }
+};
+
+struct WholeBatchAssemblyLookupKeyInfo {
+  static WholeBatchAssemblyLookupKey getEmptyKey() {
+    return {llvm::DenseMapInfo<Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
+            std::numeric_limits<ClassId>::max()};
+  }
+
+  static WholeBatchAssemblyLookupKey getTombstoneKey() {
+    return {llvm::DenseMapInfo<Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
+            std::numeric_limits<ClassId>::max()};
+  }
+
+  static unsigned getHashValue(const WholeBatchAssemblyLookupKey& key) {
+    return llvm::hash_combine(llvm::DenseMapInfo<Operation*>::getHashValue(key.sourceOp), key.resultIndex, key.classId);
+  }
+
+  static bool isEqual(const WholeBatchAssemblyLookupKey& lhs, const WholeBatchAssemblyLookupKey& rhs) {
+    return lhs == rhs;
+  }
+};
+
 using ClassSlotKey = std::pair<ClassId, SlotId>;
 
 struct MaterializedClass {
@@ -270,9 +331,36 @@ FailureOr<Value> materializeProjectedExtractReplacement(MaterializerState& state
 
 class AvailableValueStore {
 public:
-  void record(ProducerKey key, ClassId classId, Value value) { exactValues[key][classId] = value; }
+  struct ExactBatchFragmentRecord {
+    ProducerKey key;
+    Value value;
+  };
 
-  void recordPackedRun(PackedScalarRunValue run) { packedScalarRuns.push_back(std::move(run)); }
+  void record(ProducerKey key, ClassId classId, Value value) {
+    exactValues[key][classId] = value;
+
+    auto batch = dyn_cast_or_null<SpatComputeBatch>(key.instance.op);
+    if (!batch || key.instance.laneCount == 0)
+      return;
+
+    WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId};
+    SmallVector<ExactBatchFragmentRecord, 16>& bucket = exactBatchFragmentsByProducerResultClass[lookupKey];
+    for (ExactBatchFragmentRecord& record : bucket) {
+      if (!(record.key == key))
+        continue;
+      record.value = value;
+      return;
+    }
+    bucket.push_back({key, value});
+  }
+
+  void recordPackedRun(PackedScalarRunValue run) {
+    size_t runIndex = packedScalarRuns.size();
+    packedScalarRuns.push_back(std::move(run));
+    const PackedScalarRunValue& storedRun = packedScalarRuns[runIndex];
+    WholeBatchAssemblyLookupKey lookupKey {storedRun.sourceOp, storedRun.resultIndex, storedRun.targetClass};
+    packedRunsByProducerResultClass[lookupKey].push_back(runIndex);
+  }
   void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); }
 
   std::optional<Value> lookupExact(ProducerKey key, ClassId classId) const;
@@ -280,7 +368,21 @@ public:
   std::optional<Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
   IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId);
 
-  SmallVectorImpl<PackedScalarRunValue>& getPackedScalarRuns() { return packedScalarRuns; }
+  ArrayRef<size_t> getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const {
+    auto it = packedRunsByProducerResultClass.find(key);
+    if (it == packedRunsByProducerResultClass.end())
+      return {};
+    return it->second;
+  }
+
+  ArrayRef<ExactBatchFragmentRecord> getExactFragmentsForWholeBatch(WholeBatchAssemblyLookupKey key) const {
+    auto it = exactBatchFragmentsByProducerResultClass.find(key);
+    if (it == exactBatchFragmentsByProducerResultClass.end())
+      return {};
+    return it->second;
+  }
+
+  PackedScalarRunValue& getPackedRun(size_t index) { return packedScalarRuns[index]; }
 
 private:
   std::optional<Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
@@ -288,6 +390,10 @@ private:
   DenseMap<ProducerKey, DenseMap<ClassId, Value>, ProducerKeyInfo> exactValues;
   SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
   SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
+  DenseMap<WholeBatchAssemblyLookupKey, SmallVector<ExactBatchFragmentRecord, 16>, WholeBatchAssemblyLookupKeyInfo>
+    exactBatchFragmentsByProducerResultClass;
+  DenseMap<WholeBatchAssemblyLookupKey, SmallVector<size_t, 16>, WholeBatchAssemblyLookupKeyInfo>
+    packedRunsByProducerResultClass;
 };
 
 struct MaterializerState {
@@ -296,7 +402,6 @@ struct MaterializerState {
   IRRewriter rewriter;
   OperationFolder constantFolder;
   int64_t& nextChannelId;
-
   SmallVector<MaterializedClass, 8> classes;
   DenseMap<CpuId, ClassId> cpuToClass;
   DenseMap<CpuId, SmallVector<ComputeInstance, 32>> logicalInstancesByCpu;
@@ -305,7 +410,8 @@ struct MaterializerState {
   DenseSet<ClassSlotKey> materializedLogicalSlots;
 
   DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
-  DenseMap<ProducerKey, DenseSet<ClassId>, ProducerKeyInfo> sameClassConsumers;
+  DenseMap<SameClassConsumerLookupKey, SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
+    sameClassConsumerIndex;
   DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo> projectedInputMatches;
   DenseSet<ProjectedBatchInputKey, ProjectedBatchInputKeyInfo> nonProjectedInputs;
   DenseMap<Value, bool> liveExternalUseCache;
@@ -317,7 +423,9 @@ struct MaterializerState {
   DenseMap<Value, Value> hostReplacements;
   DenseSet<Operation*> oldComputeOps;
 
-  MaterializerState(func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)
+  MaterializerState(func::FuncOp func,
+                    const MergeScheduleResult& schedule,
+                    int64_t& nextChannelId)
   : func(func),
     schedule(schedule),
     rewriter(func.getContext()),
@@ -428,6 +536,14 @@ std::optional<ProducerKey> getContiguousProducerRangeForKeys(ArrayRef<ProducerKe
   return getBatchLaneProducerKey(batch, laneStart, laneCount, first.resultIndex);
 }
 
+WholeBatchAssemblyLookupKey makeWholeBatchAssemblyLookupKey(Operation* sourceOp, size_t resultIndex, ClassId classId) {
+  return {sourceOp, resultIndex, classId};
+}
+
+WholeBatchAssemblyLookupKey makeWholeBatchAssemblyLookupKey(ProducerKey key, ClassId classId) {
+  return makeWholeBatchAssemblyLookupKey(key.instance.op, key.resultIndex, classId);
+}
+
 FailureOr<RankedTensorType> getPackedBatchTensorType(Type laneType, size_t laneCount) {
   auto tensorType = dyn_cast<RankedTensorType>(laneType);
   if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
@@ -1172,14 +1288,12 @@ FailureOr<Value> materializePackedScalarRunValue(MaterializerState& state,
 
 std::optional<Value> AvailableValueStore::lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId) {
   for (PackedScalarRunValue& run : packedScalarRuns) {
-    if (run.targetClass != classId)
-      continue;
-    if (run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
+    if (run.targetClass != classId || run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
       continue;
 
     for (auto [slotIndex, slot] : llvm::enumerate(run.slots)) {
-      std::optional<ProducerKey> slotKey = getContiguousProducerRangeForKeys(slot.keys);
-      if (!slotKey || !containsProducerKey(*slotKey, key))
+      std::optional<ProducerKey> contiguousKey = getContiguousProducerRangeForKeys(slot.keys);
+      if (!contiguousKey || !containsProducerKey(*contiguousKey, key))
         continue;
 
       FailureOr<RankedTensorType> slotPackedType = getPackedBatchTensorType(run.fragmentType, slot.keys.size());
@@ -1197,12 +1311,13 @@ std::optional<Value> AvailableValueStore::lookupPackedRun(MaterializerState& sta
       Value slotPacked =
         getPackedSliceForRunIndex(state, materializedClass.op, *packed, *slotPackedType, slotIndex, (*packed).getLoc());
 
-      if (*slotKey == key) {
+      if (*contiguousKey == key) {
         record(key, classId, slotPacked);
         return slotPacked;
       }
 
-      std::optional<Value> sliced = extractPackedProducerSlice(state, materializedClass, *slotKey, slotPacked, key);
+      std::optional<Value> sliced =
+        extractPackedProducerSlice(state, materializedClass, *contiguousKey, slotPacked, key);
       if (!sliced)
         return std::nullopt;
 
@@ -1216,57 +1331,45 @@ std::optional<Value> AvailableValueStore::lookupPackedRun(MaterializerState& sta
 
 IndexedBatchRunValue* AvailableValueStore::lookupIndexedBatchRun(ProducerKey key, ClassId classId) {
   for (IndexedBatchRunValue& run : indexedBatchRuns) {
-    if (run.targetClass != classId)
+    if (run.targetClass != classId || run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
       continue;
-    if (run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
-      continue;
-
-    for (const PackedScalarRunSlot& slot : run.slots)
-      if (llvm::is_contained(slot.keys, key))
-        return &run;
+    for (const PackedScalarRunSlot& slot : run.slots) {
+      if (!llvm::is_contained(slot.keys, key))
+        continue;
+      return &run;
+    }
   }
-
   return nullptr;
 }
 
 std::optional<Value> AvailableValueStore::lookup(MaterializerState& state, ProducerKey key, ClassId classId) {
-  if (std::optional<Value> exact = lookupExact(key, classId))
+
+  if (std::optional<Value> exact = lookupExact(key, classId)) {
     return exact;
+  }
 
   if (std::optional<Value> packedRunValue = lookupPackedRun(state, key, classId))
     return packedRunValue;
 
   MaterializedClass& materializedClass = state.classes[classId];
 
-  ProducerKey containingKey;
-  Value containingValue;
-  bool foundContainingValue = false;
-
-  for (auto& entry : exactValues) {
-    ProducerKey candidateKey = entry.first;
-    if (!containsProducerKey(candidateKey, key))
+  for (const auto& [candidateKey, classValues] : exactValues) {
+    if (!sameProducerResult(candidateKey, key) || !containsProducerKey(candidateKey, key))
       continue;
 
-    auto valueIt = entry.second.find(classId);
-    if (valueIt == entry.second.end())
+    auto valueIt = classValues.find(classId);
+    if (valueIt == classValues.end())
       continue;
 
-    containingKey = candidateKey;
-    containingValue = valueIt->second;
-    foundContainingValue = true;
-    break;
+    std::optional<Value> slice =
+      extractPackedProducerSlice(state, materializedClass, candidateKey, valueIt->second, key);
+    if (!slice)
+      return std::nullopt;
+
+    record(key, classId, *slice);
+    return *slice;
   }
-
-  if (!foundContainingValue)
-    return std::nullopt;
-
-  std::optional<Value> slice =
-    extractPackedProducerSlice(state, materializedClass, containingKey, containingValue, key);
-  if (!slice)
-    return std::nullopt;
-
-  record(key, classId, *slice);
-  return *slice;
+  return std::nullopt;
 }
 
 Value createIndexTensorConstant(MaterializerState& state, Operation* anchor, ArrayRef<int64_t> values) {
@@ -1389,13 +1492,13 @@ Value createIndexedIndexValue(MaterializerState& state,
                               bool allowExhaustiveTiledSearch) {
   assert(!values.empty() && "expected at least one indexed value");
 
-  if (allEqual(values))
+  if (allEqual(values)) {
     return getOrCreateIndexConstant(state.constantFolder, anchor, values.front());
+  }
 
   if (std::optional<IndexedIndexPattern> pattern =
         getIndexedIndexPattern(values, preferredPeriod, allowExhaustiveTiledSearch))
     return createAffineIndexValue(state, *pattern, index, loc);
-
   Value table = createIndexTensorConstant(state, anchor, values);
   return tensor::ExtractOp::create(state.rewriter, loc, table, ValueRange {index}).getResult();
 }
@@ -1578,7 +1681,10 @@ LogicalResult collectProducerDestinations(MaterializerState& state) {
 
           ClassId sourceClass = state.cpuToClass.lookup(producerCpuIt->second);
           if (sourceClass == targetClass) {
-            state.sameClassConsumers[producerKey].insert(targetClass);
+            SameClassConsumerLookupKey lookupKey{producerKey.instance.op, producerKey.resultIndex, targetClass};
+            SmallVector<ProducerKey, 4>& bucket = state.sameClassConsumerIndex[lookupKey];
+            if (!llvm::is_contained(bucket, producerKey))
+              bucket.push_back(producerKey);
             continue;
           }
 
@@ -2899,11 +3005,6 @@ LogicalResult emitOutputFanout(MaterializerState& state,
   return success();
 }
 
-struct WholeBatchAssemblyRange {
-  uint32_t laneStart = 0;
-  uint32_t laneCount = 0;
-};
-
 struct DirectWholeBatchFragment {
   ProducerKey key;
   Value fragment;
@@ -2933,31 +3034,60 @@ struct WholeBatchFragmentGroup {
 struct WholeBatchAssemblyPlan {
   RankedTensorType resultType;
   int64_t rowsPerLane = 0;
+  uint32_t batchLaneCount = 0;
+  uint32_t coveredLaneCount = 0;
 
-  SmallVector<WholeBatchAssemblyRange, 16> coveredRanges;
+  SmallVector<uint8_t, 64> coveredLanes;
   SmallVector<PackedScalarRunValue*, 8> packedRuns;
   SmallVector<DirectWholeBatchFragment, 16> directFragments;
 };
 
-bool wholeBatchRangeOverlaps(ArrayRef<WholeBatchAssemblyRange> ranges, uint32_t laneStart, uint32_t laneCount) {
-  uint32_t laneEnd = laneStart + laneCount;
-  for (WholeBatchAssemblyRange range : ranges) {
-    uint32_t rangeEnd = range.laneStart + range.laneCount;
-    if (laneStart < rangeEnd && range.laneStart < laneEnd)
-      return true;
-  }
-  return false;
+bool wholeBatchLaneCovered(const WholeBatchAssemblyPlan& plan, uint32_t lane) {
+  return lane < plan.coveredLanes.size() && plan.coveredLanes[lane] != 0;
 }
 
-bool wholeBatchLaneCovered(ArrayRef<WholeBatchAssemblyRange> ranges, uint32_t lane) {
-  for (WholeBatchAssemblyRange range : ranges)
-    if (range.laneStart <= lane && lane < range.laneStart + range.laneCount)
+bool wholeBatchRangeOverlaps(const WholeBatchAssemblyPlan& plan, uint32_t laneStart, uint32_t laneCount) {
+  if (laneCount == 0)
+    return false;
+  if (laneStart >= plan.coveredLanes.size())
+    return false;
+
+  uint32_t laneEnd = std::min<uint32_t>(laneStart + laneCount, plan.coveredLanes.size());
+  for (uint32_t lane = laneStart; lane < laneEnd; ++lane)
+    if (plan.coveredLanes[lane] != 0)
       return true;
   return false;
 }
 
 void recordWholeBatchCoverage(WholeBatchAssemblyPlan& plan, uint32_t laneStart, uint32_t laneCount) {
-  plan.coveredRanges.push_back({laneStart, laneCount});
+  assert(laneCount != 0 && "cannot cover an empty whole-batch range");
+  assert(laneStart + laneCount <= plan.coveredLanes.size() && "whole-batch coverage out of bounds");
+
+  for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane) {
+    if (plan.coveredLanes[lane] != 0)
+      continue;
+    plan.coveredLanes[lane] = 1;
+    ++plan.coveredLaneCount;
+  }
+}
+
+bool localLaneRangeOverlaps(ArrayRef<uint8_t> covered, uint32_t laneStart, uint32_t laneCount) {
+  if (laneCount == 0)
+    return false;
+  if (laneStart >= covered.size())
+    return false;
+
+  uint32_t laneEnd = std::min<uint32_t>(laneStart + laneCount, covered.size());
+  for (uint32_t lane = laneStart; lane < laneEnd; ++lane)
+    if (covered[lane] != 0)
+      return true;
+  return false;
+}
+
+void markLocalLaneRangeCovered(MutableArrayRef<uint8_t> covered, uint32_t laneStart, uint32_t laneCount) {
+  assert(laneStart + laneCount <= covered.size() && "local coverage out of bounds");
+  for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane)
+    covered[lane] = 1;
 }
 
 LogicalResult
@@ -3118,13 +3248,14 @@ LogicalResult collectPackedRunsForWholeBatchInput(MaterializerState& state,
                                                   MaterializedClass& targetClass,
                                                   ProducerKey key,
                                                   WholeBatchAssemblyPlan& plan) {
-  for (PackedScalarRunValue& run : state.availableValues.getPackedScalarRuns()) {
-    if (run.targetClass != targetClass.id)
-      continue;
-    if (run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
-      continue;
+  WholeBatchAssemblyLookupKey lookupKey = makeWholeBatchAssemblyLookupKey(key, targetClass.id);
+  ArrayRef<size_t> runIndices = state.availableValues.getPackedRunIndicesForWholeBatch(lookupKey);
 
-    SmallVector<WholeBatchAssemblyRange, 16> runRanges;
+  for (size_t runIndex : runIndices) {
+    PackedScalarRunValue& run = state.availableValues.getPackedRun(runIndex);
+
+    SmallVector<ProducerKey, 16> runKeys;
+    SmallVector<uint8_t, 64> runCoveredLanes(plan.batchLaneCount, 0);
 
     for (const PackedScalarRunSlot& slot : run.slots) {
       for (ProducerKey fragmentKey : slot.keys) {
@@ -3134,23 +3265,24 @@ LogicalResult collectPackedRunsForWholeBatchInput(MaterializerState& state,
         if (fragmentKey.instance.laneCount == 0)
           return failure();
 
-        if (wholeBatchRangeOverlaps(plan.coveredRanges, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
+        if (wholeBatchRangeOverlaps(plan, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
           return failure();
 
-        if (wholeBatchRangeOverlaps(runRanges, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
+        if (localLaneRangeOverlaps(runCoveredLanes, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
           return failure();
 
-        runRanges.push_back({fragmentKey.instance.laneStart, fragmentKey.instance.laneCount});
+        markLocalLaneRangeCovered(runCoveredLanes, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount);
+        runKeys.push_back(fragmentKey);
       }
     }
 
-    if (runRanges.empty())
+    if (runKeys.empty())
       continue;
 
     plan.packedRuns.push_back(&run);
 
-    for (WholeBatchAssemblyRange range : runRanges)
-      recordWholeBatchCoverage(plan, range.laneStart, range.laneCount);
+    for (ProducerKey runKey : runKeys)
+      recordWholeBatchCoverage(plan, runKey.instance.laneStart, runKey.instance.laneCount);
   }
 
   return success();
@@ -3161,44 +3293,77 @@ LogicalResult collectDirectFragmentsForWholeBatchInput(MaterializerState& state,
                                                        SpatComputeBatch batch,
                                                        ProducerKey key,
                                                        WholeBatchAssemblyPlan& plan) {
+  struct CandidateFragment {
+    ProducerKey key;
+    Value value;
+  };
+
   uint32_t batchLaneCount = static_cast<uint32_t>(batch.getLaneCount());
-  uint32_t lane = 0;
+  if (plan.coveredLaneCount == plan.batchLaneCount) {
+    return success();
+  }
 
-  while (lane < batchLaneCount) {
-    if (wholeBatchLaneCovered(plan.coveredRanges, lane)) {
-      ++lane;
+  WholeBatchAssemblyLookupKey lookupKey = makeWholeBatchAssemblyLookupKey(key, targetClass.id);
+  ArrayRef<AvailableValueStore::ExactBatchFragmentRecord> indexedFragments =
+    state.availableValues.getExactFragmentsForWholeBatch(lookupKey);
+
+  SmallVector<CandidateFragment, 16> candidates;
+  candidates.reserve(indexedFragments.size());
+  for (const AvailableValueStore::ExactBatchFragmentRecord& record : indexedFragments) {
+    ProducerKey candidateKey = record.key;
+    if (candidateKey.instance.op != batch.getOperation() || candidateKey.resultIndex != key.resultIndex
+        || candidateKey.instance.laneCount == 0)
       continue;
+    if (wholeBatchRangeOverlaps(plan, candidateKey.instance.laneStart, candidateKey.instance.laneCount))
+      continue;
+
+    auto fragmentType = dyn_cast<RankedTensorType>(record.value.getType());
+    if (!fragmentType)
+      continue;
+
+    int64_t expectedRows = plan.rowsPerLane * static_cast<int64_t>(candidateKey.instance.laneCount);
+    if (failed(validateWholeBatchFragmentType(plan.resultType, fragmentType, expectedRows)))
+      continue;
+
+    candidates.push_back({candidateKey, record.value});
+  }
+
+  llvm::sort(candidates, [](const CandidateFragment& lhs, const CandidateFragment& rhs) {
+    if (lhs.key.instance.laneStart != rhs.key.instance.laneStart)
+      return lhs.key.instance.laneStart < rhs.key.instance.laneStart;
+    return lhs.key.instance.laneCount > rhs.key.instance.laneCount;
+  });
+
+  size_t candidateCursor = 0;
+  uint32_t lane = 0;
+  while (lane < batchLaneCount) {
+    while (lane < batchLaneCount && wholeBatchLaneCovered(plan, lane)) {
+      ++lane;
     }
 
-    bool foundFragment = false;
-
-    for (uint32_t laneCount = batchLaneCount - lane; laneCount != 0; --laneCount) {
-      if (wholeBatchRangeOverlaps(plan.coveredRanges, lane, laneCount))
-        continue;
-
-      ProducerKey candidate = getBatchLaneProducerKey(batch, lane, laneCount, key.resultIndex);
-      std::optional<Value> fragment = state.availableValues.lookupExact(candidate, targetClass.id);
-      if (!fragment)
-        continue;
-
-      auto fragmentType = dyn_cast<RankedTensorType>(fragment->getType());
-      if (!fragmentType)
-        return failure();
-
-      int64_t expectedRows = plan.rowsPerLane * static_cast<int64_t>(laneCount);
-      if (failed(validateWholeBatchFragmentType(plan.resultType, fragmentType, expectedRows)))
-        return failure();
-
-      plan.directFragments.push_back({candidate, *fragment});
-      recordWholeBatchCoverage(plan, lane, laneCount);
-
-      lane += laneCount;
-      foundFragment = true;
+    if (lane >= batchLaneCount)
       break;
+
+    while (candidateCursor < candidates.size() && candidates[candidateCursor].key.instance.laneStart < lane)
+      ++candidateCursor;
+
+    size_t candidateIndex = candidateCursor;
+    const CandidateFragment* best = nullptr;
+    while (candidateIndex < candidates.size() && candidates[candidateIndex].key.instance.laneStart == lane) {
+      const CandidateFragment& candidate = candidates[candidateIndex];
+      if (!wholeBatchRangeOverlaps(plan, lane, candidate.key.instance.laneCount)) {
+        best = &candidate;
+        break;
+      }
+      ++candidateIndex;
     }
 
-    if (!foundFragment)
+    if (!best)
       return failure();
+
+    plan.directFragments.push_back({best->key, best->value});
+    recordWholeBatchCoverage(plan, lane, best->key.instance.laneCount);
+    lane += best->key.instance.laneCount;
   }
 
   return success();
@@ -3291,11 +3456,11 @@ LogicalResult collectWholeBatchFragmentGroups(MaterializerState& state,
     }
 
     for (auto [slotIndex, slot] : llvm::enumerate(run->slots)) {
-      std::optional<ProducerKey> slotKey = getContiguousProducerRangeForKeys(slot.keys);
-      if (!slotKey)
+      std::optional<ProducerKey> contiguousKey = getContiguousProducerRangeForKeys(slot.keys);
+      if (!contiguousKey)
         return failure();
       groupIt->slotIndices.push_back(slotIndex);
-      groupIt->outputOffsets.push_back(static_cast<int64_t>(slotKey->instance.laneStart) * plan.rowsPerLane);
+      groupIt->outputOffsets.push_back(static_cast<int64_t>(contiguousKey->instance.laneStart) * plan.rowsPerLane);
     }
   }
 
@@ -3409,10 +3574,15 @@ FailureOr<WholeBatchAssemblyPlan> buildWholeBatchAssemblyPlan(MaterializerState&
   WholeBatchAssemblyPlan plan;
   plan.resultType = resultTensorType;
   plan.rowsPerLane = resultTensorType.getDimSize(0) / static_cast<int64_t>(batchLaneCount);
+  plan.batchLaneCount = batchLaneCount;
+  plan.coveredLanes.assign(batchLaneCount, 0);
 
   if (failed(collectPackedRunsForWholeBatchInput(state, targetClass, key, plan)))
     return failure();
 
+  if (plan.coveredLaneCount == plan.batchLaneCount)
+    return plan;
+
   if (failed(collectDirectFragmentsForWholeBatchInput(state, targetClass, batch, key, plan)))
     return failure();
 
@@ -4181,7 +4351,6 @@ FailureOr<SmallVector<Value, 4>> materializeBatchOutputGroupLoop(MaterializerSta
   auto sourceBatch = cast<SpatComputeBatch>(sourceOp);
   SmallVector<Type, 4>& fragmentTypes = getBatchOutputFragmentTypesCached(state, sourceBatch);
   SmallVector<Value, 4> initValues;
-
   for (size_t resultIndex : group.resultIndices) {
     if (resultIndex >= fragmentTypes.size() || !fragmentTypes[resultIndex])
       return sourceBatch.emitOpError("failed to recover per-lane output type for packed batch run");
@@ -4197,7 +4366,6 @@ FailureOr<SmallVector<Value, 4>> materializeBatchOutputGroupLoop(MaterializerSta
 
   SmallVector<int64_t, 8> logicalLanes;
   logicalLanes.reserve(run.size());
-
   for (const MaterializationRunSlot& slot : run) {
     if (slot.peers.size() != 1)
       return sourceOp->emitError("scalar batch output loop expects exactly one peer per materialization slot");
@@ -4215,34 +4383,34 @@ FailureOr<SmallVector<Value, 4>> materializeBatchOutputGroupLoop(MaterializerSta
 
   state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
   auto loop = buildNormalizedScfFor(
-    state.rewriter,
-    loc,
-    lowerBound,
-    upperBound,
-    step,
-    ValueRange(initValues),
-    [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
-      Value sourceLane = createIndexedIndexValue(state, targetClass.op, logicalLanes, loopIndex, loc);
+      state.rewriter,
+      loc,
+      lowerBound,
+      upperBound,
+      step,
+      ValueRange(initValues),
+      [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+        Value sourceLane = createIndexedIndexValue(state, targetClass.op, logicalLanes, loopIndex, loc);
 
-      FailureOr<SmallVector<Value, 4>> produced =
-        cloneBatchBodyForLane(state,
-                              targetClass,
-                              run.front().peers.front(),
-                              sourceLane,
-                              group.resultIndices,
-                              CloneIndexingContext {.runSlotIndex = loopIndex, .projectionSlotIndex = loopIndex});
-      if (failed(produced))
-        return failure();
+        FailureOr<SmallVector<Value, 4>> produced =
+          cloneBatchBodyForLane(state,
+                                targetClass,
+                                run.front().peers.front(),
+                                sourceLane,
+                                group.resultIndices,
+                                CloneIndexingContext {.runSlotIndex = loopIndex, .projectionSlotIndex = loopIndex});
+        if (failed(produced))
+          return failure();
 
-      yielded.reserve(produced->size());
-      for (auto [outputIndex, output] : llvm::enumerate(*produced)) {
-        auto fragmentType = cast<RankedTensorType>(output.getType());
-        Value acc = iterArgs[outputIndex];
-        Value firstOffset = scaleIndexByDim0Size(state, targetClass.op, loopIndex, fragmentType.getDimSize(0), loc);
-        yielded.push_back(createDim0InsertSlice(state, loc, output, acc, firstOffset));
-      }
-      return success();
-    });
+        yielded.reserve(produced->size());
+        for (auto [outputIndex, output] : llvm::enumerate(*produced)) {
+          auto fragmentType = cast<RankedTensorType>(output.getType());
+          Value acc = iterArgs[outputIndex];
+          Value firstOffset = scaleIndexByDim0Size(state, targetClass.op, loopIndex, fragmentType.getDimSize(0), loc);
+          yielded.push_back(createDim0InsertSlice(state, loc, output, acc, firstOffset));
+        }
+        return success();
+      });
   if (failed(loop))
     return failure();
 
@@ -4466,14 +4634,14 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state,
 }
 
 bool hasSameClassConsumer(MaterializerState& state, ProducerKey producerKey, ClassId classId) {
-  for (const auto& [key, consumers] : state.sameClassConsumers) {
-    if (!consumers.contains(classId))
-      continue;
-    if (!sameProducerResult(key, producerKey))
-      continue;
-    if (containsProducerKey(key, producerKey) || containsProducerKey(producerKey, key))
+  SameClassConsumerLookupKey lookupKey{producerKey.instance.op, producerKey.resultIndex, classId};
+  auto it = state.sameClassConsumerIndex.find(lookupKey);
+  if (it == state.sameClassConsumerIndex.end())
+    return false;
+
+  for (ProducerKey existing : it->second)
+    if (containsProducerKey(existing, producerKey) || containsProducerKey(producerKey, existing))
       return true;
-  }
   return false;
 }
 
@@ -4488,6 +4656,7 @@ bool canCompactBatchClassRun(MaterializerState& state,
   ArrayRef<Value> outputs = getComputeInstanceOutputValuesCached(state, run.front().peers.front());
 
   for (auto [resultIndex, ignored] : llvm::enumerate(outputs)) {
+    (void) ignored;
     for (const MaterializationRunSlot& slot : run) {
       if (slot.peers.empty())
         return false;
@@ -4533,7 +4702,8 @@ Value createBatchClassRunSourceLane(MaterializerState& state,
   SmallVector<int64_t, 16> sourceLanes;
   sourceLanes.reserve(run.size() * targetClass.cpus.size());
 
-  for (const MaterializationRunSlot& slot : run) {
+  for (auto [runSlotIndex, slot] : llvm::enumerate(run)) {
+    (void) runSlotIndex;
     assert(slot.peers.size() == targetClass.cpus.size() && "expected one peer per materialized batch lane");
     for (const ComputeInstance& peer : slot.peers)
       sourceLanes.push_back(peer.laneStart);
@@ -4577,7 +4747,6 @@ LogicalResult buildBatchRunSendPlans(MaterializerState& state,
       plan.messages.targetCoreIds.reserve(messageCount);
 
       for (size_t slotIndex = 0; slotIndex < run.size(); ++slotIndex) {
-        (void) slotIndex;
         for (auto [lane, sourceCpu] : llvm::enumerate(sourceClass.cpus)) {
           auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "batch run source core id");
           if (failed(checkedSourceCpu))
@@ -4590,6 +4759,7 @@ LogicalResult buildBatchRunSendPlans(MaterializerState& state,
             return failure();
           plan.messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu);
         }
+        (void) slotIndex;
       }
 
       plans.push_back(std::move(plan));
@@ -4773,7 +4943,8 @@ LogicalResult materializeBatchClassRun(MaterializerState& state,
   return success();
 }
 
-LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeInstance& instance) {
+LogicalResult materializeInstanceSlot(MaterializerState& state,
+                                      const ComputeInstance& instance) {
   auto cpuIt = state.schedule.computeToCpuMap.find(instance);
   if (cpuIt == state.schedule.computeToCpuMap.end())
     return instance.op->emitError("schedule materialization expected a CPU assignment for every compute instance");
@@ -4794,8 +4965,7 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns
     return success();
 
   if (isa<SpatComputeBatch>(instance.op)) {
-    FailureOr<MaterializationRun> run =
-      collectBatchMaterializationRun(state, targetClass, startLogicalSlot, instance.op);
+    FailureOr<MaterializationRun> run = collectBatchMaterializationRun(state, targetClass, startLogicalSlot, instance.op);
 
     if (succeeded(run)) {
       if (!targetClass.isBatch)
@@ -4924,6 +5094,7 @@ MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& sch
     return failure();
 
   LogicalResult _ = runRegionDCE(state.rewriter, state.func.getBody());
+  (void) _;
 
   return success();
 }

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

@@ -1,6 +1,5 @@
 #include "mlir/Analysis/TopologicalSortUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/PatternMatch.h"
@@ -14,20 +13,14 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_os_ostream.h"
-#include "llvm/Support/raw_ostream.h"
 
 #include <algorithm>
-#include <chrono>
 #include <cstddef>
 #include <cstdint>
-#include <cstdlib>
 #include <fstream>
 #include <memory>
 #include <optional>
-#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -51,83 +44,6 @@ using SpatCompute = spatial::SpatCompute;
 using SpatComputeBatch = spatial::SpatComputeBatch;
 using spatial::getProducerValueRef;
 
-bool isMergeProfilingEnabled() { return std::getenv("RAPTOR_PROFILE_MERGE") != nullptr; }
-
-class ScopedMergePhaseTimer {
-public:
-  explicit ScopedMergePhaseTimer(StringRef phaseName)
-  : enabled(isMergeProfilingEnabled()), phase(phaseName.str()) {
-    if (enabled)
-      start = std::chrono::steady_clock::now();
-  }
-
-  ~ScopedMergePhaseTimer() {
-    if (!enabled)
-      return;
-    auto elapsed = std::chrono::steady_clock::now() - start;
-    double millis = std::chrono::duration<double, std::milli>(elapsed).count();
-    llvm::errs() << "[merge-profile] " << phase << ": " << llvm::formatv("{0:F3}", millis) << " ms\n";
-  }
-
-private:
-  bool enabled = false;
-  std::string phase;
-  std::chrono::steady_clock::time_point start;
-};
-
-struct MergeIrCounts {
-  uint64_t topLevelComputeCount = 0;
-  uint64_t topLevelComputeBatchCount = 0;
-  uint64_t scalarChannelSendCount = 0;
-  uint64_t scalarChannelReceiveCount = 0;
-  uint64_t wvmmCount = 0;
-  uint64_t vaddCount = 0;
-  uint64_t scfForCount = 0;
-};
-
-MergeIrCounts collectMergeIrCounts(func::FuncOp funcOp) {
-  MergeIrCounts counts;
-
-  auto countComputeBodyOps = [&](Operation* op) {
-    op->walk([&](Operation* nestedOp) {
-      if (isa<spatial::SpatChannelSendOp>(nestedOp))
-        ++counts.scalarChannelSendCount;
-      else if (isa<spatial::SpatChannelReceiveOp>(nestedOp))
-        ++counts.scalarChannelReceiveCount;
-      else if (isa<spatial::SpatVMMOp>(nestedOp))
-        ++counts.wvmmCount;
-      else if (isa<spatial::SpatVAddOp>(nestedOp))
-        ++counts.vaddCount;
-      else if (isa<scf::ForOp>(nestedOp))
-        ++counts.scfForCount;
-    });
-  };
-
-  for (auto compute : funcOp.getOps<SpatCompute>()) {
-    ++counts.topLevelComputeCount;
-    countComputeBodyOps(compute.getOperation());
-  }
-
-  for (auto batch : funcOp.getOps<SpatComputeBatch>()) {
-    ++counts.topLevelComputeBatchCount;
-    countComputeBodyOps(batch.getOperation());
-  }
-
-  return counts;
-}
-
-void emitMergeIrCounts(StringRef phaseName, func::FuncOp funcOp) {
-  if (!isMergeProfilingEnabled())
-    return;
-
-  MergeIrCounts counts = collectMergeIrCounts(funcOp);
-  llvm::errs() << "[merge-profile] " << phaseName << " counts:"
-               << " compute=" << counts.topLevelComputeCount << " compute_batch=" << counts.topLevelComputeBatchCount
-               << " scalar_send=" << counts.scalarChannelSendCount
-               << " scalar_recv=" << counts.scalarChannelReceiveCount << " wvmm=" << counts.wvmmCount
-               << " vadd=" << counts.vaddCount << " scf_for=" << counts.scfForCount << "\n";
-}
-
 static std::optional<int32_t> getComputeCoreId(SpatCompute compute) {
   if (auto coreIdAttr = compute->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName)) {
     auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), compute, "merge compute core id");
@@ -138,16 +54,6 @@ static std::optional<int32_t> getComputeCoreId(SpatCompute compute) {
   return std::nullopt;
 }
 
-struct ComputeMotifInfo {
-  uint64_t instructionCount = 0;
-  uint64_t weightedVmmCount = 0;
-};
-
-void appendUnique(SmallVector<size_t>& values, size_t value) {
-  if (!llvm::is_contained(values, value))
-    values.push_back(value);
-}
-
 bool isTrivialSerialMergeCandidate(SpatCompute compute) {
   if (!compute->hasOneUse())
     return false;
@@ -266,212 +172,6 @@ void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
   }
 }
 
-void emitMotifProfile(func::FuncOp funcOp) {
-  if (!std::getenv("DCP_MOTIF_PROFILE"))
-    return;
-
-  SmallVector<SpatCompute> computes(funcOp.getOps<SpatCompute>());
-  DenseMap<SpatCompute, size_t> computeToIndex;
-  computeToIndex.reserve(computes.size());
-  for (auto [index, compute] : llvm::enumerate(computes))
-    computeToIndex[compute] = index;
-
-  SmallVector<ComputeMotifInfo> computeInfos(computes.size());
-  SmallVector<SmallVector<size_t>> parents(computes.size());
-  SmallVector<SmallVector<size_t>> children(computes.size());
-
-  uint64_t weightedVmmNodeCount = 0;
-  uint64_t weightedVmmOpCount = 0;
-  uint64_t edgeCount = 0;
-
-  for (auto [index, compute] : llvm::enumerate(computes)) {
-    ComputeMotifInfo& info = computeInfos[index];
-    info.instructionCount = spatial::countComputeBodyInstructions(compute.getBody());
-    compute.getBody().walk([&](spatial::SpatVMMOp) { info.weightedVmmCount++; });
-    if (info.weightedVmmCount > 0) {
-      weightedVmmNodeCount++;
-      weightedVmmOpCount += info.weightedVmmCount;
-    }
-
-    for (Value input : compute.getInputs()) {
-      auto parent = dyn_cast<SpatCompute>(input.getDefiningOp());
-      if (!parent || parent == compute)
-        continue;
-      auto parentIt = computeToIndex.find(parent);
-      if (parentIt == computeToIndex.end())
-        continue;
-
-      size_t parentIndex = parentIt->second;
-      size_t oldParentCount = parents[index].size();
-      appendUnique(parents[index], parentIndex);
-      if (parents[index].size() != oldParentCount) {
-        appendUnique(children[parentIndex], index);
-        edgeCount++;
-      }
-    }
-  }
-
-  uint64_t maxFanIn = 0;
-  uint64_t maxFanOut = 0;
-  uint64_t fanIn16 = 0;
-  uint64_t fanIn64 = 0;
-  uint64_t fanIn256 = 0;
-  uint64_t fanOut16 = 0;
-  uint64_t fanOut64 = 0;
-  uint64_t fanOut256 = 0;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    uint64_t fanIn = parents[index].size();
-    uint64_t fanOut = children[index].size();
-    maxFanIn = std::max(maxFanIn, fanIn);
-    maxFanOut = std::max(maxFanOut, fanOut);
-    fanIn16 += fanIn >= 16;
-    fanIn64 += fanIn >= 64;
-    fanIn256 += fanIn >= 256;
-    fanOut16 += fanOut >= 16;
-    fanOut64 += fanOut >= 64;
-    fanOut256 += fanOut >= 256;
-  }
-
-  uint64_t serialChainCount = 0;
-  uint64_t serialChainNodeCount = 0;
-  uint64_t maxSerialChain = 0;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    if (parents[index].size() == 1 && children[parents[index][0]].size() == 1)
-      continue;
-
-    uint64_t chainLength = 1;
-    size_t current = index;
-    while (children[current].size() == 1) {
-      size_t child = children[current][0];
-      if (parents[child].size() != 1)
-        break;
-      chainLength++;
-      current = child;
-    }
-
-    if (chainLength >= 2) {
-      serialChainCount++;
-      serialChainNodeCount += chainLength;
-      maxSerialChain = std::max(maxSerialChain, chainLength);
-    }
-  }
-
-  SmallVector<size_t> incomingEdgeCount;
-  incomingEdgeCount.reserve(parents.size());
-  for (ArrayRef<size_t> parentList : parents)
-    incomingEdgeCount.push_back(parentList.size());
-
-  SmallVector<uint64_t> level(computes.size(), 0);
-  SmallVector<size_t> readyNodes;
-  readyNodes.reserve(computes.size());
-  for (size_t index = 0; index < computes.size(); ++index)
-    if (incomingEdgeCount[index] == 0)
-      readyNodes.push_back(index);
-
-  size_t readyIndex = 0;
-  while (readyIndex != readyNodes.size()) {
-    size_t current = readyNodes[readyIndex++];
-    for (size_t child : children[current]) {
-      level[child] = std::max(level[child], level[current] + 1);
-      assert(incomingEdgeCount[child] > 0 && "incoming edge count underflow");
-      incomingEdgeCount[child]--;
-      if (incomingEdgeCount[child] == 0)
-        readyNodes.push_back(child);
-    }
-  }
-
-  SmallVector<uint64_t> weightedVmmNodesByLevel;
-  for (size_t index = 0; index < computes.size(); ++index) {
-    if (computeInfos[index].weightedVmmCount == 0)
-      continue;
-    if (weightedVmmNodesByLevel.size() <= level[index])
-      weightedVmmNodesByLevel.resize(level[index] + 1, 0);
-    weightedVmmNodesByLevel[level[index]]++;
-  }
-
-  uint64_t maxWeightedVmmLevel = 0;
-  uint64_t wideWeightedVmmLevels64 = 0;
-  uint64_t wideWeightedVmmLevels256 = 0;
-  for (uint64_t count : weightedVmmNodesByLevel) {
-    maxWeightedVmmLevel = std::max(maxWeightedVmmLevel, count);
-    wideWeightedVmmLevels64 += count >= 64;
-    wideWeightedVmmLevels256 += count >= 256;
-  }
-
-  using ShapeKey = std::tuple<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t, uint64_t>;
-  SmallVector<ShapeKey> weightedVmmShapeKeys;
-  for (auto [index, compute] : llvm::enumerate(computes)) {
-    const ComputeMotifInfo& info = computeInfos[index];
-    if (info.weightedVmmCount == 0)
-      continue;
-    weightedVmmShapeKeys.push_back({info.instructionCount,
-                                    info.weightedVmmCount,
-                                    static_cast<uint64_t>(compute.getWeights().size()),
-                                    static_cast<uint64_t>(compute.getInputs().size()),
-                                    static_cast<uint64_t>(parents[index].size()),
-                                    static_cast<uint64_t>(children[index].size())});
-  }
-
-  llvm::sort(weightedVmmShapeKeys);
-  SmallVector<std::pair<uint64_t, ShapeKey>> weightedVmmShapeCounts;
-  for (size_t index = 0; index < weightedVmmShapeKeys.size();) {
-    size_t next = index + 1;
-    while (next < weightedVmmShapeKeys.size() && weightedVmmShapeKeys[next] == weightedVmmShapeKeys[index])
-      next++;
-    weightedVmmShapeCounts.push_back({next - index, weightedVmmShapeKeys[index]});
-    index = next;
-  }
-  llvm::sort(weightedVmmShapeCounts, [](const auto& lhs, const auto& rhs) {
-    if (lhs.first != rhs.first)
-      return lhs.first > rhs.first;
-    return lhs.second < rhs.second;
-  });
-
-  llvm::errs() << llvm::formatv("[DCP-MOTIF] computes={0} edges={1} wvmmNodes={2} wvmmOps={3} "
-                                "serialChains={4} serialChainNodes={5} maxSerialChain={6} "
-                                "maxFanIn={7} maxFanOut={8} fanIn>=16/64/256={9}/{10}/{11} "
-                                "fanOut>=16/64/256={12}/{13}/{14} topoVisited={15}\n",
-                                computes.size(),
-                                edgeCount,
-                                weightedVmmNodeCount,
-                                weightedVmmOpCount,
-                                serialChainCount,
-                                serialChainNodeCount,
-                                maxSerialChain,
-                                maxFanIn,
-                                maxFanOut,
-                                fanIn16,
-                                fanIn64,
-                                fanIn256,
-                                fanOut16,
-                                fanOut64,
-                                fanOut256,
-                                readyNodes.size());
-
-  llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmLevels={0} maxWvmmLevel={1} wideWvmmLevels>=64/256={2}/{3} "
-                                "shapeGroups={4}\n",
-                                weightedVmmNodesByLevel.size(),
-                                maxWeightedVmmLevel,
-                                wideWeightedVmmLevels64,
-                                wideWeightedVmmLevels256,
-                                weightedVmmShapeCounts.size());
-
-  for (size_t rank = 0, end = std::min<size_t>(weightedVmmShapeCounts.size(), 5); rank < end; ++rank) {
-    auto [count, shape] = weightedVmmShapeCounts[rank];
-    auto [insts, vmmOps, weights, inputs, fanIn, fanOut] = shape;
-    llvm::errs() << llvm::formatv("[DCP-MOTIF] wvmmShape rank={0} count={1} insts={2} vmmOps={3} "
-                                  "weights={4} inputs={5} fanIn={6} fanOut={7}\n",
-                                  rank,
-                                  count,
-                                  insts,
-                                  vmmOps,
-                                  weights,
-                                  inputs,
-                                  fanIn,
-                                  fanOut);
-  }
-}
-
 void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) {
   std::fstream file = openReportFile(name);
   if (!file.is_open())
@@ -628,44 +328,27 @@ public:
 
   void runOnOperation() override {
     func::FuncOp func = getOperation();
-    {
-      ScopedMergePhaseTimer timer("trivial-serial-merge");
-      mergeTriviallyConnectedComputes(func);
-    }
-    if (std::getenv("DCP_MOTIF_PROFILE"))
-      emitMotifProfile(func);
+    mergeTriviallyConnectedComputes(func);
 
     const spatial::MergeScheduleResult* analysisResult = nullptr;
-    {
-      ScopedMergePhaseTimer timer("scheduling-analysis");
-      analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    }
-    {
-      ScopedMergePhaseTimer timer("schedule-materialization");
-      if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
-        signalPassFailure();
-        return;
-      }
+    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
+    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
+      signalPassFailure();
+      return;
     }
 
-    emitMergeIrCounts("after-materialization", func);
-
-    {
-      ScopedMergePhaseTimer timer("cleanup-topological-sort-report");
-      if (!sortTopologically(&func.getBody().front())) {
-        func.emitOpError("failed to topologically order merged Spatial IR");
-        signalPassFailure();
-        return;
-      }
-      if (failed(verifySpatialCommunicationInvariants(func))) {
-        func.emitOpError("merged Spatial communication invariant verification failed");
-        signalPassFailure();
-        return;
-      }
-      emitMergeIrCounts("final-post-merge", func);
-      dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
-      generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    if (!sortTopologically(&func.getBody().front())) {
+      func.emitOpError("failed to topologically order merged Spatial IR");
+      signalPassFailure();
+      return;
     }
+    if (failed(verifySpatialCommunicationInvariants(func))) {
+      func.emitOpError("merged Spatial communication invariant verification failed");
+      signalPassFailure();
+      return;
+    }
+    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
+    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
   }
 };
 

validation/.gitignore

@@ -8,3 +8,7 @@ networks/**/outputs
 networks/**/raptor
 networks/**/runner
 networks/**/simulation
+networks/**/real_image_val
+networks/**/*.png
+networks/**/*.jpg
+networks/**/*.csv

validation/gen_network_runner.py

@@ -199,7 +199,10 @@ int main(int argc, char **argv) {{
 
   // ---- Cleanup ----
   omTensorListDestroy(in_list);
-  omTensorListDestroy(out_list);
+  // Some debug-heavy models return aliased outputs. This runner is a short-
+  // lived process, so destroy only the list wrapper and let process exit
+  // reclaim the output tensors safely.
+  omTensorListDestroyShallow(out_list);
   return 0;
 }}
 """

validation/operations/gen_tests.py

@@ -1053,6 +1053,92 @@ def reducemean_large_dimension_1024():
     save_model(model, "reduce_mean/large_dimension_1024", "reduce_mean_large_dimension_1024.onnx")
 
 
+def make_legacy_reducemean_model(name, shape, output_shape, directory, filename, *, axes, keepdims=1,
+                                 noop_with_empty_axes=0):
+    """Create an opset-18 ReduceMean model that lowers to ONNXReduceMeanOp."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)
+
+    initializers = []
+    node_inputs = ["X", ""]
+    if axes is not None:
+        initializers.append(make_int64_initializer("axes", axes))
+        node_inputs = ["X", "axes"]
+
+    node = helper.make_node("ReduceMean", node_inputs, ["Y"],
+                            keepdims=keepdims, noop_with_empty_axes=noop_with_empty_axes)
+    graph = helper.make_graph([node], name, [X], [Y], initializer=initializers)
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
+    save_model(model, directory, filename)
+
+
+def reducemean_legacy_axis1_keepdims_1():
+    """Opset-18 ReduceMean over one positive axis, preserving rank."""
+    make_legacy_reducemean_model("reducemean_legacy_axis1_keepdims_1",
+                                 [2, 3, 4], [2, 1, 4],
+                                 "reduce_mean/legacy_axis1_keepdims_1",
+                                 "reduce_mean_legacy_axis1_keepdims_1.onnx",
+                                 axes=[1], keepdims=1)
+
+
+def reducemean_legacy_axis1_keepdims_0():
+    """Opset-18 ReduceMean over one positive axis, dropping the reduced axis."""
+    make_legacy_reducemean_model("reducemean_legacy_axis1_keepdims_0",
+                                 [2, 3, 4], [2, 4],
+                                 "reduce_mean/legacy_axis1_keepdims_0",
+                                 "reduce_mean_legacy_axis1_keepdims_0.onnx",
+                                 axes=[1], keepdims=0)
+
+
+def reducemean_legacy_axes_1_2_keepdims_1():
+    """Opset-18 ReduceMean over multiple positive axes."""
+    make_legacy_reducemean_model("reducemean_legacy_axes_1_2_keepdims_1",
+                                 [2, 3, 4], [2, 1, 1],
+                                 "reduce_mean/legacy_axes_1_2_keepdims_1",
+                                 "reduce_mean_legacy_axes_1_2_keepdims_1.onnx",
+                                 axes=[1, 2], keepdims=1)
+
+
+def reducemean_legacy_negative_axis():
+    """Opset-18 ReduceMean using a negative axis."""
+    make_legacy_reducemean_model("reducemean_legacy_negative_axis",
+                                 [2, 3, 4], [2, 3, 1],
+                                 "reduce_mean/legacy_negative_axis",
+                                 "reduce_mean_legacy_negative_axis.onnx",
+                                 axes=[-1], keepdims=1)
+
+
+def reducemean_legacy_reduce_all_keepdims_1():
+    """Opset-18 ReduceMean over all axes with the optional axes input omitted."""
+    make_legacy_reducemean_model("reducemean_legacy_reduce_all_keepdims_1",
+                                 [2, 3, 4], [1, 1, 1],
+                                 "reduce_mean/legacy_reduce_all_keepdims_1",
+                                 "reduce_mean_legacy_reduce_all_keepdims_1.onnx",
+                                 axes=None, keepdims=1)
+
+
+def reducemean_legacy_empty_axes_noop():
+    """Opset-18 ReduceMean with empty axes and noop_with_empty_axes enabled."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 4])
+    axes = make_int64_initializer("axes", [])
+    reduce = helper.make_node("ReduceMean", ["X", "axes"], ["R"],
+                              keepdims=1, noop_with_empty_axes=1)
+    relu = helper.make_node("Relu", ["R"], ["Y"])
+    graph = helper.make_graph([reduce, relu], "reducemean_legacy_empty_axes_noop", [X], [Y], initializer=[axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
+    save_model(model, "reduce_mean/legacy_empty_axes_noop", "reduce_mean_legacy_empty_axes_noop.onnx")
+
+
+def reducemean_legacy_nchw_spatial():
+    """Opset-18 ReduceMean over H and W on an NCHW tensor."""
+    make_legacy_reducemean_model("reducemean_legacy_nchw_spatial",
+                                 [1, 3, 5, 5], [1, 3, 1, 1],
+                                 "reduce_mean/legacy_nchw_spatial",
+                                 "reduce_mean_legacy_nchw_spatial.onnx",
+                                 axes=[2, 3], keepdims=1)
+
+
 # ---------------------------------------------------------------------------
 # Relu tests
 # ---------------------------------------------------------------------------
@@ -1340,6 +1426,118 @@ def split_uneven_channel_axis_4d():
     save_model(model, "split/uneven_channel_axis_4d", "split_uneven_channel_axis_4d.onnx")
 
 
+# ---------------------------------------------------------------------------
+# Slice tests
+# ---------------------------------------------------------------------------
+
+def slice_2d_basic():
+    """Slice a 2D tensor with explicit axes and unit steps."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4, 6])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3])
+    starts = make_int64_initializer("starts", [1, 2])
+    ends = make_int64_initializer("ends", [3, 5])
+    axes = make_int64_initializer("axes", [0, 1])
+    steps = make_int64_initializer("steps", [1, 1])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes", "steps"], ["Y"])
+    graph = helper.make_graph([node], "slice_2d_basic", [X], [Y], initializer=[starts, ends, axes, steps])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/2d_basic", "slice_2d_basic.onnx")
+
+
+def slice_default_axes():
+    """Slice with omitted axes and steps using default positional axes."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 4])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 2, 4])
+    starts = make_int64_initializer("starts", [0, 1, 0])
+    ends = make_int64_initializer("ends", [2, 3, 4])
+    node = helper.make_node("Slice", ["X", "starts", "ends"], ["Y"])
+    graph = helper.make_graph([node], "slice_default_axes", [X], [Y], initializer=[starts, ends])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/default_axes", "slice_default_axes.onnx")
+
+
+def slice_negative_axis():
+    """Slice using a negative axis."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 5])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 3])
+    starts = make_int64_initializer("starts", [1])
+    ends = make_int64_initializer("ends", [4])
+    axes = make_int64_initializer("axes", [-1])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes"], ["Y"])
+    graph = helper.make_graph([node], "slice_negative_axis", [X], [Y], initializer=[starts, ends, axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/negative_axis", "slice_negative_axis.onnx")
+
+
+def slice_negative_indices():
+    """Slice with negative indices along one axis."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 5])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 3])
+    starts = make_int64_initializer("starts", [-4])
+    ends = make_int64_initializer("ends", [-1])
+    axes = make_int64_initializer("axes", [1])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes"], ["Y"])
+    graph = helper.make_graph([node], "slice_negative_indices", [X], [Y], initializer=[starts, ends, axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/negative_indices", "slice_negative_indices.onnx")
+
+
+def slice_step2():
+    """Slice with a positive step greater than one."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 4])
+    starts = make_int64_initializer("starts", [0])
+    ends = make_int64_initializer("ends", [8])
+    axes = make_int64_initializer("axes", [2])
+    steps = make_int64_initializer("steps", [2])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes", "steps"], ["Y"])
+    graph = helper.make_graph([node], "slice_step2", [X], [Y], initializer=[starts, ends, axes, steps])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/step2", "slice_step2.onnx")
+
+
+def slice_nchw_spatial_crop():
+    """Slice an NCHW tensor across the spatial axes."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 8, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 4, 6])
+    starts = make_int64_initializer("starts", [2, 1])
+    ends = make_int64_initializer("ends", [6, 7])
+    axes = make_int64_initializer("axes", [2, 3])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes"], ["Y"])
+    graph = helper.make_graph([node], "slice_nchw_spatial_crop", [X], [Y], initializer=[starts, ends, axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/nchw_spatial_crop", "slice_nchw_spatial_crop.onnx")
+
+
+def slice_after_conv():
+    """Conv followed by a spatial crop using Slice."""
+    rng = np.random.default_rng(108)
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 8, 8])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4, 4, 4])
+    W = numpy_helper.from_array(rng.uniform(-1, 1, (4, 3, 3, 3)).astype(np.float32), name="W")
+    starts = make_int64_initializer("starts", [1, 1])
+    ends = make_int64_initializer("ends", [5, 5])
+    axes = make_int64_initializer("axes", [2, 3])
+    conv = helper.make_node("Conv", ["X", "W"], ["C"], kernel_shape=[3, 3], strides=[1, 1], pads=[0, 0, 0, 0])
+    slice_node = helper.make_node("Slice", ["C", "starts", "ends", "axes"], ["Y"])
+    graph = helper.make_graph([conv, slice_node], "slice_after_conv", [X], [Y], initializer=[W, starts, ends, axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/after_conv", "slice_after_conv.onnx")
+
+
+def slice_large_channel_1024():
+    """Slice a large channel range out of a 1024-channel tensor."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1024, 1, 1])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 512, 1, 1])
+    starts = make_int64_initializer("starts", [128])
+    ends = make_int64_initializer("ends", [640])
+    axes = make_int64_initializer("axes", [1])
+    node = helper.make_node("Slice", ["X", "starts", "ends", "axes"], ["Y"])
+    graph = helper.make_graph([node], "slice_large_channel_1024", [X], [Y], initializer=[starts, ends, axes])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
+    save_model(model, "slice/large_channel_1024", "slice_large_channel_1024.onnx")
+
+
 # ---------------------------------------------------------------------------
 # Gather tests
 # ---------------------------------------------------------------------------
@@ -1862,6 +2060,13 @@ if __name__ == "__main__":
     reducemean_4d_spatial_keepdims_0()
     reducemean_channel_axis_nchw()
     reducemean_large_dimension_1024()
+    reducemean_legacy_axis1_keepdims_1()
+    reducemean_legacy_axis1_keepdims_0()
+    reducemean_legacy_axes_1_2_keepdims_1()
+    reducemean_legacy_negative_axis()
+    reducemean_legacy_reduce_all_keepdims_1()
+    reducemean_legacy_empty_axes_noop()
+    reducemean_legacy_nchw_spatial()
 
     print("\nGenerating Relu tests:")
     relu_basic()
@@ -1880,6 +2085,16 @@ if __name__ == "__main__":
     split_negative_axis()
     split_uneven_channel_axis_4d()
 
+    print("\nGenerating Slice tests:")
+    slice_2d_basic()
+    slice_default_axes()
+    slice_negative_axis()
+    slice_negative_indices()
+    slice_step2()
+    slice_nchw_spatial_crop()
+    slice_after_conv()
+    slice_large_channel_1024()
+
     print("\nGenerating Softmax tests:")
     softmax_basic()
     softmax_3d_last_axis()

validation/raptor.py

@@ -41,7 +41,8 @@ def _format_command(cmd):
 
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
                         crossbar_size, crossbar_count, core_count=None, pim_merge_scheduler="peft",
-                        pim_memory_report="none", cwd=None, verbose=False, reporter=None, timeout_sec=None):
+                        pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False,
+                        reporter=None, timeout_sec=None):
     # Define the arguments, with the possibility to set crossbar size and count
     args = [
         network_path,
@@ -57,6 +58,8 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
         args.append(f"--core-count={core_count}")
     if pim_memory_report != "none":
         args.append(f"--pim-memory-report={pim_memory_report}")
+    if raptor_extra_args:
+        args.extend(str(arg) for arg in raptor_extra_args)
     if verbose:
         args.append("--enable-timing")
 

validation/tools/yolo_local_image_validation.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python3
+
+import argparse
+import json
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+VALIDATION_DIR = SCRIPT_DIR.parent
+if str(VALIDATION_DIR) not in sys.path:
+    sys.path.insert(0, str(VALIDATION_DIR))
+
+if sys.version_info < (3, 10):
+    raise SystemExit(
+        "yolo_local_image_validation.py requires Python 3.10+ because validation modules use modern type syntax. "
+        "Run it with a newer interpreter, for example your project venv Python."
+    )
+
+from onnx_utils import _ONNX_TO_NP, onnx_io, write_inputs_to_memory_bin
+from validate_one import MODE_COMPILE_ONLY, build_dump_ranges, run_pim_simulator, sanitize_output_name, validate_network
+from yolo_real_image_validation import (
+    IMAGE_CASES,
+    decode_yolo_output,
+    download_image,
+    draw_detections,
+    letterbox_rgb,
+    save_tensor_csv,
+    top_unique_labels,
+)
+
+
+def resolve_default_paths():
+    validation_dir = Path(__file__).resolve().parent.parent
+    repo_root = validation_dir.parent
+    return {
+        "validation_dir": validation_dir,
+        "repo_root": repo_root,
+        "network_dir": validation_dir / "networks" / "yolo11n" / "depth_51",
+        "raptor_path": repo_root / "build_release" / "Release" / "bin" / "onnx-mlir",
+        "onnx_include_dir": repo_root / "onnx-mlir" / "include",
+        "simulator_dir": repo_root / "backend-simulators" / "pim" / "pim-simulator",
+    }
+
+
+def find_network_onnx(network_dir: Path) -> Path:
+    onnx_files = sorted(network_dir.glob("*.onnx"))
+    if not onnx_files:
+        raise FileNotFoundError(f"No .onnx file found in {network_dir}")
+    if len(onnx_files) > 1:
+        names = ", ".join(path.name for path in onnx_files)
+        raise RuntimeError(f"Expected exactly one .onnx file in {network_dir}, found: {names}")
+    return onnx_files[0]
+
+
+def local_case_paths(network_dir: Path, case_name: str):
+    return {
+        "root": network_dir,
+        "runner": network_dir / "runner" / "build" / "runner",
+        "runner_build": network_dir / "runner" / "build",
+        "raptor_pim": network_dir / "raptor" / "pim",
+        "real_root": network_dir / "real_image_validation",
+        "input_csv": network_dir / "real_image_validation" / "inputs" / f"{case_name}.csv",
+        "ref_dir": network_dir / "real_image_validation" / "reference" / case_name,
+        "sim_dir": network_dir / "real_image_validation" / "simulation" / case_name,
+        "sim_bin": network_dir / "real_image_validation" / "simulation" / case_name / "out.bin",
+    }
+
+
+def ensure_local_artifacts(args, network_onnx_path: Path):
+    validate_network(
+        network_onnx_path=network_onnx_path,
+        raptor_path=args.raptor_path,
+        onnx_include_dir=args.onnx_include_dir,
+        simulator_dir=args.simulator_dir,
+        crossbar_size=args.crossbar_size,
+        crossbar_count=args.crossbar_count,
+        core_count=args.core_count,
+        command_timeout_seconds=args.command_timeout_seconds,
+        mode=MODE_COMPILE_ONLY,
+        verbose=args.verbose,
+    )
+
+
+def ensure_existing_artifacts(network_dir: Path):
+    required_paths = [
+        network_dir / "runner" / "build" / "runner",
+        network_dir / "raptor" / "pim" / "config.json",
+        network_dir / "raptor" / "pim" / "memory.bin",
+    ]
+    missing = [str(path) for path in required_paths if not path.exists()]
+    if missing:
+        raise FileNotFoundError(
+            "Missing compiled local artifacts. Re-run without --skip-compile or restore these paths:\n  "
+            + "\n  ".join(missing)
+        )
+
+
+def run_local_reference_and_simulator(args, network_dir: Path, network_onnx_path: Path, case_name: str):
+    paths = local_case_paths(network_dir, case_name)
+    paths["ref_dir"].mkdir(parents=True, exist_ok=True)
+    paths["sim_dir"].mkdir(parents=True, exist_ok=True)
+
+    output_descriptors = onnx_io(network_onnx_path)[1]
+    if len(output_descriptors) != 1:
+        raise RuntimeError(f"Expected one YOLO output tensor, found {len(output_descriptors)}")
+
+    runner_cmd = [
+        str(paths["runner"]),
+        "--in0-csv-file",
+        str(paths["input_csv"]),
+        "--in0-shape",
+        "1x3x640x640",
+        "--save-csv-dir",
+        str(paths["ref_dir"]),
+    ]
+    subprocess.run(runner_cmd, cwd=paths["runner_build"], check=True)
+
+    tensor = np.loadtxt(paths["input_csv"], delimiter=",", dtype=np.float32).reshape(1, 3, 640, 640)
+    write_inputs_to_memory_bin(paths["raptor_pim"] / "memory.bin", paths["raptor_pim"] / "config.json", [tensor])
+
+    dump_ranges = build_dump_ranges(paths["raptor_pim"] / "config.json", output_descriptors)
+    run_pim_simulator(
+        args.simulator_dir,
+        paths["raptor_pim"],
+        paths["sim_bin"],
+        dump_ranges,
+        timeout_sec=args.command_timeout_seconds,
+    )
+    return paths, output_descriptors[0]
+
+
+def analyze_case(args, network_dir: Path, network_onnx_path: Path, case, work_dir: Path):
+    image_path = work_dir / f"{case.name}{Path(case.url).suffix or '.img'}"
+    csv_path = work_dir / f"{case.name}.csv"
+    annotated_dir = args.annotated_dir
+    annotated_dir.mkdir(parents=True, exist_ok=True)
+    download_image(case.url, image_path)
+    tensor = letterbox_rgb(Image.open(image_path))
+    save_tensor_csv(tensor, csv_path)
+
+    paths = local_case_paths(network_dir, case.name)
+    paths["input_csv"].parent.mkdir(parents=True, exist_ok=True)
+    paths["input_csv"].write_bytes(csv_path.read_bytes())
+    paths, output_descriptor = run_local_reference_and_simulator(args, network_dir, network_onnx_path, case.name)
+
+    output_index, output_name, output_dtype_code, output_shape = output_descriptor
+    output_dtype = np.dtype(_ONNX_TO_NP[output_dtype_code])
+    ref_csv_path = paths["ref_dir"] / f"output{output_index}_{sanitize_output_name(output_name)}.csv"
+    ref = np.loadtxt(ref_csv_path, delimiter=",", dtype=output_dtype).reshape(output_shape)
+    sim = np.frombuffer(
+        paths["sim_bin"].read_bytes(),
+        dtype=output_dtype,
+        count=int(np.prod(output_shape)),
+    ).reshape(output_shape)
+
+    abs_diff = np.abs(sim.astype(np.float64) - ref.astype(np.float64))
+    rel_diff = abs_diff / np.maximum(np.abs(ref.astype(np.float64)), 1e-12)
+
+    ref_detections = decode_yolo_output(ref)
+    sim_detections = decode_yolo_output(sim)
+    ref_labels = top_unique_labels(ref_detections)
+    sim_labels = top_unique_labels(sim_detections)
+    ref_image_path = annotated_dir / f"{case.name}_reference.png"
+    sim_image_path = annotated_dir / f"{case.name}_simulator.png"
+    draw_detections(image_path, ref_detections, ref_image_path)
+    draw_detections(image_path, sim_detections, sim_image_path)
+
+    return {
+        "case": case.name,
+        "expected_label": case.expected_label,
+        "ref_top_labels": ref_labels,
+        "sim_top_labels": sim_labels,
+        "top1_match": bool(ref_labels and sim_labels and ref_labels[0] == sim_labels[0]),
+        "expected_in_ref": case.expected_label in ref_labels,
+        "expected_in_sim": case.expected_label in sim_labels,
+        "max_abs_diff": float(abs_diff.max()),
+        "mean_abs_diff": float(abs_diff.mean()),
+        "max_rel_diff": float(rel_diff.max()),
+        "mean_rel_diff": float(rel_diff.mean()),
+        "reference_annotated_image": str(ref_image_path),
+        "simulator_annotated_image": str(sim_image_path),
+        "ref_top_detections": ref_detections[:5],
+        "sim_top_detections": sim_detections[:5],
+    }
+
+
+def main():
+    defaults = resolve_default_paths()
+
+    parser = argparse.ArgumentParser(description="Validate YOLO detections on real images using local compilation and simulator execution.")
+    parser.add_argument("--network-dir", type=Path, default=defaults["network_dir"])
+    parser.add_argument("--network-onnx", type=Path, default=None)
+    parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"])
+    parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"])
+    parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"])
+    parser.add_argument("--crossbar-size", type=int, default=2048)
+    parser.add_argument("--crossbar-count", type=int, default=256)
+    parser.add_argument("--core-count", type=int, default=1000)
+    parser.add_argument("--command-timeout-seconds", type=float, default=7200.0)
+    parser.add_argument("--skip-compile", action="store_true")
+    parser.add_argument("--verbose", action="store_true")
+    parser.add_argument(
+        "--annotated-dir",
+        type=Path,
+        default=defaults["network_dir"] / "real_image_validation" / "annotated",
+    )
+    args = parser.parse_args()
+
+    args.network_dir = args.network_dir.resolve()
+    args.network_onnx = args.network_onnx.resolve() if args.network_onnx else find_network_onnx(args.network_dir)
+    args.raptor_path = args.raptor_path.resolve()
+    args.onnx_include_dir = args.onnx_include_dir.resolve()
+    args.simulator_dir = args.simulator_dir.resolve()
+    args.annotated_dir = args.annotated_dir.resolve()
+
+    if not args.skip_compile:
+        ensure_local_artifacts(args, args.network_onnx)
+    else:
+        ensure_existing_artifacts(args.network_dir)
+
+    reports = []
+    with tempfile.TemporaryDirectory(prefix="yolo_local_images_") as tmp_dir:
+        work_dir = Path(tmp_dir)
+        for case in IMAGE_CASES:
+            reports.append(analyze_case(args, args.network_dir, args.network_onnx, case, work_dir))
+
+    print(json.dumps({"network_dir": str(args.network_dir), "network_onnx": str(args.network_onnx), "cases": reports}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

validation/tools/yolo_real_image_validation.py

@@ -0,0 +1,425 @@
+#!/usr/bin/env python3
+
+import argparse
+import json
+import shlex
+import subprocess
+import tempfile
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+
+COCO80_CLASSES = [
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
+    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
+    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
+    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
+    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
+    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
+    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
+    "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
+    "scissors", "teddy bear", "hair drier", "toothbrush",
+]
+
+
+@dataclass(frozen=True)
+class ImageCase:
+    name: str
+    url: str
+    expected_label: str
+
+
+IMAGE_CASES = [
+    ImageCase(
+        name="cat_coco_39769",
+        url="http://images.cocodataset.org/val2017/000000039769.jpg",
+        expected_label="cat",
+    ),
+    ImageCase(
+        name="dog_pytorch_hub",
+        url="https://github.com/pytorch/hub/raw/master/images/dog.jpg",
+        expected_label="dog",
+    ),
+        ImageCase(
+        name="cute_kitty",
+        url="https://images.unsplash.com/photo-1529778873920-4da4926a72c2?q=80&w=872&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D" ,
+        expected_label="cat",
+    ),
+
+]
+
+
+def run(cmd, *, cwd=None, capture_output=False, input_bytes=None):
+    return subprocess.run(
+        cmd,
+        cwd=cwd,
+        check=True,
+        input=input_bytes,
+        capture_output=capture_output,
+    )
+
+
+def ssh_command(ssh_key: str, remote_host: str, command: str):
+    return ["ssh", "-i", ssh_key, remote_host, command]
+
+
+def remote_bash(ssh_key: str, remote_host: str, command: str, *, capture_output=False, input_bytes=None):
+    return run(
+        ssh_command(ssh_key, remote_host, command),
+        capture_output=capture_output,
+        input_bytes=input_bytes,
+    )
+
+
+def download_image(url: str, path: Path):
+    with urllib.request.urlopen(url) as response:
+        path.write_bytes(response.read())
+
+
+def letterbox_rgb(image: Image.Image, size: int = 640) -> np.ndarray:
+    image = image.convert("RGB")
+    width, height = image.size
+    scale = min(size / width, size / height)
+    resized_width = max(1, int(round(width * scale)))
+    resized_height = max(1, int(round(height * scale)))
+    resized = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
+
+    canvas = Image.new("RGB", (size, size), (114, 114, 114))
+    offset_x = (size - resized_width) // 2
+    offset_y = (size - resized_height) // 2
+    canvas.paste(resized, (offset_x, offset_y))
+
+    array = np.asarray(canvas, dtype=np.float32) / 255.0
+    chw = np.transpose(array, (2, 0, 1))
+    return np.expand_dims(chw, axis=0)
+
+
+def letterbox_params(width: int, height: int, size: int = 640):
+    scale = min(size / width, size / height)
+    resized_width = max(1, int(round(width * scale)))
+    resized_height = max(1, int(round(height * scale)))
+    offset_x = (size - resized_width) // 2
+    offset_y = (size - resized_height) // 2
+    return scale, offset_x, offset_y
+
+
+def save_tensor_csv(array: np.ndarray, path: Path):
+    flat = array.reshape(-1)
+    np.savetxt(path, flat[np.newaxis, :], delimiter=",", fmt="%.9g")
+
+
+def iou_xyxy(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+    x1 = np.maximum(box[0], boxes[:, 0])
+    y1 = np.maximum(box[1], boxes[:, 1])
+    x2 = np.minimum(box[2], boxes[:, 2])
+    y2 = np.minimum(box[3], boxes[:, 3])
+
+    inter_w = np.maximum(0.0, x2 - x1)
+    inter_h = np.maximum(0.0, y2 - y1)
+    inter = inter_w * inter_h
+
+    area_box = np.maximum(0.0, box[2] - box[0]) * np.maximum(0.0, box[3] - box[1])
+    area_boxes = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
+    union = area_box + area_boxes - inter
+    return np.divide(inter, union, out=np.zeros_like(inter), where=union > 0)
+
+
+def decode_yolo_output(
+    output: np.ndarray,
+    *,
+    conf_threshold: float = 0.25,
+    iou_threshold: float = 0.45,
+    max_detections: int = 50,
+):
+    predictions = output[0].T
+    boxes_xywh = predictions[:, :4]
+    class_scores = predictions[:, 4:]
+
+    class_ids = np.argmax(class_scores, axis=1)
+    confidences = class_scores[np.arange(class_scores.shape[0]), class_ids]
+    keep = confidences >= conf_threshold
+
+    if not np.any(keep):
+        return []
+
+    boxes_xywh = boxes_xywh[keep]
+    class_ids = class_ids[keep]
+    confidences = confidences[keep]
+
+    boxes_xyxy = np.empty_like(boxes_xywh)
+    boxes_xyxy[:, 0] = boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2.0
+    boxes_xyxy[:, 1] = boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2.0
+    boxes_xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2] / 2.0
+    boxes_xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3] / 2.0
+
+    detections = []
+    for class_id in np.unique(class_ids):
+        class_mask = class_ids == class_id
+        class_boxes = boxes_xyxy[class_mask]
+        class_scores_masked = confidences[class_mask]
+        order = np.argsort(-class_scores_masked)
+
+        while order.size > 0:
+            best = order[0]
+            detections.append({
+                "label": COCO80_CLASSES[int(class_id)],
+                "class_id": int(class_id),
+                "confidence": float(class_scores_masked[best]),
+                "box_xyxy": class_boxes[best].tolist(),
+            })
+
+            if order.size == 1:
+                break
+
+            rest = order[1:]
+            overlaps = iou_xyxy(class_boxes[best], class_boxes[rest])
+            order = rest[overlaps <= iou_threshold]
+
+    detections.sort(key=lambda det: det["confidence"], reverse=True)
+    return detections[:max_detections]
+
+
+def top_unique_labels(detections, limit: int = 5):
+    labels = []
+    seen = set()
+    for det in detections:
+        label = det["label"]
+        if label in seen:
+            continue
+        seen.add(label)
+        labels.append(label)
+        if len(labels) == limit:
+            break
+    return labels
+
+
+def clamp_box_xyxy(box_xyxy, width: int, height: int):
+    x1, y1, x2, y2 = box_xyxy
+    return [
+        max(0.0, min(float(width - 1), float(x1))),
+        max(0.0, min(float(height - 1), float(y1))),
+        max(0.0, min(float(width - 1), float(x2))),
+        max(0.0, min(float(height - 1), float(y2))),
+    ]
+
+
+def unletterbox_box_xyxy(box_xyxy, width: int, height: int, size: int = 640):
+    scale, offset_x, offset_y = letterbox_params(width, height, size=size)
+    x1, y1, x2, y2 = box_xyxy
+    return [
+        (float(x1) - offset_x) / scale,
+        (float(y1) - offset_y) / scale,
+        (float(x2) - offset_x) / scale,
+        (float(y2) - offset_y) / scale,
+    ]
+
+
+def draw_detections(image_path: Path, detections, output_path: Path, *, limit: int = 10):
+    image = Image.open(image_path).convert("RGB")
+    draw = ImageDraw.Draw(image)
+    width, height = image.size
+
+    for det in detections[:limit]:
+        box = unletterbox_box_xyxy(det["box_xyxy"], width, height)
+        box = clamp_box_xyxy(box, width, height)
+        label = f'{det["label"]} {det["confidence"]:.2f}'
+        draw.rectangle(box, outline=(255, 0, 0), width=3)
+        text_box = draw.textbbox((box[0], box[1]), label)
+        text_bg = [
+            text_box[0] - 2,
+            text_box[1] - 2,
+            text_box[2] + 2,
+            text_box[3] + 2,
+        ]
+        draw.rectangle(text_bg, fill=(255, 0, 0))
+        draw.text((box[0], box[1]), label, fill=(255, 255, 255))
+
+    image.save(output_path)
+
+
+def ensure_remote_artifacts(args):
+    remote_project = shlex.quote(args.remote_project)
+    remote_python = shlex.quote(args.remote_python)
+    validate_cmd = (
+        f"export PATH=$HOME/.cargo/bin:$PATH && "
+        f"cd {remote_project} && "
+        f"{remote_python} validation/validate.py "
+        f"--raptor-path build_release/Release/bin/onnx-mlir "
+        f"--onnx-include-dir onnx-mlir/include "
+        f"--operations-dir {shlex.quote(args.network_dir)} "
+        f"--crossbar-size {args.crossbar_size} "
+        f"--crossbar-count {args.crossbar_count} "
+        f"--core-count {args.core_count} "
+        f"--command-timeout-seconds {args.command_timeout_seconds} "
+        f"--compile-only"
+    )
+    remote_bash(args.ssh_key, args.remote_host, validate_cmd)
+
+
+def remote_case_paths(args, case_name: str):
+    network_dir = Path(args.network_dir)
+    root = Path(args.remote_project) / network_dir
+    return {
+        "root": root,
+        "runner": root / "runner" / "build" / "runner",
+        "runner_build": root / "runner" / "build",
+        "raptor_pim": root / "raptor" / "pim",
+        "real_root": root / "real_image_validation",
+        "input_csv": root / "real_image_validation" / "inputs" / f"{case_name}.csv",
+        "ref_dir": root / "real_image_validation" / "reference" / case_name,
+        "sim_dir": root / "real_image_validation" / "simulation" / case_name,
+        "sim_bin": root / "real_image_validation" / "simulation" / case_name / "out.bin",
+    }
+
+
+def write_remote_file(args, remote_path: Path, data: bytes):
+    command = (
+        f"mkdir -p {shlex.quote(str(remote_path.parent))} && "
+        f"cat > {shlex.quote(str(remote_path))}"
+    )
+    remote_bash(args.ssh_key, args.remote_host, command, input_bytes=data)
+
+
+def run_remote_reference_and_simulator(args, case_name: str):
+    paths = remote_case_paths(args, case_name)
+    quoted_project = shlex.quote(args.remote_project)
+    quoted_python = shlex.quote(args.remote_python)
+    quoted_case_csv = shlex.quote(str(paths["input_csv"]))
+    quoted_ref_dir = shlex.quote(str(paths["ref_dir"]))
+    quoted_sim_dir = shlex.quote(str(paths["sim_dir"]))
+    quoted_sim_bin = shlex.quote(str(paths["sim_bin"]))
+    quoted_runner = shlex.quote(str(paths["runner"]))
+    quoted_runner_build = shlex.quote(str(paths["runner_build"]))
+    quoted_pim = shlex.quote(str(paths["raptor_pim"]))
+
+    command = f"""
+set -e
+export PATH=$HOME/.cargo/bin:$PATH
+cd {quoted_project}
+mkdir -p {quoted_ref_dir} {quoted_sim_dir}
+cd {quoted_runner_build}
+{quoted_runner} --in0-csv-file {quoted_case_csv} --in0-shape 1x3x640x640 --save-csv-dir {quoted_ref_dir}
+cd {quoted_project}
+{quoted_python} - <<'PY'
+import json
+import numpy as np
+from pathlib import Path
+input_csv = Path({json.dumps(str(paths["input_csv"]))})
+pim_dir = Path({json.dumps(str(paths["raptor_pim"]))})
+config = json.loads((pim_dir / "config.json").read_text())
+tensor = np.loadtxt(input_csv, delimiter=",", dtype=np.float32).reshape(1, 3, 640, 640)
+with open(pim_dir / "memory.bin", "r+b") as f:
+    f.seek(config["inputs_addresses"][0])
+    f.write(tensor.tobytes(order="C"))
+output_addr = config["outputs_addresses"][0]
+output_size = 1 * 84 * 8400 * 4
+print(f"{{output_addr}},{{output_size}}")
+PY
+"""
+    result = remote_bash(args.ssh_key, args.remote_host, command, capture_output=True)
+    dump_range = result.stdout.decode().strip().splitlines()[-1]
+
+    sim_command = (
+        f"export PATH=$HOME/.cargo/bin:$PATH && "
+        f"cd {quoted_project}/backend-simulators/pim/pim-simulator && "
+        f"cargo run --no-default-features --release --package pim-simulator --bin pim-simulator -- "
+        f"-f {quoted_pim} -o {quoted_sim_bin} -d {dump_range}"
+    )
+    remote_bash(args.ssh_key, args.remote_host, sim_command)
+    return paths
+
+
+def read_remote_file(args, remote_path: Path) -> bytes:
+    result = remote_bash(
+        args.ssh_key,
+        args.remote_host,
+        f"cat {shlex.quote(str(remote_path))}",
+        capture_output=True,
+    )
+    return result.stdout
+
+
+def analyze_case(args, case: ImageCase, work_dir: Path):
+    image_path = work_dir / f"{case.name}{Path(case.url).suffix or '.img'}"
+    csv_path = work_dir / f"{case.name}.csv"
+    annotated_dir = Path(args.annotated_dir)
+    annotated_dir.mkdir(parents=True, exist_ok=True)
+    download_image(case.url, image_path)
+    tensor = letterbox_rgb(Image.open(image_path))
+    save_tensor_csv(tensor, csv_path)
+
+    remote_paths = remote_case_paths(args, case.name)
+    write_remote_file(args, remote_paths["input_csv"], csv_path.read_bytes())
+    remote_paths = run_remote_reference_and_simulator(args, case.name)
+
+    ref_csv = read_remote_file(args, remote_paths["ref_dir"] / "output0_output0.csv")
+    sim_bin = read_remote_file(args, remote_paths["sim_bin"])
+
+    ref = np.loadtxt(ref_csv.decode().splitlines(), delimiter=",", dtype=np.float32).reshape(1, 84, 8400)
+    sim = np.frombuffer(sim_bin, dtype=np.float32, count=1 * 84 * 8400).reshape(1, 84, 8400)
+    abs_diff = np.abs(sim.astype(np.float64) - ref.astype(np.float64))
+    rel_diff = abs_diff / np.maximum(np.abs(ref.astype(np.float64)), 1e-12)
+
+    ref_detections = decode_yolo_output(ref)
+    sim_detections = decode_yolo_output(sim)
+    ref_labels = top_unique_labels(ref_detections)
+    sim_labels = top_unique_labels(sim_detections)
+    ref_image_path = annotated_dir / f"{case.name}_reference.png"
+    sim_image_path = annotated_dir / f"{case.name}_simulator.png"
+    draw_detections(image_path, ref_detections, ref_image_path)
+    draw_detections(image_path, sim_detections, sim_image_path)
+
+    return {
+        "case": case.name,
+        "expected_label": case.expected_label,
+        "ref_top_labels": ref_labels,
+        "sim_top_labels": sim_labels,
+        "top1_match": bool(ref_labels and sim_labels and ref_labels[0] == sim_labels[0]),
+        "expected_in_ref": case.expected_label in ref_labels,
+        "expected_in_sim": case.expected_label in sim_labels,
+        "max_abs_diff": float(abs_diff.max()),
+        "mean_abs_diff": float(abs_diff.mean()),
+        "max_rel_diff": float(rel_diff.max()),
+        "mean_rel_diff": float(rel_diff.mean()),
+        "reference_annotated_image": str(ref_image_path),
+        "simulator_annotated_image": str(sim_image_path),
+        "ref_top_detections": ref_detections[:5],
+        "sim_top_detections": sim_detections[:5],
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Validate YOLO detections on real animal images against the simulator.")
+    parser.add_argument("--remote-host", default="gmagnani@monolith")
+    parser.add_argument("--ssh-key", default="~/.ssh/github")
+    parser.add_argument("--remote-project", default="/home/gmagnani/Project/Raptor")
+    parser.add_argument("--remote-python", default="/home/gmagnani/venv/bin/python")
+    parser.add_argument("--network-dir", default="validation/networks/yolo11n/depth_51")
+    parser.add_argument("--crossbar-size", type=int, default=2048)
+    parser.add_argument("--crossbar-count", type=int, default=256)
+    parser.add_argument("--core-count", type=int, default=1000)
+    parser.add_argument("--command-timeout-seconds", type=int, default=7200)
+    parser.add_argument("--skip-compile", action="store_true")
+    parser.add_argument("--annotated-dir", default="validation/networks/yolo11n/depth_51/real_image_validation/annotated")
+    args = parser.parse_args()
+
+    args.ssh_key = str(Path(args.ssh_key).expanduser())
+
+    if not args.skip_compile:
+        ensure_remote_artifacts(args)
+
+    reports = []
+    with tempfile.TemporaryDirectory(prefix="yolo_real_images_") as tmp_dir:
+        work_dir = Path(tmp_dir)
+        for case in IMAGE_CASES:
+            reports.append(analyze_case(args, case, work_dir))
+
+    print(json.dumps({"network_dir": args.network_dir, "cases": reports}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

validation/validate.py

@@ -67,7 +67,10 @@ def main():
     ap.add_argument("--operations-dir", default=None, help="Root of the operations tree (default: operations).")
     ap.add_argument("--simulator-dir", default=None,
                     help="Path to pim-simulator crate root (default: auto-detected relative to script).")
-    ap.add_argument("--threshold", type=float, default=1e-3, help="Max allowed diff per output element.")
+    ap.add_argument("--threshold", type=float, default=1e-3,
+                    help="Absolute tolerance for per-element output comparison.")
+    ap.add_argument("--relative-threshold", type=float, default=1e-5,
+                    help="Relative tolerance for per-element output comparison.")
     ap.add_argument("--seed", type=int, default=0, help="RNG seed for generated validation inputs.")
     ap.add_argument("--crossbar-size", type=int, default=64)
     ap.add_argument("--crossbar-count", type=int, default=8)
@@ -77,6 +80,8 @@ def main():
                     help="Scheduler used by the Spatial merge-compute-nodes pass.")
     ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none",
                     help="Emit a human-readable PIM memory planning report during codegen.")
+    ap.add_argument("--raptor-extra-arg", action="append", default=[],
+                    help="Additional argument to pass through to the Raptor compiler. Repeat as needed.")
     ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0,
                     help="Per-subprocess timeout in seconds for compiler, runner, and simulator commands.")
     ap.add_argument("--clean", action="store_true",
@@ -145,8 +150,10 @@ def main():
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
                 crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
                 pim_merge_scheduler=a.pim_merge_scheduler, pim_memory_report=a.pim_memory_report,
+                raptor_extra_args=a.raptor_extra_arg,
                 command_timeout_seconds=a.command_timeout_seconds,
                 threshold=a.threshold,
+                rtol=a.relative_threshold,
                 seed=a.seed,
                 reporter=reporter,
                 model_index=index,

validation/validate_one.py

@@ -258,14 +258,18 @@ def parse_pim_simulator_outputs(output_bin_path, outputs_descriptor):
     return arrays
 
 
-def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1e-3, verbose=False):
+def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1e-3, rtol=1e-5, verbose=False):
     all_passed = True
     rows = []
     for sim_array, (oi, name, _, shape) in zip(sim_arrays, outputs_descriptor):
         csv_name = f"output{oi}_{sanitize_output_name(name)}.csv"
         runner_array = np.loadtxt(runner_out_dir / csv_name, delimiter=',', dtype=np.float32).reshape(shape)
-        max_diff = float(np.max(np.abs(sim_array.astype(np.float64) - runner_array.astype(np.float64))))
-        passed = max_diff <= threshold
+        sim_array64 = sim_array.astype(np.float64)
+        runner_array64 = runner_array.astype(np.float64)
+        abs_diff = np.abs(sim_array64 - runner_array64)
+        allowed_diff = threshold + rtol * np.abs(runner_array64)
+        max_diff = float(np.max(abs_diff))
+        passed = bool(np.all(abs_diff <= allowed_diff))
         rows.append((name, f"{max_diff:.6e}", passed))
         if not passed:
             all_passed = False
@@ -289,7 +293,8 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                      simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None,
-                     pim_merge_scheduler="peft", pim_memory_report="none", threshold=1e-3,
+                     pim_merge_scheduler="peft", pim_memory_report="none", raptor_extra_args=None,
+                     threshold=1e-3, rtol=1e-5,
                      seed=0, reporter=None, model_index=1, model_total=1, verbose=False,
                      command_timeout_seconds=60.0, mode=MODE_FULL):
     network_onnx_path = Path(network_onnx_path).resolve()
@@ -343,7 +348,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                 pim_pass_timings = compile_with_raptor(
                     network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                     core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
-                    pim_memory_report=pim_memory_report,
+                    pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
                     cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
                 reporter.advance()
@@ -383,7 +388,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             pim_pass_timings = compile_with_raptor(
                 network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
                 core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
-                pim_memory_report=pim_memory_report,
+                pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
                 cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
             reporter.advance()
@@ -403,7 +408,7 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compare Outputs")
         sim_arrays = parse_pim_simulator_outputs(output_bin_path, outputs_descriptor)
         reporter.suspend()
-        passed = validate_outputs(sim_arrays, out_dir, outputs_descriptor, threshold, verbose=verbose)
+        passed = validate_outputs(sim_arrays, out_dir, outputs_descriptor, threshold, rtol=rtol, verbose=verbose)
         reporter.resume()
         reporter.advance()
         reporter.record_result(passed)