From aa088e2ba535a32b724675a54134fa17e937c881 Mon Sep 17 00:00:00 2001
From: ilgeco <gabriele.magnani@mail.polimi.it>
Date: Mon, 18 May 2026 17:20:40 +0200
Subject: [PATCH] Verify fix

---
 src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp  | 19 ++++++++++++++++++-
 .../Scheduling/ComputeInstanceUtils.cpp       |  1 +
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp
index 7845827..b7622d3 100644
--- a/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp
+++ b/src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp
@@ -3,6 +3,7 @@
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/LogicalResult.h"
@@ -338,6 +339,19 @@ LogicalResult SpatConcatOp::verify() {
   return success();
 }
 
+LogicalResult verifyComputeResultsUses(Operation* op) {
+  if (!isa<SpatCompute, SpatComputeBatch>(op))
+    return op->emitError("verifyComputeResultUses: Op is not a SpatCompute/SpatComputeBatch operation");
+  if (!llvm::all_of(op->getResults(), [](Value result) {
+        return llvm::all_of(result.getUsers(), [](Operation* op) {
+          return !(op->getParentOfType<SpatCompute>() || op->getParentOfType<SpatComputeBatch>());
+        });
+      })) {
+    return op->emitError("ComputeResult used directly inside another Compute" );
+  }
+  return success();
+}
+
 LogicalResult SpatCompute::verify() {
   auto& block = getBody().front();
   if (block.mightHaveTerminator()) {
@@ -375,7 +389,8 @@ LogicalResult SpatCompute::verify() {
   for (auto arg : block.getArguments())
     if (arg.use_empty())
       return emitError("ComputeOp block argument is not used");
-
+  if (failed(verifyComputeResultsUses(this->getOperation())))
+    return failure();
   return success();
 }
 
@@ -485,6 +500,8 @@ LogicalResult SpatComputeBatch::verify() {
       return emitError("body block argument type must match input type");
   }
 
+  if (failed(verifyComputeResultsUses(this->getOperation())))
+    return failure();
   return verifyBatchBody(getOperation(), block, getResultTypes(), weightsPerLane);
 }
 
diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
index cb5662d..4e53be2 100644
--- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
@@ -65,6 +65,7 @@ std::optional<ProducerValueRef> getProducerValueRef(Value value) {
   if (!op)
     return std::nullopt;
 
+  //TODO Extract Slice is not the only global non compute operation. There are other legal op
   while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
     value = extract.getSource();
     op = value.getDefiningOp();