Changed condition for op being TTSchedulable

tenstorrent · Nov 7, 2024 · 542797b · 542797b
1 parent 4ec19f1
commit 542797b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/Scheduler/Scheduler.cpp b/lib/Scheduler/Scheduler.cpp
@@ -22,7 +22,8 @@ bool isTTIROp(mlir::Operation *op) {
 }
 
 bool isTTShedulableOp(mlir::Operation *op) {
-  return (isTTNNOp(op) || isTTIROp(op)) && (not isa<func::ReturnOp>(op));
+  return (isTTNNOp(op) || isTTIROp(op)) &&
+         (not op->hasTrait<mlir::OpTrait::ReturnLike>());
 }
 
 // Init the dependencies map of all ops which are TTIR ops