Created
June 24, 2024 03:24
-
-
Save tiandiao123/b290d8baf131892bb254fd503d982fd1 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
// concurrent_queue.cpp | |
#include "concurrent_queue.h" | |
#include <iostream> | |
template <typename T> | |
void ConcurrentQueue<T>::push(T value) { | |
std::lock_guard<std::mutex> lock(mutex); | |
queue.push(std::move(value)); | |
cv.notify_one(); // Notify one waiting thread if it's waiting | |
} | |
template <typename T> | |
bool ConcurrentQueue<T>::try_pop(T& value) { | |
std::lock_guard<std::mutex> lock(mutex); | |
if (queue.empty()) { | |
return false; | |
} | |
value = std::move(queue.front()); | |
queue.pop(); | |
return true; | |
} | |
template <typename T> | |
void ConcurrentQueue<T>::wait_and_pop(T& value) { | |
std::unique_lock<std::mutex> lock(mutex); | |
cv.wait(lock, [this] { return !queue.empty(); }); // Wait until the queue is not empty | |
value = std::move(queue.front()); | |
queue.pop(); | |
} | |
template <typename T> | |
bool ConcurrentQueue<T>::empty() const { | |
std::lock_guard<std::mutex> lock(mutex); | |
return queue.empty(); | |
} | |
// Explicit template instantiation for torch::Tensor | |
template class ConcurrentQueue<torch::Tensor>; | |
void producer(ConcurrentQueue<torch::Tensor>& queue) { | |
for (int i = 0; i < 10; ++i) { | |
auto tensor = torch::rand({2, 2}); // Create a 2x2 random tensor | |
std::cout << "Producing tensor\n"; | |
queue.push(tensor); | |
std::this_thread::sleep_for(std::chrono::seconds(1)); // Simulate some delay | |
} | |
} | |
void consumer(ConcurrentQueue<torch::Tensor>& queue) { | |
for (int i = 0; i < 10; ++i) { | |
torch::Tensor tensor; | |
queue.wait_and_pop(tensor); | |
std::cout << "Consuming tensor\n"; | |
std::cout << tensor << std::endl; | |
std::this_thread::sleep_for(std::chrono::seconds(1)); // Simulate some delay | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment