Improve verification checks for broadcast op.

tenstorrent · Dec 30, 2024 · f741c50 · f741c50
1 parent 917f8fd
commit f741c50
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/lib/Dialect/TTIR/IR/TTIROps.cpp b/lib/Dialect/TTIR/IR/TTIROps.cpp
@@ -397,6 +397,11 @@ ::mlir::LogicalResult mlir::tt::ttir::BroadcastOp::verify() {
     return emitOpError("Input tensor rank should match output tensor rank.");
   }
 
+  llvm::SmallVector<int64_t, 4> broadcastedShape;
+  if (!OpTrait::util::getBroadcastedShape(
+          inputType.getShape(), outputType.getShape(), broadcastedShape)) {
+    return emitOpError("Operands are not broadcast compatible");
+  }
   return success();
 }