wehub-resource-sync
115 lines
3.3 KiB
C++
115 lines
3.3 KiB
C++
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
#include "tensorflow/lite/delegates/utils/sync_fence.h"
|
|
|
|
#include <poll.h>
|
|
|
|
#include <cerrno>
|
|
#include <cstddef>
|
|
#include <optional>
|
|
#include <variant>
|
|
#include <vector>
|
|
|
|
#include "absl/types/span.h"
|
|
#include "tensorflow/lite/delegates/utils/ret_macros.h"
|
|
#include "tensorflow/lite/minimal_logging.h"
|
|
|
|
namespace tflite::delegates::utils {
|
|
namespace {
|
|
|
|
// Returns how many file descriptors have been signalled, or an error.
|
|
// Note that the implementation is loosely based on Android's libsync.
|
|
std::optional<size_t> PollFds(absl::Span<const int> fds, bool block) {
|
|
constexpr auto kError = std::optional<size_t>();
|
|
|
|
const int timeout = block ? -1 : 0;
|
|
|
|
if (fds.empty()) {
|
|
return 0;
|
|
}
|
|
|
|
std::vector<struct pollfd> pfds;
|
|
pfds.reserve(fds.size());
|
|
for (int fd : fds) {
|
|
const struct pollfd pfd = {
|
|
fd, // .fd
|
|
POLLIN, // .events
|
|
};
|
|
pfds.push_back(pfd);
|
|
}
|
|
|
|
while (true) {
|
|
const int ret = poll(pfds.data(), pfds.size(), timeout);
|
|
|
|
// Handle redo
|
|
if (ret == -1 && (errno == EINTR || errno == EAGAIN)) {
|
|
continue;
|
|
}
|
|
|
|
// Handle error
|
|
TFLITE_RET_CHECK(ret >= 0u, "Poll failed", kError);
|
|
|
|
// Handle none ready
|
|
if (ret == 0) {
|
|
TFLITE_RET_CHECK(!block, "", kError);
|
|
return 0;
|
|
}
|
|
|
|
// Count how many fds have been signalled, setting them to -1 so they are
|
|
// not queried on subsequent passes of the loop.
|
|
size_t signalled_count = 0;
|
|
for (auto& fd : pfds) {
|
|
if (fd.revents & (POLLERR | POLLNVAL)) {
|
|
TFLITE_LOG_PROD(TFLITE_LOG_ERROR, "invalid fd to poll");
|
|
return {};
|
|
}
|
|
if (fd.fd == -1 || (fd.revents & POLLIN)) {
|
|
fd.fd = -1;
|
|
signalled_count++;
|
|
}
|
|
}
|
|
|
|
// If we are blocking and there are any fds that have not yet been
|
|
// signalled, poll again on the next loop.
|
|
if (block && signalled_count != fds.size()) {
|
|
continue;
|
|
}
|
|
|
|
return signalled_count;
|
|
}
|
|
|
|
TFLITE_ABORT_CHECK(false,
|
|
"The code should never reach this point"); // Crash OK
|
|
return {};
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::optional<std::monostate> WaitForAllFds(absl::Span<const int> fds) {
|
|
constexpr auto kError = std::optional<std::monostate>();
|
|
TFLITE_ASSIGN_OR_RETURN(const size_t signalled_count,
|
|
PollFds(fds, /*block=*/true), kError);
|
|
TFLITE_RET_CHECK(signalled_count == fds.size(), "", kError);
|
|
return std::monostate{};
|
|
}
|
|
|
|
std::optional<bool> AreAllFdsSignalled(absl::Span<const int> fds) {
|
|
TFLITE_ASSIGN_OR_RETURN(const size_t signalled_count,
|
|
PollFds(fds, /*block=*/false), std::optional<bool>());
|
|
return signalled_count == fds.size();
|
|
}
|
|
|
|
} // namespace tflite::delegates::utils
|