add verification of static weights in spatial

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -15,6 +15,7 @@
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -116,6 +117,15 @@ static bool isBatchOutputArgument(SpatComputeBatch batchOp, Value value) {
   return argNumber >= firstOutputArgNumber && argNumber < firstOutputArgNumber + batchOp.getNumResults();
 }
 
+template <typename ComputeOpTy>
+static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind) {
+  for (Value weight : computeOp.getWeights()) {
+    if (!isHostFoldableValue(weight))
+      return computeOp.emitOpError() << kind << " weights must be statically computed from constants";
+  }
+  return success();
+}
+
 static bool isConstantIndexLike(Value value) {
   APInt constantValue;
   return matchPattern(value, m_ConstantInt(&constantValue));
@@ -545,6 +555,8 @@ LogicalResult SpatCompute::verify() {
   for (unsigned inputIndex = 0; inputIndex < getInputs().size(); ++inputIndex)
     if (auto inputArg = getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
       return emitError("ComputeOp block argument is not used");
+  if (failed(verifyStaticWeights(*this, "compute")))
+    return failure();
   if (failed(verifyOnlyConstantExternalValues(this->getOperation(), getBody(), "spat.compute")))
     return failure();
   if (failed(verifyComputeResultsUses(this->getOperation())))
@@ -647,6 +659,8 @@ LogicalResult SpatComputeBatch::verify() {
 
   if (failed(verifyComputeResultsUses(this->getOperation())))
     return failure();
+  if (failed(verifyStaticWeights(*this, "compute_batch")))
+    return failure();
   if (failed(verifyOnlyConstantExternalValues(this->getOperation(), getBody(), "spat.compute_batch")))
     return failure();
   return verifyBatchBody(*this, block);

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

@@ -2350,20 +2350,29 @@ LogicalResult collectPackedRunsForWholeBatchInput(MaterializerState& state,
     if (run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
       continue;
 
-    SmallVector<WholeBatchAssemblyRange, 8> runRanges;
-    runRanges.reserve(run.slots.size());
+    SmallVector<WholeBatchAssemblyRange, 16> runRanges;
 
     for (const PackedScalarRunSlot& slot : run.slots) {
-      std::optional<ProducerKey> slotKey = getContiguousProducerKeyForKeys(slot.keys);
-      if (!slotKey)
-        return failure();
+      for (ProducerKey fragmentKey : slot.keys) {
+        if (fragmentKey.instance.op != key.instance.op || fragmentKey.resultIndex != key.resultIndex)
+          return failure();
 
-      if (wholeBatchRangeOverlaps(plan.coveredRanges, slotKey->instance.laneStart, slotKey->instance.laneCount))
-        return failure();
+        if (fragmentKey.instance.laneCount == 0)
+          return failure();
 
-      runRanges.push_back({slotKey->instance.laneStart, slotKey->instance.laneCount});
+        if (wholeBatchRangeOverlaps(plan.coveredRanges, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
+          return failure();
+
+        if (wholeBatchRangeOverlaps(runRanges, fragmentKey.instance.laneStart, fragmentKey.instance.laneCount))
+          return failure();
+
+        runRanges.push_back({fragmentKey.instance.laneStart, fragmentKey.instance.laneCount});
+      }
     }
 
+    if (runRanges.empty())
+      continue;
+
     plan.packedRuns.push_back(&run);
 
     for (WholeBatchAssemblyRange range : runRanges)