[Mosaic] Add hasVectorOperandsOrResults utility function

PiperOrigin-RevId: 820222017

Author: tlongeri

jaxlib/mosaic/dialect/tpu/util.cc

@@ -371,4 +371,13 @@ SmallVector<Operation *> getNontrivialTransitiveUsers(Value v) {
   return users;
 }
 
+bool hasVectorOperandsOrResults(Operation& op) {
+  for (Value value : llvm::concat<Value>(op.getOperands(), op.getResults())) {
+    if (isa<VectorType>(value.getType())) {
+      return true;
+    }
+  }
+  return false;
+}
+
 }  // namespace mlir::tpu

jaxlib/mosaic/dialect/tpu/util.h

@@ -299,6 +299,8 @@ std::optional<int64_t> getIntConst(Value v);
 // results.
 SmallVector<Operation *> getNontrivialTransitiveUsers(Value v);
 
+bool hasVectorOperandsOrResults(Operation& op);
+
 // Return a mod b for a, b > 0, but adjusted to return b when a mod b == 0 such
 // that the result is strictly positive.
 template <typename U, typename V>