sglang_v0.5.2/pytorch_2.8.0/c10/test/util/lazy_test.cpp

99 lines
2.3 KiB
C++

#include <atomic>
#include <thread>
#include <vector>
#include <c10/util/Lazy.h>
#include <gtest/gtest.h>
namespace c10_test {
// Long enough not to fit in typical SSO.
const std::string kLongString = "I am a long enough string";
TEST(LazyTest, OptimisticLazy) {
std::atomic<size_t> invocations = 0;
auto factory = [&] {
++invocations;
return kLongString;
};
c10::OptimisticLazy<std::string> s;
constexpr size_t kNumThreads = 16;
std::vector<std::thread> threads;
std::atomic<std::string*> address = nullptr;
threads.reserve(kNumThreads);
for (size_t i = 0; i < kNumThreads; ++i) {
threads.emplace_back([&] {
auto* p = &s.ensure(factory);
auto old = address.exchange(p);
if (old != nullptr) {
// Even racing ensure()s should return a stable reference.
EXPECT_EQ(old, p);
}
});
}
for (auto& t : threads) {
t.join();
}
EXPECT_GE(invocations.load(), 1);
EXPECT_EQ(*address.load(), kLongString);
invocations = 0;
s.reset();
s.ensure(factory);
EXPECT_EQ(invocations.load(), 1);
invocations = 0;
auto sCopy = s;
EXPECT_EQ(sCopy.ensure(factory), kLongString);
EXPECT_EQ(invocations.load(), 0);
auto sMove = std::move(s); // codespell:ignore smove
EXPECT_EQ(sMove.ensure(factory), kLongString); // codespell:ignore smove
EXPECT_EQ(invocations.load(), 0);
// NOLINTNEXTLINE(bugprone-use-after-move)
EXPECT_EQ(s.ensure(factory), kLongString);
EXPECT_EQ(invocations.load(), 1);
invocations = 0;
s = sCopy;
EXPECT_EQ(s.ensure(factory), kLongString);
EXPECT_EQ(invocations.load(), 0);
s = std::move(sCopy);
EXPECT_EQ(s.ensure(factory), kLongString);
EXPECT_EQ(invocations.load(), 0);
}
TEST(LazyTest, PrecomputedLazyValue) {
static const std::string kLongString = "I am a string";
EXPECT_EQ(
std::make_shared<c10::PrecomputedLazyValue<std::string>>(kLongString)
->get(),
kLongString);
}
TEST(LazyTest, OptimisticLazyValue) {
static const std::string kLongString = "I am a string";
class LazyString : public c10::OptimisticLazyValue<std::string> {
std::string compute() const override {
return kLongString;
}
};
auto ls = std::make_shared<LazyString>();
EXPECT_EQ(ls->get(), kLongString);
// Returned reference should be stable.
EXPECT_EQ(&ls->get(), &ls->get());
}
} // namespace c10_test